▽

Source file src/github.com/cybertec-postgresql/pgwatch/v6/internal/reaper/database_test.go

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

     1  package reaper
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"sync/atomic"
     7  	"testing"
     8  	"testing/synctest"
     9  	"time"
    10  
    11  	"github.com/cybertec-postgresql/pgwatch/v6/internal/cmdopts"
    12  	"github.com/cybertec-postgresql/pgwatch/v6/internal/db"
    13  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    14  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    15  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sinks"
    16  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
    17  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
    18  	"github.com/jackc/pgx/v5"
    19  	pgxmock "github.com/pashagolub/pgxmock/v5"
    20  	"github.com/stretchr/testify/assert"
    21  	"github.com/stretchr/testify/require"
    22  )
    23  
    24  // Helper function to create a test SourceConn with pgxmock
    25  func createTestSourceConn(t *testing.T) (*sources.DbConn, pgxmock.PgxPoolIface) {
    26  	mock, err := pgxmock.NewPool()
    27  	require.NoError(t, err)
    28  
    29  	md := &sources.DbConn{
    30  		Conn:   mock,
    31  		Source: sources.Source{Name: "testdb", Kind: sources.SourcePostgres},
    32  		RuntimeInfo: sources.RuntimeInfo{
    33  			Version:     120000,
    34  			ChangeState: make(map[string]map[string]string),
    35  		},
    36  	}
    37  	return md, mock
    38  }
    39  
    40  func TestDetectSprocChanges(t *testing.T) {
    41  	ctx := context.Background()
    42  
    43  	// Set up simple test metric (instead of complex real one)
    44  	metricDefs.MetricDefs["sproc_hashes"] = metrics.Metric{
    45  		SQLs: map[int]string{
    46  			120000: "SELECT",
    47  		},
    48  	}
    49  
    50  	// Create single connection and reaper to maintain state across calls
    51  	md, mock := createTestSourceConn(t)
    52  	defer mock.Close()
    53  
    54  	reaper := &reaper{
    55  		measurementCh: make(chan metrics.MeasurementEnvelope, 10),
    56  	}
    57  
    58  	// First run - establish baseline
    59  	initialRows := pgxmock.NewRows([]string{"tag_sproc", "tag_oid", "md5", "epoch_ns"}).
    60  		AddRow("func1", "123", "hash1", time.Now().UnixNano()).
    61  		AddRow("func2", "456", "hash2", time.Now().UnixNano())
    62  	mock.ExpectQuery("SELECT").WillReturnRows(initialRows)
    63  
    64  	result := reaper.DetectSprocChanges(ctx, md)
    65  	assert.Equal(t, 0, result.Created) // First run should not count anything as created
    66  	assert.Equal(t, 0, result.Altered)
    67  	assert.Equal(t, 0, result.Dropped)
    68  
    69  	// State should now be populated
    70  	assert.NotEmpty(t, md.ChangeState["sproc_hashes"])
    71  
    72  	// Second run - detect altered sproc (different hash for func1)
    73  	modifiedRows := pgxmock.NewRows([]string{"tag_sproc", "tag_oid", "md5", "epoch_ns"}).
    74  		AddRow("func1", "123", "new_hash", time.Now().UnixNano()).
    75  		AddRow("func2", "456", "hash2", time.Now().UnixNano())
    76  	mock.ExpectQuery("SELECT").WillReturnRows(modifiedRows)
    77  
    78  	result = reaper.DetectSprocChanges(ctx, md)
    79  	assert.Equal(t, 0, result.Created)
    80  	assert.Equal(t, 1, result.Altered) // func1 was altered
    81  	assert.Equal(t, 0, result.Dropped)
    82  
    83  	// Third run - detect new sproc (func3 added)
    84  	newSprocRows := pgxmock.NewRows([]string{"tag_sproc", "tag_oid", "md5", "epoch_ns"}).
    85  		AddRow("func1", "123", "new_hash", time.Now().UnixNano()).
    86  		AddRow("func2", "456", "hash2", time.Now().UnixNano()).
    87  		AddRow("func3", "789", "hash3", time.Now().UnixNano()) // new sproc
    88  	mock.ExpectQuery("SELECT").WillReturnRows(newSprocRows)
    89  
    90  	result = reaper.DetectSprocChanges(ctx, md)
    91  	assert.Equal(t, 1, result.Created) // func3 was created
    92  	assert.Equal(t, 0, result.Altered)
    93  	assert.Equal(t, 0, result.Dropped)
    94  
    95  	// Verify measurement is sent
    96  	select {
    97  	case <-reaper.measurementCh:
    98  		// Good, measurement was sent
    99  	default:
   100  		t.Error("Expected measurement to be sent")
   101  	}
   102  
   103  	// Fourth run - detect dropped sproc (func2 removed)
   104  	droppedSprocRows := pgxmock.NewRows([]string{"tag_sproc", "tag_oid", "md5", "epoch_ns"}).
   105  		AddRow("func1", "123", "new_hash", time.Now().UnixNano()).
   106  		AddRow("func3", "789", "hash3", time.Now().UnixNano()) // func2 dropped
   107  	mock.ExpectQuery("SELECT").WillReturnRows(droppedSprocRows)
   108  
   109  	result = reaper.DetectSprocChanges(ctx, md)
   110  	assert.Equal(t, 0, result.Created)
   111  	assert.Equal(t, 0, result.Altered)
   112  	assert.Equal(t, 1, result.Dropped) // func2 was dropped
   113  
   114  	assert.NoError(t, mock.ExpectationsWereMet())
   115  }
   116  
   117  func TestDetectTableChanges(t *testing.T) {
   118  	ctx := context.Background()
   119  
   120  	// Set up simple test metric
   121  	metricDefs.MetricDefs["table_hashes"] = metrics.Metric{
   122  		SQLs: map[int]string{
   123  			120000: "SELECT",
   124  		},
   125  	}
   126  
   127  	// Create single connection and reaper to maintain state across calls
   128  	md, mock := createTestSourceConn(t)
   129  	defer mock.Close()
   130  
   131  	reaper := &reaper{
   132  		measurementCh: make(chan metrics.MeasurementEnvelope, 10),
   133  	}
   134  
   135  	// First run - establish baseline
   136  	initialRows := pgxmock.NewRows([]string{"tag_table", "tag_oid", "md5", "epoch_ns"}).
   137  		AddRow("table1", "123", "hash1", time.Now().UnixNano()).
   138  		AddRow("table2", "456", "hash2", time.Now().UnixNano())
   139  	mock.ExpectQuery("SELECT").WillReturnRows(initialRows)
   140  
   141  	result := reaper.DetectTableChanges(ctx, md)
   142  	assert.Equal(t, 0, result.Created) // First run should not count anything as created
   143  	assert.Equal(t, 0, result.Altered)
   144  	assert.Equal(t, 0, result.Dropped)
   145  	assert.NotEmpty(t, md.ChangeState["table_hashes"])
   146  
   147  	// Second run - detect altered table (different hash for table1)
   148  	modifiedRows := pgxmock.NewRows([]string{"tag_table", "tag_oid", "md5", "epoch_ns"}).
   149  		AddRow("table1", "123", "new_hash", time.Now().UnixNano()).
   150  		AddRow("table2", "456", "hash2", time.Now().UnixNano())
   151  	mock.ExpectQuery("SELECT").WillReturnRows(modifiedRows)
   152  
   153  	result = reaper.DetectTableChanges(ctx, md)
   154  	assert.Equal(t, 0, result.Created)
   155  	assert.Equal(t, 1, result.Altered) // table1 was altered
   156  	assert.Equal(t, 0, result.Dropped)
   157  
   158  	// Third run - detect new table (table3 added)
   159  	newTableRows := pgxmock.NewRows([]string{"tag_table", "tag_oid", "md5", "epoch_ns"}).
   160  		AddRow("table1", "123", "new_hash", time.Now().UnixNano()).
   161  		AddRow("table2", "456", "hash2", time.Now().UnixNano()).
   162  		AddRow("table3", "789", "hash3", time.Now().UnixNano()) // new table
   163  	mock.ExpectQuery("SELECT").WillReturnRows(newTableRows)
   164  
   165  	result = reaper.DetectTableChanges(ctx, md)
   166  	assert.Equal(t, 1, result.Created) // table3 was created
   167  	assert.Equal(t, 0, result.Altered)
   168  	assert.Equal(t, 0, result.Dropped)
   169  
   170  	// Verify measurement is sent
   171  	select {
   172  	case msg := <-reaper.measurementCh:
   173  		assert.Equal(t, "table_changes", msg.MetricName)
   174  		assert.Equal(t, "testdb", msg.DBName)
   175  	default:
   176  		t.Error("Expected measurement to be sent")
   177  	}
   178  
   179  	// Fourth run - detect dropped table (table2 removed)
   180  	droppedTableRows := pgxmock.NewRows([]string{"tag_table", "tag_oid", "md5", "epoch_ns"}).
   181  		AddRow("table1", "123", "new_hash", time.Now().UnixNano()).
   182  		AddRow("table3", "789", "hash3", time.Now().UnixNano()) // table2 dropped
   183  	mock.ExpectQuery("SELECT").WillReturnRows(droppedTableRows)
   184  
   185  	result = reaper.DetectTableChanges(ctx, md)
   186  	assert.Equal(t, 0, result.Created)
   187  	assert.Equal(t, 0, result.Altered)
   188  	assert.Equal(t, 1, result.Dropped) // table2 was dropped
   189  
   190  	// Check that table2 was removed from state
   191  	_, exists := md.ChangeState["table_hashes"]["table2"]
   192  	assert.False(t, exists)
   193  
   194  	assert.NoError(t, mock.ExpectationsWereMet())
   195  }
   196  
   197  func TestDetectIndexChanges(t *testing.T) {
   198  	ctx := context.Background()
   199  
   200  	// Set up simple test metric
   201  	metricDefs.MetricDefs["index_hashes"] = metrics.Metric{
   202  		SQLs: map[int]string{
   203  			120000: "SELECT",
   204  		},
   205  	}
   206  
   207  	// Create single connection and reaper to maintain state across calls
   208  	md, mock := createTestSourceConn(t)
   209  	defer mock.Close()
   210  
   211  	reaper := &reaper{
   212  		measurementCh: make(chan metrics.MeasurementEnvelope, 10),
   213  	}
   214  
   215  	// First run - establish baseline
   216  	initialRows := pgxmock.NewRows([]string{"tag_index", "table", "md5", "is_valid", "epoch_ns"}).
   217  		AddRow("idx1", "table1", "hash1", "t", time.Now().UnixNano()).
   218  		AddRow("idx2", "table1", "hash2", "t", time.Now().UnixNano())
   219  	mock.ExpectQuery("SELECT").WillReturnRows(initialRows)
   220  
   221  	result := reaper.DetectIndexChanges(ctx, md)
   222  	assert.Equal(t, 0, result.Created) // First run should not count anything as created
   223  	assert.Equal(t, 0, result.Altered)
   224  	assert.Equal(t, 0, result.Dropped)
   225  	assert.NotEmpty(t, md.ChangeState["index_hashes"])
   226  
   227  	// Second run - detect altered index (is_valid changed for idx1)
   228  	modifiedRows := pgxmock.NewRows([]string{"tag_index", "table", "md5", "is_valid", "epoch_ns"}).
   229  		AddRow("idx1", "table1", "hash1", "f", time.Now().UnixNano()). // now invalid
   230  		AddRow("idx2", "table1", "hash2", "t", time.Now().UnixNano())
   231  	mock.ExpectQuery("SELECT").WillReturnRows(modifiedRows)
   232  
   233  	result = reaper.DetectIndexChanges(ctx, md)
   234  	assert.Equal(t, 0, result.Created)
   235  	assert.Equal(t, 1, result.Altered) // idx1 was altered
   236  	assert.Equal(t, 0, result.Dropped)
   237  
   238  	// Third run - detect new index (idx3 added)
   239  	newIndexRows := pgxmock.NewRows([]string{"tag_index", "table", "md5", "is_valid", "epoch_ns"}).
   240  		AddRow("idx1", "table1", "hash1", "f", time.Now().UnixNano()).
   241  		AddRow("idx2", "table1", "hash2", "t", time.Now().UnixNano()).
   242  		AddRow("idx3", "table2", "hash3", "t", time.Now().UnixNano()) // new index
   243  	mock.ExpectQuery("SELECT").WillReturnRows(newIndexRows)
   244  
   245  	result = reaper.DetectIndexChanges(ctx, md)
   246  	assert.Equal(t, 1, result.Created) // idx3 was created
   247  	assert.Equal(t, 0, result.Altered)
   248  	assert.Equal(t, 0, result.Dropped)
   249  
   250  	// Verify measurement is sent
   251  	select {
   252  	case msg := <-reaper.measurementCh:
   253  		assert.Equal(t, "index_changes", msg.MetricName)
   254  		assert.Equal(t, "testdb", msg.DBName)
   255  	default:
   256  		t.Error("Expected measurement to be sent")
   257  	}
   258  
   259  	// Fourth run - detect dropped index (idx2 removed)
   260  	droppedIndexRows := pgxmock.NewRows([]string{"tag_index", "table", "md5", "is_valid", "epoch_ns"}).
   261  		AddRow("idx1", "table1", "hash1", "f", time.Now().UnixNano()).
   262  		AddRow("idx3", "table2", "hash3", "t", time.Now().UnixNano()) // idx2 dropped
   263  	mock.ExpectQuery("SELECT").WillReturnRows(droppedIndexRows)
   264  
   265  	result = reaper.DetectIndexChanges(ctx, md)
   266  	assert.Equal(t, 0, result.Created)
   267  	assert.Equal(t, 0, result.Altered)
   268  	assert.Equal(t, 1, result.Dropped) // idx2 was dropped
   269  
   270  	// Check that idx2 was removed from state
   271  	_, exists := md.ChangeState["index_hashes"]["idx2"]
   272  	assert.False(t, exists)
   273  
   274  	assert.NoError(t, mock.ExpectationsWereMet())
   275  }
   276  
   277  func TestDetectPrivilegeChanges(t *testing.T) {
   278  	ctx := context.Background()
   279  
   280  	// Set up simple test metric
   281  	metricDefs.MetricDefs["privilege_changes"] = metrics.Metric{
   282  		SQLs: map[int]string{
   283  			120000: "SELECT",
   284  		},
   285  	}
   286  
   287  	// Create single connection and reaper to maintain state across calls
   288  	md, mock := createTestSourceConn(t)
   289  	defer mock.Close()
   290  
   291  	reaper := &reaper{
   292  		measurementCh: make(chan metrics.MeasurementEnvelope, 10),
   293  	}
   294  
   295  	// First run - establish baseline
   296  	initialRows := pgxmock.NewRows([]string{"object_type", "tag_role", "tag_object", "privilege_type", "epoch_ns"}).
   297  		AddRow("table", "user1", "table1", "SELECT", time.Now().UnixNano()).
   298  		AddRow("table", "user2", "table2", "INSERT", time.Now().UnixNano())
   299  	mock.ExpectQuery("SELECT").WillReturnRows(initialRows)
   300  
   301  	result := reaper.DetectPrivilegeChanges(ctx, md)
   302  	assert.Equal(t, 0, result.Created) // First run should not count anything as created
   303  	assert.Equal(t, 0, result.Altered)
   304  	assert.Equal(t, 0, result.Dropped)
   305  	assert.NotEmpty(t, md.ChangeState["object_privileges"])
   306  
   307  	// Second run - detect new privilege grant (user1 gets INSERT privilege on table1)
   308  	newPrivilegeRows := pgxmock.NewRows([]string{"object_type", "tag_role", "tag_object", "privilege_type", "epoch_ns"}).
   309  		AddRow("table", "user1", "table1", "SELECT", time.Now().UnixNano()).
   310  		AddRow("table", "user1", "table1", "INSERT", time.Now().UnixNano()). // new privilege
   311  		AddRow("table", "user2", "table2", "INSERT", time.Now().UnixNano())
   312  	mock.ExpectQuery("SELECT").WillReturnRows(newPrivilegeRows)
   313  
   314  	result = reaper.DetectPrivilegeChanges(ctx, md)
   315  	assert.Equal(t, 1, result.Created) // new privilege was granted
   316  	assert.Equal(t, 0, result.Altered)
   317  	assert.Equal(t, 0, result.Dropped)
   318  
   319  	// Verify measurement is sent
   320  	select {
   321  	case msg := <-reaper.measurementCh:
   322  		assert.Equal(t, "privilege_changes", msg.MetricName)
   323  		assert.Equal(t, "testdb", msg.DBName)
   324  	default:
   325  		t.Error("Expected measurement to be sent")
   326  	}
   327  
   328  	// Third run - detect privilege revoke (user1 loses INSERT privilege on table1)
   329  	revokedPrivilegeRows := pgxmock.NewRows([]string{"object_type", "tag_role", "tag_object", "privilege_type", "epoch_ns"}).
   330  		AddRow("table", "user1", "table1", "SELECT", time.Now().UnixNano()).
   331  		AddRow("table", "user2", "table2", "INSERT", time.Now().UnixNano()) // user1 INSERT privilege revoked
   332  	mock.ExpectQuery("SELECT").WillReturnRows(revokedPrivilegeRows)
   333  
   334  	result = reaper.DetectPrivilegeChanges(ctx, md)
   335  	assert.Equal(t, 0, result.Created)
   336  	assert.Equal(t, 0, result.Altered)
   337  	assert.Equal(t, 1, result.Dropped) // privilege was revoked
   338  
   339  	// Check that revoked privilege was removed from state
   340  	_, exists := md.ChangeState["object_privileges"]["table#:#user1#:#table1#:#INSERT"]
   341  	assert.False(t, exists)
   342  
   343  	assert.NoError(t, mock.ExpectationsWereMet())
   344  }
   345  
   346  func TestDetectConfigurationChanges(t *testing.T) {
   347  	ctx := context.Background()
   348  
   349  	// Set up simple test metric
   350  	metricDefs.MetricDefs["configuration_hashes"] = metrics.Metric{
   351  		SQLs: map[int]string{
   352  			120000: "SELECT",
   353  		},
   354  	}
   355  
   356  	// Create direct mock connection without transaction expectations
   357  	mock, err := pgxmock.NewPool()
   358  	require.NoError(t, err)
   359  	defer mock.Close()
   360  
   361  	md := &sources.DbConn{
   362  		Conn:   mock,
   363  		Source: sources.Source{Name: "testdb"},
   364  		RuntimeInfo: sources.RuntimeInfo{
   365  			Version:     120000,
   366  			ChangeState: make(map[string]map[string]string),
   367  		},
   368  	}
   369  
   370  	reaper := &reaper{
   371  		measurementCh: make(chan metrics.MeasurementEnvelope, 10),
   372  	}
   373  
   374  	// First run - establish baseline configuration
   375  	initialRows := pgxmock.NewRows([]string{"epoch_ns", "setting", "value"}).
   376  		AddRow(time.Now().UnixNano(), "max_connections", "100").
   377  		AddRow(time.Now().UnixNano(), "shared_buffers", "128MB")
   378  	mock.ExpectQuery("SELECT").WillReturnRows(initialRows)
   379  
   380  	result := reaper.DetectConfigurationChanges(ctx, md)
   381  	assert.Equal(t, 0, result.Created) // First run should not count anything as created
   382  	assert.Equal(t, 0, result.Altered)
   383  	assert.Equal(t, 0, result.Dropped)
   384  	assert.NotEmpty(t, md.ChangeState["configuration_hashes"])
   385  
   386  	// Second run - detect new configuration setting
   387  	newSettingRows := pgxmock.NewRows([]string{"epoch_ns", "setting", "value"}).
   388  		AddRow(time.Now().UnixNano(), "max_connections", "100").
   389  		AddRow(time.Now().UnixNano(), "shared_buffers", "128MB").
   390  		AddRow(time.Now().UnixNano(), "work_mem", "4MB") // new setting
   391  	mock.ExpectQuery("SELECT").WillReturnRows(newSettingRows)
   392  
   393  	result = reaper.DetectConfigurationChanges(ctx, md)
   394  	assert.Equal(t, 1, result.Created) // new setting was added
   395  	assert.Equal(t, 0, result.Altered)
   396  	assert.Equal(t, 0, result.Dropped)
   397  
   398  	// Verify measurement is sent
   399  	select {
   400  	case msg := <-reaper.measurementCh:
   401  		assert.Equal(t, "configuration_changes", msg.MetricName)
   402  		assert.Equal(t, "testdb", msg.DBName)
   403  	default:
   404  		t.Error("Expected measurement to be sent")
   405  	}
   406  
   407  	// Third run - detect configuration change
   408  	changedValueRows := pgxmock.NewRows([]string{"epoch_ns", "setting", "value"}).
   409  		AddRow(time.Now().UnixNano(), "max_connections", "200").  // changed value
   410  		AddRow(time.Now().UnixNano(), "shared_buffers", "256MB"). // changed value
   411  		AddRow(time.Now().UnixNano(), "work_mem", "4MB")
   412  	mock.ExpectQuery("SELECT").WillReturnRows(changedValueRows)
   413  
   414  	result = reaper.DetectConfigurationChanges(ctx, md)
   415  	assert.Equal(t, 0, result.Created)
   416  	assert.Equal(t, 2, result.Altered) // two settings were changed
   417  	assert.Equal(t, 0, result.Dropped)
   418  
   419  	// Check that state was updated
   420  	assert.Equal(t, "200", md.ChangeState["configuration_hashes"]["max_connections"])
   421  	assert.Equal(t, "256MB", md.ChangeState["configuration_hashes"]["shared_buffers"])
   422  
   423  	assert.NoError(t, mock.ExpectationsWereMet())
   424  }
   425  
   426  func TestGetInstanceUpMeasurement(t *testing.T) {
   427  	ctx := context.Background()
   428  	reaper := &reaper{}
   429  
   430  	testCases := []struct {
   431  		name            string
   432  		pingError       error
   433  		expectedUpValue int
   434  	}{
   435  		{
   436  			name:            "connection is up",
   437  			pingError:       nil,
   438  			expectedUpValue: 1,
   439  		},
   440  		{
   441  			name:            "connection is down",
   442  			pingError:       assert.AnError,
   443  			expectedUpValue: 0,
   444  		},
   445  		{
   446  			name:            "connection timeout",
   447  			pingError:       context.DeadlineExceeded,
   448  			expectedUpValue: 0,
   449  		},
   450  	}
   451  
   452  	for _, tc := range testCases {
   453  		t.Run(tc.name, func(t *testing.T) {
   454  			md, mock := createTestSourceConn(t)
   455  			defer mock.Close()
   456  
   457  			// Setup ping expectation
   458  			if tc.pingError == nil {
   459  				mock.ExpectPing()
   460  			} else {
   461  				mock.ExpectPing().WillReturnError(tc.pingError)
   462  			}
   463  
   464  			measurements, err := reaper.GetInstanceUpMeasurement(ctx, md)
   465  
   466  			// Should never return an error
   467  			assert.NoError(t, err)
   468  			require.NotNil(t, measurements)
   469  			require.Len(t, measurements, 1)
   470  
   471  			// Check instance_up metric value
   472  			measurement := measurements[0]
   473  			assert.Contains(t, measurement, "instance_up")
   474  			assert.Equal(t, tc.expectedUpValue, measurement["instance_up"])
   475  
   476  			// Verify epoch is set and valid
   477  			assert.Contains(t, measurement, metrics.EpochColumnName)
   478  			assert.Greater(t, measurement[metrics.EpochColumnName].(int64), int64(0))
   479  			assert.LessOrEqual(t, measurement[metrics.EpochColumnName].(int64), time.Now().UnixNano())
   480  
   481  			assert.NoError(t, mock.ExpectationsWereMet())
   482  		})
   483  	}
   484  }
   485  
   486  func TestGCDSlice(t *testing.T) {
   487  	tests := []struct {
   488  		name string
   489  		vals []int
   490  		want int
   491  	}{
   492  		{"empty", nil, 0},
   493  		{"single", []int{30}, 30},
   494  		{"exhaustive preset intervals", []int{30, 60, 120, 180, 300, 600, 900, 3600, 7200}, 30},
   495  		{"coprime", []int{7, 11, 13}, 1},
   496  		{"all same", []int{60, 60, 60}, 60},
   497  		{"basic preset", []int{60, 120}, 60},
   498  	}
   499  	for _, tc := range tests {
   500  		t.Run(tc.name, func(t *testing.T) {
   501  			assert.Equal(t, tc.want, GCDSlice(tc.vals))
   502  		})
   503  	}
   504  }
   505  
   506  func TestCalcTickInterval(t *testing.T) {
   507  	t.Run("exhaustive preset GCD is 30s", func(t *testing.T) {
   508  		sr := &DbConnReaper{
   509  			md: &sources.DbConn{
   510  				Source: sources.Source{
   511  					Metrics: metrics.MetricIntervals{"m1": 30, "m2": 60, "m3": 120, "m4": 300},
   512  				},
   513  			},
   514  		}
   515  		assert.Equal(t, 30*time.Second, sr.calcTickInterval())
   516  	})
   517  
   518  	t.Run("GCD floors to minimum 1s", func(t *testing.T) {
   519  		sr := &DbConnReaper{
   520  			md: &sources.DbConn{
   521  				Source: sources.Source{
   522  					Metrics: metrics.MetricIntervals{"m1": 3, "m2": 7},
   523  				},
   524  			},
   525  		}
   526  		assert.Equal(t, time.Second, sr.calcTickInterval())
   527  	})
   528  
   529  	t.Run("single metric", func(t *testing.T) {
   530  		sr := &DbConnReaper{
   531  			md: &sources.DbConn{
   532  				Source: sources.Source{
   533  					Metrics: metrics.MetricIntervals{"m1": 60},
   534  				},
   535  			},
   536  		}
   537  		assert.Equal(t, 60*time.Second, sr.calcTickInterval())
   538  	})
   539  
   540  	t.Run("empty metrics", func(t *testing.T) {
   541  		sr := &DbConnReaper{
   542  			md: &sources.DbConn{
   543  				Source: sources.Source{
   544  					Metrics: metrics.MetricIntervals{},
   545  				},
   546  			},
   547  		}
   548  		assert.Equal(t, time.Second, sr.calcTickInterval())
   549  	})
   550  
   551  	t.Run("standby metrics when in recovery", func(t *testing.T) {
   552  		sr := &DbConnReaper{
   553  			md: &sources.DbConn{
   554  				Source: sources.Source{
   555  					Metrics:        metrics.MetricIntervals{"m1": 30, "m2": 60},
   556  					MetricsStandby: metrics.MetricIntervals{"m1": 120},
   557  				},
   558  				RuntimeInfo: sources.RuntimeInfo{IsInRecovery: true},
   559  			},
   560  		}
   561  		assert.Equal(t, 120*time.Second, sr.calcTickInterval())
   562  	})
   563  }
   564  
   565  func TestNewSourceReaper(t *testing.T) {
   566  	r := &reaper{
   567  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   568  		measurementCache: NewInstanceMetricCache(),
   569  	}
   570  	md := &sources.DbConn{
   571  		Source: sources.Source{
   572  			Name:    "testdb",
   573  			Kind:    sources.SourcePostgres,
   574  			Metrics: metrics.MetricIntervals{"cpu": 30, "mem": 60, "disk": 120},
   575  		},
   576  	}
   577  	sr := NewDbConnReaper(r, md)
   578  
   579  	assert.NotNil(t, sr.lastFetch)
   580  	assert.Empty(t, sr.lastFetch)
   581  	assert.Equal(t, r, sr.reaper)
   582  	assert.Equal(t, md, sr.md)
   583  }
   584  
   585  func TestSourceReaper_ExecuteBatch(t *testing.T) {
   586  	ctx := log.WithLogger(context.Background(), log.NewNoopLogger())
   587  
   588  	metricDefs.MetricDefs["batch_metric_1"] = metrics.Metric{
   589  		SQLs: metrics.SQLs{0: "SELECT 1 as value, 100::bigint as epoch_ns"},
   590  	}
   591  	metricDefs.MetricDefs["batch_metric_2"] = metrics.Metric{
   592  		SQLs: metrics.SQLs{0: "SELECT 2 as value, 200::bigint as epoch_ns"},
   593  	}
   594  
   595  	mock, err := pgxmock.NewPool()
   596  	require.NoError(t, err)
   597  	defer mock.Close()
   598  
   599  	md := &sources.DbConn{
   600  		Source: sources.Source{
   601  			Name:    "test_source",
   602  			Kind:    sources.SourcePostgres,
   603  			Metrics: metrics.MetricIntervals{"batch_metric_1": 30, "batch_metric_2": 30},
   604  		},
   605  		Conn: mock,
   606  		RuntimeInfo: sources.RuntimeInfo{
   607  			Version:     120000,
   608  			ChangeState: make(map[string]map[string]string),
   609  		},
   610  	}
   611  
   612  	r := &reaper{
   613  		Options: &cmdopts.Options{
   614  			Metrics: metrics.CmdOpts{},
   615  			Sinks:   sinks.CmdOpts{},
   616  		},
   617  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   618  		measurementCache: NewInstanceMetricCache(),
   619  	}
   620  	sr := NewDbConnReaper(r, md)
   621  
   622  	rows1 := pgxmock.NewRows([]string{"epoch_ns", "value"}).
   623  		AddRow(time.Now().UnixNano(), int64(100))
   624  	rows2 := pgxmock.NewRows([]string{"epoch_ns", "value"}).
   625  		AddRow(time.Now().UnixNano(), int64(200))
   626  	eb := mock.ExpectBatch()
   627  	eb.ExpectQuery("SELECT 1").WillReturnRows(rows1)
   628  	eb.ExpectQuery("SELECT 2").WillReturnRows(rows2)
   629  
   630  	err = sr.executeBatch(ctx, []batchEntry{
   631  		{metricName: "batch_metric_1", metric: metricDefs.MetricDefs["batch_metric_1"], sql: "SELECT 1 as value, 100::bigint as epoch_ns"},
   632  		{metricName: "batch_metric_2", metric: metricDefs.MetricDefs["batch_metric_2"], sql: "SELECT 2 as value, 200::bigint as epoch_ns"},
   633  	})
   634  	assert.NoError(t, err)
   635  
   636  	received := 0
   637  	for {
   638  		select {
   639  		case msg := <-r.measurementCh:
   640  			assert.Equal(t, "test_source", msg.DBName)
   641  			assert.True(t, msg.MetricName == "batch_metric_1" || msg.MetricName == "batch_metric_2")
   642  			received++
   643  		default:
   644  			goto done
   645  		}
   646  	}
   647  done:
   648  	assert.Equal(t, 2, received, "should have received 2 measurement envelopes")
   649  	assert.NoError(t, mock.ExpectationsWereMet())
   650  }
   651  
   652  func TestSourceReaper_RunOneIteration(t *testing.T) {
   653  	ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
   654  
   655  	metricDefs.MetricDefs["run_test_metric"] = metrics.Metric{
   656  		SQLs: metrics.SQLs{0: "SELECT run_test"},
   657  	}
   658  
   659  	mock, err := pgxmock.NewPool()
   660  	require.NoError(t, err)
   661  	defer mock.Close()
   662  
   663  	md := &sources.DbConn{
   664  		Source: sources.Source{
   665  			Name:    "run_source",
   666  			Kind:    sources.SourcePostgres,
   667  			Metrics: metrics.MetricIntervals{"run_test_metric": 5},
   668  		},
   669  		Conn: mock,
   670  		RuntimeInfo: sources.RuntimeInfo{
   671  			Version:     120000,
   672  			ChangeState: make(map[string]map[string]string),
   673  		},
   674  	}
   675  
   676  	r := &reaper{
   677  		Options: &cmdopts.Options{
   678  			Metrics: metrics.CmdOpts{},
   679  			Sinks:   sinks.CmdOpts{},
   680  		},
   681  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   682  		measurementCache: NewInstanceMetricCache(),
   683  	}
   684  	sr := NewDbConnReaper(r, md)
   685  
   686  	// FetchRuntimeInfo sends a query
   687  	mock.ExpectQuery("select /\\* pgwatch_generated \\*/").
   688  		WillReturnError(assert.AnError)
   689  
   690  	rows := pgxmock.NewRows([]string{"epoch_ns", "value"}).
   691  		AddRow(time.Now().UnixNano(), int64(42))
   692  	eb := mock.ExpectBatch()
   693  	eb.ExpectQuery("SELECT run_test").WillReturnRows(rows)
   694  
   695  	go func() {
   696  		time.Sleep(200 * time.Millisecond)
   697  		cancel()
   698  	}()
   699  
   700  	sr.Reap(ctx)
   701  
   702  	select {
   703  	case msg := <-r.measurementCh:
   704  		assert.Equal(t, "run_source", msg.DBName)
   705  		assert.Equal(t, "run_test_metric", msg.MetricName)
   706  	case <-time.After(time.Second):
   707  		t.Error("Expected measurement but timed out")
   708  	}
   709  }
   710  
   711  func TestSourceReaper_DetectServerRestart(t *testing.T) {
   712  	sr := &DbConnReaper{
   713  		reaper: &reaper{
   714  			measurementCh: make(chan metrics.MeasurementEnvelope, 10),
   715  		},
   716  		md: &sources.DbConn{
   717  			Source: sources.Source{Name: "restart_test"},
   718  		},
   719  	}
   720  
   721  	// First observation — establish baseline
   722  	data := metrics.Measurements{
   723  		{"epoch_ns": time.Now().UnixNano(), "postmaster_uptime_s": int64(1000)},
   724  	}
   725  	sr.detectServerRestart(t.Context(), data)
   726  	assert.Equal(t, int64(1000), sr.lastUptimeS)
   727  	select {
   728  	case <-sr.reaper.measurementCh:
   729  		t.Error("should not emit restart event on first observation")
   730  	default:
   731  	}
   732  
   733  	// Second observation — uptime increased (normal)
   734  	data = metrics.Measurements{
   735  		{"epoch_ns": time.Now().UnixNano(), "postmaster_uptime_s": int64(2000)},
   736  	}
   737  	sr.detectServerRestart(t.Context(), data)
   738  	assert.Equal(t, int64(2000), sr.lastUptimeS)
   739  	select {
   740  	case <-sr.reaper.measurementCh:
   741  		t.Error("should not emit restart event when uptime increases")
   742  	default:
   743  	}
   744  
   745  	// Third observation — uptime decreased (restart!)
   746  	data = metrics.Measurements{
   747  		{"epoch_ns": time.Now().UnixNano(), "postmaster_uptime_s": int64(10)},
   748  	}
   749  	sr.detectServerRestart(t.Context(), data)
   750  	assert.Equal(t, int64(10), sr.lastUptimeS)
   751  	select {
   752  	case msg := <-sr.reaper.measurementCh:
   753  		assert.Equal(t, "object_changes", msg.MetricName)
   754  		assert.Contains(t, msg.Data[0]["details"], "restart")
   755  	default:
   756  		t.Error("expected restart event")
   757  	}
   758  }
   759  
   760  func TestSourceReaper_FetchSpecialMetric(t *testing.T) {
   761  	ctx := log.WithLogger(context.Background(), log.NewNoopLogger())
   762  
   763  	newSR := func(t *testing.T) (*DbConnReaper, *sources.DbConn, pgxmock.PgxPoolIface) {
   764  		t.Helper()
   765  		md, mock := createTestSourceConn(t)
   766  		r := &reaper{
   767  			Options: &cmdopts.Options{
   768  				Metrics: metrics.CmdOpts{},
   769  				Sinks:   sinks.CmdOpts{},
   770  			},
   771  			measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   772  			measurementCache: NewInstanceMetricCache(),
   773  		}
   774  		return NewDbConnReaper(r, md), md, mock
   775  	}
   776  
   777  	sr, _, mock := newSR(t)
   778  	defer mock.Close()
   779  
   780  	t.Run("instance_up dispatches measurement on ping success", func(t *testing.T) {
   781  		mock.ExpectPing()
   782  		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricInstanceUp, ""))
   783  		select {
   784  		case msg := <-sr.reaper.measurementCh:
   785  			assert.Equal(t, specialMetricInstanceUp, msg.MetricName)
   786  			assert.Len(t, msg.Data, 1)
   787  			assert.Equal(t, 1, msg.Data[0][specialMetricInstanceUp])
   788  		default:
   789  			t.Error("expected measurement for instance_up")
   790  		}
   791  		assert.NoError(t, mock.ExpectationsWereMet())
   792  	})
   793  
   794  	t.Run("instance_up uses storage name when set", func(t *testing.T) {
   795  		mock.ExpectPing()
   796  		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricInstanceUp, "infra_up"))
   797  		select {
   798  		case msg := <-sr.reaper.measurementCh:
   799  			assert.Equal(t, "infra_up", msg.MetricName)
   800  		default:
   801  			t.Error("expected measurement")
   802  		}
   803  		assert.NoError(t, mock.ExpectationsWereMet())
   804  	})
   805  
   806  	t.Run("change_events dispatches no measurement when no hash defs present", func(t *testing.T) {
   807  		// Doesn't contain additional defs for any of {"sproc_hashes", "table_hashes", "index_hashes", "configuration_hashes", "privilege_hashes"}
   808  		metricDefs.MetricDefs[specialMetricChangeEvents] = metrics.Metric{}
   809  		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricChangeEvents, ""))
   810  		select {
   811  		case <-sr.reaper.measurementCh:
   812  			t.Error("expected no measurement when no changes detected")
   813  		default:
   814  		}
   815  		assert.NoError(t, mock.ExpectationsWereMet())
   816  	})
   817  }
   818  
   819  func TestSourceReaper_ExecuteBatch_DegradedOnPersistentFailure(t *testing.T) {
   820  	ctx := log.WithLogger(context.Background(), log.NewNoopLogger())
   821  
   822  	metricDefs.MetricDefs["good_metric"] = metrics.Metric{
   823  		SQLs: metrics.SQLs{0: "SELECT 1 as value, 100::bigint as epoch_ns"},
   824  	}
   825  	metricDefs.MetricDefs["bad_metric"] = metrics.Metric{
   826  		SQLs: metrics.SQLs{0: "SELECT bad"},
   827  	}
   828  
   829  	mock, err := pgxmock.NewPool()
   830  	require.NoError(t, err)
   831  	defer mock.Close()
   832  
   833  	md := &sources.DbConn{
   834  		Source: sources.Source{
   835  			Name:    "degrade_test",
   836  			Kind:    sources.SourcePostgres,
   837  			Metrics: metrics.MetricIntervals{"good_metric": 30, "bad_metric": 30},
   838  		},
   839  		Conn: mock,
   840  		RuntimeInfo: sources.RuntimeInfo{
   841  			Version:     120000,
   842  			ChangeState: make(map[string]map[string]string),
   843  		},
   844  	}
   845  	r := &reaper{
   846  		Options: &cmdopts.Options{
   847  			Metrics: metrics.CmdOpts{},
   848  			Sinks:   sinks.CmdOpts{},
   849  		},
   850  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   851  		measurementCache: NewInstanceMetricCache(),
   852  	}
   853  	sr := NewDbConnReaper(r, md)
   854  
   855  	entries := []batchEntry{
   856  		{metricName: "good_metric", metric: metricDefs.MetricDefs["good_metric"], sql: "SELECT 1 as value, 100::bigint as epoch_ns"},
   857  		{metricName: "bad_metric", metric: metricDefs.MetricDefs["bad_metric"], sql: "SELECT bad"},
   858  	}
   859  
   860  	// batch: good_metric succeeds, bad_metric cascades → retry bad_metric individually → still fails
   861  	rows1 := pgxmock.NewRows([]string{"epoch_ns", "value"}).AddRow(time.Now().UnixNano(), int64(1))
   862  	eb := mock.ExpectBatch()
   863  	eb.ExpectQuery("SELECT 1").WillReturnRows(rows1)
   864  	eb.ExpectQuery("SELECT bad").WillReturnError(assert.AnError) // cascade
   865  	// individual retry of bad_metric
   866  	mock.ExpectQuery("SELECT bad").WithArgs(pgx.QueryExecModeSimpleProtocol).WillReturnError(assert.AnError)
   867  
   868  	err = sr.executeBatch(ctx, entries)
   869  	assert.Error(t, err)
   870  	assert.Contains(t, sr.degradedMetrics, "bad_metric", "bad_metric should be degraded after persistent failure")
   871  	assert.NotContains(t, sr.degradedMetrics, "good_metric", "good_metric should not be degraded")
   872  	assert.NoError(t, mock.ExpectationsWereMet())
   873  }
   874  
   875  func TestSourceReaper_ExecuteBatch_CascadeRecovery(t *testing.T) {
   876  	// A metric that errors in the batch but succeeds on individual retry must NOT be marked degraded.
   877  	ctx := log.WithLogger(context.Background(), log.NewNoopLogger())
   878  
   879  	metricDefs.MetricDefs["cascade_victim"] = metrics.Metric{
   880  		SQLs: metrics.SQLs{0: "SELECT 3 as value, 300::bigint as epoch_ns"},
   881  	}
   882  	metricDefs.MetricDefs["cascade_trigger"] = metrics.Metric{
   883  		SQLs: metrics.SQLs{0: "SELECT fail"},
   884  	}
   885  
   886  	mock, err := pgxmock.NewPool()
   887  	require.NoError(t, err)
   888  	defer mock.Close()
   889  
   890  	md := &sources.DbConn{
   891  		Source: sources.Source{
   892  			Name:    "cascade_test",
   893  			Kind:    sources.SourcePostgres,
   894  			Metrics: metrics.MetricIntervals{"cascade_trigger": 30, "cascade_victim": 30},
   895  		},
   896  		Conn: mock,
   897  		RuntimeInfo: sources.RuntimeInfo{
   898  			Version:     120000,
   899  			ChangeState: make(map[string]map[string]string),
   900  		},
   901  	}
   902  	r := &reaper{
   903  		Options: &cmdopts.Options{
   904  			Metrics: metrics.CmdOpts{},
   905  			Sinks:   sinks.CmdOpts{},
   906  		},
   907  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   908  		measurementCache: NewInstanceMetricCache(),
   909  	}
   910  	sr := NewDbConnReaper(r, md)
   911  
   912  	entries := []batchEntry{
   913  		{metricName: "cascade_trigger", metric: metricDefs.MetricDefs["cascade_trigger"], sql: "SELECT fail"},
   914  		{metricName: "cascade_victim", metric: metricDefs.MetricDefs["cascade_victim"], sql: "SELECT 3 as value, 300::bigint as epoch_ns"},
   915  	}
   916  
   917  	// batch: trigger fails, victim cascades → both retry individually
   918  	// trigger fails individually (real error), victim succeeds individually (was only a cascade)
   919  	eb := mock.ExpectBatch()
   920  	eb.ExpectQuery("SELECT fail").WillReturnError(assert.AnError)
   921  	eb.ExpectQuery("SELECT 3").WillReturnError(assert.AnError) // cascade in batch
   922  	// individual retries
   923  	mock.ExpectQuery("SELECT fail").WithArgs(pgx.QueryExecModeSimpleProtocol).WillReturnError(assert.AnError)
   924  	mock.ExpectQuery("SELECT 3").WithArgs(pgx.QueryExecModeSimpleProtocol).
   925  		WillReturnRows(pgxmock.NewRows([]string{"epoch_ns", "value"}).AddRow(time.Now().UnixNano(), int64(3)))
   926  
   927  	err = sr.executeBatch(ctx, entries)
   928  	assert.Error(t, err, "cascade_trigger error should propagate")
   929  	assert.Contains(t, sr.degradedMetrics, "cascade_trigger", "real-failure metric should be degraded")
   930  	assert.NotContains(t, sr.degradedMetrics, "cascade_victim", "cascade-only victim must not be degraded")
   931  	assert.NoError(t, mock.ExpectationsWereMet())
   932  }
   933  
   934  func TestSourceReaper_DegradedMetricRecovery(t *testing.T) {
   935  	// Uses the real Run loop (via synctest fake clock) to verify the full degraded→recovered
   936  	// lifecycle: iteration 1 the degraded metric fails individually (stays degraded),
   937  	// iteration 2 it succeeds (removed from degradedMetrics).
   938  	synctest.Test(t, func(t *testing.T) {
   939  		const (
   940  			metricName     = "recovering_metric_real"
   941  			metricInterval = 30
   942  		)
   943  
   944  		metricDefs.MetricDefs[metricName] = metrics.Metric{
   945  			SQLs: metrics.SQLs{0: "SELECT 7 as value, 700::bigint as epoch_ns"},
   946  		}
   947  
   948  		mock, err := pgxmock.NewPool()
   949  		require.NoError(t, err)
   950  		defer mock.Close()
   951  
   952  		md := &sources.DbConn{
   953  			Source: sources.Source{
   954  				Name:    "recovery_src",
   955  				Kind:    sources.SourcePostgres,
   956  				Metrics: metrics.MetricIntervals{metricName: metricInterval},
   957  			},
   958  			Conn: mock,
   959  			RuntimeInfo: sources.RuntimeInfo{
   960  				Version:     120000,
   961  				ChangeState: make(map[string]map[string]string),
   962  			},
   963  		}
   964  		r := &reaper{
   965  			Options: &cmdopts.Options{
   966  				Metrics: metrics.CmdOpts{},
   967  				Sinks:   sinks.CmdOpts{},
   968  			},
   969  			measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
   970  			measurementCache: NewInstanceMetricCache(),
   971  		}
   972  		ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   973  		sr := NewDbConnReaper(r, md)
   974  		sr.markDegraded(metricName) // pre-seed: metric already degraded
   975  
   976  		// Iteration 1: FetchRuntimeInfo + degraded individual fetch → fails → stays degraded
   977  		mock.ExpectQuery("select /\\* pgwatch_generated \\*/").WillReturnError(assert.AnError)
   978  		mock.ExpectQuery("SELECT 7").WithArgs(pgx.QueryExecModeSimpleProtocol).WillReturnError(assert.AnError)
   979  
   980  		// Iteration 2: FetchRuntimeInfo + degraded individual fetch → succeeds → recovered
   981  		mock.ExpectQuery("select /\\* pgwatch_generated \\*/").WillReturnError(assert.AnError)
   982  		mock.ExpectQuery("SELECT 7").WithArgs(pgx.QueryExecModeSimpleProtocol).
   983  			WillReturnRows(pgxmock.NewRows([]string{"epoch_ns", "value"}).AddRow(int64(700_000_000_000), int64(7)))
   984  
   985  		go sr.Reap(ctx)
   986  
   987  		// Run goroutine completes iteration 1 (pgxmock is in-memory, no real I/O) then
   988  		// blocks on time.After — the only durably-blocking operation in the loop.
   989  		synctest.Wait()
   990  		assert.True(t, sr.isDegraded(metricName), "should still be degraded after first failure")
   991  
   992  		// Advance the fake clock past the interval to trigger iteration 2.
   993  		// The Run goroutine's time.After(30s) fires first; it runs iteration 2 and
   994  		// blocks again before the test goroutine's sleep finishes.
   995  		time.Sleep(time.Duration(metricInterval)*time.Second + time.Millisecond)
   996  		synctest.Wait()
   997  		assert.False(t, sr.isDegraded(metricName), "should recover after successful fetchMetric")
   998  
   999  		assert.NoError(t, mock.ExpectationsWereMet())
  1000  	})
  1001  }
  1002  
  1003  func TestSourceReaper_NonPostgresSequential(t *testing.T) {
  1004  	ctx := log.WithLogger(context.Background(), log.NewNoopLogger())
  1005  
  1006  	metricDefs.MetricDefs["seq_metric"] = metrics.Metric{
  1007  		SQLs: metrics.SQLs{0: "SELECT seq_value"},
  1008  	}
  1009  
  1010  	mock, err := pgxmock.NewPool()
  1011  	require.NoError(t, err)
  1012  	defer mock.Close()
  1013  
  1014  	md := &sources.DbConn{
  1015  		Source: sources.Source{
  1016  			Name:    "seq_test_src",
  1017  			Kind:    sources.SourcePostgres,
  1018  			Metrics: metrics.MetricIntervals{"seq_metric": 30},
  1019  		},
  1020  		Conn: mock,
  1021  		RuntimeInfo: sources.RuntimeInfo{
  1022  			Version:     120000,
  1023  			ChangeState: make(map[string]map[string]string),
  1024  		},
  1025  	}
  1026  
  1027  	r := &reaper{
  1028  		Options: &cmdopts.Options{
  1029  			Metrics: metrics.CmdOpts{},
  1030  			Sinks:   sinks.CmdOpts{},
  1031  		},
  1032  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
  1033  		measurementCache: NewInstanceMetricCache(),
  1034  	}
  1035  	sr := NewDbConnReaper(r, md)
  1036  
  1037  	rows := pgxmock.NewRows([]string{"epoch_ns", "value"}).
  1038  		AddRow(time.Now().UnixNano(), int64(42))
  1039  	mock.ExpectQuery("SELECT seq_value").WithArgs(pgx.QueryExecModeSimpleProtocol).WillReturnRows(rows)
  1040  
  1041  	err = sr.fetchMetric(ctx, batchEntry{metricName: "seq_metric", metric: metricDefs.MetricDefs["seq_metric"], sql: "SELECT seq_value"})
  1042  	assert.NoError(t, err)
  1043  	assert.NoError(t, mock.ExpectationsWereMet())
  1044  }
  1045  
  1046  func TestSourceReaper_ExecuteBatch_BoundedByFetchDeadline(t *testing.T) {
  1047  	// SendBatch blocks forever on a BlockingPool. Without bounded fetch
  1048  	// wiring, executeBatch would only return when the parent ctx is
  1049  	// cancelled; with wiring, it returns at max(interval, MinFetchTimeout).
  1050  	origMinFetch := db.MinFetchTimeout
  1051  	t.Cleanup(func() { db.MinFetchTimeout = origMinFetch })
  1052  	db.MinFetchTimeout = 50 * time.Millisecond
  1053  
  1054  	ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
  1055  	defer cancel()
  1056  	// Safety net so a regressed implementation does not hang the test
  1057  	// indefinitely. With correct wiring the fetch deadline fires first.
  1058  	time.AfterFunc(5*time.Second, cancel)
  1059  
  1060  	metricDefs.MetricDefs["bounded_batch_metric"] = metrics.Metric{
  1061  		SQLs: metrics.SQLs{0: "SELECT 1"},
  1062  	}
  1063  
  1064  	md := &sources.DbConn{
  1065  		Source: sources.Source{
  1066  			Name:    "bounded_batch_src",
  1067  			Kind:    sources.SourcePostgres,
  1068  			Metrics: metrics.MetricIntervals{"bounded_batch_metric": 1},
  1069  		},
  1070  		Conn: testutil.BlockingPool{},
  1071  		RuntimeInfo: sources.RuntimeInfo{
  1072  			Version:     120000,
  1073  			ChangeState: make(map[string]map[string]string),
  1074  		},
  1075  	}
  1076  	r := &reaper{
  1077  		Options: &cmdopts.Options{
  1078  			Metrics: metrics.CmdOpts{},
  1079  			Sinks:   sinks.CmdOpts{},
  1080  		},
  1081  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
  1082  		measurementCache: NewInstanceMetricCache(),
  1083  	}
  1084  	sr := NewDbConnReaper(r, md)
  1085  
  1086  	start := time.Now()
  1087  	err := sr.executeBatch(ctx, []batchEntry{
  1088  		{metricName: "bounded_batch_metric", metric: metricDefs.MetricDefs["bounded_batch_metric"], sql: "SELECT 1"},
  1089  	})
  1090  	elapsed := time.Since(start)
  1091  
  1092  	require.Error(t, err)
  1093  	// Deadline error must carry the operation name so callers (and log
  1094  	// grep) can distinguish a fetch-bound timeout from ordinary query
  1095  	// errors. Fails RED today (unwired: err is plain
  1096  	// "context deadline exceeded"); goes GREEN once the call site derives
  1097  	// a "batch"-tagged ctx and surfaces its cause.
  1098  	require.ErrorContains(t, err, "batch")
  1099  	// fetch deadline = max(1s interval, 50ms floor) = 1s. executeBatch may
  1100  	// do a retry round-trip (fetchMetric) on top of the batch round-trip,
  1101  	// so the total budget is 2 * deadline + a generous scheduling margin.
  1102  	if elapsed > 2*time.Second+500*time.Millisecond {
  1103  		t.Fatalf("executeBatch took %v, want ~2 * fetch deadline", elapsed)
  1104  	}
  1105  }
  1106  
  1107  func TestSourceReaper_FetchMetric_BoundedByFetchDeadline(t *testing.T) {
  1108  	// Query blocks forever on a BlockingPool. The fetch deadline for an
  1109  	// individual metric equals max(GetMetricInterval(name), MinFetchTimeout).
  1110  	origMinFetch := db.MinFetchTimeout
  1111  	t.Cleanup(func() { db.MinFetchTimeout = origMinFetch })
  1112  	db.MinFetchTimeout = 50 * time.Millisecond
  1113  
  1114  	ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
  1115  	defer cancel()
  1116  	time.AfterFunc(5*time.Second, cancel)
  1117  
  1118  	metricDefs.MetricDefs["bounded_fetch_metric"] = metrics.Metric{
  1119  		SQLs: metrics.SQLs{0: "SELECT 1"},
  1120  	}
  1121  
  1122  	md := &sources.DbConn{
  1123  		Source: sources.Source{
  1124  			Name:    "bounded_fetch_src",
  1125  			Kind:    sources.SourcePostgres,
  1126  			Metrics: metrics.MetricIntervals{"bounded_fetch_metric": 1},
  1127  		},
  1128  		Conn: testutil.BlockingPool{},
  1129  		RuntimeInfo: sources.RuntimeInfo{
  1130  			Version:     120000,
  1131  			ChangeState: make(map[string]map[string]string),
  1132  		},
  1133  	}
  1134  	r := &reaper{
  1135  		Options: &cmdopts.Options{
  1136  			Metrics: metrics.CmdOpts{},
  1137  			Sinks:   sinks.CmdOpts{},
  1138  		},
  1139  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
  1140  		measurementCache: NewInstanceMetricCache(),
  1141  	}
  1142  	sr := NewDbConnReaper(r, md)
  1143  
  1144  	start := time.Now()
  1145  	err := sr.fetchMetric(ctx, batchEntry{
  1146  		metricName: "bounded_fetch_metric",
  1147  		metric:     metricDefs.MetricDefs["bounded_fetch_metric"],
  1148  		sql:        "SELECT 1",
  1149  	})
  1150  	elapsed := time.Since(start)
  1151  
  1152  	require.Error(t, err)
  1153  	if elapsed > 1500*time.Millisecond {
  1154  		t.Fatalf("fetchMetric took %v, want ~max(GetMetricInterval, MinFetchTimeout)", elapsed)
  1155  	}
  1156  }
  1157  
  1158  func TestSourceReaper_QueryMeasurements_BoundedByChangeDetectionTimeout(t *testing.T) {
  1159  	// Query blocks forever on a BlockingPool. QueryMeasurements feeds the
  1160  	// Detect*Changes family and must honor ChangeDetectionTimeout.
  1161  	origChange := db.ChangeDetectionTimeout
  1162  	t.Cleanup(func() { db.ChangeDetectionTimeout = origChange })
  1163  	db.ChangeDetectionTimeout = 100 * time.Millisecond
  1164  
  1165  	ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
  1166  	defer cancel()
  1167  	time.AfterFunc(5*time.Second, cancel)
  1168  
  1169  	md := &sources.DbConn{
  1170  		Source: sources.Source{
  1171  			Name: "bounded_change_src",
  1172  			Kind: sources.SourcePostgres,
  1173  		},
  1174  		Conn: testutil.BlockingPool{},
  1175  		RuntimeInfo: sources.RuntimeInfo{
  1176  			Version:     120000,
  1177  			ChangeState: make(map[string]map[string]string),
  1178  		},
  1179  	}
  1180  
  1181  	start := time.Now()
  1182  	_, err := QueryMeasurements(ctx, md, "SELECT 1")
  1183  	elapsed := time.Since(start)
  1184  
  1185  	require.Error(t, err)
  1186  	if elapsed > 1500*time.Millisecond {
  1187  		t.Fatalf("QueryMeasurements took %v, want ~ChangeDetectionTimeout (100ms)", elapsed)
  1188  	}
  1189  }
  1190  
  1191  func TestSourceReaper_DetectSprocChanges_BoundedByChangeDetectionTimeout(t *testing.T) {
  1192  	// Representative Detect*Changes function. QueryMeasurements under it
  1193  	// must honor ChangeDetectionTimeout so the whole call is bounded.
  1194  	origChange := db.ChangeDetectionTimeout
  1195  	t.Cleanup(func() { db.ChangeDetectionTimeout = origChange })
  1196  	db.ChangeDetectionTimeout = 100 * time.Millisecond
  1197  
  1198  	ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
  1199  	defer cancel()
  1200  	time.AfterFunc(5*time.Second, cancel)
  1201  
  1202  	metricDefs.MetricDefs["sproc_hashes"] = metrics.Metric{
  1203  		SQLs: map[int]string{120000: "SELECT"},
  1204  	}
  1205  
  1206  	md := &sources.DbConn{
  1207  		Source: sources.Source{Name: "bounded_detect_src"},
  1208  		Conn:   testutil.BlockingPool{},
  1209  		RuntimeInfo: sources.RuntimeInfo{
  1210  			Version:     120000,
  1211  			ChangeState: make(map[string]map[string]string),
  1212  		},
  1213  	}
  1214  	r := &reaper{
  1215  		measurementCh: make(chan metrics.MeasurementEnvelope, 10),
  1216  	}
  1217  
  1218  	start := time.Now()
  1219  	r.DetectSprocChanges(ctx, md)
  1220  	elapsed := time.Since(start)
  1221  
  1222  	if elapsed > 1500*time.Millisecond {
  1223  		t.Fatalf("DetectSprocChanges took %v, want ~ChangeDetectionTimeout (100ms)", elapsed)
  1224  	}
  1225  }
  1226  
  1227  // logSettingsCountingConn counts Query calls whose SQL contains "logging_collector"
  1228  // (the substring identifying the server_log_event_counts settings probe) so a test
  1229  // can assert the streaming log parser was started exactly once across multiple
  1230  // tick iterations.
  1231  type logSettingsCountingConn struct {
  1232  	pgxmock.PgxPoolIface
  1233  	count atomic.Int32
  1234  }
  1235  
  1236  func (c *logSettingsCountingConn) Query(ctx context.Context, query string, args ...any) (pgx.Rows, error) {
  1237  	if strings.Contains(query, "logging_collector") {
  1238  		c.count.Add(1)
  1239  	}
  1240  	return c.PgxPoolIface.Query(ctx, query, args...)
  1241  }
  1242  
  1243  // TestSourceReaper_RunStartsLogParserOnlyOnce proves the reaper spawns the
  1244  // server_log_event_counts streaming parser exactly once per worker lifetime.
  1245  // Pre-fix, the switch guard used sr.lastFetch[name] which was never assigned
  1246  // for this metric, so every tick (1s here) spawned a fresh goroutine that
  1247  // re-issued the log-settings SQL — the settings-query count would exceed 1.
  1248  // Post-fix, the dedicated logParserStarted flag ensures a single spawn.
  1249  func TestSourceReaper_RunStartsLogParserOnlyOnce(t *testing.T) {
  1250  	metricDefs.MetricDefs[specialMetricServerLogEventCounts] = metrics.Metric{
  1251  		SQLs: metrics.SQLs{0: "SELECT 1 AS value"},
  1252  	}
  1253  
  1254  	mock, err := pgxmock.NewPool()
  1255  	require.NoError(t, err)
  1256  	defer mock.Close()
  1257  	mock.MatchExpectationsInOrder(false)
  1258  
  1259  	// FetchRuntimeInfo attempt each tick; failing it keeps the loop ticking quietly.
  1260  	mock.ExpectQuery("select /\\* pgwatch_generated \\*/").WillReturnError(assert.AnError)
  1261  	// The log-settings query must be attempted exactly once; failing it makes the
  1262  	// spawned parser goroutine exit immediately, keeping the test fast.
  1263  	mock.ExpectQuery(expectedSettingsQuery).WillReturnError(assert.AnError)
  1264  
  1265  	cc := &logSettingsCountingConn{PgxPoolIface: mock}
  1266  	md := &sources.DbConn{
  1267  		Source: sources.Source{
  1268  			Name:    "log_parser_once_source",
  1269  			Kind:    sources.SourcePostgres,
  1270  			Metrics: metrics.MetricIntervals{specialMetricServerLogEventCounts: 1}, // 1s interval → 1s ticks
  1271  		},
  1272  		Conn: cc,
  1273  	}
  1274  	r := &reaper{
  1275  		Options:          &cmdopts.Options{Metrics: metrics.CmdOpts{}, Sinks: sinks.CmdOpts{}},
  1276  		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
  1277  		measurementCache: NewInstanceMetricCache(),
  1278  	}
  1279  	sr := NewDbConnReaper(r, md)
  1280  
  1281  	ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
  1282  	defer cancel()
  1283  	go sr.Reap(ctx)
  1284  
  1285  	require.Eventually(t, func() bool { return cc.count.Load() == 1 },
  1286  		2*time.Second, 10*time.Millisecond, "log parser should have been started")
  1287  	time.Sleep(2500 * time.Millisecond) // ≥2 more ticks at 1s tick interval
  1288  	assert.Equal(t, int32(1), cc.count.Load(), "log parser must be started exactly once")
  1289  }
  1290