▽

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

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

     1  package sinks
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    11  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    12  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
    13  	"github.com/jackc/pgx/v5"
    14  	jsoniter "github.com/json-iterator/go"
    15  	"github.com/pashagolub/pgxmock/v5"
    16  	"github.com/stretchr/testify/assert"
    17  	"github.com/stretchr/testify/require"
    18  )
    19  
    20  var ctx = log.WithLogger(context.Background(), log.NewNoopLogger())
    21  
    22  func TestReadMetricSchemaType(t *testing.T) {
    23  	conn, err := pgxmock.NewPool()
    24  	assert.NoError(t, err)
    25  
    26  	pgw := PostgresWriter{
    27  		ctx:    ctx,
    28  		sinkDb: conn,
    29  	}
    30  
    31  	conn.ExpectQuery("SELECT schema_type").
    32  		WillReturnError(errors.New("expected"))
    33  	assert.Error(t, pgw.ReadMetricSchemaType())
    34  
    35  	conn.ExpectQuery("SELECT schema_type").
    36  		WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
    37  	assert.NoError(t, pgw.ReadMetricSchemaType())
    38  	assert.Equal(t, DbStorageSchemaTimescale, pgw.metricSchema)
    39  }
    40  
    41  func TestNewWriterFromPostgresConn(t *testing.T) {
    42  	a := assert.New(t)
    43  	opts := &CmdOpts{
    44  		BatchingDelay:       time.Hour,
    45  		RetentionInterval:   "1 day",
    46  		MaintenanceInterval: "1 day",
    47  		PartitionInterval:   "1 hour",
    48  	}
    49  
    50  	t.Run("Success", func(*testing.T) {
    51  		conn, err := pgxmock.NewPool()
    52  		a.NoError(err)
    53  
    54  		conn.ExpectPing()
    55  		conn.ExpectQuery("SELECT extract").WithArgs("1 day", "1 day", "1 hour").WillReturnRows(
    56  			pgxmock.NewRows([]string{"col1", "col2", "col3"}).AddRow((24 * time.Hour).Seconds(), (24 * time.Hour).Seconds(), true),
    57  		)
    58  		conn.ExpectQuery("SELECT EXISTS").WithArgs("admin").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
    59  		conn.ExpectQuery("SELECT schema_type").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
    60  		for _, m := range metrics.GetDefaultBuiltInMetrics() {
    61  			conn.ExpectExec("SELECT admin.ensure_dummy_metrics_table").WithArgs(m).WillReturnResult(pgxmock.NewResult("EXECUTE", 1))
    62  		}
    63  
    64  		pgw, err := NewWriterFromPostgresConn(ctx, conn, opts)
    65  		a.NoError(err)
    66  		a.NotNil(pgw)
    67  		a.NoError(conn.ExpectationsWereMet())
    68  	})
    69  
    70  	t.Run("InitFail", func(*testing.T) {
    71  		conn, err := pgxmock.NewPool()
    72  		a.NoError(err)
    73  
    74  		conn.ExpectPing().WillReturnError(assert.AnError)
    75  
    76  		pgw, err := NewWriterFromPostgresConn(ctx, conn, opts)
    77  		a.Error(err)
    78  		a.Nil(pgw)
    79  		a.NoError(conn.ExpectationsWereMet())
    80  	})
    81  
    82  	t.Run("ReadMetricSchemaTypeFail", func(*testing.T) {
    83  		conn, err := pgxmock.NewPool()
    84  		a.NoError(err)
    85  
    86  		conn.ExpectPing()
    87  		conn.ExpectQuery("SELECT extract").WithArgs("1 day", "1 day", "1 hour").WillReturnRows(
    88  			pgxmock.NewRows([]string{"col1", "col2", "col3"}).AddRow((24 * time.Hour).Seconds(), (24 * time.Hour).Seconds(), true),
    89  		)
    90  		conn.ExpectQuery("SELECT EXISTS").WithArgs("admin").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
    91  		conn.ExpectQuery("SELECT schema_type").WillReturnError(assert.AnError)
    92  
    93  		pgw, err := NewWriterFromPostgresConn(ctx, conn, opts)
    94  		a.Error(err)
    95  		a.Nil(pgw)
    96  		a.NoError(conn.ExpectationsWereMet())
    97  	})
    98  
    99  	t.Run("EnsureBuiltinMetricDummiesFail", func(*testing.T) {
   100  		conn, err := pgxmock.NewPool()
   101  		a.NoError(err)
   102  
   103  		conn.ExpectPing()
   104  		conn.ExpectQuery("SELECT extract").WithArgs("1 day", "1 day", "1 hour").WillReturnRows(
   105  			pgxmock.NewRows([]string{"col1", "col2", "col3"}).AddRow((24 * time.Hour).Seconds(), (24 * time.Hour).Seconds(), true),
   106  		)
   107  		conn.ExpectQuery("SELECT EXISTS").WithArgs("admin").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
   108  		conn.ExpectQuery("SELECT schema_type").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
   109  		conn.ExpectExec("SELECT admin.ensure_dummy_metrics_table").WithArgs(pgxmock.AnyArg()).WillReturnError(assert.AnError)
   110  
   111  		pgw, err := NewWriterFromPostgresConn(ctx, conn, opts)
   112  		a.Error(err)
   113  		a.Nil(pgw)
   114  		a.NoError(conn.ExpectationsWereMet())
   115  	})
   116  }
   117  
   118  func TestSyncMetric(t *testing.T) {
   119  	conn, err := pgxmock.NewPool()
   120  	assert.NoError(t, err)
   121  	pgw := PostgresWriter{
   122  		ctx:    ctx,
   123  		sinkDb: conn,
   124  	}
   125  	dbUnique := "mydb"
   126  	metricName := "mymetric"
   127  	op := AddOp
   128  	conn.ExpectExec("INSERT INTO admin\\.all_distinct_dbname_metrics").WithArgs(dbUnique, metricName).WillReturnResult(pgxmock.NewResult("EXECUTE", 1))
   129  	conn.ExpectExec("SELECT admin\\.ensure_dummy_metrics_table").WithArgs(metricName).WillReturnResult(pgxmock.NewResult("EXECUTE", 1))
   130  	err = pgw.SyncMetric(dbUnique, metricName, op)
   131  	assert.NoError(t, err)
   132  	assert.NoError(t, conn.ExpectationsWereMet())
   133  
   134  	op = InvalidOp
   135  	err = pgw.SyncMetric(dbUnique, metricName, op)
   136  	assert.NoError(t, err, "ignore unknown operation")
   137  }
   138  
   139  func TestWrite(t *testing.T) {
   140  	conn, err := pgxmock.NewPool()
   141  	assert.NoError(t, err)
   142  	ctx, cancel := context.WithCancel(ctx)
   143  	pgw := PostgresWriter{
   144  		ctx:    ctx,
   145  		sinkDb: conn,
   146  	}
   147  	message := metrics.MeasurementEnvelope{
   148  		MetricName: "test_metric",
   149  		Data: metrics.Measurements{
   150  			{"number": 1, "string": "test_data"},
   151  		},
   152  		DBName:     "test_db",
   153  		CustomTags: map[string]string{"foo": "boo"},
   154  	}
   155  
   156  	highLoadTimeout = 0
   157  	err = pgw.Write(message)
   158  	assert.NoError(t, err, "messages skipped due to high load")
   159  
   160  	highLoadTimeout = time.Second * 5
   161  	pgw.input = make(chan metrics.MeasurementEnvelope, cacheLimit)
   162  	err = pgw.Write(message)
   163  	assert.NoError(t, err, "write successful")
   164  
   165  	cancel()
   166  	err = pgw.Write(message)
   167  	assert.Error(t, err, "context canceled")
   168  }
   169  
   170  func TestCopyFromMeasurements_Basic(t *testing.T) {
   171  	// Test basic iteration through single envelope with multiple measurements
   172  	data := []metrics.MeasurementEnvelope{
   173  		{
   174  			MetricName: "metric1",
   175  			DBName:     "db1",
   176  			CustomTags: map[string]string{"env": "test"},
   177  			Data: metrics.Measurements{
   178  				{"epoch_ns": int64(1000), "value": 1},
   179  				{"epoch_ns": int64(2000), "value": 2},
   180  				{"epoch_ns": int64(3000), "value": 3},
   181  			},
   182  		},
   183  	}
   184  
   185  	cfm := newCopyFromMeasurements(data)
   186  
   187  	// Test Next() and Values() for each measurement
   188  	assert.Equal(t, "metric1", cfm.MetricName()[0], "Metric name should be obtained before Next()")
   189  	assert.True(t, cfm.Next(), "Should have first measurement")
   190  	values, err := cfm.Values()
   191  	assert.NoError(t, err)
   192  	assert.Len(t, values, 4) // time, dbname, data, tag_data
   193  	assert.Equal(t, "db1", values[1])
   194  
   195  	assert.True(t, cfm.Next(), "Should have second measurement")
   196  	values, err = cfm.Values()
   197  	assert.NoError(t, err)
   198  	assert.Equal(t, "db1", values[1])
   199  
   200  	assert.True(t, cfm.Next(), "Should have third measurement")
   201  	values, err = cfm.Values()
   202  	assert.NoError(t, err)
   203  	assert.Equal(t, "db1", values[1])
   204  
   205  	assert.False(t, cfm.Next(), "Should not have more measurements")
   206  	assert.True(t, cfm.EOF(), "Should be at end")
   207  }
   208  
   209  func TestCopyFromMeasurements_MultipleEnvelopes(t *testing.T) {
   210  	// Test iteration through multiple envelopes of same metric
   211  	data := []metrics.MeasurementEnvelope{
   212  		{
   213  			MetricName: "metric1",
   214  			DBName:     "db1",
   215  			CustomTags: map[string]string{"env": "test1"},
   216  			Data: metrics.Measurements{
   217  				{"epoch_ns": int64(1000), "value": 1},
   218  				{"epoch_ns": int64(2000), "value": 2},
   219  			},
   220  		},
   221  		{
   222  			MetricName: "metric1",
   223  			DBName:     "db2",
   224  			CustomTags: map[string]string{"env": "test2"},
   225  			Data: metrics.Measurements{
   226  				{"epoch_ns": int64(3000), "value": 3},
   227  			},
   228  		},
   229  	}
   230  
   231  	cfm := newCopyFromMeasurements(data)
   232  
   233  	// First envelope, first measurement
   234  	assert.True(t, cfm.Next())
   235  	values, err := cfm.Values()
   236  	assert.NoError(t, err)
   237  	assert.Equal(t, "db1", values[1])
   238  	// First envelope, second measurement
   239  	assert.True(t, cfm.Next())
   240  	values, err = cfm.Values()
   241  	assert.NoError(t, err)
   242  	assert.Equal(t, "db1", values[1])
   243  
   244  	// Second envelope, first measurement
   245  	assert.Equal(t, "metric1", cfm.MetricName()[0])
   246  	assert.True(t, cfm.Next())
   247  	values, err = cfm.Values()
   248  	assert.NoError(t, err)
   249  	assert.Equal(t, "db2", values[1])
   250  
   251  	assert.False(t, cfm.Next())
   252  }
   253  
   254  func TestCopyFromMeasurements_MetricBoundaries(t *testing.T) {
   255  	// Test metric boundary detection with different metrics
   256  	data := []metrics.MeasurementEnvelope{
   257  		{
   258  			MetricName: "metric1",
   259  			DBName:     "db1",
   260  			CustomTags: map[string]string{},
   261  			Data: metrics.Measurements{
   262  				{"epoch_ns": int64(1000), "value": 1},
   263  				{"epoch_ns": int64(2000), "value": 2},
   264  			},
   265  		},
   266  		{
   267  			MetricName: "metric2", // Different metric
   268  			DBName:     "db1",
   269  			CustomTags: map[string]string{},
   270  			Data: metrics.Measurements{
   271  				{"epoch_ns": int64(3000), "value": 3},
   272  			},
   273  		},
   274  		{
   275  			MetricName: "metric2",
   276  			DBName:     "db2",
   277  			CustomTags: map[string]string{},
   278  			Data: metrics.Measurements{
   279  				{"epoch_ns": int64(4000), "value": 4},
   280  			},
   281  		},
   282  	}
   283  
   284  	cfm := newCopyFromMeasurements(data)
   285  
   286  	// Process metric1 completely
   287  	assert.Equal(t, "metric1", cfm.MetricName()[0])
   288  	assert.True(t, cfm.Next())
   289  	assert.True(t, cfm.Next())
   290  
   291  	// Should stop at metric boundary
   292  	assert.False(t, cfm.Next())
   293  	assert.False(t, cfm.EOF(), "Should not be at EOF yet, there's more data")
   294  
   295  	assert.Equal(t, "metric2", cfm.MetricName()[0])
   296  	assert.True(t, cfm.Next())
   297  	assert.True(t, cfm.Next())
   298  
   299  	assert.False(t, cfm.Next())
   300  	assert.True(t, cfm.EOF(), "Should be at EOF after processing all measurements")
   301  }
   302  
   303  func TestCopyFromMeasurements_EmptyData(t *testing.T) {
   304  	// Test with empty envelopes slice
   305  	cfm := newCopyFromMeasurements([]metrics.MeasurementEnvelope{})
   306  	assert.False(t, cfm.Next())
   307  	assert.True(t, cfm.EOF())
   308  }
   309  
   310  func TestCopyFromMeasurements_EmptyMeasurements(t *testing.T) {
   311  	// Test with envelope containing no measurements
   312  	data := []metrics.MeasurementEnvelope{
   313  		{
   314  			MetricName: "metric1",
   315  			DBName:     "db1",
   316  			CustomTags: map[string]string{},
   317  			Data:       metrics.Measurements{}, // Empty measurements
   318  		},
   319  		{
   320  			MetricName: "metric1",
   321  			DBName:     "db2",
   322  			CustomTags: map[string]string{},
   323  			Data: metrics.Measurements{
   324  				{"epoch_ns": int64(1000), "value": 1},
   325  			},
   326  		},
   327  	}
   328  
   329  	cfm := newCopyFromMeasurements(data)
   330  
   331  	// Should skip empty envelope and go to second one
   332  	assert.True(t, cfm.Next())
   333  	values, err := cfm.Values()
   334  	assert.NoError(t, err)
   335  	assert.Equal(t, "db2", values[1])
   336  
   337  	assert.False(t, cfm.Next())
   338  	assert.True(t, cfm.EOF())
   339  }
   340  
   341  func TestCopyFromMeasurements_TagProcessing(t *testing.T) {
   342  	// Test that tag_ prefixed fields are moved to CustomTags
   343  	data := []metrics.MeasurementEnvelope{
   344  		{
   345  			MetricName: "metric1",
   346  			DBName:     "db1",
   347  			CustomTags: map[string]string{"existing": "tag"},
   348  			Data: metrics.Measurements{
   349  				{
   350  					"epoch_ns":     int64(1000),
   351  					"value":        1,
   352  					"tag_env":      "production",
   353  					"tag_version":  "1.0",
   354  					"normal_field": "stays",
   355  				},
   356  			},
   357  		},
   358  		{
   359  			MetricName: "metric1",
   360  			DBName:     "db2",
   361  			CustomTags: nil,
   362  			Data: metrics.Measurements{
   363  				{
   364  					"epoch_ns":     int64(1000),
   365  					"value":        1,
   366  					"tag_env":      "production",
   367  					"tag_version":  "1.0",
   368  					"normal_field": "stays",
   369  				},
   370  			},
   371  		},
   372  	}
   373  
   374  	cfm := newCopyFromMeasurements(data)
   375  	assert.True(t, cfm.Next())
   376  
   377  	values, err := cfm.Values()
   378  	assert.NoError(t, err)
   379  	assert.Len(t, values, 4) // Verify structure: time, dbname, data, tag_data
   380  
   381  	// Check that custom tags were updated
   382  	// Check data JSON (should contain normal fields but not tag_ fields)
   383  	dataJSON, ok := values[2].(string)
   384  	assert.True(t, ok, "Data should be JSON string")
   385  
   386  	var dataMap map[string]any
   387  	err = jsoniter.ConfigFastest.UnmarshalFromString(dataJSON, &dataMap)
   388  	assert.NoError(t, err)
   389  	assert.Contains(t, dataMap, "normal_field")
   390  	assert.NotContains(t, dataMap, "tag_env", "tag_env should not be in data")
   391  	assert.NotContains(t, dataMap, "tag_version", "tag_version should not be in data")
   392  
   393  	// Check tag JSON (should contain converted tags)
   394  	tagJSON, ok := values[3].(string)
   395  	assert.True(t, ok, "Tag data should be JSON string")
   396  
   397  	var tagMap map[string]string
   398  	err = jsoniter.ConfigFastest.UnmarshalFromString(tagJSON, &tagMap)
   399  	assert.NoError(t, err)
   400  	assert.Contains(t, tagMap, "existing")
   401  	assert.Contains(t, tagMap, "env", "tag_env should be converted to env")
   402  	assert.Contains(t, tagMap, "version", "tag_version should be converted to version")
   403  	assert.Equal(t, "production", tagMap["env"])
   404  	assert.Equal(t, "1.0", tagMap["version"])
   405  
   406  	assert.True(t, cfm.Next())
   407  	_, err = cfm.Values()
   408  	assert.NoError(t, err, "should process nil CustomTags without error")
   409  }
   410  
   411  func TestCopyFromMeasurements_JsonMarshaling(t *testing.T) {
   412  	// Test that JSON marshaling works correctly
   413  	data := []metrics.MeasurementEnvelope{
   414  		{
   415  			MetricName: "metric1",
   416  			DBName:     "db1",
   417  			CustomTags: map[string]string{"env": "test"},
   418  			Data: metrics.Measurements{
   419  				{
   420  					"epoch_ns": int64(1000),
   421  					"value":    42,
   422  					"name":     "test_measurement",
   423  				},
   424  				{
   425  					"epoch_ns": int64(1000),
   426  					"value": func() string {
   427  						return "should produce error while marshaled"
   428  					},
   429  					"name": "test_measurement",
   430  				},
   431  			},
   432  		},
   433  	}
   434  
   435  	cfm := newCopyFromMeasurements(data)
   436  	assert.True(t, cfm.Next())
   437  
   438  	values, err := cfm.Values()
   439  	assert.NoError(t, err)
   440  	assert.Len(t, values, 4)
   441  
   442  	// Values should be: [time, dbname, data_json, tag_data_json]
   443  	assert.Equal(t, "db1", values[1])
   444  
   445  	// Check that JSON strings are valid
   446  	dataJSON, ok := values[2].(string)
   447  	assert.True(t, ok, "Data should be JSON string")
   448  	assert.Contains(t, dataJSON, `"value":42`)
   449  	assert.Contains(t, dataJSON, `"name":"test_measurement"`)
   450  
   451  	tagJSON, ok := values[3].(string)
   452  	assert.True(t, ok, "Tag data should be JSON string")
   453  	assert.Contains(t, tagJSON, `"env":"test"`)
   454  
   455  	assert.True(t, cfm.Next())
   456  	_, err = cfm.Values()
   457  	assert.Error(t, err, "cannot marshal function value to JSON")
   458  
   459  	cfm.NextEnvelope()
   460  	assert.NotPanics(t, func() { _ = cfm.MetricName() })
   461  }
   462  
   463  func TestCopyFromMeasurements_ErrorHandling(t *testing.T) {
   464  	// Test Err() method
   465  	cfm := newCopyFromMeasurements([]metrics.MeasurementEnvelope{})
   466  	assert.NoError(t, cfm.Err(), "Err() should always return nil")
   467  }
   468  
   469  func TestCopyFromMeasurements_StateManagement(t *testing.T) {
   470  	// Test that internal state is managed correctly during iteration
   471  	data := []metrics.MeasurementEnvelope{
   472  		{
   473  			MetricName: "metric1",
   474  			DBName:     "db1",
   475  			CustomTags: map[string]string{},
   476  			Data: metrics.Measurements{
   477  				{"epoch_ns": int64(1000), "value": 1},
   478  			},
   479  		},
   480  		{
   481  			MetricName: "metric2", // Different metric
   482  			DBName:     "db1",
   483  			CustomTags: map[string]string{},
   484  			Data: metrics.Measurements{
   485  				{"epoch_ns": int64(2000), "value": 2},
   486  			},
   487  		},
   488  	}
   489  
   490  	cfm := newCopyFromMeasurements(data)
   491  
   492  	// Initial state
   493  	assert.Equal(t, -1, cfm.envelopeIdx)
   494  	assert.Equal(t, -1, cfm.measurementIdx)
   495  	assert.Equal(t, "", cfm.metricName)
   496  
   497  	// After first Next()
   498  	assert.True(t, cfm.Next())
   499  	assert.Equal(t, 0, cfm.envelopeIdx)
   500  	assert.Equal(t, 0, cfm.measurementIdx)
   501  	assert.Equal(t, "metric1", cfm.metricName)
   502  
   503  	// After hitting metric boundary
   504  	assert.False(t, cfm.Next())
   505  	// State should be positioned to restart on next metric
   506  	assert.Equal(t, "", cfm.metricName)
   507  }
   508  
   509  func TestCopyFromMeasurements_CopyFail(t *testing.T) {
   510  	a := assert.New(t)
   511  	r := require.New(t)
   512  
   513  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   514  	r.NoError(err)
   515  	defer pgTearDown()
   516  
   517  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   518  	r.NoError(err)
   519  	conn, err := pgx.Connect(ctx, connStr)
   520  	r.NoError(err)
   521  
   522  	_, err = conn.Exec(ctx, `CREATE TABLE IF NOT EXISTS test_metric (
   523  		time timestamptz not null default now(),
   524  		dbname text NOT NULL,
   525  		data jsonb,
   526  		tag_data jsonb)`)
   527  	r.NoError(err)
   528  
   529  	msgs := []metrics.MeasurementEnvelope{
   530  		{
   531  			MetricName: "test_metric",
   532  			Data: metrics.Measurements{
   533  				{"epoch_ns": int64(2000), "value": func() {}},
   534  				{"epoch_ns": int64(2000), "value": struct{}{}},
   535  			},
   536  			DBName: "test_db",
   537  		},
   538  	}
   539  
   540  	cfm := newCopyFromMeasurements(msgs)
   541  
   542  	for !cfm.EOF() {
   543  		_, err = conn.CopyFrom(context.Background(), cfm.MetricName(), targetColumns[:], cfm)
   544  		a.Error(err)
   545  		if err != nil {
   546  			if !cfm.NextEnvelope() {
   547  				break
   548  			}
   549  		}
   550  	}
   551  
   552  }
   553  
   554  // tests interval string validation for all
   555  // cli flags that expect a PostgreSQL interval string
   556  func TestIntervalValidation(t *testing.T) {
   557  	a := assert.New(t)
   558  
   559  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   560  	a.NoError(err)
   561  	defer pgTearDown()
   562  
   563  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   564  	a.NoError(err)
   565  
   566  	opts := &CmdOpts{
   567  		PartitionInterval:   "1 minute",
   568  		MaintenanceInterval: "-1 hours",
   569  		RetentionInterval:   "00:01:30",
   570  		BatchingDelay:       time.Second,
   571  	}
   572  
   573  	_, err = NewPostgresWriter(ctx, connStr, opts)
   574  	a.EqualError(err, "--partition-interval must be at least 1 hour, got: 1 minute")
   575  	opts.PartitionInterval = "1 hour"
   576  
   577  	_, err = NewPostgresWriter(ctx, connStr, opts)
   578  	a.EqualError(err, "--maintenance-interval must be a positive PostgreSQL interval or 0 to disable it")
   579  	opts.MaintenanceInterval = "0 hours"
   580  
   581  	_, err = NewPostgresWriter(ctx, connStr, opts)
   582  	a.Error(err)
   583  
   584  	invalidIntervals := []string{
   585  		"not an interval", "3 dayss",
   586  		"four hours",
   587  	}
   588  
   589  	for _, interval := range invalidIntervals {
   590  		opts.PartitionInterval = interval
   591  		_, err = NewPostgresWriter(ctx, connStr, opts)
   592  		a.Error(err)
   593  		opts.PartitionInterval = "1 hour"
   594  
   595  		opts.MaintenanceInterval = interval
   596  		_, err = NewPostgresWriter(ctx, connStr, opts)
   597  		a.Error(err)
   598  		opts.MaintenanceInterval = "1 hour"
   599  
   600  		opts.RetentionInterval = interval
   601  		_, err = NewPostgresWriter(ctx, connStr, opts)
   602  		a.Error(err)
   603  		opts.RetentionInterval = "1 hour"
   604  	}
   605  
   606  	validIntervals := []string{
   607  		"3 days 4 hours", "1 year",
   608  		"P3D", "PT3H", "0-02", "1 00:00:00",
   609  		"P0-02", "P1", "2 weeks",
   610  	}
   611  
   612  	for _, interval := range validIntervals {
   613  		opts.PartitionInterval = interval
   614  		opts.MaintenanceInterval = interval
   615  		opts.RetentionInterval = interval
   616  
   617  		_, err = NewPostgresWriter(ctx, connStr, opts)
   618  		a.NoError(err)
   619  	}
   620  }
   621  
   622  func TestPartitionInterval(t *testing.T) {
   623  	a := assert.New(t)
   624  	r := require.New(t)
   625  
   626  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   627  	r.NoError(err)
   628  	defer pgTearDown()
   629  
   630  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   631  	r.NoError(err)
   632  
   633  	opts := &CmdOpts{
   634  		PartitionInterval:   "3 weeks",
   635  		RetentionInterval:   "14 days",
   636  		MaintenanceInterval: "12 hours",
   637  		BatchingDelay:       time.Second,
   638  	}
   639  
   640  	pgw, err := NewPostgresWriter(ctx, connStr, opts)
   641  	r.NoError(err)
   642  
   643  	conn, err := pgx.Connect(ctx, connStr)
   644  	r.NoError(err)
   645  
   646  	m := map[string]ExistingPartitionInfo{
   647  		"test_metric": {
   648  			time.Now(), time.Now().Add(time.Hour),
   649  		},
   650  	}
   651  	err = pgw.EnsureMetricTimePartsExist(m)
   652  	r.NoError(err)
   653  
   654  	var partitionsNum int
   655  	err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric');").Scan(&partitionsNum)
   656  	a.NoError(err)
   657  	// 1 the metric table itself + 4 time partitions (1 we asked for + 3 precreated)
   658  	a.Equal(5, partitionsNum)
   659  
   660  	part := pgw.partitionMapMetric["test_metric"]
   661  	// partition bounds should have a difference of 3 weeks
   662  	a.Equal(part.StartTime.Add(3*7*24*time.Hour), part.EndTime)
   663  }
   664  
   665  // TestPartitionForwardTimeJump checks #1529
   666  func TestPartitionForwardTimeJump(t *testing.T) {
   667  	a := assert.New(t)
   668  	r := require.New(t)
   669  
   670  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   671  	r.NoError(err)
   672  	defer pgTearDown()
   673  
   674  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   675  	r.NoError(err)
   676  
   677  	opts := &CmdOpts{
   678  		PartitionInterval:   "1 day",
   679  		RetentionInterval:   "14 days",
   680  		MaintenanceInterval: "12 hours",
   681  		BatchingDelay:       time.Second,
   682  	}
   683  
   684  	pgw, err := NewPostgresWriter(ctx, connStr, opts)
   685  	r.NoError(err)
   686  
   687  	conn, err := pgx.Connect(ctx, connStr)
   688  	r.NoError(err)
   689  	defer conn.Close(ctx)
   690  
   691  	// 1. Create partitions for an early period.
   692  	early := time.Date(2026, 2, 5, 10, 0, 0, 0, time.UTC)
   693  	r.NoError(pgw.EnsureMetricTimePartsExist(map[string]ExistingPartitionInfo{
   694  		"jump_metric": {StartTime: early, EndTime: early},
   695  	}))
   696  
   697  	// 2. Simulate a restart: the in-memory partition cache is lost, while the
   698  	//    existing partitions remain in the database.
   699  	pgw.partitionMapMetric = make(map[string]ExistingPartitionInfo)
   700  
   701  	// 3. New data arrives weeks later, well beyond the pre-created partitions'
   702  	//    upper bound (the "forward time jump" / gap scenario).
   703  	late := time.Date(2026, 8, 19, 0, 49, 18, 0, time.UTC)
   704  	err = pgw.EnsureMetricTimePartsExist(map[string]ExistingPartitionInfo{
   705  		"jump_metric": {StartTime: late, EndTime: late},
   706  	})
   707  	// Must not error (previously failed with "cannot scan NULL into *time.Time").
   708  	r.NoError(err)
   709  
   710  	// The returned/cached range must be non-zero and must contain the new timestamp.
   711  	part := pgw.partitionMapMetric["jump_metric"]
   712  	a.False(part.StartTime.IsZero(), "part_available_from must not be NULL")
   713  	a.False(part.EndTime.IsZero(), "part_available_to must not be NULL")
   714  	a.False(late.Before(part.StartTime), "new timestamp must be >= partition start")
   715  	a.True(late.Before(part.EndTime), "new timestamp must be < partition end")
   716  
   717  	// A subpartition physically covering the new timestamp must exist, otherwise
   718  	// the CopyFrom insert would fail with SQLSTATE 23514 (no partition for row).
   719  	var covering int
   720  	r.NoError(conn.QueryRow(ctx, `
   721  		SELECT count(*)
   722  		FROM pg_catalog.pg_class c
   723  		JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid
   724  		JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent
   725  		WHERE c.relispartition
   726  		  AND c.relnamespace = 'subpartitions'::regnamespace
   727  		  AND parent.relname = 'jump_metric'
   728  		  AND $1::timestamptz >= substring(pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) from 'FOR VALUES FROM \(''([^'']+)''')::timestamptz
   729  		  AND $1::timestamptz <  substring(pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) from 'TO \(''([^'']+)''')::timestamptz
   730  	`, late).Scan(&covering))
   731  	a.Equal(1, covering, "exactly one subpartition must cover the new timestamp")
   732  }
   733  
   734  func Test_Maintain(t *testing.T) {
   735  	a := assert.New(t)
   736  	r := require.New(t)
   737  
   738  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
   739  	r.NoError(err)
   740  	defer pgTearDown()
   741  
   742  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   743  	r.NoError(err)
   744  	conn, err := pgx.Connect(ctx, connStr)
   745  	r.NoError(err)
   746  
   747  	opts := &CmdOpts{
   748  		PartitionInterval:   "1 hour",
   749  		RetentionInterval:   "1 hour",
   750  		MaintenanceInterval: "0 days",
   751  		BatchingDelay:       time.Hour,
   752  	}
   753  
   754  	pgw, err := NewPostgresWriter(ctx, connStr, opts)
   755  	r.NoError(err)
   756  
   757  	t.Run("MaintainUniqueSources", func(_ *testing.T) {
   758  		// adds an entry to `admin.all_distinct_dbname_metrics`
   759  		err = pgw.SyncMetric("test", "test_metric_1", AddOp)
   760  		r.NoError(err)
   761  
   762  		var numOfEntries int
   763  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
   764  		a.NoError(err)
   765  		a.Equal(1, numOfEntries)
   766  
   767  		// manually call the maintenance routine
   768  		pgw.MaintainUniqueSources()
   769  
   770  		// entry should have been deleted, because it has no corresponding entries in `test_metric_1` table.
   771  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
   772  		a.NoError(err)
   773  		a.Equal(0, numOfEntries)
   774  
   775  		message := []metrics.MeasurementEnvelope{
   776  			{
   777  				MetricName: "test_metric_1",
   778  				Data: metrics.Measurements{
   779  					{"number": 1, "string": "test_data"},
   780  				},
   781  				DBName: "test_db",
   782  			},
   783  		}
   784  		pgw.flush(message)
   785  
   786  		// manually call the maintenance routine
   787  		pgw.MaintainUniqueSources()
   788  
   789  		// entry should have been added, because there is a corresponding entry in `test_metric_1` table just written.
   790  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
   791  		a.NoError(err)
   792  		a.Equal(1, numOfEntries)
   793  
   794  		_, err = conn.Exec(ctx, "DROP TABLE test_metric_1;")
   795  		r.NoError(err)
   796  
   797  		// the corresponding entry should be deleted
   798  		pgw.MaintainUniqueSources()
   799  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
   800  		a.NoError(err)
   801  		a.Equal(0, numOfEntries)
   802  	})
   803  
   804  	t.Run("MaintainUniqueSources_MultipleMetricsAndSources", func(_ *testing.T) {
   805  		// Create metric tables with partitions
   806  		err = pgw.EnsureMetricDummy("test_metric_a")
   807  		r.NoError(err)
   808  		err = pgw.EnsureMetricDummy("test_metric_b")
   809  		r.NoError(err)
   810  
   811  		// Create time-based partitions for each metric
   812  		_, err = conn.Exec(ctx, `
   813  			CREATE TABLE subpartitions.test_metric_a_2024w01 PARTITION OF public.test_metric_a FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
   814  			CREATE TABLE subpartitions.test_metric_b_2024w01 PARTITION OF public.test_metric_b FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
   815  		`)
   816  		r.NoError(err)
   817  
   818  		// Directly insert test data with different dbnames
   819  		_, err = conn.Exec(ctx, `
   820  			INSERT INTO test_metric_a (time, dbname, data) VALUES
   821  				('2024-01-03', 'db1', '{}'::jsonb),
   822  				('2024-01-03', 'db2', '{}'::jsonb),
   823  				('2024-01-03', 'db3', '{}'::jsonb)
   824  		`)
   825  		r.NoError(err)
   826  
   827  		_, err = conn.Exec(ctx, `
   828  			INSERT INTO test_metric_b (time, dbname, data) VALUES
   829  				('2024-01-03', 'db1', '{}'::jsonb),
   830  				('2024-01-03', 'db2', '{}'::jsonb)
   831  		`)
   832  		r.NoError(err)
   833  
   834  		// Run maintenance
   835  		pgw.MaintainUniqueSources()
   836  
   837  		// Should have 3 entries for test_metric_a and 2 for test_metric_b
   838  		var count int
   839  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_a';").Scan(&count)
   840  		a.NoError(err)
   841  		a.Equal(3, count)
   842  
   843  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_b';").Scan(&count)
   844  		a.NoError(err)
   845  		a.Equal(2, count)
   846  
   847  		// Cleanup
   848  		_, err = conn.Exec(ctx, "DROP TABLE test_metric_a, test_metric_b;")
   849  		r.NoError(err)
   850  		pgw.MaintainUniqueSources()
   851  	})
   852  
   853  	t.Run("MaintainUniqueSources_StaleEntriesCleanup", func(_ *testing.T) {
   854  		// Create metric table
   855  		err = pgw.EnsureMetricDummy("test_metric_c")
   856  		r.NoError(err)
   857  
   858  		// Create time partition for the metric
   859  		_, err = conn.Exec(ctx, `
   860  			CREATE TABLE subpartitions.test_metric_c_2024w01 PARTITION OF public.test_metric_c FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
   861  		`)
   862  		r.NoError(err)
   863  
   864  		// Directly insert test data with one active dbname
   865  		_, err = conn.Exec(ctx, `
   866  			INSERT INTO test_metric_c (time, dbname, data) VALUES
   867  				('2024-01-03', 'db_active', '{}'::jsonb)
   868  		`)
   869  		r.NoError(err)
   870  
   871  		// Manually add the active entry and stale entries to the listing table
   872  		_, err = conn.Exec(ctx, "INSERT INTO admin.all_distinct_dbname_metrics (dbname, metric) VALUES ('db_active', 'test_metric_c'), ('db_stale1', 'test_metric_c'), ('db_stale2', 'test_metric_c');")
   873  		r.NoError(err)
   874  
   875  		var count int
   876  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_c';").Scan(&count)
   877  		a.NoError(err)
   878  		a.Equal(3, count) // 1 active + 2 stale
   879  
   880  		// Run maintenance - should remove stale entries
   881  		pgw.MaintainUniqueSources()
   882  
   883  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_c';").Scan(&count)
   884  		a.NoError(err)
   885  		a.Equal(1, count) // only active one remains
   886  
   887  		var dbname string
   888  		err = conn.QueryRow(ctx, "SELECT dbname FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_c';").Scan(&dbname)
   889  		a.NoError(err)
   890  		a.Equal("db_active", dbname)
   891  
   892  		// Cleanup
   893  		_, err = conn.Exec(ctx, "DROP TABLE test_metric_c;")
   894  		r.NoError(err)
   895  		pgw.MaintainUniqueSources()
   896  	})
   897  
   898  	t.Run("MaintainUniqueSources_AdvisoryLock", func(_ *testing.T) {
   899  		// Create a second connection to simulate concurrent maintenance
   900  		conn2, err := pgx.Connect(ctx, connStr)
   901  		r.NoError(err)
   902  		defer conn2.Close(ctx)
   903  
   904  		// Create metric table and partition
   905  		err = pgw.EnsureMetricDummy("test_metric_d")
   906  		r.NoError(err)
   907  
   908  		_, err = conn.Exec(ctx, `
   909  			CREATE TABLE subpartitions.test_metric_d_2024w01 PARTITION OF public.test_metric_d FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
   910  		`)
   911  		r.NoError(err)
   912  
   913  		// Directly insert test data for only db1
   914  		_, err = conn.Exec(ctx, `
   915  			INSERT INTO test_metric_d (time, dbname, data) VALUES
   916  				('2024-01-03', 'db1', '{}'::jsonb)
   917  		`)
   918  		r.NoError(err)
   919  
   920  		// Add both active and stale entries to the listing table
   921  		_, err = conn.Exec(ctx, "INSERT INTO admin.all_distinct_dbname_metrics (dbname, metric) VALUES ('db1', 'test_metric_d'), ('db_stale', 'test_metric_d');")
   922  		r.NoError(err)
   923  
   924  		var count int
   925  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_d';").Scan(&count)
   926  		a.NoError(err)
   927  		a.Equal(2, count, "Should have 2 entries initially (1 active + 1 stale)")
   928  
   929  		// Acquire the advisory lock using session-level lock in conn2
   930  		// This will block transaction-level locks from the same lock ID
   931  		var lockAcquired bool
   932  		err = conn2.QueryRow(ctx, "SELECT pg_try_advisory_lock(1571543679778230000);").Scan(&lockAcquired)
   933  		r.NoError(err)
   934  		a.True(lockAcquired, "Should acquire advisory lock")
   935  
   936  		// Try to run maintenance - should skip because lock is held by conn2
   937  		pgw.MaintainUniqueSources()
   938  
   939  		// Stale entry should still exist because maintenance was skipped
   940  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_d';").Scan(&count)
   941  		a.NoError(err)
   942  		a.Equal(2, count, "Stale entry should remain because maintenance was skipped due to lock")
   943  
   944  		// Release lock from conn2
   945  		_, err = conn2.Exec(ctx, "SELECT pg_advisory_unlock(1571543679778230000);")
   946  		r.NoError(err)
   947  
   948  		// Now maintenance should work and clean up stale entry
   949  		pgw.MaintainUniqueSources()
   950  
   951  		// Should only have the active entry, stale one removed
   952  		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_d';").Scan(&count)
   953  		a.NoError(err)
   954  		a.Equal(1, count, "Only active entry should remain after maintenance runs")
   955  
   956  		var dbname string
   957  		err = conn.QueryRow(ctx, "SELECT dbname FROM admin.all_distinct_dbname_metrics WHERE metric = 'test_metric_d';").Scan(&dbname)
   958  		a.NoError(err)
   959  		a.Equal("db1", dbname, "Remaining entry should be the active one")
   960  
   961  		// Cleanup
   962  		_, err = conn.Exec(ctx, "DROP TABLE test_metric_d;")
   963  		r.NoError(err)
   964  		pgw.MaintainUniqueSources()
   965  	})
   966  
   967  	t.Run("DeleteOldPartitions", func(_ *testing.T) {
   968  		// Creates a new top level table for `test_metric_2`
   969  		err = pgw.SyncMetric("test", "test_metric_2", AddOp)
   970  		r.NoError(err)
   971  
   972  		boundStart := time.Now().Add(-1 * 2 * 24 * time.Hour).Format("2006-01-02")
   973  		boundEnd := time.Now().Add(-1 * 24 * time.Hour).Format("2006-01-02")
   974  
   975  		// create the time partition with end bound yesterday
   976  		_, err = conn.Exec(ctx,
   977  			fmt.Sprintf(
   978  				`CREATE TABLE subpartitions.test_metric_2_yesterday
   979  			PARTITION OF public.test_metric_2
   980  			FOR VALUES FROM ('%s') TO ('%s')`,
   981  				boundStart, boundEnd),
   982  		)
   983  		a.NoError(err)
   984  		_, err = conn.Exec(ctx, "COMMENT ON TABLE subpartitions.test_metric_2_yesterday IS $$pgwatch-generated-metric-time-lvl$$")
   985  		a.NoError(err)
   986  
   987  		var partitionsNum int
   988  		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric_2');").Scan(&partitionsNum)
   989  		a.NoError(err)
   990  		a.Equal(2, partitionsNum)
   991  
   992  		pgw.opts.RetentionInterval = "2 days"
   993  		pgw.DeleteOldPartitions() // 1 day < 2 days, shouldn't delete anything
   994  
   995  		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric_2');").Scan(&partitionsNum)
   996  		a.NoError(err)
   997  		a.Equal(2, partitionsNum)
   998  
   999  		pgw.opts.RetentionInterval = "1 hour"
  1000  		pgw.DeleteOldPartitions() // 1 day > 1 hour, should delete the partition
  1001  
  1002  		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric_2');").Scan(&partitionsNum)
  1003  		a.NoError(err)
  1004  		a.Equal(1, partitionsNum)
  1005  	})
  1006  
  1007  	t.Run("Epoch to Duration Conversion", func(_ *testing.T) {
  1008  		table := map[string]time.Duration{
  1009  			"1 hour":   time.Hour,
  1010  			"2 hours":  2 * time.Hour,
  1011  			"4 days":   4 * 24 * time.Hour,
  1012  			"1 day":    24 * time.Hour,
  1013  			"1 year":   365.25 * 24 * time.Hour,
  1014  			"1 week":   7 * 24 * time.Hour,
  1015  			"3 weeks":  3 * 7 * 24 * time.Hour,
  1016  			"2 months": 2 * 30 * 24 * time.Hour,
  1017  			"1 month":  30 * 24 * time.Hour,
  1018  		}
  1019  
  1020  		for k, v := range table {
  1021  			opts := &CmdOpts{
  1022  				PartitionInterval:   "1 hour",
  1023  				RetentionInterval:   k,
  1024  				MaintenanceInterval: k,
  1025  				BatchingDelay:       time.Hour,
  1026  			}
  1027  
  1028  			pgw, err := NewPostgresWriter(ctx, connStr, opts)
  1029  			a.NoError(err)
  1030  			a.Equal(pgw.retentionInterval, v)
  1031  			a.Equal(pgw.maintenanceInterval, v)
  1032  		}
  1033  	})
  1034  }
  1035  
  1036  // TestEnsureMetricTimePartsExist_SpecialMetricNames verifies that metric names with special characters
  1037  // (dots, uppercase, hyphens, underscores) are accepted by the partition functions.
  1038  func TestEnsureMetricTimePartsExist_SpecialMetricNames(t *testing.T) {
  1039  	r := require.New(t)
  1040  	a := assert.New(t)
  1041  
  1042  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
  1043  	r.NoError(err)
  1044  	defer pgTearDown()
  1045  
  1046  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
  1047  	r.NoError(err)
  1048  
  1049  	opts := &CmdOpts{
  1050  		PartitionInterval:   "1 day",
  1051  		RetentionInterval:   "7 days",
  1052  		MaintenanceInterval: "12 hours",
  1053  		BatchingDelay:       time.Second,
  1054  	}
  1055  	pgw, err := NewPostgresWriter(ctx, connStr, opts)
  1056  	r.NoError(err)
  1057  
  1058  	specialNames := []string{
  1059  		"metric.new",
  1060  		"Metric.With.Dots",
  1061  		"UPPERCASE_METRIC",
  1062  		"MixedCase_Metric-123",
  1063  		"metric-with-hyphens",
  1064  		"metric_with_underscores",
  1065  		"Metric123.Test_Name-456",
  1066  	}
  1067  
  1068  	m := make(map[string]ExistingPartitionInfo)
  1069  	for _, name := range specialNames {
  1070  		m[name] = ExistingPartitionInfo{
  1071  			StartTime: time.Now(),
  1072  			EndTime:   time.Now().Add(time.Hour),
  1073  		}
  1074  	}
  1075  
  1076  	err = pgw.EnsureMetricTimePartsExist(m)
  1077  	r.NoError(err, "EnsureMetricTimePartsExist should handle special metric names")
  1078  
  1079  	conn, err := pgx.Connect(ctx, connStr)
  1080  	r.NoError(err)
  1081  	defer conn.Close(ctx)
  1082  
  1083  	var partitionCount int
  1084  	err = conn.QueryRow(ctx, `SELECT COUNT(*) FROM pg_partition_tree('"metric.new"') WHERE level = 1`).Scan(&partitionCount)
  1085  	r.NoError(err)
  1086  	// 4 time partitions (1 requested + 3 precreated) per metric
  1087  	a.Equal(4, partitionCount)
  1088  }
  1089  
  1090  // TestEnsureMetricTimePartsExist_IdempotentAcrossRestarts verifies that repeated calls to
  1091  // EnsureMetricTimePartsExist with fresh writer instances (simulating process restarts)
  1092  // do not create duplicate partitions.
  1093  func TestEnsureMetricTimePartsExist_IdempotentAcrossRestarts(t *testing.T) {
  1094  	r := require.New(t)
  1095  	a := assert.New(t)
  1096  
  1097  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
  1098  	r.NoError(err)
  1099  	defer pgTearDown()
  1100  
  1101  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
  1102  	r.NoError(err)
  1103  
  1104  	opts := &CmdOpts{
  1105  		PartitionInterval:   "1 day",
  1106  		RetentionInterval:   "7 days",
  1107  		MaintenanceInterval: "12 hours",
  1108  		BatchingDelay:       time.Second,
  1109  	}
  1110  
  1111  	m := map[string]ExistingPartitionInfo{
  1112  		"restart_test_metric": {
  1113  			StartTime: time.Now(),
  1114  			EndTime:   time.Now().Add(time.Hour),
  1115  		},
  1116  	}
  1117  
  1118  	for i := range 5 {
  1119  		pgw, err := NewPostgresWriter(ctx, connStr, opts)
  1120  		r.NoError(err)
  1121  		r.NoError(pgw.EnsureMetricTimePartsExist(m))
  1122  
  1123  		conn, err := pgx.Connect(ctx, connStr)
  1124  		r.NoError(err)
  1125  		var count int
  1126  		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('restart_test_metric') WHERE isleaf").Scan(&count)
  1127  		conn.Close(ctx)
  1128  		r.NoError(err)
  1129  		a.Equal(4, count, "partition count should not grow on restart %d", i+1)
  1130  	}
  1131  }
  1132  
  1133  func TestFlush_SinkDBIsDown(t *testing.T) {
  1134  	r := require.New(t)
  1135  
  1136  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
  1137  	r.NoError(err)
  1138  
  1139  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
  1140  	r.NoError(err)
  1141  
  1142  	opts := &CmdOpts{
  1143  		PartitionInterval:   "1 day",
  1144  		RetentionInterval:   "7 days",
  1145  		MaintenanceInterval: "12 hours",
  1146  		BatchingDelay:       time.Second,
  1147  	}
  1148  
  1149  	pgw, err := NewPostgresWriter(ctx, connStr, opts)
  1150  	r.NoError(err)
  1151  	// Bring the sink db down
  1152  	pgTearDown()
  1153  
  1154  	msgs := []metrics.MeasurementEnvelope{
  1155  		{
  1156  			MetricName: "test_metric",
  1157  			Data: metrics.Measurements{
  1158  				{"epoch_ns": time.Now().UnixNano(), "value": 1},
  1159  			},
  1160  			DBName: "test_db",
  1161  		},
  1162  	}
  1163  
  1164  	// It should return, Issue:#1426 will keep it spinning forever
  1165  	pgw.flush(msgs)
  1166  }
  1167  
  1168  // TestDropAllMetricTables verifies that admin.drop_all_metric_tables() is a procedure that
  1169  // drops every top-level metric table partition by partition and cleans up the listing table.
  1170  // The partition-by-partition drop with COMMIT between drops is the contract that lets large
  1171  // installs avoid blowing past max_locks_per_transaction. See issue #1474.
  1172  func TestDropAllMetricTables(t *testing.T) {
  1173  	a := assert.New(t)
  1174  	r := require.New(t)
  1175  
  1176  	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
  1177  	r.NoError(err)
  1178  	defer pgTearDown()
  1179  
  1180  	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
  1181  	r.NoError(err)
  1182  	conn, err := pgx.Connect(ctx, connStr)
  1183  	r.NoError(err)
  1184  	defer conn.Close(ctx)
  1185  
  1186  	opts := &CmdOpts{
  1187  		PartitionInterval:   "1 hour",
  1188  		RetentionInterval:   "1 hour",
  1189  		MaintenanceInterval: "0 days",
  1190  		BatchingDelay:       time.Hour,
  1191  	}
  1192  
  1193  	pgw, err := NewPostgresWriter(ctx, connStr, opts)
  1194  	r.NoError(err)
  1195  
  1196  	// Create two dummy metric tables and seed a listing row for each.
  1197  	r.NoError(pgw.SyncMetric("test", "test_metric_a", AddOp))
  1198  	r.NoError(pgw.SyncMetric("test", "test_metric_b", AddOp))
  1199  
  1200  	// Attach a single weekly partition to each so the partition-by-partition drop branch
  1201  	// is exercised.
  1202  	_, err = conn.Exec(ctx, `
  1203  		CREATE TABLE subpartitions.test_metric_a_2024w01 PARTITION OF public.test_metric_a FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
  1204  		COMMENT ON TABLE subpartitions.test_metric_a_2024w01 IS 'pgwatch-generated-metric-time-lvl';
  1205  		CREATE TABLE subpartitions.test_metric_b_2024w01 PARTITION OF public.test_metric_b FOR VALUES FROM ('2024-01-01') TO ('2024-01-08');
  1206  		COMMENT ON TABLE subpartitions.test_metric_b_2024w01 IS 'pgwatch-generated-metric-time-lvl';
  1207  	`)
  1208  	r.NoError(err)
  1209  
  1210  	// Sanity: the two test top-level metric tables exist before the call.
  1211  	var regA0, regB0 *string
  1212  	r.NoError(conn.QueryRow(ctx, "SELECT to_regclass('public.test_metric_a')::text, to_regclass('public.test_metric_b')::text").Scan(&regA0, &regB0))
  1213  	r.NotNil(regA0, "public.test_metric_a must exist after SyncMetric")
  1214  	r.NotNil(regB0, "public.test_metric_b must exist after SyncMetric")
  1215  
  1216  	var topLevelCount int
  1217  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM admin.get_top_level_metric_tables() WHERE table_name IN ('test_metric_a', 'test_metric_b')").Scan(&topLevelCount))
  1218  	a.Equal(2, topLevelCount)
  1219  
  1220  	// Invoke the procedure on a bare connection (CALL is rejected inside an explicit transaction).
  1221  	_, err = conn.Exec(ctx, "CALL admin.drop_all_metric_tables();")
  1222  	r.NoError(err)
  1223  
  1224  	// Both top-level tables must be gone.
  1225  	var regA, regB *string
  1226  	r.NoError(conn.QueryRow(ctx, "SELECT to_regclass('public.test_metric_a')::text, to_regclass('public.test_metric_b')::text").Scan(&regA, &regB))
  1227  	a.Nil(regA, "public.test_metric_a should be dropped")
  1228  	a.Nil(regB, "public.test_metric_b should be dropped")
  1229  
  1230  	// The leaf partitions must be gone too (not left behind as detached tables).
  1231  	var subpartCount int
  1232  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM pg_class WHERE relnamespace = 'subpartitions'::regnamespace").Scan(&subpartCount))
  1233  	a.Equal(0, subpartCount, "no tables should remain in the subpartitions schema")
  1234  
  1235  	// The listing table must be truncated.
  1236  	var listingCount int
  1237  	r.NoError(conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics").Scan(&listingCount))
  1238  	a.Equal(0, listingCount)
  1239  
  1240  	// The routine must exist as a procedure (prokind = 'p'), not a function.
  1241  	var prokind string
  1242  	r.NoError(conn.QueryRow(ctx, `
  1243  		SELECT p.prokind FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
  1244  		WHERE n.nspname = 'admin' AND p.proname = 'drop_all_metric_tables'
  1245  	`).Scan(&prokind))
  1246  	a.Equal("p", prokind, "admin.drop_all_metric_tables must be a procedure")
  1247  }
  1248