From 0d6e31713ff8431331392eafd2e14ca702e66988 Mon Sep 17 00:00:00 2001
From: Michael Hohn <mhhohn@gmail.com>
Date: Mon, 14 Jul 2025 13:26:23 -0700
Subject: [PATCH] Create a new PGState, reading connection info from env vars

---
 pkg/agent/agent.go          | 10 +++++++-
 pkg/state/state_postgres.go | 46 +++++++++++++++++++++++++++++--------
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
index 632678e..434ac70 100644
--- a/pkg/agent/agent.go
+++ b/pkg/agent/agent.go
@@ -124,7 +124,15 @@ func RunAnalysisJob(
 
 	databaseData, err := dbs.GetDatabase(job.Spec.NameWithOwner)
 	if err != nil {
-		return result, fmt.Errorf("failed to get database: %w", err)
+		slog.Error("Failed to get database",
+			slog.String("owner", job.Spec.Owner),
+			slog.String("repo", job.Spec.Repo),
+			slog.Int("session_id", job.Spec.SessionID),
+			slog.String("operation", "GetDatabase"),
+			slog.Any("error", err),
+		)
+		return result, fmt.Errorf("failed to get database for %s/%s: %w",
+			job.Spec.Owner, job.Spec.Repo, err)
 	}
 
 	// Write the CodeQL database data to the filesystem
diff --git a/pkg/state/state_postgres.go b/pkg/state/state_postgres.go
index 7cf15e0..1cf0894 100644
--- a/pkg/state/state_postgres.go
+++ b/pkg/state/state_postgres.go
@@ -17,25 +17,52 @@ type PGState struct {
 	pool *pgxpool.Pool
 }
 
-// ----- Create a new PGState, reading connection info from env vars
+func validateEnvVars(requiredEnvVars []string) {
+	missing := false
+
+	for _, envVar := range requiredEnvVars {
+		if _, ok := os.LookupEnv(envVar); !ok {
+			slog.Error("Missing required environment variable", "key", envVar)
+			missing = true
+		}
+	}
+
+	if missing {
+		os.Exit(1)
+	}
+}
+
 func NewPGState() *PGState {
 	ctx := context.Background()
 
-	// Get connection URL from environment (standard: PGHOST, PGPORT, etc.)
-	dbURL := os.Getenv("DATABASE_URL") // canonical way: use DATABASE_URL
-	if dbURL == "" {
-		slog.Error("DATABASE_URL environment variable must be set")
-		os.Exit(1)
+	required := []string{
+		"POSTGRES_USER",
+		"POSTGRES_PASSWORD",
+		"POSTGRES_DB",
+		// Host & port may be omitted if you rely on Docker DNS, but list
+		// them here to make the requirement explicit:
+		"POSTGRES_HOST",
+		"POSTGRES_PORT",
 	}
 
-	// Create a pool config
+	validateEnvVars(required)
+
+	// Assemble from vars
+	user := os.Getenv("POSTGRES_USER")
+	pass := os.Getenv("POSTGRES_PASSWORD")
+	host := os.Getenv("POSTGRES_HOST")
+	port := os.Getenv("POSTGRES_PORT")
+	db := os.Getenv("POSTGRES_DB")
+
+	dbURL := fmt.Sprintf("postgres://%s:%s@%s:%s/%s", user, pass, host, port, db)
+	slog.Info("Assembled Postgres connection URL from POSTGRES_* variables", "url", dbURL)
+
 	config, err := pgxpool.ParseConfig(dbURL)
 	if err != nil {
-		slog.Error("Failed to parse DATABASE_URL", "error", err)
+		slog.Error("Failed to parse connection URL", "url", dbURL, "error", err)
 		os.Exit(1)
 	}
 
-	// Optionally set pool tuning here
 	config.MaxConns = 10
 
 	pool, err := pgxpool.NewWithConfig(ctx, config)
@@ -46,7 +73,6 @@ func NewPGState() *PGState {
 
 	slog.Info("Connected to Postgres", "max_conns", config.MaxConns)
 
-	// schema initialization
 	SetupSchemas(pool)
 
 	return &PGState{pool: pool}