Fix: Adding test for testing newsysdb with custom pool

Author: vr-varad

dbos/dbos.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/robfig/cron/v3"
 )
 
@@ -42,6 +43,7 @@ type Config struct {
 	ApplicationVersion string          // Application version (optional, overridden by DBOS__APPVERSION env var)
 	ExecutorID         string          // Executor ID (optional, overridden by DBOS__VMID env var)
 	Context            context.Context // User Context
+	Pool               *pgxpool.Pool   // Custom Pool
 }
 
 // processConfig enforces mandatory fields and applies defaults.
@@ -67,6 +69,7 @@ func processConfig(inputConfig *Config) (*Config, error) {
 		ConductorAPIKey:    inputConfig.ConductorAPIKey,
 		ApplicationVersion: inputConfig.ApplicationVersion,
 		ExecutorID:         inputConfig.ExecutorID,
+		Pool:               inputConfig.Pool,
 	}
 
 	// Load defaults
@@ -327,7 +330,7 @@ func NewDBOSContext(ctx context.Context, inputConfig Config) (DBOSContext, error
 
 	newSystemDatabaseInputs := newSystemDatabaseInput{
 		databaseURL: config.DatabaseURL,
-		custom_pool: nil,
+		custom_pool: config.Pool,
 		logger:      initExecutor.logger,
 	}
 

dbos/dbos_test.go

@@ -1,10 +1,13 @@
 package dbos
 
 import (
+	"bytes"
 	"context"
+	"log/slog"
 	"testing"
 	"time"
 
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -58,6 +61,55 @@ func TestConfig(t *testing.T) {
 		assert.Equal(t, expectedMsg, dbosErr.Message)
 	})
 
+	t.Run("NewSystemDatabaseWithCustomPool", func(t *testing.T) {
+
+		// Logger
+		var buf bytes.Buffer
+		slogLogger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{
+			Level: slog.LevelDebug,
+		}))
+
+		slogLogger = slogLogger.With("service", "dbos-test", "environment", "test")
+
+		// Custom Pool
+		poolConfig, _ := pgxpool.ParseConfig(databaseURL)
+
+		poolConfig.MaxConns = 10
+		poolConfig.MinConns = 5
+		poolConfig.MaxConnLifetime = 2 * time.Hour
+		poolConfig.MaxConnIdleTime = time.Minute * 2
+
+		poolConfig.ConnConfig.ConnectTimeout = 10 * time.Second
+
+		pool, err := pgxpool.NewWithConfig(context.Background(), poolConfig)
+		require.NoError(t, err)
+		defer pool.Close()
+
+		config := Config{
+			DatabaseURL: databaseURL,
+			AppName:     "test-custom-pool",
+			Logger:      slogLogger,
+			Pool:        pool,
+		}
+
+		customdbosContext, err := NewDBOSContext(context.Background(), config)
+		require.NoError(t, err)
+		require.NotNil(t, customdbosContext)
+
+		dbosCtx, ok := customdbosContext.(*dbosContext)
+		require.True(t, ok)
+		sysDB, ok := dbosCtx.systemDB.(*sysDB)
+		require.True(t, ok)
+		assert.Same(t, pool, sysDB.pool, "The pool in dbosContext should be the same as the custom pool provided")
+
+		stats := sysDB.pool.Stat()
+		assert.Equal(t, int32(10), stats.MaxConns(), "MaxConns should match custom pool config")
+
+		customdbosContext.Shutdown(1 * time.Minute)
+		err = pool.Ping(context.Background())
+		assert.NoError(t, err, "Custom pool should not be closed by Shutdown and should still be usable")
+	})
+
 	t.Run("FailsWithoutDatabaseURL", func(t *testing.T) {
 		config := Config{
 			AppName: "test-app",