feat: shared connection pool for pg (#1584)

* **New Features**
* Worker-level PostgreSQL connection pool configuration to manage
connections across multiple DAGs in a single worker process.
* Centralized connection pool management for shared-nothing worker mode
with reference-counted pool lifecycle.

* **Improvements**
* Executor resource cleanup to ensure proper connection and resource
release after execution.
* Fixed connection pool defaults for improved stability and consistency.

Author: yottahmd

internal/cmn/config/config.go

@@ -338,6 +338,10 @@ type Worker struct {
 	MaxActiveRuns int               // Maximum number of active runs (default: 100)
 	Labels        map[string]string // Worker labels for capability matching
 	Coordinators  []string          // Coordinator addresses for static discovery (host:port)
+
+	// PostgresPool holds connection pool settings for shared-nothing mode.
+	// When multiple DAGs run concurrently in a worker, they share this pool.
+	PostgresPool PostgresPoolConfig
 }
 
 // Scheduler represents the scheduler configuration
@@ -359,6 +363,28 @@ type Scheduler struct {
 	ZombieDetectionInterval time.Duration
 }
 
+// PostgresPoolConfig holds PostgreSQL connection pool settings for workers.
+// Used in shared-nothing worker mode to prevent connection exhaustion
+// when multiple DAGs run concurrently in a single worker process.
+type PostgresPoolConfig struct {
+	// MaxOpenConns is the maximum total open connections across ALL PostgreSQL DSNs.
+	// This is the hard limit shared across all database connections.
+	// Default: 25
+	MaxOpenConns int
+
+	// MaxIdleConns is the maximum number of idle connections per DSN.
+	// Default: 5
+	MaxIdleConns int
+
+	// ConnMaxLifetime is the maximum lifetime of a connection in seconds.
+	// Default: 300 (5 minutes)
+	ConnMaxLifetime int
+
+	// ConnMaxIdleTime is the maximum idle time for a connection in seconds.
+	// Default: 60 (1 minute)
+	ConnMaxIdleTime int
+}
+
 // Peer holds the certificate and TLS configuration for peer connections over gRPC.
 type Peer struct {
 	// CertFile is the path to the server's TLS certificate file.

internal/cmn/config/definition.go

@@ -386,6 +386,9 @@ type WorkerDef struct {
 	// When specified, the worker will connect directly to these coordinators
 	// instead of using the file-based service registry.
 	Coordinators interface{} `mapstructure:"coordinators"`
+
+	// PostgresPool holds connection pool settings for shared-nothing mode.
+	PostgresPool *PostgresPoolDef `mapstructure:"postgresPool"`
 }
 
 // SchedulerDef holds the configuration for the scheduler.
@@ -405,3 +408,23 @@ type SchedulerDef struct {
 	// Default is 45 seconds. Set to 0 to disable.
 	ZombieDetectionInterval string `mapstructure:"zombieDetectionInterval"`
 }
+
+// PostgresPoolDef holds the definition for PostgreSQL connection pool configuration.
+// Used in shared-nothing worker mode to prevent connection exhaustion.
+type PostgresPoolDef struct {
+	// MaxOpenConns is the maximum total open connections across ALL PostgreSQL DSNs.
+	// Default: 25
+	MaxOpenConns int `mapstructure:"maxOpenConns"`
+
+	// MaxIdleConns is the maximum number of idle connections per DSN.
+	// Default: 5
+	MaxIdleConns int `mapstructure:"maxIdleConns"`
+
+	// ConnMaxLifetime is the maximum lifetime of a connection in seconds.
+	// Default: 300 (5 minutes)
+	ConnMaxLifetime int `mapstructure:"connMaxLifetime"`
+
+	// ConnMaxIdleTime is the maximum idle time for a connection in seconds.
+	// Default: 60 (1 minute)
+	ConnMaxIdleTime int `mapstructure:"connMaxIdleTime"`
+}

internal/cmn/config/loader.go

@@ -669,6 +669,37 @@ func (l *ConfigLoader) loadWorkerConfig(cfg *Config, def Definition) {
 			cfg.Worker.Coordinators = addresses
 			l.warnings = append(l.warnings, addrWarnings...)
 		}
+
+		// Load PostgresPool config
+		if def.Worker.PostgresPool != nil {
+			pp := def.Worker.PostgresPool
+			if pp.MaxOpenConns > 0 {
+				cfg.Worker.PostgresPool.MaxOpenConns = pp.MaxOpenConns
+			}
+			if pp.MaxIdleConns > 0 {
+				cfg.Worker.PostgresPool.MaxIdleConns = pp.MaxIdleConns
+			}
+			if pp.ConnMaxLifetime > 0 {
+				cfg.Worker.PostgresPool.ConnMaxLifetime = pp.ConnMaxLifetime
+			}
+			if pp.ConnMaxIdleTime > 0 {
+				cfg.Worker.PostgresPool.ConnMaxIdleTime = pp.ConnMaxIdleTime
+			}
+		}
+	}
+
+	// Set PostgresPool defaults if not configured
+	if cfg.Worker.PostgresPool.MaxOpenConns == 0 {
+		cfg.Worker.PostgresPool.MaxOpenConns = 25
+	}
+	if cfg.Worker.PostgresPool.MaxIdleConns == 0 {
+		cfg.Worker.PostgresPool.MaxIdleConns = 5
+	}
+	if cfg.Worker.PostgresPool.ConnMaxLifetime == 0 {
+		cfg.Worker.PostgresPool.ConnMaxLifetime = 300
+	}
+	if cfg.Worker.PostgresPool.ConnMaxIdleTime == 0 {
+		cfg.Worker.PostgresPool.ConnMaxIdleTime = 60
 	}
 }
 
