[management] Add configurable statement timeout to postgres store (#8381)

management/internals/modules/reverseproxy/accesslogs/manager/repository.go

@@ -14,18 +14,21 @@ import (
 	"github.com/netbirdio/netbird/shared/management/status"
 )
 
+const deleteBatchSize = 10000
+
 type sqlRepository struct {
-	conn *db.Conn
-	db   *gorm.DB
+	conn            *db.Conn
+	db              *gorm.DB
+	deleteBatchSize int
 }
 
 // NewRepository returns the access log repository backed by conn.
 func NewRepository(conn *db.Conn) accesslogs.Repository {
-	return &sqlRepository{conn: conn, db: conn.DB(nil)}
+	return &sqlRepository{conn: conn, db: conn.DB(nil), deleteBatchSize: deleteBatchSize}
 }
 
 func (r *sqlRepository) WithTx(tx *db.Tx) accesslogs.Repository {
-	return &sqlRepository{conn: r.conn, db: r.conn.DB(tx)}
+	return &sqlRepository{conn: r.conn, db: r.conn.DB(tx), deleteBatchSize: r.deleteBatchSize}
 }
 
 func (r *sqlRepository) Create(ctx context.Context, entry *accesslogs.AccessLogEntry) error {
@@ -73,12 +76,28 @@ func (r *sqlRepository) ListByAccount(ctx context.Context, lockStrength db.Locki
 }
 
 func (r *sqlRepository) DeleteOlderThan(ctx context.Context, olderThan time.Time) (int64, error) {
-	result := r.db.Where("timestamp < ?", olderThan).Delete(&accesslogs.AccessLogEntry{})
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
-		return 0, status.Errorf(status.Internal, "failed to delete old access logs")
+	var deleted int64
+	for {
+		if err := ctx.Err(); err != nil {
+			return deleted, err
+		}
+
+		batch := r.db.Model(&accesslogs.AccessLogEntry{}).
+			Select("id").
+			Where("timestamp < ?", olderThan).
+			Order("timestamp").
+			Limit(r.deleteBatchSize)
+		result := r.db.Where("id IN (?)", batch).Delete(&accesslogs.AccessLogEntry{})
+		if result.Error != nil {
+			log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
+			return deleted, status.Errorf(status.Internal, "failed to delete old access logs")
+		}
+
+		deleted += result.RowsAffected
+		if result.RowsAffected < int64(r.deleteBatchSize) {
+			return deleted, nil
+		}
 	}
-	return result.RowsAffected, nil
 }
 
 func applyFilters(query *gorm.DB, filter accesslogs.AccessLogFilter) *gorm.DB {

management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go

@@ -3,6 +3,8 @@ package manager
 import (
 	"context"
 	"errors"
+	"fmt"
+	"os"
 	"testing"
 	"time"
 
@@ -12,6 +14,7 @@ import (
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
 	"github.com/netbirdio/netbird/management/internals/shared/db"
 	"github.com/netbirdio/netbird/management/internals/shared/db/dbtest"
+	"github.com/netbirdio/netbird/management/server/testutil"
 )
 
 func newTestRepository(t *testing.T) (accesslogs.Repository, *db.Conn) {
@@ -157,3 +160,70 @@ func TestSqlRepository_ListByAccount_Search(t *testing.T) {
 	assert.Equal(t, []string{"req-api", "req-web"}, search("example.com"), "search should list every match newest first")
 	assert.Empty(t, search("nothing-matches"))
 }
+
+func TestSqlRepository_DeleteOlderThan_DeletesInBatches(t *testing.T) {
+	cases := map[string]struct{ old, recent int }{
+		"partial last batch":      {old: 5, recent: 2},
+		"exact multiple of batch": {old: 4, recent: 1},
+		"nothing to delete":       {old: 0, recent: 3},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			repo, _ := newTestRepository(t)
+			assertBatchedDelete(t, repo, tc.old, tc.recent)
+		})
+	}
+}
+
+func TestSqlRepository_DeleteOlderThan_DeletesInBatchesOnPostgres(t *testing.T) {
+	if os.Getenv("NETBIRD_STORE_ENGINE") != string(db.PostgresStoreEngine) {
+		t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
+	}
+	cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+	conn, err := db.OpenPostgres(context.Background(), dsn, db.DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	require.NoError(t, conn.AutoMigrate(&accesslogs.AccessLogEntry{}))
+	require.NoError(t, conn.DB(nil).Exec("DELETE FROM access_log_entries").Error)
+
+	assertBatchedDelete(t, NewRepository(conn), 5, 2)
+}
+
+func TestSqlRepository_DeleteOlderThan_StopsWhenContextDone(t *testing.T) {
+	repo, _ := newTestRepository(t)
+	require.NoError(t, repo.Create(context.Background(), newEntry("old", "acc", "GET", 48*time.Hour)))
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour))
+
+	require.ErrorIs(t, err, context.Canceled)
+	assert.Zero(t, deleted)
+	assert.EqualValues(t, 1, countEntries(t, repo), "nothing should be deleted once the context is done")
+}
+
+func assertBatchedDelete(t *testing.T, repo accesslogs.Repository, old, recent int) {
+	t.Helper()
+	repo.(*sqlRepository).deleteBatchSize = 2
+	ctx := context.Background()
+	for i := 0; i < old; i++ {
+		require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("old-%d", i), "acc", "GET", 48*time.Hour)))
+	}
+	for i := 0; i < recent; i++ {
+		require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("new-%d", i), "acc", "GET", time.Hour)))
+	}
+
+	deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour))
+	require.NoError(t, err)
+	assert.EqualValues(t, old, deleted, "every old entry should be deleted across batches")
+	assert.EqualValues(t, recent, countEntries(t, repo), "recent entries should be kept")
+}
+
+func countEntries(t *testing.T, repo accesslogs.Repository) int64 {
+	t.Helper()
+	_, total, err := repo.ListByAccount(context.Background(), db.LockingStrengthNone, "acc", accesslogs.AccessLogFilter{Page: 1, PageSize: 10})
+	require.NoError(t, err)
+	return total
+}

