<commit_before>// Copyright 2015-2017 trivago GmbH
//
// Licensed 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 producer

import (
	"context"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core"
	"github.com/trivago/tgo"
	"github.com/trivago/tgo/tcontainer"
	"gopkg.in/olivere/elastic.v5"
	"net/http"
	"sync"
	"syscall"
	"time"
)

// ElasticSearch producer plugin
//
// The ElasticSearch producer sends messages to elastic search using the bulk
// http API. The producer expects a json payload.
//
// Parameters
//
// - Retry/Count: Set the amount of retries before a Elasticsearch request
// fail finally.
// By default this parameter is set to "3".
//
// - Retry/TimeToWaitSec: This value denotes the time in seconds after which a
// failed dataset will be  transmitted again.
// By default this parameter is set to "3".
//
// - SetGzip: This value enables or disables gzip compression for Elasticsearch
// requests (disabled by default). This option is used one to one for the library
// package. See http://godoc.org/gopkg.in/olivere/elastic.v5#SetGzip
// By default this parameter is set to "false".
//
// - Servers: This value defines a list of servers to connect to.
//
// - User: This value used as the username for the elasticsearch server.
// By default this parameter is set to "".
//
// - Password: This value used as the password for the elasticsearch server.
// By default this parameter is set to "".
//
// - StreamProperties: This value defines the mapping and settings for each stream.
// As index use the stream name here.
//
// - StreamProperties/<streamName>/Index: The value defines the Elasticsearch
// index used for the stream.
//
// - StreamProperties/<streamName>/Type: This value defines the document type
// used for the stream.
//
// - StreamProperties/<streamName>/DayBasedIndex: This value can be set to "true"
// to append the date of the message to the index as in "<index>_YYYY-MM-DD".
// NOTE: This setting incurs a performance penalty because it is necessary to
// check if an index exists for each message!
// By default this parameter is set to "false".
//
// - StreamProperties/<streamName>/Mapping: This value is a map which is used
// for the document field mapping. As document type, the already defined type is
// reused for the field mapping. See
// https://www.elastic.co/guide/en/elasticsearch/reference/5.4/indices-create-index.html#mappings
//
// - StreamProperties/<streamName>/Settings: This value is a map which is used
// for the index settings. See
// https://www.elastic.co/guide/en/elasticsearch/reference/5.4/indices-create-index.html#mappings
//
// Examples
//
// This example starts a simple twitter example producer for local running ElasticSearch:
//
//  producerElasticSearch:
//    Type: producer.ElasticSearch
//    Streams: tweets_stream
//    SetGzip: true
//    Servers:
//      - http://127.0.0.1:9200
//    StreamProperties:
//      tweets_stream:
//        Index: twitter
//        DayBasedIndex: true
//        Type: tweet
//        Mapping:
//          # index mapping for payload
//          user: keyword
//          message: text
//        Settings:
//          number_of_shards: 1
//          number_of_replicas: 1
type ElasticSearch struct {
	core.BatchedProducer `gollumdoc:"embed_type"`
	connection           elasticConnection
	indexMap             map[core.MessageStreamID]*indexMapItem
}

type indexMapItem struct {
	name        string
	typeName    string
	settings    *elasticIndex
	useDayIndex bool
}

func newIndexMapItem() *indexMapItem {
	return &indexMapItem{
		"",
		"",
		nil,
		false,
	}
}

type elasticIndex struct {
	Settings map[string]interface{}    `json:"settings,omitempty"`
	Mappings map[string]elasticMapping `json:"mappings,omitempty"`
}

type elasticMapping struct {
	Properties map[string]elasticType `json:"properties,omitempty"`
}

type elasticType struct {
	TypeName interface{} `json:"type"`
}

func newElasticIndex(property tcontainer.MarshalMap) *elasticIndex {
	elType, _ := property.String("type")
	mapping, _ := property.MarshalMap("mapping")
	settings, _ := property.MarshalMap("settings")

	elMappings := elasticMapping{
		Properties: make(map[string]elasticType),
	}

	for fieldName := range mapping {
		typeName, _ := mapping.String(fieldName)
		elMappings.Properties[fieldName] = elasticType{TypeName: typeName}
	}

	// init elasticIndex instance
	elIndex := elasticIndex{
		Settings: make(map[string]interface{}),
		Mappings: make(map[string]elasticMapping),
	}

	elIndex.Mappings[elType] = elMappings
	elIndex.Settings = settings

	return &elIndex
}

