Test PgBouncer against our code-level partitioning logic (#3178)

* test partition drops with pgbouncer

* dedicated DDLPool to bypass pgbouncer

* fix generate

* add env var

* min env var

* update documentation

* no go mod redundancy

Author: mnafees

examples/go/guides/go.mod

@@ -7,7 +7,7 @@ The Hatchet server and engine can be configured via environment variables using
 Hatchet uses the following environment variable prefixes:
 
 - **`SERVER_`** (173 variables) - Main server configuration including runtime, authentication, encryption, monitoring, and integrations
-- **`DATABASE_`** (15 variables) - PostgreSQL database connection and configuration
+- **`DATABASE_`** (19 variables) - PostgreSQL database connection and configuration
 - **`READ_REPLICA_`** (4 variables) - Read replica database configuration
 - **`ADMIN_`** (3 variables) - Administrator user setup for initial seeding
 - **`DEFAULT_`** (3 variables) - Default tenant configuration
@@ -140,36 +140,40 @@ Variables marked with ⚠️ are conditionally required when specific features a
 
 ## Database Configuration
 
-| Variable                      | Description                                           | Default Value       |
-| ----------------------------- | ----------------------------------------------------- | ------------------- |
-| `DATABASE_URL`                | PostgreSQL connection string                          | `127.0.0.1`         |
-| `DATABASE_POSTGRES_HOST`      | PostgreSQL host                                       | `127.0.0.1`         |
-| `DATABASE_POSTGRES_PORT`      | PostgreSQL port                                       | `5431`              |
-| `DATABASE_POSTGRES_USERNAME`  | PostgreSQL username                                   | `hatchet`           |
-| `DATABASE_POSTGRES_PASSWORD`  | PostgreSQL password                                   | `hatchet`           |
-| `DATABASE_POSTGRES_DB_NAME`   | PostgreSQL database name                              | `hatchet`           |
-| `DATABASE_POSTGRES_SSL_MODE`  | PostgreSQL SSL mode                                   | `disable`           |
-| `DATABASE_MAX_CONNS`          | Max database connections                              | `50`                |
-| `DATABASE_MIN_CONNS`          | Min database connections                              | `10`                |
-| `DATABASE_MAX_QUEUE_CONNS`    | Max queue connections                                 | `50`                |
-| `DATABASE_MIN_QUEUE_CONNS`    | Min queue connections                                 | `10`                |
-| `DATABASE_MAX_CONN_LIFETIME`  | Max lifetime of a connection                          | `15m`               |
-| `DATABASE_MAX_CONN_IDLE_TIME` | Max time a connection can be idle before being closed | `1m`                |
-| `DATABASE_LOG_QUERIES`        | Log database queries                                  | `false`             |
-| `CACHE_DURATION`              | Cache duration                                        | `5s`                |
-| `ADMIN_EMAIL`                 | Admin email for seeding                               | `admin@example.com` |
-| `ADMIN_PASSWORD`              | Admin password for seeding                            | `Admin123!!`        |
-| `ADMIN_NAME`                  | Admin name for seeding                                | `Admin`             |
-| `DEFAULT_TENANT_NAME`         | Default tenant name                                   | `Default`           |
-| `DEFAULT_TENANT_SLUG`         | Default tenant slug                                   | `default`           |
-| `DEFAULT_TENANT_ID`           | Default tenant ID                                     |                     |
-| `SEED_DEVELOPMENT`            | Development seeding flag                              |                     |
-| `READ_REPLICA_ENABLED`        | Enable read replica                                   | `false`             |
-| `READ_REPLICA_DATABASE_URL`   | Read replica database URL                             |                     |
-| `READ_REPLICA_MAX_CONNS`      | Read replica max connections                          | `50`                |
-| `READ_REPLICA_MIN_CONNS`      | Read replica min connections                          | `10`                |
-| `DATABASE_LOGGER_LEVEL`       | Database logger level                                 |                     |
-| `DATABASE_LOGGER_FORMAT`      | Database logger format                                |                     |
+| Variable                      | Description                                                                | Default Value       |
+| ----------------------------- | -------------------------------------------------------------------------- | ------------------- |
+| `DATABASE_URL`                | PostgreSQL connection string                                               | `127.0.0.1`         |
+| `DATABASE_POSTGRES_HOST`      | PostgreSQL host                                                            | `127.0.0.1`         |
+| `DATABASE_POSTGRES_PORT`      | PostgreSQL port                                                            | `5431`              |
+| `DATABASE_POSTGRES_USERNAME`  | PostgreSQL username                                                        | `hatchet`           |
+| `DATABASE_POSTGRES_PASSWORD`  | PostgreSQL password                                                        | `hatchet`           |
+| `DATABASE_POSTGRES_DB_NAME`   | PostgreSQL database name                                                   | `hatchet`           |
+| `DATABASE_POSTGRES_SSL_MODE`  | PostgreSQL SSL mode                                                        | `disable`           |
+| `DATABASE_MAX_CONNS`          | Max database connections                                                   | `50`                |
+| `DATABASE_MIN_CONNS`          | Min database connections                                                   | `10`                |
+| `DATABASE_MAX_QUEUE_CONNS`    | Max queue connections                                                      | `50`                |
+| `DATABASE_MIN_QUEUE_CONNS`    | Min queue connections                                                      | `10`                |
+| `DATABASE_MAX_CONN_LIFETIME`  | Max lifetime of a connection                                               | `15m`               |
+| `DATABASE_MAX_CONN_IDLE_TIME` | Max time a connection can be idle before being closed                      | `1m`                |
+| `DATABASE_LOG_QUERIES`        | Log database queries                                                       | `false`             |
+| `DATABASE_PGBOUNCER_ENABLED`  | Enable pgbouncer support; requires `DATABASE_DIRECT_URL` to be set         | `false`             |
+| `DATABASE_DIRECT_URL`         | Direct PostgreSQL connection string bypassing pgbouncer for DDL operations |                     |
+| `DATABASE_DIRECT_MAX_CONNS`   | Max connections for the direct (non-pgbouncer) pool                        | `2`                 |
+| `DATABASE_DIRECT_MIN_CONNS`   | Min connections for the direct (non-pgbouncer) pool                        | `1`                 |
+| `CACHE_DURATION`              | Cache duration                                                             | `5s`                |
+| `ADMIN_EMAIL`                 | Admin email for seeding                                                    | `admin@example.com` |
+| `ADMIN_PASSWORD`              | Admin password for seeding                                                 | `Admin123!!`        |
+| `ADMIN_NAME`                  | Admin name for seeding                                                     | `Admin`             |
+| `DEFAULT_TENANT_NAME`         | Default tenant name                                                        | `Default`           |
+| `DEFAULT_TENANT_SLUG`         | Default tenant slug                                                        | `default`           |
+| `DEFAULT_TENANT_ID`           | Default tenant ID                                                          |                     |
+| `SEED_DEVELOPMENT`            | Development seeding flag                                                   |                     |
+| `READ_REPLICA_ENABLED`        | Enable read replica                                                        | `false`             |
+| `READ_REPLICA_DATABASE_URL`   | Read replica database URL                                                  |                     |
+| `READ_REPLICA_MAX_CONNS`      | Read replica max connections                                               | `50`                |
+| `READ_REPLICA_MIN_CONNS`      | Read replica min connections                                               | `10`                |
+| `DATABASE_LOGGER_LEVEL`       | Database logger level                                                      |                     |
+| `DATABASE_LOGGER_FORMAT`      | Database logger format                                                     |                     |
 
 ## Security Check Configuration
 

frontend/docs/pages/self-hosting/configuration-options.mdx

@@ -38,6 +38,7 @@ require (
 	github.com/pingcap/errors v0.11.4
 	github.com/posthog/posthog-go v1.10.0
 	github.com/pressly/goose/v3 v3.27.0
+	github.com/sashabaranov/go-openai v1.41.2
 	github.com/sethvargo/go-retry v0.3.0
 	github.com/spf13/cobra v1.10.2
 	github.com/spf13/viper v1.21.0

go.mod

@@ -407,6 +407,8 @@ github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sagikazarmark/locafero v0.11.0 h1:1iurJgmM9G3PA/I+wWYIOw/5SyBtxapeHDcg+AAIFXc=
 github.com/sagikazarmark/locafero v0.11.0/go.mod h1:nVIGvgyzw595SUSUE6tvCp3YYTeHs15MvlmU87WwIik=
+github.com/sashabaranov/go-openai v1.41.2 h1:vfPRBZNMpnqu8ELsclWcAvF19lDNgh1t6TVfFFOPiSM=
+github.com/sashabaranov/go-openai v1.41.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/sethvargo/go-retry v0.3.0 h1:EEt31A35QhrcRZtrYFDTBg91cqZVnFL2navjDrah2SE=
 github.com/sethvargo/go-retry v0.3.0/go.mod h1:mNX17F0C/HguQMyMyJxcnU471gOZGxCLyYaFyAZraas=
 github.com/shirou/gopsutil/v4 v4.25.7 h1:bNb2JuqKuAu3tRlPv5piSmBZyMfecwQ+t/ILq+1JqVM=

go.sum

@@ -29,6 +29,18 @@ type ConfigFile struct {
 	MaxQueueConns int `mapstructure:"maxQueueConns" json:"maxQueueConns,omitempty" default:"50"`
 	MinQueueConns int `mapstructure:"minQueueConns" json:"minQueueConns,omitempty" default:"10"`
 
+	// PgBouncerEnabled indicates that the main DATABASE_URL connects through pgbouncer.
+	// When true, DATABASE_DIRECT_URL must also be set so that DDL operations like
+	// DETACH PARTITION CONCURRENTLY can bypass pgbouncer.
+	PgBouncerEnabled bool `mapstructure:"pgbouncerEnabled" json:"pgbouncerEnabled,omitempty" default:"false"`
+
+	// DirectDatabaseURL is a connection string that bypasses pgbouncer and connects directly
+	// to PostgreSQL. This is used for DDL operations like DETACH PARTITION CONCURRENTLY that
+	// cannot run inside a transaction block. Required when PgBouncerEnabled is true.
+	DirectDatabaseURL      string `mapstructure:"directDatabaseUrl" json:"directDatabaseUrl,omitempty" default:""`
+	DirectDatabaseMaxConns int    `mapstructure:"directDatabaseMaxConns" json:"directDatabaseMaxConns,omitempty" default:"2"`
+	DirectDatabaseMinConns int    `mapstructure:"directDatabaseMinConns" json:"directDatabaseMinConns,omitempty" default:"1"`
+
 	MaxConnLifetime time.Duration `mapstructure:"maxConnLifetime" json:"maxConnLifetime,omitempty" default:"15m"`
 	MaxConnIdleTime time.Duration `mapstructure:"maxConnIdleTime" json:"maxConnIdleTime,omitempty" default:"1m"`
 
@@ -67,6 +79,11 @@ type Layer struct {
 
 	QueuePool *pgxpool.Pool
 
+	// DirectPool is a small pool (max 2 connections) that bypasses pgbouncer and connects
+	// directly to PostgreSQL for DDL operations like DETACH PARTITION CONCURRENTLY.
+	// If pgbouncer is not used, this may be nil and callers should fall back to Pool.
+	DirectPool *pgxpool.Pool
+
 	V1 v1.Repository
 
 	Seed SeedConfigFile
@@ -87,6 +104,11 @@ func BindAllEnv(v *viper.Viper) {
 	_ = v.BindEnv("maxConnLifetime", "DATABASE_MAX_CONN_LIFETIME")
 	_ = v.BindEnv("maxConnIdleTime", "DATABASE_MAX_CONN_IDLE_TIME")
 
+	_ = v.BindEnv("pgbouncerEnabled", "DATABASE_PGBOUNCER_ENABLED")
+	_ = v.BindEnv("directDatabaseUrl", "DATABASE_DIRECT_URL")
+	_ = v.BindEnv("directDatabaseMaxConns", "DATABASE_DIRECT_MAX_CONNS")
+	_ = v.BindEnv("directDatabaseMinConns", "DATABASE_DIRECT_MIN_CONNS")
+
 	_ = v.BindEnv("readReplicaEnabled", "READ_REPLICA_ENABLED")
 	_ = v.BindEnv("readReplicaDatabaseUrl", "READ_REPLICA_DATABASE_URL")
 	_ = v.BindEnv("readReplicaMaxConns", "READ_REPLICA_MAX_CONNS")

pkg/config/database/config.go

@@ -263,6 +263,40 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 		}
 	}
 
+	// A small direct pool for DDL operations that cannot go through pgbouncer
+	// (e.g. DETACH PARTITION CONCURRENTLY which cannot run inside a transaction block).
+	var directPool *pgxpool.Pool
+
+	if cf.PgBouncerEnabled && cf.DirectDatabaseURL == "" {
+		return nil, fmt.Errorf("DATABASE_PGBOUNCER_ENABLED is set but DATABASE_DIRECT_URL is not; " +
+			"a direct PostgreSQL connection is required for DDL operations like DETACH PARTITION CONCURRENTLY")
+	}
+
+	if cf.DirectDatabaseURL != "" {
+		directConfig, err := pgxpool.ParseConfig(cf.DirectDatabaseURL)
+		if err != nil {
+			return nil, fmt.Errorf("could not parse direct database url: %w", err)
+		}
+
+		if cf.DirectDatabaseMaxConns != 0 {
+			directConfig.MaxConns = int32(cf.DirectDatabaseMaxConns) // nolint: gosec
+		}
+
+		if cf.DirectDatabaseMinConns != 0 {
+			directConfig.MinConns = int32(cf.DirectDatabaseMinConns) // nolint: gosec
+		}
+
+		directConfig.MaxConnLifetime = cf.MaxConnLifetime
+		directConfig.MaxConnIdleTime = cf.MaxConnIdleTime
+		directConfig.AfterConnect = pgxpoolConnAfterConnect
+		directConfig.ConnConfig.Tracer = otelpgx.NewTracer()
+
+		directPool, err = pgxpool.NewWithConfig(context.Background(), directConfig)
+		if err != nil {
+			return nil, fmt.Errorf("could not connect to direct database: %w", err)
+		}
+	}
+
 	ch := cache.New(cf.CacheDuration)
 
 	retentionPeriod, err := time.ParseDuration(scf.Runtime.Limits.DefaultTenantRetentionPeriod)
@@ -305,6 +339,7 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 
 	v1, cleanupV1 := repov1.NewRepository(
 		pool,
+		directPool,
 		&l,
 		cf.CacheDuration,
 		retentionPeriod,
@@ -328,10 +363,11 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 
 			return cleanupV1()
 		},
-		Pool:      pool,
-		QueuePool: pool,
-		V1:        v1,
-		Seed:      cf.Seed,
+		Pool:       pool,
+		QueuePool:  pool,
+		DirectPool: directPool,
+		V1:         v1,
+		Seed:       cf.Seed,
 	}, nil
 
 }

pkg/config/loader/loader.go

@@ -309,7 +309,7 @@ func NewOLAPRepositoryFromPool(
 ) (OLAPRepository, func() error) {
 	v := validator.NewDefaultValidator()
 
-	shared, cleanupShared := newSharedRepository(pool, v, l, payloadStoreOpts, tenantLimitConfig, enforceLimits, cacheDuration, enableDurableUserEventLog)
+	shared, cleanupShared := newSharedRepository(pool, nil, v, l, payloadStoreOpts, tenantLimitConfig, enforceLimits, cacheDuration, enableDurableUserEventLog)
 
 	return newOLAPRepository(shared, olapRetentionPeriod, shouldPartitionEventsTables, statusUpdateBatchSizeLimits), cleanupShared
 }
@@ -400,10 +400,14 @@ func (r *OLAPRepositoryImpl) UpdateTablePartitions(ctx context.Context) error {
 		r.l.Warn().Msgf("removing partitions before %s using retention period of %s", removeBefore.Format(time.RFC3339), r.olapRetentionPeriod)
 	}
 
+	// Use the direct pool (bypasses pgbouncer) for DDL operations because
+	// DETACH PARTITION CONCURRENTLY cannot run inside a transaction block.
+	ddlPool := r.DDLPool()
+
 	for _, partition := range partitions {
 		r.l.Debug().Msgf("detaching partition %s", partition.PartitionName)
 
-		conn, release, err := sqlchelpers.AcquireConnectionWithStatementTimeout(ctx, r.pool, r.l, 30*60*1000) // 30 minutes
+		conn, release, err := sqlchelpers.AcquireConnectionWithStatementTimeout(ctx, ddlPool, r.l, 30*60*1000) // 30 minutes
 
 		if err != nil {
 			return err

pkg/repository/olap.go

@@ -92,6 +92,7 @@ type repositoryImpl struct {
 
 func NewRepository(
 	pool *pgxpool.Pool,
+	directPool *pgxpool.Pool,
 	l *zerolog.Logger,
 	cacheDuration time.Duration,
 	taskRetentionPeriod, olapRetentionPeriod time.Duration,
@@ -105,7 +106,7 @@ func NewRepository(
 ) (Repository, func() error) {
 	v := validator.NewDefaultValidator()
 
-	shared, cleanupShared := newSharedRepository(pool, v, l, payloadStoreOpts, tenantLimitConfig, enforceLimits, cacheDuration, enableDurableUserEventLog)
+	shared, cleanupShared := newSharedRepository(pool, directPool, v, l, payloadStoreOpts, tenantLimitConfig, enforceLimits, cacheDuration, enableDurableUserEventLog)
 
 	mq, cleanupMq := newMessageQueueRepository(shared)