management/internals/shared/db/conn.go

@@ -11,6 +11,8 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 	log "github.com/sirupsen/logrus"
 	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/util"
 )
 
 const (
@@ -64,13 +66,11 @@ func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.
 		opt(&options)
 	}
 
-	txTimeout := defaultTransactionTimeout
-	if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
-		if parsed, err := time.ParseDuration(v); err == nil {
-			txTimeout = parsed
-		}
-	}
+	txTimeout := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout)
 	log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)
+	if d := util.DurationFromEnv(statementTimeoutEnv, 0); engine == PostgresStoreEngine && d > 0 {
+		log.WithContext(ctx).Infof("Setting statement timeout to %v", d)
+	}
 
 	conns := runtime.NumCPU()
 	configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))

management/internals/shared/db/open.go

@@ -10,7 +10,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/jackc/pgx/v5/stdlib"
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
@@ -100,9 +102,10 @@ func OpenSqliteFile(ctx context.Context, dataDir, storeFile string) (*Conn, erro
 	return conn, nil
 }
 
-// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by pool.
+// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by
+// pool, applying the configured timeouts to every session.
 func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Option) (*Conn, error) {
-	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	gormDB, err := openGormPostgres(ctx, dsn)
 	if err != nil {
 		return nil, err
 	}
@@ -123,7 +126,7 @@ func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Opti
 // OpenPostgresWithoutPool opens a Postgres database through gorm alone, for
 // repositories that never take the pgx read paths.
 func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*Conn, error) {
-	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	gormDB, err := openGormPostgres(ctx, dsn)
 	if err != nil {
 		return nil, err
 	}
@@ -135,13 +138,29 @@ func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*
 	return conn, nil
 }
 
-// OpenPgxPool opens a pgx pool to dsn sized by cfg and verifies it with a ping.
+func openGormPostgres(ctx context.Context, dsn string) (*gorm.DB, error) {
+	config, err := pgx.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("parse database config: %w", err)
+	}
+	sqlDB := stdlib.OpenDB(*config, stdlib.OptionAfterConnect(sessionTimeouts()))
+	gormDB, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDB}), GormConfig())
+	if err != nil {
+		_ = sqlDB.Close()
+		return nil, err
+	}
+	return gormDB, nil
+}
+
+// OpenPgxPool opens a pgx pool to dsn sized by cfg, applying the configured
+// timeouts to every session, and verifies it with a ping.
 func OpenPgxPool(ctx context.Context, dsn string, cfg PoolConfig) (*pgxpool.Pool, error) {
 	config, err := pgxpool.ParseConfig(dsn)
 	if err != nil {
 		return nil, fmt.Errorf("unable to parse database config: %w", err)
 	}
 
+	config.AfterConnect = sessionTimeouts()
 	config.MaxConns = cfg.MaxConns
 	config.MinConns = cfg.MinConns
 	config.MaxConnLifetime = cfg.MaxConnLifetime

management/internals/shared/db/timeout.go

@@ -0,0 +1,52 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+
+	"github.com/netbirdio/netbird/util"
+)
+
+const (
+	transactionTimeoutEnv = "NB_STORE_TRANSACTION_TIMEOUT"
+	statementTimeoutEnv   = "NB_STORE_STATEMENT_TIMEOUT"
+)
+
+// sessionTimeouts sets the configured timeouts on each new Postgres session with
+// SET, since poolers such as PgBouncer reject them as startup parameters.
+func sessionTimeouts() func(context.Context, *pgx.Conn) error {
+	var settings []string
+	if d := util.DurationFromEnv(statementTimeoutEnv, 0); d > 0 {
+		settings = append(settings, "SET statement_timeout = "+postgresMillis(d))
+	}
+	if d := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout); d > 0 {
+		settings = append(settings, "SET idle_in_transaction_session_timeout = "+postgresMillis(d))
+	}
+	query := strings.Join(settings, "; ")
+
+	return func(ctx context.Context, conn *pgx.Conn) error {
+		if query == "" {
+			return nil
+		}
+		if _, err := conn.Exec(ctx, query); err != nil {
+			return fmt.Errorf("set session timeouts: %w", err)
+		}
+		return nil
+	}
+}
+
+// postgresMillis converts d to milliseconds for a Postgres timeout, rounding up
+// so a positive d never becomes 0 (disabled) and capping at the server maximum.
+func postgresMillis(d time.Duration) string {
+	ms := d.Milliseconds()
+	if d%time.Millisecond != 0 {
+		ms++
+	}
+	return strconv.FormatInt(min(ms, math.MaxInt32), 10)
+}

