youtubebeat/vendor/github.com/elastic/beats/libbeat/outputs/elasticsearch/bulkapi.go

166 lines
3.8 KiB
Go

// Licensed to Elasticsearch B.V. under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Elasticsearch B.V. licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
package elasticsearch
import (
"bytes"
"io"
"io/ioutil"
"net/http"
"strings"
)
// MetaBuilder creates meta data for bulk requests
type MetaBuilder func(interface{}) interface{}
type bulkRequest struct {
requ *http.Request
}
type bulkResult struct {
raw []byte
}
// Bulk performs many index/delete operations in a single API call.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
func (conn *Connection) Bulk(
index, docType string,
params map[string]string, body []interface{},
) (*QueryResult, error) {
return conn.BulkWith(index, docType, params, nil, body)
}
// BulkWith creates a HTTP request containing a bunch of operations and send
// them to Elasticsearch. The request is retransmitted up to max_retries before
// returning an error.
func (conn *Connection) BulkWith(
index string,
docType string,
params map[string]string,
metaBuilder MetaBuilder,
body []interface{},
) (*QueryResult, error) {
if len(body) == 0 {
return nil, nil
}
enc := conn.encoder
enc.Reset()
if err := bulkEncode(enc, metaBuilder, body); err != nil {
return nil, err
}
requ, err := newBulkRequest(conn.URL, index, docType, params, enc)
if err != nil {
return nil, err
}
_, result, err := conn.sendBulkRequest(requ)
if err != nil {
return nil, err
}
return readQueryResult(result.raw)
}
func newBulkRequest(
urlStr string,
index, docType string,
params map[string]string,
body bodyEncoder,
) (*bulkRequest, error) {
path, err := makePath(index, docType, "_bulk")
if err != nil {
return nil, err
}
url := addToURL(urlStr, path, "", params)
var reader io.Reader
if body != nil {
reader = body.Reader()
}
requ, err := http.NewRequest("POST", url, reader)
if err != nil {
return nil, err
}
if body != nil {
body.AddHeader(&requ.Header)
}
return &bulkRequest{
requ: requ,
}, nil
}
func (r *bulkRequest) Reset(body bodyEncoder) {
bdy := body.Reader()
rc, ok := bdy.(io.ReadCloser)
if !ok && body != nil {
rc = ioutil.NopCloser(bdy)
}
switch v := bdy.(type) {
case *bytes.Buffer:
r.requ.ContentLength = int64(v.Len())
case *bytes.Reader:
r.requ.ContentLength = int64(v.Len())
case *strings.Reader:
r.requ.ContentLength = int64(v.Len())
}
r.requ.Header = http.Header{}
r.requ.Body = rc
body.AddHeader(&r.requ.Header)
}
func (conn *Connection) sendBulkRequest(requ *bulkRequest) (int, bulkResult, error) {
status, resp, err := conn.execHTTPRequest(requ.requ)
if err != nil {
return status, bulkResult{}, err
}
result, err := readBulkResult(resp)
return status, result, err
}
func readBulkResult(obj []byte) (bulkResult, error) {
return bulkResult{obj}, nil
}
func bulkEncode(out bulkWriter, metaBuilder MetaBuilder, body []interface{}) error {
if metaBuilder == nil {
for _, obj := range body {
if err := out.AddRaw(obj); err != nil {
debugf("Failed to encode message: %s", err)
return err
}
}
} else {
for _, obj := range body {
meta := metaBuilder(obj)
if err := out.Add(meta, obj); err != nil {
debugf("Failed to encode event (dropping event): %s", err)
}
}
}
return nil
}