@@ -1075,6 +1106,12 @@ func (l *ConfigLoader) setViperDefaultValues(paths Paths) {
 	// Monitoring settings
 	l.v.SetDefault("monitoring.retention", "24h")
 	l.v.SetDefault("monitoring.interval", "5s")
+
+	// Worker PostgreSQL pool settings (for shared-nothing mode)
+	l.v.SetDefault("worker.postgresPool.maxOpenConns", 25)
+	l.v.SetDefault("worker.postgresPool.maxIdleConns", 5)
+	l.v.SetDefault("worker.postgresPool.connMaxLifetime", 300)
+	l.v.SetDefault("worker.postgresPool.connMaxIdleTime", 60)
 }
 
 // envBinding defines a mapping between a config key and its environment variable.
@@ -1211,6 +1248,12 @@ var envBindings = []envBinding{
 	// Monitoring configuration
 	{key: "monitoring.retention", env: "MONITORING_RETENTION"},
 	{key: "monitoring.interval", env: "MONITORING_INTERVAL"},
+
+	// Worker PostgreSQL pool configuration (for shared-nothing mode)
+	{key: "worker.postgresPool.maxOpenConns", env: "WORKER_POSTGRES_POOL_MAX_OPEN_CONNS"},
+	{key: "worker.postgresPool.maxIdleConns", env: "WORKER_POSTGRES_POOL_MAX_IDLE_CONNS"},
+	{key: "worker.postgresPool.connMaxLifetime", env: "WORKER_POSTGRES_POOL_CONN_MAX_LIFETIME"},
+	{key: "worker.postgresPool.connMaxIdleTime", env: "WORKER_POSTGRES_POOL_CONN_MAX_IDLE_TIME"},
 }
 
 // bindEnvironmentVariables binds configuration keys to environment variables using the loader's viper instance.

internal/cmn/config/loader_test.go

