github.com/abayer/test-infra@v0.0.5/velodrome/transform/transform.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  	"flag"
    21  	"fmt"
    22  	"os"
    23  	"path/filepath"
    24  	"strings"
    25  	"time"
    26  
    27  	"k8s.io/test-infra/velodrome/sql"
    28  	"k8s.io/test-infra/velodrome/transform/plugins"
    29  
    30  	"github.com/golang/glog"
    31  	"github.com/spf13/cobra"
    32  )
    33  
    34  type transformConfig struct {
    35  	InfluxConfig
    36  	sql.MySQLConfig
    37  
    38  	repository string
    39  	once       bool
    40  	frequency  int
    41  	metricName string
    42  }
    43  
    44  func (config *transformConfig) CheckRootFlags() error {
    45  	if config.repository == "" {
    46  		return fmt.Errorf("repository must be set")
    47  	}
    48  	config.repository = strings.ToLower(config.repository)
    49  
    50  	if config.metricName == "" {
    51  		return fmt.Errorf("metric name must be set")
    52  	}
    53  
    54  	return nil
    55  }
    56  
    57  func (config *transformConfig) AddFlags(cmd *cobra.Command) {
    58  	cmd.PersistentFlags().IntVar(&config.frequency, "frequency", 2, "Number of iterations per hour")
    59  	cmd.PersistentFlags().BoolVar(&config.once, "once", false, "Run once and then leave")
    60  	cmd.PersistentFlags().StringVar(&config.repository, "repository", "", "Repository to use for metrics")
    61  	cmd.PersistentFlags().StringVar(&config.metricName, "name", "", "Name of the metric")
    62  	cmd.PersistentFlags().AddGoFlagSet(flag.CommandLine)
    63  }
    64  
    65  // Dispatch receives channels to each type of events, and dispatch them to each plugins.
    66  func Dispatch(plugin plugins.Plugin, DB *InfluxDB, issues chan sql.Issue, eventsCommentsChannel chan interface{}) {
    67  	for {
    68  		var points []plugins.Point
    69  		select {
    70  		case issue, ok := <-issues:
    71  			if !ok {
    72  				return
    73  			}
    74  			points = plugin.ReceiveIssue(issue)
    75  		case event, ok := <-eventsCommentsChannel:
    76  			if !ok {
    77  				return
    78  			}
    79  			switch event := event.(type) {
    80  			case sql.IssueEvent:
    81  				points = plugin.ReceiveIssueEvent(event)
    82  			case sql.Comment:
    83  				points = plugin.ReceiveComment(event)
    84  			default:
    85  				glog.Fatal("Received invalid object: ", event)
    86  			}
    87  		}
    88  
    89  		for _, point := range points {
    90  			if err := DB.Push(point.Tags, point.Values, point.Date); err != nil {
    91  				glog.Fatal("Failed to push point: ", err)
    92  			}
    93  		}
    94  	}
    95  }
    96  
    97  // Plugins constantly wait for new issues/events/comments
    98  func (config *transformConfig) run(plugin plugins.Plugin) error {
    99  	if err := config.CheckRootFlags(); err != nil {
   100  		return err
   101  	}
   102  
   103  	mysqldb, err := config.MySQLConfig.CreateDatabase()
   104  	if err != nil {
   105  		return err
   106  	}
   107  
   108  	influxdb, err := config.InfluxConfig.CreateDatabase(
   109  		map[string]string{"repository": config.repository},
   110  		config.metricName)
   111  	if err != nil {
   112  		return err
   113  	}
   114  
   115  	fetcher := NewFetcher(config.repository)
   116  
   117  	// Plugins constantly wait for new issues/events/comments
   118  	go Dispatch(plugin, influxdb, fetcher.IssuesChannel,
   119  		fetcher.EventsCommentsChannel)
   120  
   121  	ticker := time.Tick(time.Hour / time.Duration(config.frequency))
   122  	for {
   123  		// Fetch new events from MySQL, push it to plugins
   124  		if err := fetcher.Fetch(mysqldb); err != nil {
   125  			return err
   126  		}
   127  		if err := influxdb.PushBatchPoints(); err != nil {
   128  			return err
   129  		}
   130  
   131  		if config.once {
   132  			break
   133  		}
   134  		<-ticker
   135  	}
   136  
   137  	return nil
   138  }
   139  
   140  func main() {
   141  	config := &transformConfig{}
   142  	root := &cobra.Command{
   143  		Use:   filepath.Base(os.Args[0]),
   144  		Short: "Transform sql database info into influx stats",
   145  	}
   146  	config.AddFlags(root)
   147  	config.MySQLConfig.AddFlags(root)
   148  	config.InfluxConfig.AddFlags(root)
   149  
   150  	root.AddCommand(plugins.NewCountPlugin(config.run))
   151  
   152  	if err := root.Execute(); err != nil {
   153  		glog.Fatalf("%v\n", err)
   154  	}
   155  }