Create user in envtool for PostgreSQL

Taskfile.yml

@@ -419,7 +419,7 @@ tasks:
         --proxy-addr=127.0.0.1:47017
         --mode=diff-normal
         --handler=pg
-        --postgresql-url='postgres://127.0.0.1:5433/ferretdb?search_path='
+        --postgresql-url='postgres://username:password@127.0.0.1:5433/ferretdb?search_path='
         --test-records-dir=tmp/records
 
   run-proxy:

cmd/envtool/envtool.go

@@ -33,6 +33,8 @@ import (
 	"time"
 
 	"github.com/alecthomas/kong"
+	"github.com/jackc/pgerrcode"
+	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
@@ -52,6 +54,12 @@ var (
 	//go:embed error.tmpl
 	errorTemplateB []byte
 
+	//go:embed test_user.json
+	testUser string
+
+	//go:embed test_system_users.json
+	testSystemUsers string
+
 	// Parsed error template.
 	errorTemplate = template.Must(template.New("error").Option("missingkey=error").Parse(string(errorTemplateB)))
 )
@@ -128,22 +136,50 @@ func setupAnyPostgres(ctx context.Context, logger *zap.SugaredLogger, uri string
 		ctxutil.SleepWithJitter(ctx, time.Second, retry)
 	}
 
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
-	return nil
+	return ctx.Err()
 }
 
 // setupPostgres configures `postgres` container.
+//
+// It also creates a user in ferretdb and template1 databases so all subsequent
+// databases created from template1 have the user.
+// See also https://www.postgresql.org/docs/current/manage-ag-templatedbs.html.
 func setupPostgres(ctx context.Context, logger *zap.SugaredLogger) error {
+	l := logger.Named("postgres")
+
 	// user `username` must exist, but password may be any, even empty
-	return setupAnyPostgres(ctx, logger.Named("postgres"), "postgres://username@127.0.0.1:5432/ferretdb")
+	err := setupAnyPostgres(ctx, l, "postgres://username@127.0.0.1:5432/ferretdb")
+	if err != nil {
+		return err
+	}
+
+	err = setupUserInPostgres(ctx, logger, "postgres://username@127.0.0.1:5432/ferretdb")
+	if err != nil {
+		return err
+	}
+
+	return setupUserInPostgres(ctx, logger, "postgres://username@127.0.0.1:5432/template1")
 }
 
 // setupPostgresSecured configures `postgres_secured` container.
+//
+// It also creates a user in ferretdb and template1 databases so all subsequent
+// databases created from template1 have the user.
+// See also https://www.postgresql.org/docs/current/manage-ag-templatedbs.html.
 func setupPostgresSecured(ctx context.Context, logger *zap.SugaredLogger) error {
-	return setupAnyPostgres(ctx, logger.Named("postgres_secured"), "postgres://username:password@127.0.0.1:5433/ferretdb")
+	l := logger.Named("postgres_secured")
+
+	err := setupAnyPostgres(ctx, l, "postgres://username:password@127.0.0.1:5433/ferretdb")
+	if err != nil {
+		return err
+	}
+
+	err = setupUserInPostgres(ctx, logger, "postgres://username:password@127.0.0.1:5433/ferretdb")
+	if err != nil {
+		return err
+	}
+
+	return setupUserInPostgres(ctx, logger, "postgres://username:password@127.0.0.1:5433/template1")
 }
 
 // setupMySQL configures `mysql` container.
@@ -219,6 +255,92 @@ func setupMongodb(ctx context.Context, logger *zap.SugaredLogger) error {
 	return ctx.Err()
 }
 
