<commit_before>package influxdb_output

import (
	"context"
	"encoding/json"
	"log"
	"os"
	"time"

	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/influxdata/influxdb-client-go/v2/api"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	defaultURL               = "http://localhost:8086"
	defaultBatchSize         = 1000
	defaultFlushTimer        = 10 * time.Second
	defaultHealthCheckPeriod = 30 * time.Second
)

func init() {
	outputs.Register("influxdb", func() outputs.Output {
		return &InfluxDBOutput{
			Cfg: &Config{},
		}
	})
}

type InfluxDBOutput struct {
	Cfg      *Config
	client   influxdb2.Client
	writer   api.WriteAPI
	metrics  []prometheus.Collector
	logger   *log.Logger
	cancelFn context.CancelFunc
}
type Config struct {
	URL               string        `mapstructure:"url,omitempty"`
	Org               string        `mapstructure:"org,omitempty"`
	Bucket            string        `mapstructure:"bucket,omitempty"`
	Token             string        `mapstructure:"token,omitempty"`
	BatchSize         uint          `mapstructure:"batch_size,omitempty"`
	FlushTimer        time.Duration `mapstructure:"flush_timer,omitempty"`
	UseGzip           bool          `mapstructure:"use_gzip,omitempty"`
	HealthCheckPeriod time.Duration `mapstructure:"health_check_period,omitempty"`
	Debug             bool          `mapstructure:"debug,omitempty"`
}

func (k *InfluxDBOutput) String() string {
	b, err := json.Marshal(k)
	if err != nil {
		return ""
	}
	return string(b)
}
func (i *InfluxDBOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	var ctx context.Context
	ctx, i.cancelFn = context.WithCancel(context.Background())
	err := mapstructure.Decode(cfg, i.Cfg)
	if err != nil {
		return err
	}
	if i.Cfg.URL == "" {
		i.Cfg.URL = defaultURL
	}
	if i.Cfg.BatchSize == 0 {
		i.Cfg.BatchSize = defaultBatchSize
	}
	if i.Cfg.FlushTimer == 0 {
		i.Cfg.FlushTimer = defaultFlushTimer
	}
	if i.Cfg.HealthCheckPeriod == 0 {
		i.Cfg.HealthCheckPeriod = defaultHealthCheckPeriod
	}
	i.logger = log.New(os.Stderr, "influxdb_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		i.logger.SetOutput(logger.Writer())
		i.logger.SetFlags(logger.Flags())
	}
	opts := influxdb2.DefaultOptions().
		SetUseGZip(i.Cfg.UseGzip).
		SetBatchSize(i.Cfg.BatchSize).
		SetFlushInterval(uint(i.Cfg.FlushTimer.Milliseconds()))
	if i.Cfg.Debug {
		opts.SetLogLevel(3)
	}
	i.client = influxdb2.NewClientWithOptions(i.Cfg.URL, i.Cfg.Token, opts)
	// start influx health check
	go i.healthCheck(ctx)
	i.writer = i.client.WriteAPI(i.Cfg.Org, i.Cfg.Bucket)
	i.logger.Printf("initialized influxdb write API: %s", i.String())
	// start influxdb error logs
	go func() {
		for err = range i.writer.Errors() {
			i.logger.Printf("writeAPI error: %v", err)
		}
	}()
	return nil
}

func (i *InfluxDBOutput) Write(rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := collector.ResponseToEventMsgs(measName, rsp, meta)
		if err != nil {
			i.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			i.writer.WritePoint(influxdb2.NewPoint(ev.Name, ev.Tags, ev.Values, time.Unix(0, ev.Timestamp)))
		}
	}
}

func (i *InfluxDBOutput) Close() error {
	i.logger.Printf("flushing data...")
	i.writer.Flush()
	i.logger.Printf("closing client...")
	i.client.Close()
	i.cancelFn()
	i.logger.Printf("closed.")
	return nil
}
func (i *InfluxDBOutput) Metrics() []prometheus.Collector { return i.metrics }

func (i *InfluxDBOutput) healthCheck(ctx context.Context) {
	i.health(ctx)
	ticker := time.NewTicker(i.Cfg.HealthCheckPeriod)
	for {
		select {
		case <-ticker.C:
			i.health(ctx)
		case <-ctx.Done():
			return
		}
	}
}

func (i *InfluxDBOutput) health(ctx context.Context) {
	res, err := i.client.Health(ctx)
	if err != nil {
		i.logger.Printf("failed to initialize influxdb output: %v", err)
		return
	}
	if res != nil {
		i.logger.Printf("health check result: name=%s, version=%s, status=%s, message=%s, commit=%s, checks=%v",
			res.Name, *res.Version, res.Status, *res.Message, *res.Commit, *res.Checks)
		return
	}
	i.logger.Print("health check result is nil")
}<commit_msg>add enable_tls | change healthcheck log message<commit_after>package influxdb_output

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"log"
	"os"
	"time"

	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/influxdata/influxdb-client-go/v2/api"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	defaultURL               = "http://localhost:8086"
	defaultBatchSize         = 1000
	defaultFlushTimer        = 10 * time.Second
	defaultHealthCheckPeriod = 30 * time.Second
)