management/internals/shared/db/timeout_test.go

@@ -0,0 +1,146 @@
+package db
+
+import (
+	"context"
+	"errors"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/server/testutil"
+)
+
+// queryCanceledCode is the SQLSTATE Postgres reports when statement_timeout fires.
+const queryCanceledCode = "57014"
+
+func TestPostgresMillis(t *testing.T) {
+	cases := map[time.Duration]string{
+		time.Minute:             "60000",
+		1500 * time.Microsecond: "2",
+		500 * time.Microsecond:  "1",
+		720 * time.Hour:         "2147483647",
+	}
+	for d, want := range cases {
+		assert.Equal(t, want, postgresMillis(d), "milliseconds for %v", d)
+	}
+}
+
+func TestSessionTimeouts_NothingToApply(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "0")
+	t.Setenv(transactionTimeoutEnv, "0")
+	ctx := context.Background()
+
+	assert.NoError(t, sessionTimeouts()(ctx, nil), "the hook should not touch the session when both timeouts are disabled")
+}
+
+func TestOpenPostgres_BoundsEverySession(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	t.Setenv(transactionTimeoutEnv, "3m")
+	conn := openPostgresTestConn(t)
+	ctx := context.Background()
+
+	requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+
+	_, err := conn.Pool(nil).Exec(ctx, "SELECT pg_sleep(1)")
+	requireStatementTimeout(t, err)
+
+	var inTx string
+	require.NoError(t, conn.RunInTx(ctx, func(tx *Tx) error {
+		return conn.DB(tx).Raw("SHOW statement_timeout").Scan(&inTx).Error
+	}))
+	assert.Equal(t, "1min", inTx, "transactions should keep their own 1 minute bound")
+
+	assert.Equal(t, "3min", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+		"the idle transaction timeout should follow the transaction timeout")
+}
+
+func TestOpenPostgres_NoStatementTimeoutKeepsServerSetting(t *testing.T) {
+	for name, value := range map[string]string{"unset": "", "zero": "0"} {
+		t.Run(name, func(t *testing.T) {
+			t.Setenv(statementTimeoutEnv, value)
+			if value == "" {
+				require.NoError(t, os.Unsetenv(statementTimeoutEnv))
+			}
+			ctx := context.Background()
+			dsn := postgresTestDSN(t)
+			admin, err := pgx.Connect(ctx, dsn)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = admin.Close(ctx) })
+			_, err = admin.Exec(ctx, "ALTER ROLE CURRENT_USER SET statement_timeout = '7s'")
+			require.NoError(t, err)
+
+			conn, err := OpenPostgres(ctx, dsn, DefaultPoolConfig)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = conn.Close() })
+
+			assert.Equal(t, "7s", showSetting(t, conn.DB(nil), "statement_timeout"), "sessions should keep the server setting")
+		})
+	}
+}
+
+func TestOpenPostgres_CapsTimeoutsAtServerMaximum(t *testing.T) {
+	t.Setenv(transactionTimeoutEnv, "720h")
+	conn := openPostgresTestConn(t)
+
+	assert.Equal(t, "2147483647ms", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+		"a timeout above the server maximum should be capped instead of failing the connection")
+}
+
+func TestOpenPostgresWithoutPool_BoundsSessions(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	conn, err := OpenPostgresWithoutPool(context.Background(), postgresTestDSN(t))
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+}
+
+func TestOpenPgxPool_BoundsSessions(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	ctx := context.Background()
+	pool, err := OpenPgxPool(ctx, postgresTestDSN(t), DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(pool.Close)
+
+	_, err = pool.Exec(ctx, "SELECT pg_sleep(1)")
+	requireStatementTimeout(t, err)
+}
+
+func postgresTestDSN(t *testing.T) string {
+	t.Helper()
+	if os.Getenv("NETBIRD_STORE_ENGINE") != string(PostgresStoreEngine) {
+		t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
+	}
+	cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+	return dsn
+}
+
+func openPostgresTestConn(t *testing.T) *Conn {
+	t.Helper()
+	conn, err := OpenPostgres(context.Background(), postgresTestDSN(t), DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	return conn
+}
+
+func requireStatementTimeout(t *testing.T, err error) {
+	t.Helper()
+	var pgErr *pgconn.PgError
+	require.True(t, errors.As(err, &pgErr), "expected a Postgres error, got %v", err)
+	assert.Equal(t, queryCanceledCode, pgErr.Code, "the statement should be cancelled by statement_timeout")
+}
+
+func showSetting(t *testing.T, db *gorm.DB, name string) string {
+	t.Helper()
+	var value string
+	require.NoError(t, db.Raw("SHOW "+name).Scan(&value).Error)
+	return value
+}