stillbox/pkg/database/partman/partman_test.go

package partman_test

import (
	"context"
	"testing"
	"time"

	"dynatron.me/x/stillbox/internal/common"
	"dynatron.me/x/stillbox/pkg/config"
	"dynatron.me/x/stillbox/pkg/database"
	"dynatron.me/x/stillbox/pkg/database/mocks"
	"dynatron.me/x/stillbox/pkg/database/partman"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

var mctx = mock.Anything

func inTx(s *mocks.Store) {
	s.EXPECT().InTx(mctx, mock.AnythingOfType("func(database.Store) error"), mock.AnythingOfType("pgx.TxOptions")).RunAndReturn(func(ctx context.Context, f func(db database.Store) error, po pgx.TxOptions) error {
		return f(s)
	})

}

type timeRange struct {
	start time.Time
	end   time.Time
}

func TestPartman(t *testing.T) {
	ctx := context.Background()

	timeInUTC := func(s string) time.Time {
		t, err := time.ParseInLocation("2006-01-02 15:04:05", s, time.UTC)
		if err != nil {
			panic(err)
		}

		return t
	}

	dateInUTC := func(s string) time.Time {
		t, err := time.ParseInLocation("2006-01-02", s, time.UTC)
		if err != nil {
			panic(err)
		}

		return t
	}

	etb := func(name, low, up string) database.PartitionResult {
		return database.PartitionResult{
			ParentTable: "calls",
			Schema:      "public",
			Name:        name,
			LowerBound:  low,
			UpperBound:  up,
		}
	}

	tests := []struct {
		name          string
		now           time.Time
		cfg           config.Partition
		extant        []database.PartitionResult
		expectCreate  []string
		expectDrop    []string
		expectDetach  []string
		expectSweep   []timeRange
		expectCleanup []timeRange
	}{
		{
			name: "monthly base",
			now:  timeInUTC("2024-11-28 11:37:04"),
			cfg: config.Partition{
				Enabled:      true,
				Schema:       "public",
				Interval:     "monthly",
				Retain:       2,
				Drop:         true,
				PreProvision: common.PtrTo(2),
			},
			extant: []database.PartitionResult{
				etb("calls_p_2024_10", "2024-10-01", "2024-11-01"),
				etb("calls_p_2024_09", "2024-09-01", "2024-10-01"),
				etb("calls_p_2024_08", "2024-08-01", "2024-09-01"),
				etb("calls_p_2024_07", "2024-07-01", "2024-08-01"),
			},
			expectCreate: []string{
				"calls_p_2024_11",
				"calls_p_2024_12",
				"calls_p_2025_01",
			},
			expectDrop: []string{
				"public.calls_p_2024_07",
				"public.calls_p_2024_08",
			},
			expectSweep: []timeRange{
				timeRange{start: dateInUTC("2024-07-01"), end: dateInUTC("2024-08-01")},
				timeRange{start: dateInUTC("2024-08-01"), end: dateInUTC("2024-09-01")},
			},
			expectCleanup: []timeRange{
				timeRange{start: dateInUTC("2024-07-01"), end: dateInUTC("2024-08-01")},
				timeRange{start: dateInUTC("2024-08-01"), end: dateInUTC("2024-09-01")},
			},
			expectDetach: []string{
				"public.calls_p_2024_07",
				"public.calls_p_2024_08",
			},
		},
		{
			name: "weekly base",
			now:  timeInUTC("2024-11-28 11:37:04"),
			cfg: config.Partition{
				Enabled:      true,
				Schema:       "public",
				Interval:     "weekly",
				Retain:       2,
				Drop:         false,
				PreProvision: common.PtrTo(2),
			},
			extant: []database.PartitionResult{
				etb("calls_p_2024_w44", "2024-10-28", "2024-11-04"),
				etb("calls_p_2024_w45", "2024-11-04", "2024-11-11"),
				etb("calls_p_2024_w46", "2024-11-11", "2024-11-18"),
			},
			expectCreate: []string{
				"calls_p_2024_w47",
				"calls_p_2024_w48",
				"calls_p_2024_w49",
				"calls_p_2024_w50",
			},
			expectSweep: []timeRange{
				timeRange{start: dateInUTC("2024-10-28"), end: dateInUTC("2024-11-04")},
				timeRange{start: dateInUTC("2024-11-04"), end: dateInUTC("2024-11-11")},
			},
			expectCleanup: []timeRange{
				timeRange{start: dateInUTC("2024-10-28"), end: dateInUTC("2024-11-04")},
				timeRange{start: dateInUTC("2024-11-04"), end: dateInUTC("2024-11-11")},
			},
			expectDetach: []string{
				"public.calls_p_2024_w44",
				"public.calls_p_2024_w45",
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			db := mocks.NewStore(t)
			createdPartitions := make([]string, 0, len(tc.expectCreate))
			sweptRanges := make([]timeRange, 0, len(tc.expectSweep))
			droppedPartitions := make([]string, 0, len(tc.expectDrop))
			cleanupRanges := make([]timeRange, 0, len(tc.expectCleanup))
			detachedPartitions := make([]string, 0, len(tc.expectDetach))

			db.EXPECT().
				CreatePartition(
					mctx, mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("time.Time"),
					mock.AnythingOfType("time.Time"),
				).
				Run(func(ctx context.Context, tableName, partitionName string, start, end time.Time) {
					createdPartitions = append(createdPartitions, partitionName)
				}).Return(nil)

			db.EXPECT().
				SweepCalls(
					mctx, mock.AnythingOfType("pgtype.Timestamptz"), mock.AnythingOfType("pgtype.Timestamptz"),
				).
				Run(func(ctx context.Context, start, end pgtype.Timestamptz) {
					sweptRanges = append(sweptRanges, timeRange{start: start.Time, end: end.Time})
				}).Return(nil)

			db.EXPECT().
				CleanupSweptCalls(
					mctx, mock.AnythingOfType("pgtype.Timestamptz"), mock.AnythingOfType("pgtype.Timestamptz"),
				).Run(func(ctx context.Context, start, end pgtype.Timestamptz) {
				cleanupRanges = append(cleanupRanges, timeRange{start: start.Time, end: end.Time})
			}).Return(nil)

			if tc.cfg.Drop {
				db.EXPECT().
					DropPartition(mctx, mock.AnythingOfType("string")).
					Run(func(ctx context.Context, partName string) {
						droppedPartitions = append(droppedPartitions, partName)
					}).Return(nil)
			}

			db.EXPECT().
				DetachPartition(
					mctx, mock.AnythingOfType("string")).
				Run(func(ctx context.Context, partName string) {
					detachedPartitions = append(detachedPartitions, partName)
				}).Return(nil)

			inTx(db)

			db.EXPECT().GetTablePartitions(mctx, "public", "calls").Return(tc.extant, nil)

			pm, err := partman.New(db, tc.cfg)
			require.NoError(t, err)

			err = pm.Check(ctx, tc.now)
			require.NoError(t, err)

			assert.ElementsMatch(t, tc.expectCreate, createdPartitions, "created partitions")
			assert.ElementsMatch(t, tc.expectSweep, sweptRanges, "swept ranges")
			assert.ElementsMatch(t, tc.expectDrop, droppedPartitions, "dropped partitions")
			assert.ElementsMatch(t, tc.expectCleanup, cleanupRanges, "cleaned up ranges")
			assert.ElementsMatch(t, tc.expectDetach, detachedPartitions, "detached partitions")
		})
	}
}