type indexExists uint8

const (
	indexExistsFalse indexExists = iota
	indexExistsTrue
	indexExistsAuto
)

func init() {
	core.TypeRegistry.Register(ElasticSearch{})
}

// Configure initializes this producer with values from a plugin config.
func (prod *ElasticSearch) Configure(conf core.PluginConfigReader) {
	prod.connection.servers = conf.GetStringArray("Servers", []string{"http://127.0.0.1:9200"})
	prod.connection.user = conf.GetString("User", "")
	prod.connection.password = conf.GetString("Password", "")
	prod.connection.setGzip = conf.GetBool("SetGzip", false)
	prod.connection.isConnectedStatus = false

	prod.configureIndexSettings(conf.GetMap("StreamProperties", tcontainer.NewMarshalMap()), conf.Errors)
	prod.configureRetrySettings(conf.GetInt("Retry/Count", 3), conf.GetInt("Retry/TimeToWaitSec", 3))
}

// Produce starts the producer
func (prod *ElasticSearch) Produce(workers *sync.WaitGroup) {
	defer prod.WorkerDone()

	prod.initIndex()
	prod.AddMainWorker(workers)
	prod.BatchMessageLoop(workers, prod.submitMessages)
}

func (prod *ElasticSearch) configureRetrySettings(retry, timeToWaitSec int64) {
	prod.connection.retrier.retry = int(retry)
	prod.connection.retrier.backoff = elastic.NewConstantBackoff(time.Duration(timeToWaitSec) * time.Second)
	prod.connection.retrier.logger = prod.Logger.WithField("Scope", "connection.retrier")

	prod.Logger.Debugf("Using retrier with a retry count of '%d' and a time to wait of '%d' sec", retry, timeToWaitSec)
}

func (prod *ElasticSearch) configureIndexSettings(properties tcontainer.MarshalMap, errors *tgo.ErrorStack) {
	prod.indexMap = map[core.MessageStreamID]*indexMapItem{}

	if len(properties) <= 0 {
		prod.Logger.Error("No stream configuration found. Please check your config.")
		return
	}

	for streamName := range properties {
		streamID := core.GetStreamID(streamName)
		indexMapItem := newIndexMapItem()

		property, err := properties.MarshalMap(streamName)
		if err != nil {
			prod.Logger.Errorf("no configuration found for stream '%s'. Please check your config.", streamName)
			errors.Push(err)
			continue
		}

		indexName, err := property.String("index")
		if err != nil {
			prod.Logger.Errorf("no index configured for stream '%s'. Please check your config.", streamName)
			errors.Push(err)
			continue
		}

		dayBasedIndex, _ := property.Bool("daybasedindex")
		if dayBasedIndex {
			currentTime := time.Now()
			indexName += currentTime.Format("_2006-01-02")
			indexMapItem.useDayIndex = true
		}

		typeName, err := property.String("type")
		if err != nil {
			prod.Logger.Errorf("no data type configured for stream '%s'. Please check your config.", streamName)
		}

		indexMapItem.settings = newElasticIndex(property)
		indexMapItem.typeName = typeName
		indexMapItem.name = indexName

		prod.indexMap[streamID] = indexMapItem
	}
}

func (prod *ElasticSearch) getClient() (*elastic.Client, error) {
	if !prod.connection.isConnected() {
		err := prod.connection.connect()
		if err != nil {
			prod.Logger.Error("Error during connection: ", err)
			return nil, err
		}
	}

	return prod.connection.client, nil
}

func (prod *ElasticSearch) initIndex() {
	for _, indexMapItem := range prod.indexMap {
		prod.createIndex(indexMapItem.name, indexMapItem.settings, indexExistsAuto)
	}
}

func (prod *ElasticSearch) createIndex(indexName string, elIndex *elasticIndex, indexExistsCheck indexExists) error {
	client, err := prod.getClient()
	if err != nil {
		return err
	}

	ctx := context.Background()

	// avoid unnecessary request for checking index (see submitMessage())
	var exists bool
	switch indexExistsCheck {
	case indexExistsTrue:
		exists = true
		break
	case indexExistsFalse:
		exists = false
		break
	case indexExistsAuto:
		exists, err = prod.isIndexExists(ctx, indexName, client)
		if err != nil {
			return err
		}
	}

	if !exists {
		_, err = client.CreateIndex(indexName).Do(ctx)
		if err != nil {
			prod.Logger.Errorln("Issue during creating index: ", err)
			return err
		}

		prod.Logger.Debugf("Created index '%s'\n", indexName)
	}

	for typeName, properties := range elIndex.Mappings {
		mapping := client.PutMapping()
		mapping.Index(indexName)
		mapping.Type(typeName)

		json := map[string]interface{}{typeName: properties}
		mapping.BodyJson(json)

		_, err := mapping.Do(ctx)
		if err != nil {
			prod.Logger.Errorln("Issue during creating index mapping: ", err)
		}
	}

	return err
}

