[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
+}