▽

Source file src/github.com/cybertec-postgresql/pgwatch/v6/internal/sources/resolver_test.go

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

     1  package sources_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"strings"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/pashagolub/pgxmock/v5"
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  	client "go.etcd.io/etcd/client/v3"
    16  
    17  	"github.com/cybertec-postgresql/pgwatch/v6/internal/db"
    18  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
    19  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
    20  )
    21  
    22  func TestMonitoredDatabase_ResolveDatabasesFromPostgres(t *testing.T) {
    23  	pgContainer, pgTeardown, err := testutil.SetupPostgresContainer()
    24  	require.NoError(t, err)
    25  	defer pgTeardown()
    26  
    27  	// Create a new MonitoredDatabase instance
    28  	md := sources.Source{}
    29  	md.Name = "continuous"
    30  	md.Kind = sources.SourcePostgresDiscovery
    31  	md.ConnStr, err = pgContainer.ConnectionString(ctx, "sslmode=disable")
    32  	assert.NoError(t, err)
    33  
    34  	// Call the ResolveDatabasesFromPostgres method
    35  	dbs, err := md.ResolveDatabases()
    36  	assert.NoError(t, err)
    37  	assert.True(t, len(dbs) == 2) //postgres and mydatabase
    38  
    39  	// check the "continuous_mydatabase"
    40  	db := dbs.GetMonitoredDatabase(md.Name + "_mydatabase")
    41  	assert.NotNil(t, db)
    42  	dbConn, ok := db.(*sources.DbConn)
    43  	assert.True(t, ok)
    44  	assert.Equal(t, "mydatabase", dbConn.GetDatabaseName())
    45  
    46  	//check unexpected database
    47  	db = dbs.GetMonitoredDatabase(md.Name + "_unexpected")
    48  	assert.Nil(t, db)
    49  }
    50  
    51  func TestResolveDatabasesFromPostgres_ResolverTimeout(t *testing.T) {
    52  	// Shrink the resolver timeout so the test completes quickly.
    53  	orig := db.ResolverTimeout
    54  	db.ResolverTimeout = 150 * time.Millisecond
    55  	t.Cleanup(func() { db.ResolverTimeout = orig })
    56  
    57  	// BlackholeListener accepts TCP but never completes the pgwire handshake,
    58  	// causing NewConn (or the subsequent Query) to stall indefinitely when
    59  	// called with context.Background() / context.TODO().
    60  	addr, closeFn := testutil.BlackholeListener(t)
    61  	defer closeFn()
    62  
    63  	// connect_timeout is set longer than ResolverTimeout so that only the
    64  	// ResolverTimeout deadline (once wired) terminates the call.
    65  	connStr := fmt.Sprintf("postgres://postgres@%s/postgres?connect_timeout=10&sslmode=disable", addr)
    66  
    67  	md := sources.Source{}
    68  	md.Name = "stall_test"
    69  	md.Kind = sources.SourcePostgresDiscovery
    70  	md.ConnStr = connStr
    71  
    72  	start := time.Now()
    73  	_, err := sources.NewResolver().ResolveDatabasesFromPostgres(md)
    74  	elapsed := time.Since(start)
    75  
    76  	require.Error(t, err, "expected an error from a stalled resolver")
    77  
    78  	// Must return within the deadline (with generous test-scheduling slack).
    79  	assert.Less(t, elapsed, db.ResolverTimeout+2*time.Second,
    80  		"ResolveDatabasesFromPostgres took too long: %v", elapsed)
    81  
    82  	// The error must name the resolver operation (surfaced via WithOpTimeout cause).
    83  	assert.Contains(t, err.Error(), "resolve stall_test",
    84  		"error should name the resolver operation; got: %v", err)
    85  }
    86  
    87  func TestMonitoredDatabase_ResolveDatabasesFromPatroni(t *testing.T) {
    88  	etcdContainer, etcdTeardown, err := testutil.SetupEtcdContainer()
    89  	require.NoError(t, err)
    90  	defer etcdTeardown()
    91  
    92  	endpoint, err := etcdContainer.ClientEndpoint(ctx)
    93  	require.NoError(t, err)
    94  
    95  	cli, err := client.New(client.Config{
    96  		Endpoints:   []string{endpoint},
    97  		DialTimeout: 10 * time.Second,
    98  	})
    99  	require.NoError(t, err, "failed to create etcd client")
   100  	defer cli.Close()
   101  
   102  	// Start postgres server for testing
   103  	pgContainer, pgTeardown, err := testutil.SetupPostgresContainerWithInitScripts("../../docker/bootstrap/create_role_db.sql")
   104  	require.NoError(t, err)
   105  	defer pgTeardown()
   106  
   107  	pgConnStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
   108  	require.NoError(t, err)
   109  	// Put values to etcd server
   110  	kv := map[string]string{
   111  		`/service/demo/config`:           `{"ttl":30,"loop_wait":10,"retry_timeout":10,"maximum_lag_on_failover":1048576,"postgresql":{"use_pg_rewind":true,"pg_hba":["local all all trust","host replication replicator all md5","host all all all md5"],"parameters":{"max_connections":100}}}`,
   112  		`/service/demo/initialize`:       `7553211779477532695`,
   113  		`/service/demo/leader`:           `patroni3`,
   114  		`/service/demo/members/patroni1`: `{"conn_url":"postgres://172.18.0.8:5432/postgres","api_url":"http://172.18.0.8:8008/patroni","state":"running","role":"replica","version":"4.0.7","xlog_location":67108960,"replay_lsn":67108960,"receive_lsn":67108960,"replication_state":"streaming","timeline":1}`,
   115  		`/service/demo/members/patroni2`: `{"conn_url":"postgres://172.18.0.4:5432/postgres","api_url":"http://172.18.0.4:8008/patroni","state":"running","role":"replica","version":"4.0.7","xlog_location":67108960,"replay_lsn":67108960,"receive_lsn":67108960,"replication_state":"streaming","timeline":1}`,
   116  		`/service/demo/members/patroni3`: `{"conn_url":"` + pgConnStr + `","api_url":"http://172.18.0.3:8008/patroni","state":"running","role":"primary","version":"4.0.7","xlog_location":67108960,"timeline":1}`,
   117  		`/service/demo/status`:           `{"optime":67108960,"slots":{"patroni1":67108960,"patroni2":67108960,"patroni3":67108960},"retain_slots":["patroni1","patroni2","patroni3"]}}`}
   118  
   119  	cancelCtx, cancel := context.WithTimeout(context.Background(), time.Second*5)
   120  	for k, v := range kv {
   121  		_, err = cli.Put(cancelCtx, k, v)
   122  		require.NoError(t, err, "failed to put key %s to etcd", k)
   123  	}
   124  	cancel()
   125  
   126  	md := sources.Source{}
   127  	md.Name = "continuous"
   128  	md.OnlyIfMaster = true
   129  
   130  	t.Run("simple patroni discovery", func(t *testing.T) {
   131  		md.Kind = sources.SourcePatroniDiscovery
   132  		md.ConnStr = "etcd://" + strings.TrimPrefix(endpoint, "http://")
   133  		md.ConnStr += "/service"
   134  		md.ConnStr += "/demo"
   135  
   136  		// Run ResolveDatabasesFromPatroni
   137  		dbs, err := md.ResolveDatabases()
   138  		assert.NoError(t, err)
   139  		assert.NotNil(t, dbs)
   140  		assert.Len(t, dbs, 4) // postgres, mydatrabase, pgwatch, pgwatch_metrics}
   141  	})
   142  
   143  	t.Run("several endpoints patroni discovery", func(t *testing.T) {
   144  		md.Kind = sources.SourcePatroniDiscovery
   145  		e := strings.TrimPrefix(endpoint, "http://")
   146  		md.ConnStr = "etcd://" + strings.Join([]string{e, e, e}, ",")
   147  		md.ConnStr += "/service"
   148  		md.ConnStr += "/demo"
   149  
   150  		// Run ResolveDatabasesFromPatroni
   151  		dbs, err := md.ResolveDatabases()
   152  		assert.NoError(t, err)
   153  		assert.NotNil(t, dbs)
   154  		assert.Len(t, dbs, 4) // postgres, mydatrabase, pgwatch, pgwatch_metrics}
   155  	})
   156  
   157  	t.Run("namespace patroni discovery", func(t *testing.T) {
   158  		md.Kind = sources.SourcePatroniDiscovery
   159  		md.ConnStr = "etcd://" + strings.TrimPrefix(endpoint, "http://")
   160  
   161  		// Run ResolveDatabasesFromPatroni
   162  		dbs, err := md.ResolveDatabases()
   163  		assert.NoError(t, err)
   164  		assert.NotNil(t, dbs)
   165  		assert.Len(t, dbs, 4) // postgres, mydatrabase, pgwatch, pgwatch_metrics}
   166  	})
   167  }
   168  
   169  func TestMonitoredDatabase_UnsupportedDCS(t *testing.T) {
   170  	md := sources.Source{}
   171  	md.Name = "continuous"
   172  	md.Kind = sources.SourcePatroniDiscovery
   173  
   174  	md.ConnStr = "consul://foo"
   175  	_, err := md.ResolveDatabases()
   176  	assert.ErrorIs(t, err, errors.ErrUnsupported)
   177  
   178  	md.ConnStr = "zookeeper://foo"
   179  	_, err = md.ResolveDatabases()
   180  	assert.ErrorIs(t, err, errors.ErrUnsupported)
   181  
   182  	md.ConnStr = "unknown://foo"
   183  	_, err = md.ResolveDatabases()
   184  	assert.EqualError(t, err, "unsupported DCS type: unknown")
   185  
   186  }
   187  
   188  func TestNewHostConfig_BasicParsing(t *testing.T) {
   189  	tests := []struct {
   190  		name     string
   191  		uri      string
   192  		expected sources.HostConfig
   193  		wantErr  bool
   194  	}{
   195  		{
   196  			name: "simple etcd URI",
   197  			uri:  "etcd://localhost:2379/service/demo",
   198  			expected: sources.HostConfig{
   199  				DcsType:      "etcd",
   200  				DcsEndpoints: []string{"http://localhost:2379"},
   201  				Path:         "/service/demo",
   202  			},
   203  		},
   204  		{
   205  			name: "etcd with multiple hosts",
   206  			uri:  "etcd://host1:2379,host2:2379,host3:2379/service/demo",
   207  			expected: sources.HostConfig{
   208  				DcsType:      "etcd",
   209  				DcsEndpoints: []string{"http://host1:2379", "http://host2:2379", "http://host3:2379"},
   210  				Path:         "/service/demo",
   211  			},
   212  		},
   213  		{
   214  			name: "zookeeper URI",
   215  			uri:  "zookeeper://localhost:2181/patroni",
   216  			expected: sources.HostConfig{
   217  				DcsType:      "zookeeper",
   218  				DcsEndpoints: []string{"localhost:2181"},
   219  				Path:         "/patroni",
   220  			},
   221  		},
   222  		{
   223  			name: "consul URI",
   224  			uri:  "consul://localhost:8500/service",
   225  			expected: sources.HostConfig{
   226  				DcsType:      "consul",
   227  				DcsEndpoints: []string{"localhost:8500"},
   228  				Path:         "/service",
   229  			},
   230  		},
   231  		{
   232  			name:    "invalid URI - no scheme",
   233  			uri:     "localhost:2379/service",
   234  			wantErr: true,
   235  		},
   236  		{
   237  			name:    "unsupported scheme",
   238  			uri:     "redis://localhost:6379",
   239  			wantErr: true,
   240  		},
   241  	}
   242  
   243  	for _, tt := range tests {
   244  		t.Run(tt.name, func(t *testing.T) {
   245  			hc, err := sources.NewHostConfig(tt.uri)
   246  			if tt.wantErr {
   247  				assert.Error(t, err)
   248  				return
   249  			}
   250  			require.NoError(t, err)
   251  			assert.Equal(t, tt.expected.DcsType, hc.DcsType)
   252  			assert.Equal(t, tt.expected.DcsEndpoints, hc.DcsEndpoints)
   253  			assert.Equal(t, tt.expected.Path, hc.Path)
   254  		})
   255  	}
   256  }
   257  
   258  func TestNewHostConfig_WithUserInfo(t *testing.T) {
   259  	tests := []struct {
   260  		name     string
   261  		uri      string
   262  		username string
   263  		password string
   264  	}{
   265  		{
   266  			name:     "username only",
   267  			uri:      "etcd://admin@localhost:2379/service",
   268  			username: "admin",
   269  			password: "",
   270  		},
   271  		{
   272  			name:     "username and password",
   273  			uri:      "etcd://admin:secret@localhost:2379/service",
   274  			username: "admin",
   275  			password: "secret",
   276  		},
   277  		{
   278  			name:     "multiple hosts with auth",
   279  			uri:      "etcd://user:pass@host1:2379,host2:2379/service",
   280  			username: "user",
   281  			password: "pass",
   282  		},
   283  	}
   284  
   285  	for _, tt := range tests {
   286  		t.Run(tt.name, func(t *testing.T) {
   287  			hc, err := sources.NewHostConfig(tt.uri)
   288  			require.NoError(t, err)
   289  			assert.Equal(t, tt.username, hc.Username)
   290  			assert.Equal(t, tt.password, hc.Password)
   291  		})
   292  	}
   293  }
   294  
   295  func TestNewHostConfig_WithQueryParameters(t *testing.T) {
   296  	tests := []struct {
   297  		name     string
   298  		uri      string
   299  		caFile   string
   300  		certFile string
   301  		keyFile  string
   302  	}{
   303  		{
   304  			name:     "all TLS parameters",
   305  			uri:      "etcd://localhost:2379/service?ca_file=/path/to/ca.crt&cert_file=/path/to/cert.crt&key_file=/path/to/key.key",
   306  			caFile:   "/path/to/ca.crt",
   307  			certFile: "/path/to/cert.crt",
   308  			keyFile:  "/path/to/key.key",
   309  		},
   310  		{
   311  			name:   "only ca_file",
   312  			uri:    "etcd://localhost:2379/service?ca_file=/ca.crt",
   313  			caFile: "/ca.crt",
   314  		},
   315  		{
   316  			name:     "cert and key only",
   317  			uri:      "etcd://localhost:2379/service?cert_file=/cert.crt&key_file=/key.key",
   318  			certFile: "/cert.crt",
   319  			keyFile:  "/key.key",
   320  		},
   321  		{
   322  			name: "no TLS parameters",
   323  			uri:  "etcd://localhost:2379/service",
   324  		},
   325  		{
   326  			name:     "TLS params with multiple hosts",
   327  			uri:      "etcd://host1:2379,host2:2379/service?ca_file=/ca.crt&cert_file=/cert.crt",
   328  			caFile:   "/ca.crt",
   329  			certFile: "/cert.crt",
   330  		},
   331  	}
   332  
   333  	for _, tt := range tests {
   334  		t.Run(tt.name, func(t *testing.T) {
   335  			hc, err := sources.NewHostConfig(tt.uri)
   336  			require.NoError(t, err)
   337  			assert.Equal(t, tt.caFile, hc.CAFile)
   338  			assert.Equal(t, tt.certFile, hc.CertFile)
   339  			assert.Equal(t, tt.keyFile, hc.KeyFile)
   340  		})
   341  	}
   342  }
   343  
   344  func TestNewHostConfig_WithAuthAndTLS(t *testing.T) {
   345  	uri := "etcd://admin:secret@host1:2379,host2:2379/service/demo?ca_file=/ca.crt&cert_file=/cert.crt&key_file=/key.key"
   346  	hc, err := sources.NewHostConfig(uri)
   347  	require.NoError(t, err)
   348  
   349  	assert.Equal(t, "etcd", hc.DcsType)
   350  	assert.Equal(t, []string{"http://host1:2379", "http://host2:2379"}, hc.DcsEndpoints)
   351  	assert.Equal(t, "/service/demo", hc.Path)
   352  	assert.Equal(t, "admin", hc.Username)
   353  	assert.Equal(t, "secret", hc.Password)
   354  	assert.Equal(t, "/ca.crt", hc.CAFile)
   355  	assert.Equal(t, "/cert.crt", hc.CertFile)
   356  	assert.Equal(t, "/key.key", hc.KeyFile)
   357  }
   358  
   359  func TestNewHostConfig_PathVariations(t *testing.T) {
   360  	tests := []struct {
   361  		name  string
   362  		uri   string
   363  		path  string
   364  		scope bool
   365  	}{
   366  		{
   367  			name:  "namespace only",
   368  			uri:   "etcd://localhost:2379/service",
   369  			path:  "/service",
   370  			scope: false,
   371  		},
   372  		{
   373  			name:  "namespace and scope",
   374  			uri:   "etcd://localhost:2379/service/demo",
   375  			path:  "/service/demo",
   376  			scope: true,
   377  		},
   378  		{
   379  			name:  "deep path",
   380  			uri:   "etcd://localhost:2379/service/demo/v1",
   381  			path:  "/service/demo/v1",
   382  			scope: true,
   383  		},
   384  		{
   385  			name:  "no path",
   386  			uri:   "etcd://localhost:2379",
   387  			path:  "",
   388  			scope: false,
   389  		},
   390  	}
   391  
   392  	for _, tt := range tests {
   393  		t.Run(tt.name, func(t *testing.T) {
   394  			hc, err := sources.NewHostConfig(tt.uri)
   395  			require.NoError(t, err)
   396  			assert.Equal(t, tt.path, hc.Path)
   397  			assert.Equal(t, tt.scope, hc.IsScopeSpecified())
   398  		})
   399  	}
   400  }
   401  
   402  func TestNewHostConfig_EdgeCases(t *testing.T) {
   403  	tests := []struct {
   404  		name    string
   405  		uri     string
   406  		wantErr bool
   407  	}{
   408  		{
   409  			name:    "empty URI",
   410  			uri:     "",
   411  			wantErr: true,
   412  		},
   413  		{
   414  			name:    "URI without scheme separator",
   415  			uri:     "etcdlocalhost:2379",
   416  			wantErr: true,
   417  		},
   418  		{
   419  			name:    "URI with invalid host format",
   420  			uri:     "etcd://[::1:2379/service", // malformed IPv6
   421  			wantErr: true,
   422  		},
   423  	}
   424  
   425  	for _, tt := range tests {
   426  		t.Run(tt.name, func(t *testing.T) {
   427  			_, err := sources.NewHostConfig(tt.uri)
   428  			if tt.wantErr {
   429  				assert.Error(t, err)
   430  			} else {
   431  				assert.NoError(t, err)
   432  			}
   433  		})
   434  	}
   435  }
   436  
   437  // stubNewConnWithDatnames replaces sources.NewConn with a stub that builds a fresh
   438  // pgxmock pool per call so each invocation gets its own mocked discovery query.
   439  // The returned rows closure is invoked inside the stub to construct the row set
   440  // for the current call; this lets a single stub serve concurrent goroutines.
   441  func stubNewConnWithDatnames(t *testing.T, rowsFn func() []string) {
   442  	t.Helper()
   443  	orig := sources.NewConn
   444  	sources.NewConn = func(_ context.Context, _ string, _ ...db.ConnConfigCallback) (db.PgxPoolIface, error) {
   445  		mock, err := pgxmock.NewPool()
   446  		if err != nil {
   447  			return nil, err
   448  		}
   449  		rows := pgxmock.NewRows([]string{"datname"})
   450  		for _, n := range rowsFn() {
   451  			rows.AddRow(n)
   452  		}
   453  		mock.ExpectQuery("pg_database").
   454  			WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg()).
   455  			WillReturnRows(rows)
   456  		return mock, nil
   457  	}
   458  	t.Cleanup(func() { sources.NewConn = orig })
   459  }
   460  
   461  // stubNewConnWithError makes sources.NewConn return the given error verbatim.
   462  func stubNewConnWithError(t *testing.T, err error) {
   463  	t.Helper()
   464  	orig := sources.NewConn
   465  	sources.NewConn = func(_ context.Context, _ string, _ ...db.ConnConfigCallback) (db.PgxPoolIface, error) {
   466  		return nil, err
   467  	}
   468  	t.Cleanup(func() { sources.NewConn = orig })
   469  }
   470  
   471  func TestResolveDatabasesFromPostgres_LKGFallbackOnFailure(t *testing.T) {
   472  	resolver := sources.NewResolver()
   473  
   474  	md := sources.Source{
   475  		Name:    "lkg_failure",
   476  		Kind:    sources.SourcePostgresDiscovery,
   477  		ConnStr: "postgres://user:pw@a:5432/postgres?sslmode=disable",
   478  	}
   479  
   480  	// Wave 1: success populates the cache.
   481  	stubNewConnWithDatnames(t, func() []string { return []string{"db1", "db2", "db3"} })
   482  	dbs, err := resolver.ResolveDatabasesFromPostgres(md)
   483  	require.NoError(t, err)
   484  	require.Len(t, dbs, 3)
   485  	firstNames := []string{dbs[0].GetSource().Name, dbs[1].GetSource().Name, dbs[2].GetSource().Name}
   486  
   487  	// Wave 2: discovery fails; cache MUST be served with nil error.
   488  	sentinel := errors.New("boom")
   489  	stubNewConnWithError(t, sentinel)
   490  	dbs, err = resolver.ResolveDatabasesFromPostgres(md)
   491  	require.NoError(t, err, "expected cached fallback to swallow the error")
   492  	require.Len(t, dbs, 3, "expected the previously cached list to be returned")
   493  	gotNames := []string{dbs[0].GetSource().Name, dbs[1].GetSource().Name, dbs[2].GetSource().Name}
   494  	assert.Equal(t, firstNames, gotNames)
   495  }
   496  
   497  func TestResolveDatabasesFromPostgres_LKGReplacementOnSuccess(t *testing.T) {
   498  	resolver := sources.NewResolver()
   499  
   500  	md := sources.Source{
   501  		Name:    "lkg_replace",
   502  		Kind:    sources.SourcePostgresDiscovery,
   503  		ConnStr: "postgres://user:pw@b:5432/postgres?sslmode=disable",
   504  	}
   505  
   506  	// Seed cache with the first list.
   507  	stubNewConnWithDatnames(t, func() []string { return []string{"alpha", "beta"} })
   508  	dbs, err := resolver.ResolveDatabasesFromPostgres(md)
   509  	require.NoError(t, err)
   510  	require.Len(t, dbs, 2)
   511  
   512  	// Replace with a new, different list via a successful resolution.
   513  	stubNewConnWithDatnames(t, func() []string { return []string{"gamma", "delta", "epsilon"} })
   514  	dbs, err = resolver.ResolveDatabasesFromPostgres(md)
   515  	require.NoError(t, err)
   516  	require.Len(t, dbs, 3)
   517  	newNames := []string{dbs[0].GetSource().Name, dbs[1].GetSource().Name, dbs[2].GetSource().Name}
   518  	assert.Contains(t, newNames, "lkg_replace_gamma")
   519  	assert.Contains(t, newNames, "lkg_replace_delta")
   520  	assert.Contains(t, newNames, "lkg_replace_epsilon")
   521  
   522  	// Discovery fails now: the NEW list must be served, not the old one.
   523  	stubNewConnWithError(t, errors.New("still down"))
   524  	dbs, err = resolver.ResolveDatabasesFromPostgres(md)
   525  	require.NoError(t, err)
   526  	require.Len(t, dbs, 3)
   527  	for _, d := range dbs {
   528  		n := d.GetSource().Name
   529  		assert.NotContains(t, []string{"lkg_replace_alpha", "lkg_replace_beta"}, n,
   530  			"stale entry %q was served from cache", n)
   531  	}
   532  }
   533  
   534  func TestResolveDatabasesFromPostgres_EmptyCacheErrorPropagates(t *testing.T) {
   535  	resolver := sources.NewResolver()
   536  
   537  	sentinel := errors.New("no cache yet")
   538  	stubNewConnWithError(t, sentinel)
   539  
   540  	md := sources.Source{
   541  		Name:    "lkg_empty",
   542  		Kind:    sources.SourcePostgresDiscovery,
   543  		ConnStr: "postgres://user:pw@c:5432/postgres?sslmode=disable",
   544  	}
   545  
   546  	dbs, err := resolver.ResolveDatabasesFromPostgres(md)
   547  	require.Error(t, err)
   548  	assert.ErrorIs(t, err, sentinel)
   549  	assert.Empty(t, dbs, "no cached entry should have been served")
   550  }
   551  
   552  func TestResolveDatabasesFromPostgres_CacheKeyIdentity(t *testing.T) {
   553  	resolver := sources.NewResolver()
   554  
   555  	base := sources.Source{
   556  		Name:    "lkg_key",
   557  		Kind:    sources.SourcePostgresDiscovery,
   558  		ConnStr: "postgres://user:pw@d:5432/postgres?sslmode=disable",
   559  	}
   560  
   561  	stubNewConnWithDatnames(t, func() []string { return []string{"one", "two"} })
   562  	dbs, err := resolver.ResolveDatabasesFromPostgres(base)
   563  	require.NoError(t, err)
   564  	require.Len(t, dbs, 2)
   565  
   566  	// Same Name, DIFFERENT ConnStr: failing must not serve the cached entry.
   567  	sentinelA := errors.New("connstr-A down")
   568  	stubNewConnWithError(t, sentinelA)
   569  	other := sources.Source{
   570  		Name:    base.Name,
   571  		Kind:    sources.SourcePostgresDiscovery,
   572  		ConnStr: "postgres://user:pw@d-other:5432/postgres?sslmode=disable",
   573  	}
   574  	dbs, err = resolver.ResolveDatabasesFromPostgres(other)
   575  	require.Error(t, err)
   576  	assert.ErrorIs(t, err, sentinelA)
   577  	assert.Empty(t, dbs, "a reconfigured ConnStr must not serve the previous target's list")
   578  
   579  	// Same ConnStr, changed IncludePattern: failing must not serve the cached entry either.
   580  	sentinelB := errors.New("include-pattern changed")
   581  	stubNewConnWithError(t, sentinelB)
   582  	repattern := sources.Source{
   583  		Name:           base.Name,
   584  		Kind:           sources.SourcePostgresDiscovery,
   585  		ConnStr:        base.ConnStr,
   586  		IncludePattern: "^foo_",
   587  	}
   588  	dbs, err = resolver.ResolveDatabasesFromPostgres(repattern)
   589  	require.Error(t, err)
   590  	assert.ErrorIs(t, err, sentinelB)
   591  	assert.Empty(t, dbs, "a reconfigured include_pattern must not serve the previous result set")
   592  }
   593  
   594  func TestResolveDatabasesFromPostgres_ConcurrentFallback(t *testing.T) {
   595  	resolver := sources.NewResolver()
   596  
   597  	const n = 8
   598  	srcs := make(sources.Sources, n)
   599  	for i := range n {
   600  		srcs[i] = sources.Source{
   601  			Name:    fmt.Sprintf("concurrent_%d", i),
   602  			Kind:    sources.SourcePostgresDiscovery,
   603  			ConnStr: fmt.Sprintf("postgres://user:pw@h%d:5432/postgres?sslmode=disable", i),
   604  		}
   605  	}
   606  
   607  	// Wave 1: every source resolves successfully, each to a single distinct datname.
   608  	stubNewConnWithDatnames(t, func() []string {
   609  		// The stub doesn't know which source is calling; that's fine - the source
   610  		// name is taken from the Source struct, not from the mocked rows. We just
   611  		// need at least one row so resolution does not error.
   612  		return []string{"only"}
   613  	})
   614  	dbs, err := resolver.ResolveDatabases(srcs, func(string) {})
   615  	require.NoError(t, err)
   616  	require.GreaterOrEqual(t, len(dbs), n, "each source should produce one resolved DB")
   617  
   618  	// Wave 2: discovery fails for every source. Cached lists must be served.
   619  	stubNewConnWithError(t, errors.New("discovery down"))
   620  	// Add one brand-new source that has no cached entry.
   621  	extra := sources.Source{
   622  		Name:    "concurrent_extra",
   623  		Kind:    sources.SourcePostgresDiscovery,
   624  		ConnStr: "postgres://user:pw@extra:5432/postgres?sslmode=disable",
   625  	}
   626  	srcs2 := append(sources.Sources{}, srcs...)
   627  	srcs2 = append(srcs2, extra)
   628  
   629  	var onErrorNames sync.Map
   630  	dbs2, err := resolver.ResolveDatabases(srcs2, func(name string) {
   631  		onErrorNames.Store(name, struct{}{})
   632  	})
   633  	require.Error(t, err, "the never-cached extra source must propagate its error")
   634  	// Every original (cached) source still contributes its last-known DBs.
   635  	for i := range n {
   636  		require.NotNil(t, dbs2.GetMonitoredDatabase(fmt.Sprintf("concurrent_%d_only", i)),
   637  			"source %d should be served from cache", i)
   638  	}
   639  	// Extra source never cached: its name must appear in onError.
   640  	if _, ok := onErrorNames.Load(extra.Name); !ok {
   641  		t.Fatalf("expected onError to fire for never-cached source %q", extra.Name)
   642  	}
   643  }
   644