▽

Source file src/github.com/cybertec-postgresql/pgwatch/v6/internal/sinks/postgres_migration_test.go

Documentation: github.com/cybertec-postgresql/pgwatch/v6/internal/sinks

     1  package sinks
     2  
     3  import (
     4  	"os"
     5  	"regexp"
     6  	"testing"
     7  
     8  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
     9  
    10  	"github.com/jackc/pgx/v5"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  // TestMigrationsCountInvariant is a fast, dependency-free unit test that guards against
    16  // the "seed drift" class of bug. The source of truth for "how many migrations exist" is
    17  // migrations() itself, exposed via registeredMigrationsCount(). This test verifies that
    18  // the rows seeded into admin.migration in admin_schema.sql stay in lock-step with the
    19  // registered migrations: both must be updated together whenever a migration is added or
    20  // removed.
    21  func TestMigrationsCountInvariant(t *testing.T) {
    22  	a := assert.New(t)
    23  
    24  	registered := registeredMigrationsCount()
    25  
    26  	// Count rows seeded into admin.migration in admin_schema.sql. The seed looks like:
    27  	//     INSERT INTO admin.migration (id, version) VALUES
    28  	//         (0, '...'),
    29  	//         (1, '...'),
    30  	//         (2, '...');
    31  	seeded := len(regexp.MustCompile(`(?m)^\s*\(\d+,\s*'`).FindAllString(sqlMetricAdminSchema, -1))
    32  	a.Equal(registered, seeded,
    33  		"registeredMigrationsCount (%d) must equal the number of rows seeded into admin.migration in admin_schema.sql (%d)",
    34  		registered, seeded)
    35  }
    36  
    37  // simulatePreV6Migration wipes every row from admin.migration so that the next
    38  // mig.Migrate() replays the full migration chain (01110 → 01180 → 01409 → 01474)
    39  // against the freshly-bootstrapped database. A fresh NewPostgresSinkMigrator
    40  // bootstrap seeds all migrations as applied (the current install path); here
    41  // we emulate an older database that must be upgraded from scratch so every
    42  // migration is exercised in every test that calls this helper.
    43  //
    44  // Callers that need a specific pre-migration object state (e.g. the legacy
    45  // function form of admin.drop_all_metric_tables) must recreate that object
    46  // after this call returns.
    47  func simulatePreV6Migration(t *testing.T, conn *pgx.Conn) {
    48  	t.Helper()
    49  	_, err := conn.Exec(ctx, `DELETE FROM admin.migration`)
    50  	require.NoError(t, err)
    51  
    52  	var count int
    53  	require.NoError(t, conn.QueryRow(ctx, `SELECT count(*) FROM admin.migration`).Scan(&count))
    54  	require.Equal(t, 0, count, "all migration rows should be wiped after rollback")
    55  }
    56  
    57  // oldSchemaMetricTable creates a metric table using the pre-v6 (dbname -> time) two-level
    58  // partitioning layout that the "01409 Switch to time-only partitioning" migration converts from:
    59  //
    60  //	public.<metric>                      PARTITION BY LIST (dbname)      -- top level
    61  //	  subpartitions.<metric>_<dbname>    PARTITION BY RANGE (time)       -- dbname level
    62  //	    subpartitions.<metric>_<...>_<w> leaf partition                  -- time level
    63  //
    64  // It seeds a couple of rows so the test can assert data survives the migration.
    65  func oldSchemaMetricTable(t *testing.T, conn *pgx.Conn, metric string) {
    66  	t.Helper()
    67  	_, err := conn.Exec(ctx, `
    68  		CREATE TABLE public.`+metric+` (LIKE admin.metrics_template INCLUDING INDEXES) PARTITION BY LIST (dbname);
    69  		COMMENT ON TABLE public.`+metric+` IS 'pgwatch-generated-metric-lvl';
    70  
    71  		CREATE TABLE subpartitions.`+metric+`_db1 PARTITION OF public.`+metric+`
    72  			FOR VALUES IN ('db1') PARTITION BY RANGE (time);
    73  
    74  		CREATE TABLE subpartitions.`+metric+`_db1_2024w01 PARTITION OF subpartitions.`+metric+`_db1
    75  			FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
    76  		COMMENT ON TABLE subpartitions.`+metric+`_db1_2024w01 IS 'pgwatch-generated-metric-time-lvl';
    77  
    78  		INSERT INTO public.`+metric+` (time, dbname, data) VALUES
    79  			('2024-01-03 10:00:00+00', 'db1', '{"x": 1}'::jsonb),
    80  			('2024-01-04 11:00:00+00', 'db1', '{"x": 2}'::jsonb);
    81  	`)
    82  	require.NoError(t, err)
    83  }
    84  
    85  // isRangePartitioned reports whether the given relation is now top-level RANGE partitioned.
    86  func isRangePartitioned(t *testing.T, conn *pgx.Conn, metric string) bool {
    87  	t.Helper()
    88  	var ok bool
    89  	err := conn.QueryRow(ctx,
    90  		`SELECT EXISTS (SELECT 1 FROM pg_partitioned_table WHERE partrelid = to_regclass($1) AND partstrat = 'r')`,
    91  		metric).Scan(&ok)
    92  	require.NoError(t, err)
    93  	return ok
    94  }
    95  
    96  // rowCount returns the number of rows currently stored in the (partitioned) metric table.
    97  func rowCount(t *testing.T, conn *pgx.Conn, metric string) int {
    98  	t.Helper()
    99  	var n int
   100  	require.NoError(t, conn.QueryRow(ctx, `SELECT count(*) FROM public.`+metric).Scan(&n))
   101  	return n
   102  }
   103  
   104  // TestMigration01409_TimeOnlyPartitioning is an end-to-end test against a real PostgreSQL
   105  // container. It builds the old (dbname -> time) partitioned layout, runs the sink migrations,
   106  // and asserts that the table is converted to time-only RANGE partitioning while preserving data.
   107  // It also runs the migration a second time to verify idempotency / re-run safety.
   108  func TestMigration01409_TimeOnlyPartitioning(t *testing.T) {
   109  	if os.Getenv("PGWATCH_TEST_SKIP_MIGRATION") != "" {
   110  		t.Skip("migration integration test skipped via PGWATCH_TEST_SKIP_MIGRATION")
   111  	}
   112  	r := require.New(t)
   113  	a := assert.New(t)
   114  
   115  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   116  	r.NoError(err)
   117  	defer pgTearDown()
   118  
   119  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   120  	r.NoError(err)
   121  
   122  	// Bootstrap the admin schema (creates admin.*, subpartitions, metrics_template, etc.).
   123  	// A fresh bootstrap seeds admin.migration with ALL migrations already applied, so we
   124  	// delete the 01409 row to simulate an older database that predates this migration.
   125  	mig, err := NewPostgresSinkMigrator(ctx, connStr)
   126  	r.NoError(err)
   127  
   128  	conn, err := pgx.Connect(ctx, connStr)
   129  	r.NoError(err)
   130  	defer conn.Close(ctx)
   131  
   132  	simulatePreV6Migration(t, conn)
   133  
   134  	const metric = "old_style_metric"
   135  	oldSchemaMetricTable(t, conn, metric)
   136  
   137  	// Sanity check: before migration the table is LIST (dbname) partitioned, not RANGE.
   138  	a.False(isRangePartitioned(t, conn, metric), "table should start as LIST(dbname) partitioned")
   139  	a.Equal(2, rowCount(t, conn, metric), "seeded rows should be present before migration")
   140  
   141  	// Run the migrations (this executes 01110, 01180 and the 01409 conversion).
   142  	r.NoError(mig.Migrate())
   143  
   144  	// After migration the top-level table must be RANGE(time) partitioned and keep its data.
   145  	a.True(isRangePartitioned(t, conn, metric), "table should be converted to RANGE(time) partitioning")
   146  	a.Equal(2, rowCount(t, conn, metric), "data must be preserved across the migration")
   147  
   148  	// The temporary *_before_v6_migration table must have been cleaned up.
   149  	var leftover bool
   150  	r.NoError(conn.QueryRow(ctx,
   151  		`SELECT to_regclass($1) IS NOT NULL`, metric+"_before_v6_migration").Scan(&leftover))
   152  	a.False(leftover, "the *_before_v6_migration scratch table should be dropped")
   153  
   154  	// Idempotency: running the migrations again must not error and must not lose data.
   155  	needs, err := mig.NeedsMigration()
   156  	r.NoError(err)
   157  	a.False(needs, "no migrations should be pending immediately after a successful migrate")
   158  
   159  	r.NoError(mig.Migrate(), "re-running Migrate() must be a no-op and not error")
   160  	a.True(isRangePartitioned(t, conn, metric))
   161  	a.Equal(2, rowCount(t, conn, metric), "data must remain intact after a second migrate")
   162  }
   163  
   164  // TestMigration01409_EmptyTable verifies the migration handles a metric table with no rows:
   165  // MIN(time) is NULL server-side, so it should create a single empty time partition without error.
   166  func TestMigration01409_EmptyTable(t *testing.T) {
   167  	if os.Getenv("PGWATCH_TEST_SKIP_MIGRATION") != "" {
   168  		t.Skip("migration integration test skipped via PGWATCH_TEST_SKIP_MIGRATION")
   169  	}
   170  	r := require.New(t)
   171  	a := assert.New(t)
   172  
   173  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   174  	r.NoError(err)
   175  	defer pgTearDown()
   176  
   177  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   178  	r.NoError(err)
   179  
   180  	mig, err := NewPostgresSinkMigrator(ctx, connStr)
   181  	r.NoError(err)
   182  
   183  	conn, err := pgx.Connect(ctx, connStr)
   184  	r.NoError(err)
   185  	defer conn.Close(ctx)
   186  
   187  	simulatePreV6Migration(t, conn)
   188  
   189  	const metric = "empty_old_metric"
   190  	_, err = conn.Exec(ctx, `
   191  		CREATE TABLE public.`+metric+` (LIKE admin.metrics_template INCLUDING INDEXES) PARTITION BY LIST (dbname);
   192  		COMMENT ON TABLE public.`+metric+` IS 'pgwatch-generated-metric-lvl';
   193  		CREATE TABLE subpartitions.`+metric+`_db1 PARTITION OF public.`+metric+`
   194  			FOR VALUES IN ('db1') PARTITION BY RANGE (time);
   195  		CREATE TABLE subpartitions.`+metric+`_db1_2024w01 PARTITION OF subpartitions.`+metric+`_db1
   196  			FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
   197  		COMMENT ON TABLE subpartitions.`+metric+`_db1_2024w01 IS 'pgwatch-generated-metric-time-lvl';
   198  	`)
   199  	r.NoError(err)
   200  
   201  	r.NoError(mig.Migrate())
   202  
   203  	a.True(isRangePartitioned(t, conn, metric), "empty table should still be converted to RANGE(time)")
   204  	a.Equal(0, rowCount(t, conn, metric))
   205  
   206  	// at least one (empty) time partition should have been created
   207  	var leaves int
   208  	r.NoError(conn.QueryRow(ctx,
   209  		`SELECT count(*) FROM pg_partition_tree($1) WHERE isleaf`, metric).Scan(&leaves))
   210  	a.GreaterOrEqual(leaves, 1, "a single empty time partition should exist")
   211  }
   212  
   213  // TestMigration01474_DropAllMetricTablesProcedure verifies that an older database that
   214  // still has admin.drop_all_metric_tables as a function gets upgraded in place to the
   215  // new procedure by migration 01474, and the upgraded routine actually drops partitions
   216  // partition-by-partition (not in one bulk DROP).
   217  func TestMigration01474_DropAllMetricTablesProcedure(t *testing.T) {
   218  	if os.Getenv("PGWATCH_TEST_SKIP_MIGRATION") != "" {
   219  		t.Skip("migration integration test skipped via PGWATCH_TEST_SKIP_MIGRATION")
   220  	}
   221  	r := require.New(t)
   222  	a := assert.New(t)
   223  
   224  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   225  	r.NoError(err)
   226  	defer pgTearDown()
   227  
   228  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   229  	r.NoError(err)
   230  
   231  	mig, err := NewPostgresSinkMigrator(ctx, connStr)
   232  	r.NoError(err)
   233  
   234  	conn, err := pgx.Connect(ctx, connStr)
   235  	r.NoError(err)
   236  	defer conn.Close(ctx)
   237  
   238  	// Simulate a pre-01474 database: drop the procedure seeded by bootstrap and recreate
   239  	// the legacy function; delete the highest migration row so the migrator picks 01474 up.
   240  	_, err = conn.Exec(ctx, `
   241  		DROP PROCEDURE IF EXISTS admin.drop_all_metric_tables();
   242  		CREATE FUNCTION admin.drop_all_metric_tables() RETURNS int AS $$ BEGIN RETURN 0; END $$ LANGUAGE plpgsql;
   243  		DELETE FROM admin.migration WHERE id >= 3;
   244  	`)
   245  	r.NoError(err)
   246  
   247  	// Before the migration the routine must be a function (prokind = 'f').
   248  	var prokindBefore string
   249  	r.NoError(conn.QueryRow(ctx, `
   250  		SELECT p.prokind FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
   251  		WHERE n.nspname = 'admin' AND p.proname = 'drop_all_metric_tables'
   252  	`).Scan(&prokindBefore))
   253  	a.Equal("f", prokindBefore, "legacy routine must start as a function")
   254  
   255  	// Build a tiny dummy metric with one partition so the upgraded procedure has something
   256  	_, err = conn.Exec(ctx, `
   257  		CREATE TABLE public.upgraded_metric (LIKE admin.metrics_template INCLUDING INDEXES) PARTITION BY RANGE (time);
   258  		COMMENT ON TABLE public.upgraded_metric IS 'pgwatch-generated-metric-lvl';
   259  		INSERT INTO admin.all_distinct_dbname_metrics (dbname, metric) VALUES ('db1', 'upgraded_metric');
   260  		CREATE TABLE subpartitions.upgraded_metric_2024w01 PARTITION OF public.upgraded_metric FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
   261  		COMMENT ON TABLE subpartitions.upgraded_metric_2024w01 IS 'pgwatch-generated-metric-time-lvl';
   262  	`)
   263  	r.NoError(err)
   264  
   265  	r.NoError(mig.Migrate())
   266  
   267  	// After migration the routine must be a procedure (prokind = 'p').
   268  	var prokindAfter string
   269  	r.NoError(conn.QueryRow(ctx, `
   270  		SELECT p.prokind FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
   271  		WHERE n.nspname = 'admin' AND p.proname = 'drop_all_metric_tables'
   272  	`).Scan(&prokindAfter))
   273  	a.Equal("p", prokindAfter, "admin.drop_all_metric_tables must be upgraded to a procedure")
   274  
   275  	// The upgraded procedure must actually drop partition-by-partition. The bare
   276  	// conn (not a transaction-wrapped one) is required: CALL ... COMMIT is rejected
   277  	// inside an explicit transaction block.
   278  	_, err = conn.Exec(ctx, "CALL admin.drop_all_metric_tables();")
   279  	r.NoError(err, "upgraded procedure must execute end-to-end")
   280  
   281  	var leftover *string
   282  	r.NoError(conn.QueryRow(ctx, "SELECT to_regclass('public.upgraded_metric')::text").Scan(&leftover))
   283  	a.Nil(leftover, "top-level metric table must be dropped by the upgraded procedure")
   284  	var subpartCount int
   285  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM pg_class WHERE relnamespace = 'subpartitions'::regnamespace").Scan(&subpartCount))
   286  	a.Equal(0, subpartCount, "partition must be dropped by the upgraded procedure")
   287  	var listingCount int
   288  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics").Scan(&listingCount))
   289  	a.Equal(0, listingCount, "listing table must be truncated by the upgraded procedure")
   290  
   291  	needs, err := mig.NeedsMigration()
   292  	r.NoError(err)
   293  	a.False(needs, "no migrations should be pending after the upgrade")
   294  }
   295  
   296  // TestMigration_AllMigrationsRunFromEmpty is the explicit all-migrations exercise.
   297  // It wipes every row from admin.migration, runs mig.Migrate() against the fresh
   298  // bootstrap, and asserts that every registered migration in the migrator chain
   299  // actually executed and recorded itself in admin.migration — including the new
   300  // 01474 routine. The migrator's "did it run" signal is purely the row count in
   301  // admin.migration, so this test guards against two failure modes at once:
   302  //   - a registered migration that fails silently (count stays below registeredMigrationsCount())
   303  //   - registeredMigrationsCount() drifting above the actual number of registered migrations
   304  func TestMigration_AllMigrationsRunFromEmpty(t *testing.T) {
   305  	if os.Getenv("PGWATCH_TEST_SKIP_MIGRATION") != "" {
   306  		t.Skip("migration integration test skipped via PGWATCH_TEST_SKIP_MIGRATION")
   307  	}
   308  	r := require.New(t)
   309  	a := assert.New(t)
   310  
   311  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   312  	r.NoError(err)
   313  	defer pgTearDown()
   314  
   315  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   316  	r.NoError(err)
   317  
   318  	mig, err := NewPostgresSinkMigrator(ctx, connStr)
   319  	r.NoError(err)
   320  
   321  	conn, err := pgx.Connect(ctx, connStr)
   322  	r.NoError(err)
   323  	defer conn.Close(ctx)
   324  
   325  	simulatePreV6Migration(t, conn)
   326  
   327  	// After wipe: exactly zero applied migrations.
   328  	var before int
   329  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM admin.migration").Scan(&before))
   330  	a.Equal(0, before)
   331  
   332  	r.NoError(mig.Migrate(), "migrate from empty must run the entire chain without error")
   333  
   334  	// After migrate: every registered migration row must be present.
   335  	var after int
   336  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM admin.migration").Scan(&after))
   337  	a.Equal(registeredMigrationsCount(), after,
   338  		"every registered migration (incl. 01474) must record itself in admin.migration after a wipe + migrate")
   339  
   340  	// The new routine from 01474 must have been created as a procedure.
   341  	var prokind string
   342  	r.NoError(conn.QueryRow(ctx, `
   343  		SELECT p.prokind FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
   344  		WHERE n.nspname = 'admin' AND p.proname = 'drop_all_metric_tables'
   345  	`).Scan(&prokind))
   346  	a.Equal("p", prokind, "01474 must install the routine as a procedure")
   347  
   348  	// ensure_partition_metric_time must have been (re)installed by 01180.
   349  	var fnExists bool
   350  	r.NoError(conn.QueryRow(ctx, `
   351  		SELECT EXISTS (SELECT 1 FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
   352  		               WHERE n.nspname = 'admin' AND p.proname = 'ensure_partition_metric_time')
   353  	`).Scan(&fnExists))
   354  	a.True(fnExists, "01180 must install admin.ensure_partition_metric_time")
   355  
   356  	// needs-migration must report clean.
   357  	needs, err := mig.NeedsMigration()
   358  	r.NoError(err)
   359  	a.False(needs, "no migrations should be pending after a full replay")
   360  
   361  	// Idempotency: a second migrate must be a clean no-op.
   362  	r.NoError(mig.Migrate(), "second migrate after full replay must be a no-op")
   363  }
   364