+// setupUserInPostgres creates a user with username/password credential in admin database
+// with supported mechanisms.
+// It creates admin database (PostgreSQL admin schema), if it does not exist.
+// It also creates system.users collection, if it does not exist.
+//
+// Without this, once the first user is created, the authentication fails
+// as username/password does not exist in admin.system.users collection.
+func setupUserInPostgres(ctx context.Context, logger *zap.SugaredLogger, uri string) error {
+	sp, err := state.NewProvider("")
+	if err != nil {
+		return err
+	}
+
+	u, err := url.Parse(uri)
+	if err != nil {
+		return err
+	}
+
+	p, err := pool.New(uri, logger.Desugar(), sp)
+	if err != nil {
+		return err
+	}
+
+	defer p.Close()
+
+	username := u.User.Username()
+	password, _ := u.User.Password()
+
+	dbPool, err := p.Get(username, password)
+	if err != nil {
+		return err
+	}
+
+	defer dbPool.Close()
+
+	var pgErr *pgconn.PgError
+
+	q := `CREATE SCHEMA admin`
+	if _, err = dbPool.Exec(ctx, q); err != nil && (!errors.As(err, &pgErr) || pgErr.Code != pgerrcode.DuplicateSchema) {
+		return err
+	}
+
+	if err == nil {
+		q = `CREATE TABLE admin._ferretdb_database_metadata (_jsonb jsonb)`
+		if _, err = dbPool.Exec(ctx, q); err != nil {
+			return err
+		}
+
+		q = `CREATE UNIQUE INDEX _ferretdb_database_metadata_id_idx ON admin._ferretdb_database_metadata (((_jsonb->'_id')))`
+		if _, err = dbPool.Exec(ctx, q); err != nil {
+			return err
+		}
+
+		q = `CREATE UNIQUE INDEX _ferretdb_database_metadata_table_idx ON admin._ferretdb_database_metadata (((_jsonb->'table')))`
+		if _, err = dbPool.Exec(ctx, q); err != nil {
+			return err
+		}
+	}
+
+	q = `CREATE TABLE admin.system_users_aff2f7ce (_jsonb jsonb)`
+	if _, err = dbPool.Exec(ctx, q); err != nil && (!errors.As(err, &pgErr) || pgErr.Code != pgerrcode.DuplicateTable) {
+		return err
+	}
+
+	if err == nil {
+		q = `CREATE UNIQUE INDEX "system_users_aff2f7ce__id__67399184_idx" ON admin.system_users_aff2f7ce (((_jsonb->'_id')))`
+		if _, err = dbPool.Exec(ctx, q); err != nil {
+			return err
+		}
+
+		q = `INSERT INTO admin._ferretdb_database_metadata (_jsonb) VALUES ($1)`
+		if _, err = dbPool.Exec(ctx, q, testSystemUsers); err != nil {
+			return err
+		}
+	}
+
+	q = `INSERT INTO admin.system_users_aff2f7ce (_jsonb) VALUES ($1)`
+
+	_, err = dbPool.Exec(ctx, q, testUser)
+	if err != nil && (!errors.As(err, &pgErr) || pgErr.Code != pgerrcode.UniqueViolation) {
+		return err
+	}
+
+	return nil
+}
+
 // setupMongodbSecured configures `mongodb_secured` container.
 func setupMongodbSecured(ctx context.Context, logger *zap.SugaredLogger) error {
 	if err := waitForPort(ctx, logger.Named("mongodb_secured"), 47018); err != nil {

cmd/envtool/envtool_test.go

@@ -16,9 +16,15 @@ package main
 
 import (
 	"bytes"
+	"context"
+	"fmt"
 	"os"
 	"testing"
 
+	zapadapter "github.com/jackc/pgx-zap"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/jackc/pgx/v5/tracelog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -92,3 +98,53 @@ func TestPackageVersion(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "1.0.0", output.String())
 }
+
+func TestSetupUserInPostgres(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping in -short mode")
+	}
+
+	t.Parallel()
+
+	baseURI := "postgres://username@127.0.0.1:5432/ferretdb?search_path="
+	cfg, err := pgxpool.ParseConfig(baseURI)
+	require.NoError(t, err)
+
+	l := testutil.Logger(t)
+	cfg.MinConns = 0
+	cfg.MaxConns = 1
+	cfg.ConnConfig.Tracer = &tracelog.TraceLog{
+		Logger:   zapadapter.NewLogger(l),
+		LogLevel: tracelog.LogLevelTrace,
+	}
+
+	ctx := testutil.Ctx(t)
+	p, err := pgxpool.NewWithConfig(ctx, cfg)
+	require.NoError(t, err)
+
+	dbName := testutil.DatabaseName(t)
+	sanitizedName := pgx.Identifier{dbName}.Sanitize()
+
+	_, err = p.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", sanitizedName))
+	require.NoError(t, err)
+
+	// use template0 because template1 may already have the user created
+	_, err = p.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s TEMPLATE template0", sanitizedName))
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		defer p.Close()
+
+		_, err = p.Exec(context.Background(), fmt.Sprintf("DROP DATABASE %s", sanitizedName))
+		require.NoError(t, err)
+	})
+
+	uri := fmt.Sprintf("postgres://username@127.0.0.1:5432/%s", dbName)
+
+	err = setupUserInPostgres(ctx, l.Sugar(), uri)
+	require.NoError(t, err)
+
+	// if the user already exists, it should not fail
+	err = setupUserInPostgres(ctx, l.Sugar(), uri)
+	require.NoError(t, err)
+}

cmd/envtool/test_system_users.json

@@ -0,0 +1,84 @@
+{
+  "$s": {
+    "p": {
+      "_id": {
+        "t": "string"
+      },
+      "uuid": {
+        "t": "string"
+      },
+      "table": {
+        "t": "string"
+      },
+      "indexes": {
+        "t": "array",
+        "i": [
+          {
+            "t": "object",
+            "$s": {
+              "p": {
+                "pgindex": {
+                  "t": "string"
+                },
+                "name": {
+                  "t": "string"
+                },
+                "key": {
+                  "t": "object",
+                  "$s": {
+                    "p": {
+                      "_id": {
+                        "t": "int"
+                      }
+                    },
+                    "$k": [
+                      "_id"
+                    ]
+                  }
+                },
+                "unique": {
+                  "t": "bool"
+                }
+              },
+              "$k": [
+                "pgindex",
+                "name",
+                "key",
+                "unique"
+              ]
+            }
+          }
+        ]
+      },
+      "cappedSize": {
+        "t": "long"
+      },
+      "cappedDocs": {
+        "t": "long"
+      }
+    },
+    "$k": [
+      "_id",
+      "uuid",
+      "table",
+      "indexes",
+      "cappedSize",
+      "cappedDocs"
+    ]
+  },
+  "_id": "system.users",
+  "uuid": "b34c2094-2086-4184-8761-f5692fadb2f2",
+  "table": "system_users_aff2f7ce",
+  "indexes": [
+    {
+      "pgindex": "system_users_aff2f7ce__id__67399184_idx",
+      "name": "_id_",
+      "key": {
+        "_id": 1
+      },
+      "unique": true
+    }
+  ],
+  "cappedSize": 0,
+  "cappedDocs": 0
+}