diff --git a/config.yaml.sample b/config.yaml.sample index a3c4f57..0628d6a 100644 --- a/config.yaml.sample +++ b/config.yaml.sample @@ -9,5 +9,7 @@ dhis2_endpoints: sendgrid_accounts: - account_name: "Account1" api_key: "key" + time_zone: "timezone" - account_name: "Account2" api_key: "key" + time_zone: "timezone" diff --git a/main.go b/main.go index bd10358..5f33ab4 100644 --- a/main.go +++ b/main.go @@ -2,28 +2,24 @@ package main import ( "context" - "io/ioutil" - "log" - "net/http" - "os" - "os/signal" - "time" - "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/simpledotorg/rtsl_exporter/alphasms" "github.com/simpledotorg/rtsl_exporter/dhis2" "github.com/simpledotorg/rtsl_exporter/sendgrid" "gopkg.in/yaml.v2" + "io/ioutil" + "log" + "net/http" + "os" + "os/signal" + "time" ) type Config struct { - ALPHASMSAPIKey string `yaml:"alphasms_api_key"` - SendGridAccounts []struct { - AccountName string `yaml:"account_name"` - APIKey string `yaml:"api_key"` - } `yaml:"sendgrid_accounts"` - DHIS2Endpoints []struct { + ALPHASMSAPIKey string `yaml:"alphasms_api_key"` + SendGridAccounts []sendgrid.AccountConfig `yaml:"sendgrid_accounts"` + DHIS2Endpoints []struct { BaseURL string `yaml:"base_url"` Username string `yaml:"username"` Password string `yaml:"password"` @@ -42,12 +38,10 @@ func readConfig(configPath string) (*Config, error) { } return config, nil } - func gracefulShutdown(server *http.Server) { // Create a context with a timeout for the shutdown ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - // Notify when the shutdown process is complete idleConnectionsClosed := make(chan struct{}) go func() { @@ -56,7 +50,6 @@ func gracefulShutdown(server *http.Server) { log.Printf("HTTP Server Shutdown Error: %v", err) } }() - // Wait for the server to shut down select { case <-ctx.Done(): @@ -65,15 +58,12 @@ func gracefulShutdown(server *http.Server) { log.Println("HTTP Server Shutdown Complete") } } - func main() { log.SetFlags(0) - config, err := readConfig("config.yaml") if err != nil { log.Fatalf("Error reading config file: %v", err) } - // Alphasms if config.ALPHASMSAPIKey == "" { log.Fatalf("ALPHASMS_API_KEY not provided in config file") @@ -81,7 +71,6 @@ func main() { alphasmsClient := alphasms.Client{APIKey: config.ALPHASMSAPIKey} alphasmsExporter := alphasms.NewExporter(&alphasmsClient) prometheus.MustRegister(alphasmsExporter) - // DHIS2 dhis2Clients := []*dhis2.Client{} for _, endpoint := range config.DHIS2Endpoints { @@ -95,22 +84,22 @@ func main() { } dhis2Exporter := dhis2.NewExporter(dhis2Clients) prometheus.MustRegister(dhis2Exporter) - - // Register SendGrid exporters - apiKeys := make(map[string]string) + // Register SendGrid exporters with time zones + sendGridConfigMap := make(map[string]sendgrid.AccountConfig) for _, account := range config.SendGridAccounts { - apiKeys[account.AccountName] = account.APIKey + sendGridConfigMap[account.AccountName] = sendgrid.AccountConfig{ + AccountName: account.AccountName, + APIKey: account.APIKey, + TimeZone: account.TimeZone, + } } - sendgridExporter := sendgrid.NewExporter(apiKeys) + sendgridExporter := sendgrid.NewExporter(sendGridConfigMap) prometheus.MustRegister(sendgridExporter) - http.Handle("/metrics", promhttp.Handler()) log.Println("Starting server on :8080") - httpServer := &http.Server{ Addr: ":8080", } - go func() { sigint := make(chan os.Signal, 1) signal.Notify(sigint, os.Interrupt) @@ -118,10 +107,8 @@ func main() { log.Println("Shutdown signal received") gracefulShutdown(httpServer) }() - if err := httpServer.ListenAndServe(); err != http.ErrServerClosed { log.Fatalf("HTTP server ListenAndServe Error: %v", err) } - log.Println("Bye bye") } diff --git a/sendgrid/exporter.go b/sendgrid/exporter.go index 9522d95..99883a5 100755 --- a/sendgrid/exporter.go +++ b/sendgrid/exporter.go @@ -1,25 +1,42 @@ package sendgrid import ( + "github.com/prometheus/client_golang/prometheus" "log" "time" - - "github.com/prometheus/client_golang/prometheus" ) type Exporter struct { - client *Client - emailLimit *prometheus.GaugeVec - emailRemaining *prometheus.GaugeVec - emailUsed *prometheus.GaugeVec - planExpiration *prometheus.GaugeVec - httpReturnCode *prometheus.GaugeVec + client *Client + timeZones map[string]*time.Location // Map of time locations + emailLimit *prometheus.GaugeVec + emailRemaining *prometheus.GaugeVec + emailUsed *prometheus.GaugeVec + planExpiration *prometheus.GaugeVec + httpReturnCode *prometheus.GaugeVec httpResponseTime *prometheus.GaugeVec } +type AccountConfig struct { + AccountName string `yaml:"account_name"` + APIKey string `yaml:"api_key"` + TimeZone string `yaml:"time_zone"` +} -func NewExporter(apiKeys map[string]string) *Exporter { +func NewExporter(accounts map[string]AccountConfig) *Exporter { + apiKeys := make(map[string]string) + timeZones := make(map[string]*time.Location) + for accountName, accountConfig := range accounts { + apiKeys[accountName] = accountConfig.APIKey + loc, err := time.LoadLocation(accountConfig.TimeZone) + if err != nil { + log.Printf("Error loading time zone for account %s: %v", accountName, err) + loc = time.UTC // Default to UTC if time zone cannot be loaded + } + timeZones[accountName] = loc + } return &Exporter{ - client: NewClient(apiKeys), + client: NewClient(apiKeys), + timeZones: timeZones, emailLimit: prometheus.NewGaugeVec(prometheus.GaugeOpts{ Namespace: "sendgrid", Name: "email_limit_count", @@ -52,9 +69,6 @@ func NewExporter(apiKeys map[string]string) *Exporter { }, []string{"account_name"}), } } - - -// Describe sends the descriptions of the metrics to Prometheus. func (e *Exporter) Describe(ch chan<- *prometheus.Desc) { e.emailLimit.Describe(ch) e.emailRemaining.Describe(ch) @@ -63,42 +77,29 @@ func (e *Exporter) Describe(ch chan<- *prometheus.Desc) { e.httpReturnCode.Describe(ch) e.httpResponseTime.Describe(ch) } - -// Collect retrieves metrics and sends them to Prometheus. func (e *Exporter) Collect(ch chan<- prometheus.Metric) { - // Collect metrics for each account for accountName := range e.client.APIKeys { metrics, statusCode, responseTime, err := e.client.FetchMetrics(accountName) if err != nil { log.Printf("Failed to get metrics for account %s: %v", accountName, err) continue } - - // Set metrics values for each account e.emailLimit.WithLabelValues(accountName).Set(metrics.Total) e.emailRemaining.WithLabelValues(accountName).Set(metrics.Remaining) e.emailUsed.WithLabelValues(accountName).Set(metrics.Used) - - // Parse the plan expiration date + timeZone := e.timeZones[accountName] dateFormat := "2006-01-02" - planResetDate, parseErr := time.Parse(dateFormat, metrics.NextReset) + planResetDate, parseErr := time.ParseInLocation(dateFormat, metrics.NextReset, timeZone) if parseErr != nil { - log.Printf("Failed to parse plan reset date: %v", parseErr) + log.Printf("Failed to parse plan reset date for account %s: %v", accountName, parseErr) continue } - - // Calculate time until expiration - timeUntilExpiration := planResetDate.Sub(time.Now()).Seconds() - if timeUntilExpiration < 0 { - timeUntilExpiration = 0 - } - + currentTime := time.Now().In(timeZone) + timeUntilExpiration := planResetDate.Sub(currentTime).Seconds() e.planExpiration.WithLabelValues(accountName).Set(timeUntilExpiration) - e.httpReturnCode.WithLabelValues(accountName).Set(float64(statusCode)) + e.httpReturnCode.WithLabelValues(accountName).Set(float64(statusCode)) e.httpResponseTime.WithLabelValues(accountName).Set(responseTime.Seconds()) } - - // Collect all metrics once e.emailLimit.Collect(ch) e.emailRemaining.Collect(ch) e.emailUsed.Collect(ch) diff --git a/sendgrid/exporter_test.go b/sendgrid/exporter_test.go index 87d3362..2b284b3 100755 --- a/sendgrid/exporter_test.go +++ b/sendgrid/exporter_test.go @@ -1,9 +1,11 @@ package sendgrid + import ( - "testing" "github.com/jarcoal/httpmock" "github.com/prometheus/client_golang/prometheus/testutil" + "testing" ) + func TestExporterCollect(t *testing.T) { // Activate the HTTP mock httpmock.Activate() @@ -17,10 +19,14 @@ func TestExporterCollect(t *testing.T) { "next_reset": "2024-02-20" }`)) // Created a new Exporter - accountNames := map[string]string{ - "mockAccount": "mockAPIKey", + accountConfigs := map[string]AccountConfig{ + "mockAccount": { + AccountName: "mockAccount", + APIKey: "mockAPIKey", + TimeZone: "UTC", + }, } - exporter := NewExporter(accountNames) + exporter := NewExporter(accountConfigs) t.Run("Successful metrics collection", func(t *testing.T) { expectedMetrics := []string{ "sendgrid_email_limit_count",