netbird/management/internals/shared/db/timeout_test.go

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
}