Use UTC for all pgx connections and check for database TZ (#2398)

* set utc for all pgx sessions

* helper func

* also accept IANA Etc/UTC

pkg/config/database/config.go

@@ -37,6 +37,11 @@ type ConfigFile struct {
 	LogQueries bool `mapstructure:"logQueries" json:"logQueries,omitempty" default:"false"`
 
 	CacheDuration time.Duration `mapstructure:"cacheDuration" json:"cacheDuration,omitempty" default:"5s"`
+
+	// EnforceUTCTimezone enforces that the database instance timezone is set to UTC.
+	// If enabled and the database timezone is not UTC, the server will panic on startup.
+	// To disable this check, set DATABASE_ENFORCE_UTC_TIMEZONE=false
+	EnforceUTCTimezone bool `mapstructure:"enforceUtcTimezone" json:"enforceUtcTimezone,omitempty" default:"true"`
 }
 
 type SeedConfigFile struct {
@@ -90,6 +95,7 @@ func BindAllEnv(v *viper.Viper) {
 	_ = v.BindEnv("readReplicaMinConns", "READ_REPLICA_MIN_CONNS")
 
 	_ = v.BindEnv("cacheDuration", "CACHE_DURATION")
+	_ = v.BindEnv("enforceUtcTimezone", "DATABASE_ENFORCE_UTC_TIMEZONE")
 
 	_ = v.BindEnv("seed.adminEmail", "ADMIN_EMAIL")
 	_ = v.BindEnv("seed.adminPassword", "ADMIN_PASSWORD")

pkg/config/loader/loader.go

@@ -141,8 +141,14 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 		return nil, err
 	}
 
-	// ref: https://github.com/jackc/pgx/issues/1549
 	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+		// Set timezone to UTC for all connections
+		if _, err := conn.Exec(ctx, "SET TIME ZONE 'UTC'"); err != nil {
+			return err
+		}
+
+		// ref: https://github.com/jackc/pgx/issues/1549
+
 		t, err := conn.LoadType(ctx, "v1_readable_status_olap")
 		if err != nil {
 			return err
@@ -188,6 +194,13 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 		config.BeforeAcquire = debugger.beforeAcquire
 	}
 
+	// Check database instance timezone if enforcement is enabled
+	if cf.EnforceUTCTimezone {
+		if err := checkDatabaseTimezone(config.ConnConfig, cf.PostgresDbName, "primary database", &l); err != nil {
+			return nil, err
+		}
+	}
+
 	pool, err := pgxpool.NewWithConfig(context.Background(), config)
 
 	if err != nil {
@@ -219,6 +232,19 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 		readReplicaConfig.MaxConnLifetime = 15 * 60 * time.Second
 		readReplicaConfig.ConnConfig.Tracer = otelpgx.NewTracer()
 
+		// Check read replica database instance timezone if enforcement is enabled
+		if cf.EnforceUTCTimezone {
+			if err := checkDatabaseTimezone(readReplicaConfig.ConnConfig, "", "read replica database", &l); err != nil {
+				return nil, err
+			}
+		}
+
+		// Set timezone to UTC for read replica connections
+		readReplicaConfig.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+			_, err := conn.Exec(ctx, "SET TIME ZONE 'UTC'")
+			return err
+		}
+
 		readReplicaPool, err = pgxpool.NewWithConfig(context.Background(), readReplicaConfig)
 
 		if err != nil {
@@ -832,3 +858,37 @@ func loadInternalClient(l *zerolog.Logger, conf *server.InternalClientTLSConfigF
 		clientv1.WithLogger(l),
 	), nil
 }
+
+// checkDatabaseTimezone validates that the database instance timezone is set to UTC.
+// It creates a temporary connection to check the timezone without using the AfterConnect hook.
+func checkDatabaseTimezone(connConfig *pgx.ConnConfig, dbName string, dbLabel string, l *zerolog.Logger) error {
+	tempConn, err := pgx.ConnectConfig(context.Background(), connConfig)
+	if err != nil {
+		return fmt.Errorf("could not create temporary connection to %s to check timezone: %w", dbLabel, err)
+	}
+	defer tempConn.Close(context.Background())
+
+	var dbTimezone string
+	if err := tempConn.QueryRow(context.Background(), "SHOW timezone").Scan(&dbTimezone); err != nil {
+		return fmt.Errorf("could not query %s timezone: %w", dbLabel, err)
+	}
+
+	// Accept both "UTC" and "Etc/UTC" as valid UTC timezones
+	if dbTimezone != "UTC" && dbTimezone != "Etc/UTC" {
+		if dbName == "" {
+			dbName = "<your_database_name>"
+		}
+		return fmt.Errorf(
+			"%s instance timezone is set to '%s' but must be 'UTC' or 'Etc/UTC'\n"+
+				"This check ensures time-based operations work correctly across all sessions\n"+
+				"To fix this issue, you have two options:\n"+
+				"  1. Set your PostgreSQL instance timezone to UTC by running: ALTER DATABASE %s SET TIMEZONE='UTC'\n"+
+				"  2. Disable this check by setting the environment variable: DATABASE_ENFORCE_UTC_TIMEZONE=false\n"+
+				"Note: Disabling this check is not recommended as it may lead to timezone-related issues",
+			dbLabel, dbTimezone, dbName,
+		)
+	}
+
+	l.Info().Msgf("%s instance timezone verified: %s", dbLabel, dbTimezone)
+	return nil
+}

pkg/repository/v1/task_partition_test.go

@@ -53,7 +53,16 @@ func setupPostgresWithMigration(t *testing.T) (*pgxpool.Pool, func()) {
 	migrate.RunMigrations(ctx)
 	t.Log("Migration completed successfully")
 
-	pool, err := pgxpool.New(ctx, connStr)
+	config, err := pgxpool.ParseConfig(connStr)
+	require.NoError(t, err)
+
+	// Set timezone to UTC for all test connections
+	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+		_, err := conn.Exec(ctx, "SET TIME ZONE 'UTC'")
+		return err
+	}
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
 	require.NoError(t, err)
 
 	err = pool.Ping(ctx)