From c47c0967fc167aaf928ec8edc0c0288ac1a0d80c Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Thu, 6 Aug 2026 17:43:30 -0400 Subject: [PATCH 1/4] add option to set pgx connection config Signed-off-by: Jet Chiang --- go/core/internal/database/connect.go | 36 +++++++++++--- go/core/internal/database/connect_test.go | 38 +++++++++++++- go/core/pkg/app/app.go | 49 +++++++++++++++---- go/core/pkg/app/app_test.go | 42 ++++++++++++++++ .../templates/controller-configmap.yaml | 14 ++++++ .../tests/controller-deployment_test.yaml | 36 ++++++++++++++ helm/kagent/values.yaml | 11 +++++ 7 files changed, 210 insertions(+), 16 deletions(-) diff --git a/go/core/internal/database/connect.go b/go/core/internal/database/connect.go index ff02138be..6df8ee11f 100644 --- a/go/core/internal/database/connect.go +++ b/go/core/internal/database/connect.go @@ -16,9 +16,16 @@ import ( // PostgresConfig holds the connection parameters for a Postgres database. // URL must be a resolved connection string — use ResolveURL to resolve from // a file path before constructing this config. +// +// Pool fields are optional: nil leaves the corresponding pgxpool.Config value +// from ParseConfig unchanged (pgx library defaults). type PostgresConfig struct { - URL string - VectorEnabled bool + URL string + VectorEnabled bool + MaxConns *int32 + MinConns *int32 + MaxConnIdleTime *time.Duration + MaxConnLifetime *time.Duration } const ( @@ -30,21 +37,38 @@ const ( // Connect opens a Postgres connection pool using cfg and retries Ping with // exponential backoff until the connection succeeds or defaultMaxTimeout elapses. func Connect(ctx context.Context, cfg *PostgresConfig) (*pgxpool.Pool, error) { - return retryDBConnection(ctx, cfg.URL, cfg.VectorEnabled) + return retryDBConnection(ctx, cfg) +} + +// applyPoolConfig copies non-nil pool settings from cfg onto config. +func applyPoolConfig(config *pgxpool.Config, cfg *PostgresConfig) { + if cfg.MaxConns != nil { + config.MaxConns = *cfg.MaxConns + } + if cfg.MinConns != nil { + config.MinConns = *cfg.MinConns + } + if cfg.MaxConnIdleTime != nil { + config.MaxConnIdleTime = *cfg.MaxConnIdleTime + } + if cfg.MaxConnLifetime != nil { + config.MaxConnLifetime = *cfg.MaxConnLifetime + } } // retryDBConnection opens a pgxpool connection, registering pgvector types when // vectorEnabled is true, and retries Ping with exponential backoff until the // connection succeeds or defaultMaxTimeout elapses. -func retryDBConnection(ctx context.Context, url string, vectorEnabled bool) (*pgxpool.Pool, error) { +func retryDBConnection(ctx context.Context, cfg *PostgresConfig) (*pgxpool.Pool, error) { ctx, cancel := context.WithTimeout(ctx, defaultMaxTimeout) defer cancel() - config, err := pgxpool.ParseConfig(url) + config, err := pgxpool.ParseConfig(cfg.URL) if err != nil { return nil, fmt.Errorf("failed to parse database URL: %w", err) } - if vectorEnabled { + applyPoolConfig(config, cfg) + if cfg.VectorEnabled { config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { return pgvectorpgx.RegisterTypes(ctx, conn) } diff --git a/go/core/internal/database/connect_test.go b/go/core/internal/database/connect_test.go index 681e58f32..cb9e7c85b 100644 --- a/go/core/internal/database/connect_test.go +++ b/go/core/internal/database/connect_test.go @@ -7,17 +7,53 @@ import ( "testing" "time" + "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRetryDBConnection_DeadlineExceeded(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() - _, err := retryDBConnection(ctx, "postgres://user:pass@localhost:1/nodb?connect_timeout=1", false) + _, err := retryDBConnection(ctx, &PostgresConfig{ + URL: "postgres://user:pass@localhost:1/nodb?connect_timeout=1", + }) assert.ErrorIs(t, err, context.DeadlineExceeded) } +func TestApplyPoolConfig(t *testing.T) { + base, err := pgxpool.ParseConfig("postgres://user:pass@localhost:5432/db") + require.NoError(t, err) + + t.Run("unset leaves pgx defaults", func(t *testing.T) { + config := base.Copy() + applyPoolConfig(config, &PostgresConfig{}) + assert.Equal(t, base.MaxConns, config.MaxConns) + assert.Equal(t, int32(0), config.MinConns) + assert.Equal(t, 30*time.Minute, config.MaxConnIdleTime) + assert.Equal(t, time.Hour, config.MaxConnLifetime) + }) + + t.Run("set fields override", func(t *testing.T) { + config := base.Copy() + maxConns := int32(8) + minConns := int32(0) + idle := time.Minute + lifetime := 10 * time.Minute + applyPoolConfig(config, &PostgresConfig{ + MaxConns: &maxConns, + MinConns: &minConns, + MaxConnIdleTime: &idle, + MaxConnLifetime: &lifetime, + }) + assert.Equal(t, int32(8), config.MaxConns) + assert.Equal(t, int32(0), config.MinConns) + assert.Equal(t, time.Minute, config.MaxConnIdleTime) + assert.Equal(t, 10*time.Minute, config.MaxConnLifetime) + }) +} + func TestResolveURLFile(t *testing.T) { tests := []struct { name string diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index c8c6efd35..8b04f25d9 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -136,11 +136,15 @@ type Config struct { // http://host: so traffic egresses in plaintext to a proxy // that originates TLS upstream. Off by default; MCPEgressPlaintext bool - Database struct { - Url string - UrlFile string - VectorEnabled bool - SkipMigrations bool + Database struct { + Url string + UrlFile string + VectorEnabled bool + SkipMigrations bool + MaxConns int // 0 = unset (pgx default) + MinConns int // -1 = unset (pgx default); 0 is a valid value + MaxConnIdleTime time.Duration // 0 = unset (pgx default) + MaxConnLifetime time.Duration // 0 = unset (pgx default) } Substrate struct { AteAPIEndpoint string @@ -183,6 +187,10 @@ func (cfg *Config) SetFlags(commandLine *flag.FlagSet) { commandLine.StringVar(&cfg.Database.UrlFile, "postgres-database-url-file", "", "Path to a file containing the PostgreSQL database URL. Takes precedence over --postgres-database-url.") commandLine.BoolVar(&cfg.Database.VectorEnabled, "database-vector-enabled", true, "Enable pgvector extension and memory table. Requires pgvector to be installed on the PostgreSQL server.") commandLine.BoolVar(&cfg.Database.SkipMigrations, "skip-migrations", false, "Do not run database migrations at startup; instead verify the database is already migrated and fail if it is not. Migrations must be applied out-of-band (e.g. from a pipeline or pre-upgrade hook). Settable via the SKIP_MIGRATIONS env var.") + commandLine.IntVar(&cfg.Database.MaxConns, "db-max-conns", 0, "Maximum number of connections in the Postgres pool. 0 leaves the pgx default.") + commandLine.IntVar(&cfg.Database.MinConns, "db-min-conns", -1, "Minimum number of connections in the Postgres pool. -1 leaves the pgx default; 0 is valid (recommended for serverless).") + commandLine.DurationVar(&cfg.Database.MaxConnIdleTime, "db-max-conn-idle-time", 0, "Maximum idle time before a Postgres pool connection is closed. 0 leaves the pgx default (30m).") + commandLine.DurationVar(&cfg.Database.MaxConnLifetime, "db-max-conn-lifetime", 0, "Maximum lifetime of a Postgres pool connection. 0 leaves the pgx default (1h).") commandLine.StringVar(&cfg.WatchNamespaces, "watch-namespaces", "", "The namespaces to watch for .") @@ -230,6 +238,32 @@ func (cfg *Config) SetFlags(commandLine *flag.FlagSet) { commandLine.StringVar(&agent_translator.DefaultAgentBindHost, "default-agent-bind-host", agent_translator.DefaultAgentBindHost, "Default host address for agent pods to bind to. Use '0.0.0.0' for IPv4 only or '::' for dual-stack (IPv4+IPv6).") } +// postgresConfigFromApp builds a database.PostgresConfig from app flags. +// Zero/unset flag values leave the corresponding pool field nil so pgx defaults apply. +func postgresConfigFromApp(dbURL string, cfg *Config) *database.PostgresConfig { + pgCfg := &database.PostgresConfig{ + URL: dbURL, + VectorEnabled: cfg.Database.VectorEnabled, + } + if cfg.Database.MaxConns > 0 { + v := int32(cfg.Database.MaxConns) + pgCfg.MaxConns = &v + } + if cfg.Database.MinConns >= 0 { + v := int32(cfg.Database.MinConns) + pgCfg.MinConns = &v + } + if cfg.Database.MaxConnIdleTime > 0 { + v := cfg.Database.MaxConnIdleTime + pgCfg.MaxConnIdleTime = &v + } + if cfg.Database.MaxConnLifetime > 0 { + v := cfg.Database.MaxConnLifetime + pgCfg.MaxConnLifetime = &v + } + return pgCfg +} + // LoadFromEnv loads configuration values from environment variables. // Flag names are converted to uppercase with underscores (e.g., metrics-bind-address -> METRICS_BIND_ADDRESS). func LoadFromEnv(fs *flag.FlagSet) error { @@ -512,10 +546,7 @@ func Start(getExtensionConfig GetExtensionConfig, extraSources []migrations.Sour } // Connect to database - db, err := database.Connect(ctx, &database.PostgresConfig{ - URL: dbURL, - VectorEnabled: cfg.Database.VectorEnabled, - }) + db, err := database.Connect(ctx, postgresConfigFromApp(dbURL, &cfg)) if err != nil { setupLog.Error(err, "unable to connect to database") os.Exit(1) diff --git a/go/core/pkg/app/app_test.go b/go/core/pkg/app/app_test.go index 78d076ca1..b06caf596 100644 --- a/go/core/pkg/app/app_test.go +++ b/go/core/pkg/app/app_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestFilterValidNamespaces(t *testing.T) { @@ -264,6 +265,47 @@ func TestDatabaseUrlFileFlag(t *testing.T) { assert.Equal(t, "/etc/credentials/db-url", cfg.Database.UrlFile) } +func TestDatabasePoolFlags(t *testing.T) { + fs := flag.NewFlagSet("test", flag.ContinueOnError) + cfg := Config{} + cfg.SetFlags(fs) + + assert.Equal(t, "0", fs.Lookup("db-max-conns").DefValue) + assert.Equal(t, "-1", fs.Lookup("db-min-conns").DefValue) + assert.Equal(t, "0s", fs.Lookup("db-max-conn-idle-time").DefValue) + assert.Equal(t, "0s", fs.Lookup("db-max-conn-lifetime").DefValue) + + t.Setenv("DB_MAX_CONNS", "4") + t.Setenv("DB_MIN_CONNS", "0") + t.Setenv("DB_MAX_CONN_IDLE_TIME", "1m") + t.Setenv("DB_MAX_CONN_LIFETIME", "10m") + assert.NoError(t, LoadFromEnv(fs)) + assert.Equal(t, 4, cfg.Database.MaxConns) + assert.Equal(t, 0, cfg.Database.MinConns) + assert.Equal(t, time.Minute, cfg.Database.MaxConnIdleTime) + assert.Equal(t, 10*time.Minute, cfg.Database.MaxConnLifetime) + + pgCfg := postgresConfigFromApp("postgres://localhost/db", &cfg) + require.NotNil(t, pgCfg.MaxConns) + assert.Equal(t, int32(4), *pgCfg.MaxConns) + require.NotNil(t, pgCfg.MinConns) + assert.Equal(t, int32(0), *pgCfg.MinConns) + require.NotNil(t, pgCfg.MaxConnIdleTime) + assert.Equal(t, time.Minute, *pgCfg.MaxConnIdleTime) + require.NotNil(t, pgCfg.MaxConnLifetime) + assert.Equal(t, 10*time.Minute, *pgCfg.MaxConnLifetime) +} + +func TestPostgresConfigFromAppUnset(t *testing.T) { + cfg := Config{} + cfg.Database.MinConns = -1 // flag default + pgCfg := postgresConfigFromApp("postgres://localhost/db", &cfg) + assert.Nil(t, pgCfg.MaxConns) + assert.Nil(t, pgCfg.MinConns) + assert.Nil(t, pgCfg.MaxConnIdleTime) + assert.Nil(t, pgCfg.MaxConnLifetime) +} + func TestSubstrateAteAPITokenFileFlag(t *testing.T) { fs := flag.NewFlagSet("test", flag.ContinueOnError) cfg := Config{} diff --git a/helm/kagent/templates/controller-configmap.yaml b/helm/kagent/templates/controller-configmap.yaml index 110902341..1d7f2d56c 100644 --- a/helm/kagent/templates/controller-configmap.yaml +++ b/helm/kagent/templates/controller-configmap.yaml @@ -60,6 +60,20 @@ data: {{- end }} DATABASE_VECTOR_ENABLED: {{ .Values.database.postgres.vectorEnabled | quote }} SKIP_MIGRATIONS: {{ .Values.database.postgres.skipMigrations | default false | quote }} + {{- with .Values.database.postgres.pool }} + {{- if and (hasKey . "maxConns") (ne .maxConns nil) }} + DB_MAX_CONNS: {{ .maxConns | quote }} + {{- end }} + {{- if and (hasKey . "minConns") (ne .minConns nil) }} + DB_MIN_CONNS: {{ .minConns | quote }} + {{- end }} + {{- if .maxConnIdleTime }} + DB_MAX_CONN_IDLE_TIME: {{ .maxConnIdleTime | quote }} + {{- end }} + {{- if .maxConnLifetime }} + DB_MAX_CONN_LIFETIME: {{ .maxConnLifetime | quote }} + {{- end }} + {{- end }} WATCH_NAMESPACES: {{ include "kagent.watchNamespaces" . | quote }} MCP_EGRESS_PLAINTEXT: {{ .Values.controller.mcpEgressPlaintext | default false | quote }} {{- if .Values.controller.a2aClientTimeout }} diff --git a/helm/kagent/tests/controller-deployment_test.yaml b/helm/kagent/tests/controller-deployment_test.yaml index c3eaeec27..00f19e721 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -596,6 +596,42 @@ tests: path: data.DATABASE_VECTOR_ENABLED value: "true" + - it: should not set DB pool env vars by default + template: controller-configmap.yaml + asserts: + - notExists: + path: data.DB_MAX_CONNS + - notExists: + path: data.DB_MIN_CONNS + - notExists: + path: data.DB_MAX_CONN_IDLE_TIME + - notExists: + path: data.DB_MAX_CONN_LIFETIME + + - it: should set DB pool env vars when pool is configured + template: controller-configmap.yaml + set: + database: + postgres: + pool: + maxConns: 4 + minConns: 0 + maxConnIdleTime: 1m + maxConnLifetime: 10m + asserts: + - equal: + path: data.DB_MAX_CONNS + value: "4" + - equal: + path: data.DB_MIN_CONNS + value: "0" + - equal: + path: data.DB_MAX_CONN_IDLE_TIME + value: "1m" + - equal: + path: data.DB_MAX_CONN_LIFETIME + value: "10m" + - it: should not set POSTGRES_DATABASE_URL in configmap template: controller-configmap.yaml asserts: diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index d6a3204c8..acad2722d 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -88,6 +88,17 @@ database: # The controller instead verifies the database is already migrated and fails if it is not. # Migrations must be applied out-of-band (e.g. from a CI/CD pipeline) before install/upgrade. skipMigrations: false + # -- Optional pgxpool settings. Leave unset/null to keep pgx library defaults + # (MaxConns≈max(4,NumCPU), MinConns=0, MaxConnIdleTime=30m, MaxConnLifetime=1h). + # For Aurora Serverless / scale-to-zero, recommended: + # minConns: 0 + # maxConnIdleTime: 1m + # maxConns: 4 + pool: + maxConns: null + minConns: null + maxConnIdleTime: "" + maxConnLifetime: "" # -- Bundled PostgreSQL instance — for development and evaluation only. # Not suitable for production. Deployed when enabled is true and url/urlFile are not set. bundled: From 1f1ac2ad95d934550a19733e2fb15abe85111499 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Fri, 7 Aug 2026 10:30:56 -0400 Subject: [PATCH 2/4] empty commit to retrigger runs Signed-off-by: Jet Chiang From 6a001700c7ebbadc899e38cd82093e8037c99f4d Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Fri, 7 Aug 2026 10:47:52 -0400 Subject: [PATCH 3/4] go format Signed-off-by: Jet Chiang --- go/core/internal/database/connect.go | 2 +- go/core/internal/database/connect_test.go | 2 +- go/core/pkg/app/app.go | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/go/core/internal/database/connect.go b/go/core/internal/database/connect.go index 6df8ee11f..8eedb63c5 100644 --- a/go/core/internal/database/connect.go +++ b/go/core/internal/database/connect.go @@ -25,7 +25,7 @@ type PostgresConfig struct { MaxConns *int32 MinConns *int32 MaxConnIdleTime *time.Duration - MaxConnLifetime *time.Duration + MaxConnLifetime *time.Duration } const ( diff --git a/go/core/internal/database/connect_test.go b/go/core/internal/database/connect_test.go index cb9e7c85b..0525301e6 100644 --- a/go/core/internal/database/connect_test.go +++ b/go/core/internal/database/connect_test.go @@ -45,7 +45,7 @@ func TestApplyPoolConfig(t *testing.T) { MaxConns: &maxConns, MinConns: &minConns, MaxConnIdleTime: &idle, - MaxConnLifetime: &lifetime, + MaxConnLifetime: &lifetime, }) assert.Equal(t, int32(8), config.MaxConns) assert.Equal(t, int32(0), config.MinConns) diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 8b04f25d9..d25d963d5 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -136,7 +136,7 @@ type Config struct { // http://host: so traffic egresses in plaintext to a proxy // that originates TLS upstream. Off by default; MCPEgressPlaintext bool - Database struct { + Database struct { Url string UrlFile string VectorEnabled bool @@ -144,7 +144,7 @@ type Config struct { MaxConns int // 0 = unset (pgx default) MinConns int // -1 = unset (pgx default); 0 is a valid value MaxConnIdleTime time.Duration // 0 = unset (pgx default) - MaxConnLifetime time.Duration // 0 = unset (pgx default) + MaxConnLifetime time.Duration // 0 = unset (pgx default) } Substrate struct { AteAPIEndpoint string From 1d332b704ca22fb20150d9259e85dcea207691fd Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Fri, 7 Aug 2026 18:42:14 -0400 Subject: [PATCH 4/4] review comments Signed-off-by: Jet Chiang --- go/core/internal/database/connect.go | 16 +++++++++++++--- go/core/pkg/app/app.go | 2 +- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/go/core/internal/database/connect.go b/go/core/internal/database/connect.go index 8eedb63c5..47daec917 100644 --- a/go/core/internal/database/connect.go +++ b/go/core/internal/database/connect.go @@ -40,8 +40,9 @@ func Connect(ctx context.Context, cfg *PostgresConfig) (*pgxpool.Pool, error) { return retryDBConnection(ctx, cfg) } -// applyPoolConfig copies non-nil pool settings from cfg onto config. -func applyPoolConfig(config *pgxpool.Config, cfg *PostgresConfig) { +// applyPoolConfig copies non-nil pool settings from cfg onto config and +// validates the resulting pool bounds. +func applyPoolConfig(config *pgxpool.Config, cfg *PostgresConfig) error { if cfg.MaxConns != nil { config.MaxConns = *cfg.MaxConns } @@ -54,6 +55,13 @@ func applyPoolConfig(config *pgxpool.Config, cfg *PostgresConfig) { if cfg.MaxConnLifetime != nil { config.MaxConnLifetime = *cfg.MaxConnLifetime } + if config.MaxConns < 1 { + return fmt.Errorf("db maxConns must be >= 1, got %d", config.MaxConns) + } + if config.MinConns > config.MaxConns { + return fmt.Errorf("db minConns (%d) cannot be greater than maxConns (%d)", config.MinConns, config.MaxConns) + } + return nil } // retryDBConnection opens a pgxpool connection, registering pgvector types when @@ -67,7 +75,9 @@ func retryDBConnection(ctx context.Context, cfg *PostgresConfig) (*pgxpool.Pool, if err != nil { return nil, fmt.Errorf("failed to parse database URL: %w", err) } - applyPoolConfig(config, cfg) + if err := applyPoolConfig(config, cfg); err != nil { + return nil, err + } if cfg.VectorEnabled { config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { return pgvectorpgx.RegisterTypes(ctx, conn) diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index d25d963d5..f7235eef2 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -188,7 +188,7 @@ func (cfg *Config) SetFlags(commandLine *flag.FlagSet) { commandLine.BoolVar(&cfg.Database.VectorEnabled, "database-vector-enabled", true, "Enable pgvector extension and memory table. Requires pgvector to be installed on the PostgreSQL server.") commandLine.BoolVar(&cfg.Database.SkipMigrations, "skip-migrations", false, "Do not run database migrations at startup; instead verify the database is already migrated and fail if it is not. Migrations must be applied out-of-band (e.g. from a pipeline or pre-upgrade hook). Settable via the SKIP_MIGRATIONS env var.") commandLine.IntVar(&cfg.Database.MaxConns, "db-max-conns", 0, "Maximum number of connections in the Postgres pool. 0 leaves the pgx default.") - commandLine.IntVar(&cfg.Database.MinConns, "db-min-conns", -1, "Minimum number of connections in the Postgres pool. -1 leaves the pgx default; 0 is valid (recommended for serverless).") + commandLine.IntVar(&cfg.Database.MinConns, "db-min-conns", -1, "Minimum number of connections in the Postgres pool. -1 leaves the pgx default.") commandLine.DurationVar(&cfg.Database.MaxConnIdleTime, "db-max-conn-idle-time", 0, "Maximum idle time before a Postgres pool connection is closed. 0 leaves the pgx default (30m).") commandLine.DurationVar(&cfg.Database.MaxConnLifetime, "db-max-conn-lifetime", 0, "Maximum lifetime of a Postgres pool connection. 0 leaves the pgx default (1h).")