github.com/abayer/test-infra@v0.0.5/velodrome/transform/influx.go (about)

     1  /*
     2  Copyright 2016 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package main
    18  
    19  import (
    20  	"fmt"
    21  	"sort"
    22  	"strings"
    23  	"time"
    24  
    25  	"github.com/golang/glog"
    26  	influxdb "github.com/influxdata/influxdb/client/v2"
    27  	_ "github.com/jinzhu/gorm/dialects/mysql"
    28  	"github.com/spf13/cobra"
    29  )
    30  
    31  // InfluxConfig creates an InfluxDB
    32  type InfluxConfig struct {
    33  	Host     string
    34  	DB       string
    35  	User     string
    36  	Password string
    37  }
    38  
    39  // AddFlags parses options for database configuration
    40  func (config *InfluxConfig) AddFlags(cmd *cobra.Command) {
    41  	cmd.PersistentFlags().StringVar(&config.User, "influx-user", "root", "InfluxDB user")
    42  	cmd.PersistentFlags().StringVar(&config.Password, "influx-password", "", "InfluxDB password")
    43  	cmd.PersistentFlags().StringVar(&config.Host, "influx-host", "http://localhost:8086", "InfluxDB http server")
    44  	cmd.PersistentFlags().StringVar(&config.DB, "influx-database", "github", "InfluxDB database name")
    45  }
    46  
    47  func dropSeries(client influxdb.Client, measurement, database string, tags map[string]string) error {
    48  	query := influxdb.Query{
    49  		Command:  fmt.Sprintf(`DROP SERIES FROM %s %s`, measurement, tagsToWhere(tags)),
    50  		Database: database,
    51  	}
    52  	_, err := client.Query(query)
    53  	return err
    54  }
    55  
    56  // CreateDatabase creates and connects a new instance of an InfluxDB
    57  // It is created based on the fields set in the configuration.
    58  func (config *InfluxConfig) CreateDatabase(tags map[string]string, measurement string) (*InfluxDB, error) {
    59  	client, err := influxdb.NewHTTPClient(influxdb.HTTPConfig{
    60  		Addr:     config.Host,
    61  		Username: config.User,
    62  		Password: config.Password,
    63  	})
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  
    68  	err = dropSeries(client, measurement, config.DB, tags)
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  
    73  	bp, err := influxdb.NewBatchPoints(influxdb.BatchPointsConfig{
    74  		Database:  config.DB,
    75  		Precision: "s",
    76  	})
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  
    81  	return &InfluxDB{
    82  		client:      client,
    83  		database:    config.DB,
    84  		batch:       bp,
    85  		tags:        tags,
    86  		measurement: measurement,
    87  	}, err
    88  }
    89  
    90  // InfluxDB is a connection handler to a Influx database
    91  type InfluxDB struct {
    92  	client      influxdb.Client
    93  	database    string
    94  	measurement string
    95  	batch       influxdb.BatchPoints
    96  	batchSize   int
    97  	tags        map[string]string
    98  }
    99  
   100  // mergeTags merges the default tags with the exta tags. Default will be overridden if it conflicts.
   101  func mergeTags(defaultTags, extraTags map[string]string) map[string]string {
   102  	newTags := map[string]string{}
   103  
   104  	for k, v := range defaultTags {
   105  		newTags[k] = v
   106  	}
   107  	for k, v := range extraTags {
   108  		newTags[k] = v
   109  	}
   110  
   111  	return newTags
   112  }
   113  
   114  // tagsToWhere creates a where query to match tags element
   115  func tagsToWhere(tags map[string]string) string {
   116  	if len(tags) == 0 {
   117  		return ""
   118  	}
   119  
   120  	sortedKeys := []string{}
   121  	for k := range tags {
   122  		sortedKeys = append(sortedKeys, k)
   123  	}
   124  	sort.Strings(sortedKeys)
   125  
   126  	conditions := []string{}
   127  	for _, key := range sortedKeys {
   128  		conditions = append(conditions, fmt.Sprintf(`"%s" = '%v'`, key, tags[key]))
   129  	}
   130  	return "WHERE " + strings.Join(conditions, " AND ")
   131  }
   132  
   133  // Push a point to the database. This appends to current batchpoint
   134  func (i *InfluxDB) Push(tags map[string]string, fields map[string]interface{}, date time.Time) error {
   135  	pt, err := influxdb.NewPoint(i.measurement, mergeTags(i.tags, tags), fields, date)
   136  	if err != nil {
   137  		return err
   138  	}
   139  
   140  	i.batch.AddPoint(pt)
   141  	i.batchSize++
   142  
   143  	return nil
   144  }
   145  
   146  // PushBatchPoints pushes the batch points (for real)
   147  func (i *InfluxDB) PushBatchPoints() error {
   148  	// Push
   149  	err := i.client.Write(i.batch)
   150  	if err != nil {
   151  		return err
   152  	}
   153  	glog.Infof("Sent to influx: %d points", i.batchSize)
   154  
   155  	// Recreate new batch
   156  	i.batch, err = influxdb.NewBatchPoints(influxdb.BatchPointsConfig{
   157  		Database:  i.database,
   158  		Precision: "s",
   159  	})
   160  	i.batchSize = 0
   161  
   162  	if err != nil {
   163  		return err
   164  	}
   165  
   166  	return nil
   167  }