Merge branch 'pks-testhelper-db-setup' into 'master'

glsql: Unify paths to retrieve database configuration

See merge request gitlab-org/gitaly!3910

3 files changed, 54 insertions, 61 deletions

diff --git a/internal/praefect/datastore/glsql/postgres_test.go b/internal/praefect/datastore/glsql/postgres_test.go
index 5616274bd..b91db2b58 100644
--- a/internal/praefect/datastore/glsql/postgres_test.go
+++ b/internal/praefect/datastore/glsql/postgres_test.go
@@ -1,28 +1,13 @@
 package glsql
 
 import (
-	"os"
-	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
 )
 
 func TestOpenDB(t *testing.T) {
-	getEnvFromGDK(t)
-
-	dbCfg := config.DB{
-		Host: os.Getenv("PGHOST"),
-		Port: func() int {
-			pgPort := os.Getenv("PGPORT")
-			port, err := strconv.Atoi(pgPort)
-			require.NoError(t, err, "failed to parse PGPORT %q", pgPort)
-			return port
-		}(),
-		DBName:  "postgres",
-		SSLMode: "disable",
-	}
+	dbCfg := GetDBConfig(t, "postgres")
 
 	t.Run("failed to ping because of incorrect config", func(t *testing.T) {
 		badCfg := dbCfg
diff --git a/internal/praefect/datastore/glsql/testing.go b/internal/praefect/datastore/glsql/testing.go
index bb9430b1f..3ca634bff 100644
--- a/internal/praefect/datastore/glsql/testing.go
+++ b/internal/praefect/datastore/glsql/testing.go
@@ -8,6 +8,7 @@ import (
 	"os/exec"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/google/uuid"
@@ -131,39 +132,34 @@ func NewDB(t testing.TB) DB {
 // GetDBConfig returns the database configuration determined by
 // environment variables.  See NewDB() for the list of variables.
 func GetDBConfig(t testing.TB, database string) config.DB {
-	getEnvFromGDK(t)
+	env := getDatabaseEnvironment(t)
 
-	host, hostFound := os.LookupEnv("PGHOST")
-	require.True(t, hostFound, "PGHOST env var expected to be provided to connect to Postgres database")
+	require.Contains(t, env, "PGHOST", "PGHOST env var expected to be provided to connect to Postgres database")
+	require.Contains(t, env, "PGPORT", "PGHOST env var expected to be provided to connect to Postgres database")
 
-	port, portFound := os.LookupEnv("PGPORT")
-	require.True(t, portFound, "PGPORT env var expected to be provided to connect to Postgres database")
-	portNumber, pErr := strconv.Atoi(port)
-	require.NoError(t, pErr, "PGPORT must be a port number of the Postgres database listens for incoming connections")
+	portNumber, err := strconv.Atoi(env["PGPORT"])
+	require.NoError(t, err, "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
 	conf := config.DB{
-		Host:    host,
+		Host:    env["PGHOST"],
 		Port:    portNumber,
 		DBName:  database,
 		SSLMode: "disable",
-		User:    os.Getenv("PGUSER"),
+		User:    env["PGUSER"],
 		SessionPooled: config.DBConnection{
-			Host: host,
+			Host: env["PGHOST"],
 			Port: portNumber,
 		},
 	}
 
-	bouncerHost, bouncerHostFound := os.LookupEnv("PGHOST_PGBOUNCER")
-	if bouncerHostFound {
+	if bouncerHost, ok := env["PGHOST_PGBOUNCER"]; ok {
 		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")
-
+	if bouncerPort, ok := env["PGPORT_PGBOUNCER"]; ok {
+		bouncerPortNumber, err := strconv.Atoi(bouncerPort)
+		require.NoError(t, err, "PGPORT_PGBOUNCER must be a port number of the PgBouncer")
 		conf.Port = bouncerPortNumber
 	}
 
@@ -287,25 +283,48 @@ func scanSingleBool(t testing.TB, db *sql.DB, query string, args ...interface{})
 	return flag
 }
 
-func getEnvFromGDK(t testing.TB) {
-	gdkEnv, err := exec.Command("gdk", "env").Output()
-	if err != nil {
-		// Assume we are not in a GDK setup; this is not an error so just return.
-		return
-	}
+var (
+	// Running `gdk env` takes about 250ms on my system and is thus comparatively slow. When
+	// running with Praefect as proxy, this time adds up and may thus slow down tests by quite a
+	// margin. We thus amortize these costs by only running it once.
+	databaseEnvOnce sync.Once
+	databaseEnv     map[string]string
+)
 
-	for _, line := range strings.Split(string(gdkEnv), "\n") {
-		const prefix = "export "
-		if !strings.HasPrefix(line, prefix) {
-			continue
+func getDatabaseEnvironment(t testing.TB) map[string]string {
+	databaseEnvOnce.Do(func() {
+		envvars := map[string]string{}
+
+		// We only process output if `gdk env` returned success. If it didn't, we simply assume that
+		// we are not running in a GDK environment and will try to extract variables from the
+		// environment instead.
+		if output, err := exec.Command("gdk", "env").Output(); err == nil {
+			for _, line := range strings.Split(string(output), "\n") {
+				const prefix = "export "
+				if !strings.HasPrefix(line, prefix) {
+					continue
+				}
+
+				split := strings.SplitN(strings.TrimPrefix(line, prefix), "=", 2)
+				if len(split) != 2 {
+					continue
+				}
+
+				envvars[split[0]] = split[1]
+			}
 		}
 
-		split := strings.SplitN(strings.TrimPrefix(line, prefix), "=", 2)
-		if len(split) != 2 {
-			continue
+		for _, key := range []string{"PGHOST", "PGPORT", "PGUSER", "PGHOST_PGBOUNCER", "PGPORT_PGBOUNCER"} {
+			if _, ok := envvars[key]; !ok {
+				value, ok := os.LookupEnv(key)
+				if ok {
+					envvars[key] = value
+				}
+			}
 		}
-		key, value := split[0], split[1]
 
-		require.NoError(t, os.Setenv(key, value), "set env var %v", key)
-	}
+		databaseEnv = envvars
+	})
+
+	return databaseEnv
 }
diff --git a/internal/testhelper/testserver/gitaly.go b/internal/testhelper/testserver/gitaly.go
index da8ddd92f..0c7ed2645 100644
--- a/internal/testhelper/testserver/gitaly.go
+++ b/internal/testhelper/testserver/gitaly.go
@@ -7,7 +7,6 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"strconv"
 	"testing"
 	"time"
 
@@ -73,10 +72,6 @@ func runPraefectProxy(t testing.TB, cfg config.Cfg, gitalyAddr, praefectBinPath
 
 	praefectServerSocketPath := "unix://" + testhelper.GetTemporaryGitalySocketFileName(t)
 
-	pgport := os.Getenv("PGPORT")
-	port, err := strconv.Atoi(pgport)
-	require.NoError(t, err)
-
 	dbName := createDatabase(t)
 
 	conf := praefectconfig.Config{
@@ -85,13 +80,7 @@ func runPraefectProxy(t testing.TB, cfg config.Cfg, gitalyAddr, praefectBinPath
 		Auth: auth.Config{
 			Token: cfg.Auth.Token,
 		},
-		DB: praefectconfig.DB{
-			Host:    os.Getenv("PGHOST"),
-			Port:    port,
-			User:    os.Getenv("PGUSER"),
-			DBName:  dbName,
-			SSLMode: "disable",
-		},
+		DB: glsql.GetDBConfig(t, dbName),
 		Failover: praefectconfig.Failover{
 			Enabled:           true,
 			ElectionStrategy:  praefectconfig.ElectionStrategyLocal,