func init() {
	outputs.Register("influxdb", func() outputs.Output {
		return &InfluxDBOutput{
			Cfg: &Config{},
		}
	})
}

type InfluxDBOutput struct {
	Cfg      *Config
	client   influxdb2.Client
	writer   api.WriteAPI
	metrics  []prometheus.Collector
	logger   *log.Logger
	cancelFn context.CancelFunc
}
type Config struct {
	URL               string        `mapstructure:"url,omitempty"`
	Org               string        `mapstructure:"org,omitempty"`
	Bucket            string        `mapstructure:"bucket,omitempty"`
	Token             string        `mapstructure:"token,omitempty"`
	BatchSize         uint          `mapstructure:"batch_size,omitempty"`
	FlushTimer        time.Duration `mapstructure:"flush_timer,omitempty"`
	UseGzip           bool          `mapstructure:"use_gzip,omitempty"`
	EnableTLS         bool          `mapstructure:"enable_tls,omitempty"`
	HealthCheckPeriod time.Duration `mapstructure:"health_check_period,omitempty"`
	Debug             bool          `mapstructure:"debug,omitempty"`
}

func (k *InfluxDBOutput) String() string {
	b, err := json.Marshal(k)
	if err != nil {
		return ""
	}
	return string(b)
}
func (i *InfluxDBOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	var ctx context.Context
	ctx, i.cancelFn = context.WithCancel(context.Background())
	err := mapstructure.Decode(cfg, i.Cfg)
	if err != nil {
		return err
	}
	if i.Cfg.URL == "" {
		i.Cfg.URL = defaultURL
	}
	if i.Cfg.BatchSize == 0 {
		i.Cfg.BatchSize = defaultBatchSize
	}
	if i.Cfg.FlushTimer == 0 {
		i.Cfg.FlushTimer = defaultFlushTimer
	}
	if i.Cfg.HealthCheckPeriod == 0 {
		i.Cfg.HealthCheckPeriod = defaultHealthCheckPeriod
	}
	i.logger = log.New(os.Stderr, "influxdb_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		i.logger.SetOutput(logger.Writer())
		i.logger.SetFlags(logger.Flags())
	}
	opts := influxdb2.DefaultOptions().
		SetUseGZip(i.Cfg.UseGzip).
		SetBatchSize(i.Cfg.BatchSize).
		SetFlushInterval(uint(i.Cfg.FlushTimer.Milliseconds()))
	if i.Cfg.EnableTLS {
		opts.SetTLSConfig(&tls.Config{
			InsecureSkipVerify: true,
		})
	}
	if i.Cfg.Debug {
		opts.SetLogLevel(3)
	}
	i.client = influxdb2.NewClientWithOptions(i.Cfg.URL, i.Cfg.Token, opts)
	// start influx health check
	go i.healthCheck(ctx)
	i.writer = i.client.WriteAPI(i.Cfg.Org, i.Cfg.Bucket)
	i.logger.Printf("initialized influxdb write API: %s", i.String())
	// start influxdb error logs
	go func() {
		for err = range i.writer.Errors() {
			i.logger.Printf("writeAPI error: %v", err)
		}
	}()
	return nil
}

func (i *InfluxDBOutput) Write(rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := collector.ResponseToEventMsgs(measName, rsp, meta)
		if err != nil {
			i.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			i.writer.WritePoint(influxdb2.NewPoint(ev.Name, ev.Tags, ev.Values, time.Unix(0, ev.Timestamp)))
		}
	}
}

func (i *InfluxDBOutput) Close() error {
	i.logger.Printf("flushing data...")
	i.writer.Flush()
	i.logger.Printf("closing client...")
	i.client.Close()
	i.cancelFn()
	i.logger.Printf("closed.")
	return nil
}
func (i *InfluxDBOutput) Metrics() []prometheus.Collector { return i.metrics }

func (i *InfluxDBOutput) healthCheck(ctx context.Context) {
	i.health(ctx)
	ticker := time.NewTicker(i.Cfg.HealthCheckPeriod)
	for {
		select {
		case <-ticker.C:
			i.health(ctx)
		case <-ctx.Done():
			return
		}
	}
}

func (i *InfluxDBOutput) health(ctx context.Context) {
	res, err := i.client.Health(ctx)
	if err != nil {
		i.logger.Printf("failed health check: %v", err)
		return
	}
	if res != nil {
		b, err := json.Marshal(res)
		if err != nil {
			i.logger.Printf("failed to marshal health check result: %v", err)
			i.logger.Printf("health check result: %+v", res)
			return
		}
		i.logger.Printf("health check result: %s", string(b))
		return
	}
	i.logger.Print("health check result is nil")
}
{'fraction_non_alphanumeric': 0.1413264759455695, 'fraction_numerical': 0.005143040822886531, 'mean_word_length': 3.546517291768144, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, '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': '921717', 'n_tokens_mistral': 3514, 'n_tokens_neox': 3326, 'n_words': 664}