github.com/grafana/synthetic-monitoring-agent@v0.64.2/cmd/test-api/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"flag"
     7  	"fmt"
     8  	"io"
     9  	"log"
    10  	"os"
    11  	"path/filepath"
    12  
    13  	"github.com/grafana/synthetic-monitoring-agent/cmd/test-api/internal/db"
    14  	"github.com/grafana/synthetic-monitoring-agent/cmd/test-api/internal/grpc"
    15  	"github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
    16  	"github.com/rs/zerolog"
    17  	"golang.org/x/sync/errgroup"
    18  )
    19  
    20  func main() {
    21  	if err := run(os.Args, os.Stdout); err != nil {
    22  		log.Fatal(err)
    23  	}
    24  }
    25  
    26  func run(args []string, stdout io.Writer) error {
    27  	flags := flag.NewFlagSet(filepath.Base(args[0]), flag.ExitOnError)
    28  
    29  	var (
    30  		debug          = flags.Bool("debug", false, "debug output (enables verbose)")
    31  		verbose        = flags.Bool("verbose", false, "verbose logging")
    32  		grpcListenAddr = flags.String("grpc-listen-address", ":4031", "GRPC listen address")
    33  		dataFn         = flags.String("load-data", "data.json", "filename with data to load")
    34  	)
    35  
    36  	if err := flags.Parse(args[1:]); err != nil {
    37  		return err
    38  	}
    39  
    40  	zl := getLogger(filepath.Base(args[0]), *verbose, *debug, stdout)
    41  
    42  	ctx, cancel := context.WithCancel(context.Background())
    43  	defer cancel()
    44  
    45  	eg, ctx := errgroup.WithContext(ctx)
    46  
    47  	db := db.New()
    48  
    49  	if err := loadData(*dataFn, db); err != nil {
    50  		return err
    51  	}
    52  
    53  	checksServer, err := grpc.NewChecksServer(grpc.ChecksServerOpts{
    54  		Logger: zl.With().Str("subsystem", "checks_server").Logger(),
    55  		Db:     db,
    56  	})
    57  	if err != nil {
    58  		zl.Error().Err(err).Msg("cannot create checks server")
    59  		return err
    60  	}
    61  
    62  	tenantsServer, err := grpc.NewTenantsServer(grpc.TenantsServerOpts{
    63  		Logger: zl.With().Str("subsystem", "tenants_server").Logger(),
    64  		Db:     db,
    65  	})
    66  	if err != nil {
    67  		zl.Error().Err(err).Msg("cannot create tenants server")
    68  		return err
    69  	}
    70  
    71  	grpcServer, err := grpc.NewServer(ctx, &grpc.Opts{
    72  		Logger:        zl.With().Str("subsystem", "grpc").Logger(),
    73  		ListenAddr:    *grpcListenAddr,
    74  		ChecksServer:  checksServer,
    75  		TenantsServer: tenantsServer,
    76  		Db:            db,
    77  	})
    78  	if err != nil {
    79  		zl.Error().Err(err).Msg("cannot create GRPC server")
    80  		return err
    81  	}
    82  
    83  	eg.Go(func() error { return grpcServer.Run(ctx) })
    84  
    85  	if err := eg.Wait(); err != nil {
    86  		zl.Error().Err(err).Send()
    87  		return err
    88  	}
    89  
    90  	return nil
    91  }
    92  
    93  func getLogger(name string, verbose, debug bool, w io.Writer) zerolog.Logger {
    94  	zerolog.TimeFieldFormat = zerolog.TimeFormatUnixMs
    95  
    96  	zl := zerolog.New(w).With().Timestamp().Str("program", name).Logger()
    97  
    98  	switch {
    99  	case debug:
   100  		zerolog.SetGlobalLevel(zerolog.DebugLevel)
   101  
   102  		zl = zl.With().Caller().Logger()
   103  
   104  	case verbose:
   105  		zerolog.SetGlobalLevel(zerolog.InfoLevel)
   106  
   107  	default:
   108  		zerolog.SetGlobalLevel(zerolog.ErrorLevel)
   109  	}
   110  
   111  	return zl
   112  }
   113  
   114  func loadData(fn string, db *db.Db) error {
   115  	var data struct {
   116  		Tenants     []synthetic_monitoring.Tenant `json:"tenants"`
   117  		Probes      []synthetic_monitoring.Probe  `json:"probes"`
   118  		ProbeTokens map[int64]string              `json:"probeTokens"`
   119  		Checks      []synthetic_monitoring.Check  `json:"checks"`
   120  	}
   121  
   122  	fh, err := os.Open(fn) //#nosec -- yes, I want to read whatever file the user tells me.
   123  	if err != nil {
   124  		return err
   125  	}
   126  
   127  	defer fh.Close() //#nosec -- file is open for reading.
   128  
   129  	dec := json.NewDecoder(fh)
   130  
   131  	if err := dec.Decode(&data); err != nil {
   132  		return err
   133  	}
   134  
   135  	ctx := context.Background()
   136  
   137  	for _, tenant := range data.Tenants {
   138  		if err := db.AddTenant(ctx, &tenant); err != nil {
   139  			return err
   140  		}
   141  	}
   142  
   143  	for _, probe := range data.Probes {
   144  		token, found := data.ProbeTokens[probe.Id]
   145  		if !found {
   146  			return fmt.Errorf("token not found for probe %d", probe.Id)
   147  		}
   148  
   149  		if err := db.AddProbe(ctx, &probe, []byte(token)); err != nil {
   150  			return err
   151  		}
   152  	}
   153  
   154  	for _, check := range data.Checks {
   155  		if err := db.AddCheck(ctx, &check); err != nil {
   156  			return err
   157  		}
   158  	}
   159  
   160  	return nil
   161  }