@@ -235,6 +235,12 @@ func TestLoad_Env(t *testing.T) {
 		Worker: Worker{
 			ID:            "test-worker-123",
 			MaxActiveRuns: 200,
+			PostgresPool: PostgresPoolConfig{
+				MaxOpenConns:    25,
+				MaxIdleConns:    5,
+				ConnMaxLifetime: 300,
+				ConnMaxIdleTime: 60,
+			},
 		},
 		Scheduler: Scheduler{
 			Port:                    9999,
@@ -484,6 +490,12 @@ scheduler:
 				"env":    "production",
 				"region": "us-west-2",
 			},
+			PostgresPool: PostgresPoolConfig{
+				MaxOpenConns:    25,
+				MaxIdleConns:    5,
+				ConnMaxLifetime: 300,
+				ConnMaxIdleTime: 60,
+			},
 		},
 		Scheduler: Scheduler{
 			Port:                    7890,

internal/runtime/builtin/sql/config.go

@@ -17,11 +17,6 @@ type Config struct {
 	// - SQLite: "file:./data.db?mode=rw" or ":memory:"
 	DSN string `mapstructure:"dsn"`
 
-	// Connection pool settings
-	MaxOpenConns    int `mapstructure:"maxOpenConns"`    // Maximum open connections (default: 5)
-	MaxIdleConns    int `mapstructure:"maxIdleConns"`    // Maximum idle connections (default: 2)
-	ConnMaxLifetime int `mapstructure:"connMaxLifetime"` // Connection max lifetime in seconds (default: 300)
-
 	// Parameterized queries (SQL injection prevention)
 	// Can be map[string]any for named params or []any for positional params
 	Params any `mapstructure:"params"`
@@ -86,15 +81,11 @@ type ImportConfig struct {
 }
 
 // DefaultConfig returns a Config with default values.
-// These defaults match the JSON schema documentation.
 func DefaultConfig() *Config {
 	return &Config{
-		MaxOpenConns:    5,   // Match schema default
-		MaxIdleConns:    2,   // Match schema default
-		ConnMaxLifetime: 300, // Match schema default (seconds)
-		Timeout:         60,
-		OutputFormat:    "jsonl",
-		NullString:      "null",
+		Timeout:      60,
+		OutputFormat: "jsonl",
+		NullString:   "null",
 	}
 }
 
@@ -138,19 +129,6 @@ func ParseConfig(_ context.Context, mapCfg map[string]any) (*Config, error) {
 		}
 	}
 
-	// Validate connection pool settings
-	if cfg.MaxOpenConns < 0 {
-		return nil, fmt.Errorf("maxOpenConns must be non-negative")
-	}
-	if cfg.MaxIdleConns < 0 {
-		return nil, fmt.Errorf("maxIdleConns must be non-negative")
-	}
-	if cfg.MaxIdleConns > cfg.MaxOpenConns && cfg.MaxOpenConns > 0 {
-		return nil, fmt.Errorf("maxIdleConns (%d) cannot exceed maxOpenConns (%d)", cfg.MaxIdleConns, cfg.MaxOpenConns)
-	}
-	if cfg.ConnMaxLifetime < 0 {
-		return nil, fmt.Errorf("connMaxLifetime must be non-negative")
-	}
 	if cfg.Timeout < 0 {
 		return nil, fmt.Errorf("timeout must be non-negative")
 	}
@@ -277,10 +255,7 @@ var importConfigSchema = &jsonschema.Schema{
 var postgresConfigSchema = &jsonschema.Schema{
 	Type: "object",
 	Properties: map[string]*jsonschema.Schema{
-		"dsn":             {Type: "string", Description: "PostgreSQL connection string (DSN)"},
-		"maxOpenConns":    {Type: "integer", Description: "Maximum open connections"},
-		"maxIdleConns":    {Type: "integer", Description: "Maximum idle connections"},
-		"connMaxLifetime": {Type: "integer", Description: "Connection max lifetime in seconds"},
+		"dsn": {Type: "string", Description: "PostgreSQL connection string (DSN)"},
 		"params": {
 			Description: "Query parameters (map for named, array for positional)",
 			OneOf: []*jsonschema.Schema{
@@ -306,9 +281,7 @@ var postgresConfigSchema = &jsonschema.Schema{
 var sqliteConfigSchema = &jsonschema.Schema{
 	Type: "object",
 	Properties: map[string]*jsonschema.Schema{
-		"dsn":          {Type: "string", Description: "SQLite connection string (file path or :memory:)"},
-		"maxOpenConns": {Type: "integer", Description: "Maximum open connections"},
-		"maxIdleConns": {Type: "integer", Description: "Maximum idle connections"},
+		"dsn": {Type: "string", Description: "SQLite connection string (file path or :memory:)"},
 		"params": {
 			Description: "Query parameters (map for named, array for positional)",
 			OneOf: []*jsonschema.Schema{

internal/runtime/builtin/sql/connection.go

@@ -19,13 +19,20 @@ type ConnectionManager struct {
 	closed   bool
 }
 
-// Connection retry settings
-// These are generous to handle slow container startup in CI environments
+// Connection settings
 const (
+	// Retry settings - generous to handle slow container startup in CI environments
 	maxConnRetries    = 30
 	initialRetryDelay = 500 * time.Millisecond
 	maxRetryDelay     = 2 * time.Second
 	pingTimeout       = 5 * time.Second
+
+	// Default connection pool settings for non-worker mode
+	// These are conservative values suitable for single-step execution.
+	// In worker mode, the global pool manager handles pooling.
+	defaultMaxOpenConns    = 1 // Single connection per step
+	defaultMaxIdleConns    = 1 // Keep the connection ready
+	defaultConnMaxLifetime = 5 * time.Minute
 )
 
 // NewConnectionManager creates a new connection manager.
@@ -56,10 +63,12 @@ func NewConnectionManager(ctx context.Context, driver Driver, cfg *Config) (*Con
 			return nil, fmt.Errorf("failed to connect to database after %d attempts: %w", attempt, lastErr)
 		}
 
-		// Configure connection pool
-		db.SetMaxOpenConns(cfg.MaxOpenConns)
-		db.SetMaxIdleConns(cfg.MaxIdleConns)
-		db.SetConnMaxLifetime(time.Duration(cfg.ConnMaxLifetime) * time.Second)
+		// Configure connection pool with fixed defaults
+		// In non-worker mode, each step has its own isolated connection.
+		// In worker mode, the global pool manager handles pooling.
+		db.SetMaxOpenConns(defaultMaxOpenConns)
+		db.SetMaxIdleConns(defaultMaxIdleConns)
+		db.SetConnMaxLifetime(defaultConnMaxLifetime)
 
 		// Verify connection with ping
 		pingCtx, cancel := context.WithTimeout(ctx, pingTimeout)