func (prod *ElasticSearch) isIndexExists(ctx context.Context, indexName string, client *elastic.Client) (bool, error) {
	exists, err := client.IndexExists(indexName).Do(ctx)
	prod.Logger.Debugf("Index %s exists: %v", indexName, exists)
	if err != nil {
		prod.Logger.Errorln("Issue during checking index: ", err)
		return false, err
	}

	return exists, nil
}

func (prod *ElasticSearch) submitMessages() core.AssemblyFunc {
	return func(messages []*core.Message) {
		client, err := prod.getClient()
		if err != nil {
			prod.Logger.Error("Sending messages failed: ", err)
			return
		}

		bulkRequest := client.Bulk()
		ctx := context.Background()

		for _, msg := range messages {
			indexMapItem := prod.indexMap[msg.GetStreamID()]

			if indexMapItem.useDayIndex {
				// create daily index ...
				indexExists, _ := prod.isIndexExists(ctx, indexMapItem.name, client)
				if !indexExists {
					prod.createIndex(indexMapItem.name, indexMapItem.settings, indexExistsFalse)
				}
			}

			bulkIndexRequest := elastic.NewBulkIndexRequest()
			bulkIndexRequest.Index(indexMapItem.name).
				Type(indexMapItem.typeName).
				Doc(msg.String())

			bulkRequest.Add(bulkIndexRequest)
		}

		// NumberOfActions contains the number of requests in a bulk<commit_msg>fixed nil pointer panic on unmapped streams<commit_after>// Copyright 2015-2017 trivago GmbH
//
// Licensed 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 producer

import (
	"context"
	"net/http"
	"sync"
	"syscall"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core"
	"github.com/trivago/tgo"
	"github.com/trivago/tgo/tcontainer"
	"gopkg.in/olivere/elastic.v5"
)

// ElasticSearch producer plugin
//
// The ElasticSearch producer sends messages to elastic search using the bulk
// http API. The producer expects a json payload.
//
// Parameters
//
// - Retry/Count: Set the amount of retries before a Elasticsearch request
// fail finally.
// By default this parameter is set to "3".
//
// - Retry/TimeToWaitSec: This value denotes the time in seconds after which a
// failed dataset will be  transmitted again.
// By default this parameter is set to "3".
//
// - SetGzip: This value enables or disables gzip compression for Elasticsearch
// requests (disabled by default). This option is used one to one for the library
// package. See http://godoc.org/gopkg.in/olivere/elastic.v5#SetGzip
// By default this parameter is set to "false".
//
// - Servers: This value defines a list of servers to connect to.
//
// - User: This value used as the username for the elasticsearch server.
// By default this parameter is set to "".
//
// - Password: This value used as the password for the elasticsearch server.
// By default this parameter is set to "".
//
// - StreamProperties: This value defines the mapping and settings for each stream.
// As index use the stream name here.
//
// - StreamProperties/<streamName>/Index: The value defines the Elasticsearch
// index used for the stream.
//
// - StreamProperties/<streamName>/Type: This value defines the document type
// used for the stream.
//
// - StreamProperties/<streamName>/DayBasedIndex: This value can be set to "true"
// to append the date of the message to the index as in "<index>_YYYY-MM-DD".
// NOTE: This setting incurs a performance penalty because it is necessary to
// check if an index exists for each message!
// By default this parameter is set to "false".
//
// - StreamProperties/<streamName>/Mapping: This value is a map which is used
// for the document field mapping. As document type, the already defined type is
// reused for the field mapping. See
// https://www.elastic.co/guide/en/elasticsearch/reference/5.4/indices-create-index.html#mappings
//
// - StreamProperties/<streamName>/Settings: This value is a map which is used
// for the index settings. See
// https://www.elastic.co/guide/en/elasticsearch/reference/5.4/indices-create-index.html#mappings
//
// Examples
//
// This example starts a simple twitter example producer for local running ElasticSearch:
//
//  producerElasticSearch:
//    Type: producer.ElasticSearch
//    Streams: tweets_stream
//    SetGzip: true
//    Servers:
//      - http://127.0.0.1:9200
//    StreamProperties:
//      tweets_stream:
//        Index: twitter
//        DayBasedIndex: true
//        Type: tweet
//        Mapping:
//          # index mapping for payload
//          user: keyword
//          message: text
//        Settings:
//          number_of_shards: 1
//          number_of_replicas: 1
type ElasticSearch struct {
	core.BatchedProducer `gollumdoc:"embed_type"`
	connection           elasticConnection
	indexMap             map[core.MessageStreamID]*indexMapItem
}

type indexMapItem struct {
	name        string
	typeName    string
	settings    *elasticIndex
	useDayIndex bool
}

func newIndexMapItem() *indexMapItem {
	return &indexMapItem{
		"",
		"",
		nil,
		false,
	}
}

type elasticIndex struct {
	Settings map[string]interface{}    `json:"settings,omitempty"`
	Mappings map[string]elasticMapping `json:"mappings,omitempty"`
}

type elasticMapping struct {
	Properties map[string]elasticType `json:"properties,omitempty"`
}

type elasticType struct {
	TypeName interface{} `json:"type"`
}

func newElasticIndex(property tcontainer.MarshalMap) *elasticIndex {
	elType, _ := property.String("type")
	mapping, _ := property.MarshalMap("mapping")
	settings, _ := property.MarshalMap("settings")

	elMappings := elasticMapping{
		Properties: make(map[string]elasticType),
	}

	for fieldName := range mapping {
		typeName, _ := mapping.String(fieldName)
		elMappings.Properties[fieldName] = elasticType{TypeName: typeName}
	}

	// init elasticIndex instance
	elIndex := elasticIndex{
		Settings: make(map[string]interface{}),
		Mappings: make(map[string]elasticMapping),
	}

	elIndex.Mappings[elType] = elMappings
	elIndex.Settings = settings

	return &elIndex
}

type indexExists uint8

const (
	indexExistsFalse indexExists = iota
	indexExistsTrue
	indexExistsAuto
)

func init() {
	core.TypeRegistry.Register(ElasticSearch{})
}

// Configure initializes this producer with values from a plugin config.
func (prod *ElasticSearch) Configure(conf core.PluginConfigReader) {
	prod.connection.servers = conf.GetStringArray("Servers", []string{"http://127.0.0.1:9200"})
	prod.connection.user = conf.GetString("User", "")
	prod.connection.password = conf.GetString("Password", "")
	prod.connection.setGzip = conf.GetBool("SetGzip", false)
	prod.connection.isConnectedStatus = false

	prod.configureIndexSettings(conf.GetMap("StreamProperties", tcontainer.NewMarshalMap()), conf.Errors)
	prod.configureRetrySettings(conf.GetInt("Retry/Count", 3), conf.GetInt("Retry/TimeToWaitSec", 3))
}

// Produce starts the producer
func (prod *ElasticSearch) Produce(workers *sync.WaitGroup) {
	defer prod.WorkerDone()

	prod.initIndex()
	prod.AddMainWorker(workers)
	prod.BatchMessageLoop(workers, prod.submitMessages)
}

func (prod *ElasticSearch) configureRetrySettings(retry, timeToWaitSec int64) {
	prod.connection.retrier.retry = int(retry)
	prod.connection.retrier.backoff = elastic.NewConstantBackoff(time.Duration(timeToWaitSec) * time.Second)
	prod.connection.retrier.logger = prod.Logger.WithField("Scope", "connection.retrier")

	prod.Logger.Debugf("Using retrier with a retry count of '%d' and a time to wait of '%d' sec", retry, timeToWaitSec)
}

func (prod *ElasticSearch) configureIndexSettings(properties tcontainer.MarshalMap, errors *tgo.ErrorStack) {
	prod.indexMap = map[core.MessageStreamID]*indexMapItem{}

	if len(properties) <= 0 {
		prod.Logger.Error("No stream configuration found. Please check your config.")
		return
	}

	for streamName := range properties {
		streamID := core.GetStreamID(streamName)
		indexMapItem := newIndexMapItem()

		property, err := properties.MarshalMap(streamName)
		if err != nil {
			prod.Logger.Errorf("no configuration found for stream '%s'. Please check your config.", streamName)
			errors.Push(err)
			continue
		}

		indexName, err := property.String("index")
		if err != nil {
			prod.Logger.Errorf("no index configured for stream '%s'. Please check your config.", streamName)
			errors.Push(err)
			continue
		}

		dayBasedIndex, _ := property.Bool("daybasedindex")
		if dayBasedIndex {
			currentTime := time.Now()
			indexName += currentTime.Format("_2006-01-02")
			indexMapItem.useDayIndex = true
		}

		typeName, err := property.String("type")
		if err != nil {
			prod.Logger.Errorf("no data type configured for stream '%s'. Please check your config.", streamName)
		}

		indexMapItem.settings = newElasticIndex(property)
		indexMapItem.typeName = typeName
		indexMapItem.name = indexName

		prod.indexMap[streamID] = indexMapItem
	}
}

func (prod *ElasticSearch) getClient() (*elastic.Client, error) {
	if !prod.connection.isConnected() {
		err := prod.connection.connect()
		if err != nil {
			prod.Logger.Error("Error during connection: ", err)
			return nil, err
		}
	}

	return prod.connection.client, nil
}

func (prod *ElasticSearch) initIndex() {
	for _, indexMapItem := range prod.indexMap {
		prod.createIndex(indexMapItem.name, indexMapItem.settings, indexExistsAuto)
	}
}

func (prod *ElasticSearch) createIndex(indexName string, elIndex *elasticIndex, indexExistsCheck indexExists) error {
	client, err := prod.getClient()
	if err != nil {
		return err
	}

	ctx := context.Background()

	// avoid unnecessary request for checking index (see submitMessage())
	var exists bool
	switch indexExistsCheck {
	case indexExistsTrue:
		exists = true
		break
	case indexExistsFalse:
		exists = false
		break
	case indexExistsAuto:
		exists, err = prod.isIndexExists(ctx, indexName, client)
		if err != nil {
			return err
		}
	}

	if !exists {
		_, err = client.CreateIndex(indexName).Do(ctx)
		if err != nil {
			prod.Logger.Errorln("Issue during creating index: ", err)
			return err
		}

		prod.Logger.Debugf("Created index '%s'\n", indexName)
	}

	for typeName, properties := range elIndex.Mappings {
		mapping := client.PutMapping()
		mapping.Index(indexName)
		mapping.Type(typeName)

		json := map[string]interface{}{typeName: properties}
		mapping.BodyJson(json)

		_, err := mapping.Do(ctx)
		if err != nil {
			prod.Logger.Errorln("Issue during creating index mapping: ", err)
		}
	}

	return err
}

func (prod *ElasticSearch) isIndexExists(ctx context.Context, indexName string, client *elastic.Client) (bool, error) {
	exists, err := client.IndexExists(indexName).Do(ctx)
	prod.Logger.Debugf("Index %s exists: %v", indexName, exists)
	if err != nil {
		prod.Logger.Errorln("Issue during checking index: ", err)
		return false, err
	}

	return exists, nil
}

func (prod *ElasticSearch) submitMessages() core.AssemblyFunc {
	return func(messages []*core.Message) {
		client, err := prod.getClient()
		if err != nil {
			prod.Logger.Error("Sending messages failed: ", err)
			return
		}

		bulkRequest := client.Bulk()
		ctx := context.Background()

		for _, msg := range messages {
			indexMapItem, isSet := prod.indexMap[msg.GetStreamID()]

			if !isSet {
				prod.Logger.Warningf("No index setting for stream %s", msg.GetStreamID().GetName())
				continue
			}

			if indexMapItem.useDayIndex {
				// create daily index ...
				indexExists, _ := prod.isIndexExists(ctx, indexMapItem.name, client)
				if !indexExists {
					prod.createIndex(indexMapItem.name, indexMapItem.settings, indexExistsFalse)
				}
			}

			bulkIndexRequest := elastic.NewBulkIndexRequest()
			bulkIndexRequest.Index(indexMapItem.name).
				Type(indexMapItem.typeName).
				Doc(msg.String())

			bulkRequest.Add(bulkIndexRequest)
		}

		// NumberOfActions contains the number of requests in a bulk
{'fraction_non_alphanumeric': 0.11712126796903796, 'fraction_numerical': 0.005160339107998526, 'mean_word_length': 3.8775280898876403, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 15, 'https://': 4, 'lorem ipsum': 0, 'www.': 6, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5772555', 'n_tokens_mistral': 6948, 'n_tokens_neox': 6508, 'n_words': 1878}