1 files changed, 59 insertions, 26 deletions
diff --git a/pgruntime/connector.go b/pgruntime/connector.go
index 60e0de2..92780f6 100644
--- a/pgruntime/connector.go
+++ b/pgruntime/connector.go
@@ -5,7 +5,9 @@ package pgruntime
 
 import (
 	"context"
+	"crypto/sha256"
 	"database/sql"
+	"encoding/base32"
 	"fmt"
 	"regexp"
 	"strings"
@@ -19,13 +21,17 @@ import (
 // This type acts as a dependency injection surface, abstracting the different ways in which different database drivers and libraries perform connection.
 //
 // [libpq-style connection URI]: https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING-URIS
-type Connector[T gsql.ConnectionHandle] func(ctx context.Context, dbURL string) (T, error)
+type Connector[T gsql.ConnectionHandle] func(context.Context, ConnectionTarget) (T, error)
 
 // StdConnector returns a [Connector] for database/sql drivers.
 // When used with [lib/pq], the driver name must be "postgres".
 func StdConnector(driverName string) Connector[*gsql.DB] {
-	return func(ctx context.Context, dbURL string) (*gsql.DB, error) {
-		db, err := sql.Open(driverName, dbURL)
+	return func(ctx context.Context, target ConnectionTarget) (*gsql.DB, error) {
+		u, err := target.IntoURL()
+		if err != nil {
+			return nil, err
+		}
+		db, err := sql.Open(driverName, u.String())
 		if err != nil {
 			return nil, err
 		}
@@ -37,11 +43,7 @@ func StdConnector(driverName string) Connector[*gsql.DB] {
 func (c Connector[T]) Connect(ctx context.Context, target ConnectionTarget, behavior ConnectionBehavior) (T, error) {
 	var none T // shorthand for error return paths
 
-	u, err := target.IntoURL()
-	if err != nil {
-		return none, err
-	}
-	db, err := c(ctx, u.String())
+	db, err := c(ctx, target)
 	if err != nil {
 		return none, err
 	}
@@ -90,15 +92,16 @@ func (c Connector[T]) ConnectForTest(t assert.TestingTB, behavior ConnectionBeha
 
 	// normalize t.Name() into an acceptable database name for PostgreSQL
 	//   - only alphanumerics and underscore -> replace all other symbols with _
-	//   - max 63 chars -> reject longer names
+	//   - max 63 chars -> if overflown, truncate and append a short digest to hopefully make it unique
 	dbName := strings.ToLower(params.DatabaseName)
 	dbName = regexp.MustCompile(`[^a-z_]`).ReplaceAllString(dbName, "_")
 	if len(dbName) > 63 {
-		t.Fatalf("cannot use t.Name() = %q (normalized to %q) as a database name because it is longer than 63 chars", params.DatabaseName, dbName)
+		digest := sha256.Sum256([]byte(params.DatabaseName))
+		encoded := base32.HexEncoding.EncodeToString(digest[:])
+		dbName = dbName[0:53] + "__" + strings.ToLower(encoded[0:8])
 	}
 
-	// connect to "postgres" database for the DROP/CREATE DATABASE queries
-	// TODO: DROP/CREATE DATABASE turns out to be very slow (in a real-world scenario: 100ms per test instead of ~10ms to wipe just the DB contents and reset sequences)
+	// connect to "postgres" database for the CREATE DATABASE query (if necessary)
 	target := ConnectionTarget{
 		HostName:          "127.0.0.1",
 		Port:              testdbPort,
@@ -106,36 +109,66 @@ func (c Connector[T]) ConnectForTest(t assert.TestingTB, behavior ConnectionBeha
 		DatabaseName:      "postgres",
 		ConnectionOptions: "sslmode=disable",
 	}
-	err := c.prepareTestDatabase(ctx, target, dbName)
+	adminDB, err := c(ctx, target)
+	if err != nil {
+		t.Fatal(err.Error() + " (if this error is about the database server not running, check if your TestMain() calls pgruntime.WithTestDB())")
+	}
+	err = createDatabaseIfMissing(ctx, adminDB, dbName)
+	err = errext.WithCleanup(err, "db.Close", adminDB.GSQLClose(ctx))
 	if err != nil {
 		t.Fatal(err.Error())
 	}
 
 	// connect to actual test database
 	target.DatabaseName = dbName
-	handle, err := c.Connect(ctx, target, behavior)
+	testDB, err := c.Connect(ctx, target, behavior)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	err = resetTestDatabase(ctx, testDB, params, behavior)
 	if err != nil {
+		err = errext.WithCleanup(err, "db.Close", testDB.GSQLClose(ctx))
 		t.Fatal(err.Error())
 	}
-	return handle, target
+	return testDB, target
 }
 
-func (c Connector[T]) prepareTestDatabase(ctx context.Context, target ConnectionTarget, dbName string) error {
-	u, err := target.IntoURL()
+func createDatabaseIfMissing(ctx context.Context, db gsql.Handle, dbName string) error {
+	// check if database exists
+	exists, err := selectOneValue[bool](ctx, db, `SELECT COUNT(*) > 0 FROM pg_catalog.pg_database WHERE datname = $1`, dbName)
 	if err != nil {
-		return err
+		return fmt.Errorf("while reading from pg_catalog.pg_database: %w", err)
 	}
-	db, err := c(ctx, u.String())
-	if err != nil {
-		return fmt.Errorf("%w (if this error is about the database server not running, check if your TestMain() calls pgruntime.WithTestDB())", err)
+
+	// create database if necessary
+	if !exists {
+		_, err = execQuery(ctx, db, "CREATE DATABASE "+quoteIdentifier(dbName), nil)
+		if err != nil {
+			return fmt.Errorf("during CREATE DATABASE: %w", err)
+		}
+	}
+
+	return nil
+}
+
+func resetTestDatabase(ctx context.Context, db gsql.Handle, params testSetupParams, behavior ConnectionBehavior) error {
+	// enumerate all tables that need to be truncated (all tables that are not managed by pgruntime)
+	condition := `table_schema = 'public' AND table_type = 'BASE TABLE'`
+	if len(behavior.Migrations) > 0 {
+		condition += ` AND table_name != 'schema_migrations'`
 	}
-	_, err = execQuery(ctx, db, "DROP DATABASE IF EXISTS "+quoteIdentifier(dbName), nil)
+	query := fmt.Sprintf(`SELECT quote_ident(table_name) FROM information_schema.tables WHERE %s ORDER BY table_name`, condition)
+	quotedTableNames, err := selectSeveralValues[string](ctx, db, query)
 	if err != nil {
-		return errext.WithCleanup(fmt.Errorf("during DROP DATABASE: %w", err), "db.Close", db.GSQLClose(ctx))
+		return fmt.Errorf("while listing tables to truncate: %w", err)
 	}
-	_, err = execQuery(ctx, db, "CREATE DATABASE "+quoteIdentifier(dbName), nil)
+
+	// truncate all tables at once
+	query = fmt.Sprintf(`TRUNCATE %s RESTART IDENTITY CASCADE`, strings.Join(quotedTableNames, ", "))
+	_, err = execQuery(ctx, db, query, nil)
 	if err != nil {
-		return errext.WithCleanup(fmt.Errorf("during CREATE DATABASE: %w", err), "db.Close", db.GSQLClose(ctx))
+		return fmt.Errorf("during %s: %w", query, err)
 	}
-	return errext.WithCleanup(nil, "db.Close", db.GSQLClose(ctx))
+
+	return nil
 }