Merge branch 'ps-direct-postgres-connection' into 'master'

Praefect should be able to connect directly to the PostgreSQL

Closes #3119

See merge request gitlab-org/gitaly!2586

7 files changed, 71 insertions, 19 deletions

diff --git a/internal/praefect/config/config.go b/internal/praefect/config/config.go
index 9343f489e..a06acf914 100644
--- a/internal/praefect/config/config.go
+++ b/internal/praefect/config/config.go
@@ -258,17 +258,30 @@ type DB struct {
 	SSLCert     string `toml:"sslcert"`
 	SSLKey      string `toml:"sslkey"`
 	SSLRootCert string `toml:"sslrootcert"`
+	HostNoProxy string `toml:"host_no_proxy"`
+	PortNoProxy int    `toml:"port_no_proxy"`
 }
 
 // ToPQString returns a connection string that can be passed to github.com/lib/pq.
-func (db DB) ToPQString() string {
+func (db DB) ToPQString(direct bool) string {
+	var hostVal string
+	var portVal int
+
+	if direct {
+		hostVal = db.HostNoProxy
+		portVal = db.PortNoProxy
+	} else {
+		hostVal = db.Host
+		portVal = db.Port
+	}
+
 	var fields []string
-	if db.Port > 0 {
-		fields = append(fields, fmt.Sprintf("port=%d", db.Port))
+	if portVal > 0 {
+		fields = append(fields, fmt.Sprintf("port=%d", portVal))
 	}
 
 	for _, kv := range []struct{ key, value string }{
-		{"host", db.Host},
+		{"host", hostVal},
 		{"user", db.User},
 		{"password", db.Password},
 		{"dbname", db.DBName},
diff --git a/internal/praefect/config/config_test.go b/internal/praefect/config/config_test.go
index 9bd6b0090..be87c9f89 100644
--- a/internal/praefect/config/config_test.go
+++ b/internal/praefect/config/config_test.go
@@ -271,6 +271,8 @@ func TestConfigParsing(t *testing.T) {
 					SSLCert:     "/path/to/cert",
 					SSLKey:      "/path/to/key",
 					SSLRootCert: "/path/to/root-cert",
+					HostNoProxy: "2.3.4.5",
+					PortNoProxy: 6432,
 				},
 				MemoryQueueEnabled:  true,
 				GracefulStopTimeout: config.Duration(30 * time.Second),
@@ -358,13 +360,14 @@ func TestStorageNames(t *testing.T) {
 
 func TestToPQString(t *testing.T) {
 	testCases := []struct {
-		desc string
-		in   DB
-		out  string
+		desc   string
+		in     DB
+		direct bool
+		out    string
 	}{
 		{desc: "empty", in: DB{}, out: "binary_parameters=yes"},
 		{
-			desc: "basic example",
+			desc: "proxy connection",
 			in: DB{
 				Host:        "1.2.3.4",
 				Port:        2345,
@@ -376,7 +379,24 @@ func TestToPQString(t *testing.T) {
 				SSLKey:      "/path/to/key",
 				SSLRootCert: "/path/to/root-cert",
 			},
-			out: `port=2345 host=1.2.3.4 user=praefect-user password=secret dbname=praefect_production sslmode=require sslcert=/path/to/cert sslkey=/path/to/key sslrootcert=/path/to/root-cert binary_parameters=yes`,
+			direct: false,
+			out:    `port=2345 host=1.2.3.4 user=praefect-user password=secret dbname=praefect_production sslmode=require sslcert=/path/to/cert sslkey=/path/to/key sslrootcert=/path/to/root-cert binary_parameters=yes`,
+		},
+		{
+			desc: "direct connection",
+			in: DB{
+				HostNoProxy: "1.2.3.4",
+				PortNoProxy: 2345,
+				User:        "praefect-user",
+				Password:    "secret",
+				DBName:      "praefect_production",
+				SSLMode:     "require",
+				SSLCert:     "/path/to/cert",
+				SSLKey:      "/path/to/key",
+				SSLRootCert: "/path/to/root-cert",
+			},
+			direct: true,
+			out:    `port=2345 host=1.2.3.4 user=praefect-user password=secret dbname=praefect_production sslmode=require sslcert=/path/to/cert sslkey=/path/to/key sslrootcert=/path/to/root-cert binary_parameters=yes`,
 		},
 		{
 			desc: "with spaces and quotes",
@@ -389,7 +409,7 @@ func TestToPQString(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			require.Equal(t, tc.out, tc.in.ToPQString())
+			require.Equal(t, tc.out, tc.in.ToPQString(tc.direct))
 		})
 	}
 }
diff --git a/internal/praefect/config/testdata/config.toml b/internal/praefect/config/testdata/config.toml
index 430e7fdd0..4fe025293 100644
--- a/internal/praefect/config/testdata/config.toml
+++ b/internal/praefect/config/testdata/config.toml
@@ -53,6 +53,8 @@ sslmode = "require"
 sslcert = "/path/to/cert"
 sslkey = "/path/to/key"
 sslrootcert = "/path/to/root-cert"
+host_no_proxy = "2.3.4.5"
+port_no_proxy = 6432
 
 [failover]
 error_threshold_window = "20s"
diff --git a/internal/praefect/datastore/glsql/postgres.go b/internal/praefect/datastore/glsql/postgres.go
index 125de5956..4b474c179 100644
--- a/internal/praefect/datastore/glsql/postgres.go
+++ b/internal/praefect/datastore/glsql/postgres.go
@@ -15,7 +15,7 @@ import (
 
 // OpenDB returns connection pool to the database.
 func OpenDB(conf config.DB) (*sql.DB, error) {
-	db, err := sql.Open("postgres", conf.ToPQString())
+	db, err := sql.Open("postgres", conf.ToPQString(false))
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/praefect/datastore/glsql/testdata/pgbouncer.ini b/internal/praefect/datastore/glsql/testdata/pgbouncer.ini
index 59580510e..17395e9da 100644
--- a/internal/praefect/datastore/glsql/testdata/pgbouncer.ini
+++ b/internal/praefect/datastore/glsql/testdata/pgbouncer.ini
@@ -1,6 +1,6 @@
 [pgbouncer]
 listen_addr = *
-listen_port = 5432
+listen_port = 6432
 auth_type = trust
 pool_mode = transaction
 ignore_startup_parameters = extra_float_digits
diff --git a/internal/praefect/datastore/glsql/testing.go b/internal/praefect/datastore/glsql/testing.go
index 0fb396fe4..44a1e02f9 100644
--- a/internal/praefect/datastore/glsql/testing.go
+++ b/internal/praefect/datastore/glsql/testing.go
@@ -116,13 +116,30 @@ func GetDBConfig(t testing.TB, database string) config.DB {
 	require.NoError(t, pErr, "PGPORT must be a port number of the Postgres database listens for incoming connections")
 
 	// connect to 'postgres' database first to re-create testing database from scratch
-	return config.DB{
-		Host:    host,
-		Port:    portNumber,
-		DBName:  database,
-		SSLMode: "disable",
-		User:    os.Getenv("PGUSER"),
+	conf := config.DB{
+		Host:        host,
+		HostNoProxy: host,
+		Port:        portNumber,
+		PortNoProxy: portNumber,
+		DBName:      database,
+		SSLMode:     "disable",
+		User:        os.Getenv("PGUSER"),
 	}
+
+	bouncerHost, bouncerHostFound := os.LookupEnv("PGHOST_PGBOUNCER")
+	if bouncerHostFound {
+		conf.Host = bouncerHost
+	}
+
+	bouncerPort, bouncerPortFound := os.LookupEnv("PGPORT_PGBOUNCER")
+	if bouncerPortFound {
+		bouncerPortNumber, pErr := strconv.Atoi(bouncerPort)
+		require.NoError(t, pErr, "PGPORT_PGBOUNCER must be a port number of the PgBouncer")
+
+		conf.Port = bouncerPortNumber
+	}
+
+	return conf
 }
 
 func initGitalyTestDB(t testing.TB, database string) *sql.DB {
diff --git a/internal/praefect/datastore/queue_test.go b/internal/praefect/datastore/queue_test.go
index 17d605ab9..de0a53af6 100644
--- a/internal/praefect/datastore/queue_test.go
+++ b/internal/praefect/datastore/queue_test.go
@@ -756,7 +756,7 @@ func TestPostgresReplicationEventQueue_StartHealthUpdate(t *testing.T) {
 			if updatedJobLocks[i].JobID == dequeuedEventsUntriggered[0].ID {
 				require.Equal(t, initialJobLocks[i].TriggeredAt, updatedJobLocks[i].TriggeredAt, "no update expected as it was not submitted")
 			} else {
-				require.True(t, updatedJobLocks[i].TriggeredAt.After(initialJobLocks[i].TriggeredAt))
+				require.GreaterOrEqual(t, updatedJobLocks[i].TriggeredAt.UnixNano(), initialJobLocks[i].TriggeredAt.UnixNano())
 			}
 		}