Merge pull request #4 from simpledotorg/sc-13403

sc-13403 Added timezone parameter

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"

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,23 +58,19 @@ 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")
 	}
 	alphasmsClient := alphasms.Client{APIKey: config.ALPHASMSAPIKey}
 	alphasmsExporter := alphasms.NewExporter(&alphasmsClient)
 	prometheus.MustRegister(alphasmsExporter)
-
 	// DHIS2
 	dhis2Clients := []*dhis2.Client{}
 	for _, endpoint := range config.DHIS2Endpoints {
@@ -95,33 +84,31 @@ 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)
 		<-sigint
 		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")
 }

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)

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",