pgruntime: in ConnectForTest, only truncate tables instead of recreating the database

DROP + CREATE DATABASE was measured at 100 ms per ConnectForTest(),
which is prohibitively slow in large test suites. This should be more
in the ballpark of 10 ms per test.

1 files changed, 35 insertions, 0 deletions

diff --git a/pgruntime/helpers.go b/pgruntime/helpers.go
index 11bc84e..1019e6e 100644
--- a/pgruntime/helpers.go
+++ b/pgruntime/helpers.go
@@ -13,6 +13,7 @@ import (
 )
 
 // Convenience function for executing a one-off SQL query returning no rows.
+// TODO: move to gsql
 func execQuery(ctx context.Context, db gsql.Handle, query string, args []any) (sql.Result, error) {
 	stmt, err := db.GSQLPrepare(ctx, query, false)
 	if err != nil {
@@ -32,6 +33,40 @@ func queryRow(ctx context.Context, db gsql.Handle, query string, args, slots []a
 	return errext.WithCleanup(err, "stmt.Close", stmt.Close())
 }
 
+// Convenience function for executing a one-off SQL query returning one value.
+// TODO: move to gsql
+func selectOneValue[T any](ctx context.Context, db gsql.Handle, query string, args ...any) (T, error) {
+	stmt, err := db.GSQLPrepare(ctx, query, false)
+	if err != nil {
+		var none T
+		return none, err
+	}
+
+	var result T
+	err = stmt.QueryRow(ctx, args, []any{&result})
+	return result, errext.WithCleanup(err, "stmt.Close", stmt.Close())
+}
+
+// Convenience function for executing a one-off SQL query returning several single-column rows.
+// TODO: move to gsql (and also add ForeachValue with a callback instead of a slice return, maybe even ForeachPair and ForeachTriple)
+func selectSeveralValues[T any](ctx context.Context, db gsql.Handle, query string, args ...any) ([]T, error) {
+	rows, err := db.GSQLQuery(ctx, query, args)
+	if err != nil {
+		return nil, err
+	}
+	var result []T
+	for rows.Next() {
+		// TODO: this should share growRecordSlice() from Oblast to optimize allocations
+		var value T
+		err := rows.Scan(&value)
+		if err != nil {
+			return nil, errext.WithCleanup(err, "rows.Close", rows.Close())
+		}
+		result = append(result, value)
+	}
+	return result, errext.WithCleanup(nil, "rows.Err", rows.Err())
+}
+
 // Convenience function for preparing an identifier that needs to be inserted into a query verbatim
 // (e.g. a database name for CREATE DATABASE).
 func quoteIdentifier(name string) string {