pgtest: scan topologypgtest

3 files changed, 193 insertions, 1 deletions

diff --git a/pgtest/tracker.go b/pgtest/pgtest.go
index ebd53cd..4461e5e 100644
--- a/pgtest/tracker.go
+++ b/pgtest/pgtest.go
@@ -5,6 +5,11 @@
 package pgtest
 
 import (
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
 	"go.xyrillian.de/gg/assert"
 	"go.xyrillian.de/gg/gsql"
 )
@@ -14,6 +19,8 @@ type Tracker struct {
 	t   assert.TestingTB
 	dbh gsql.Handle
 	s   Snapshot
+
+	topo topology
 }
 
 // NewTracker creates a new Tracker.
@@ -22,7 +29,18 @@ type Tracker struct {
 // This is an optimization, since it is often desirable to assert on the full DB contents when creating the tracker.
 // Calling [Tracker.DBContent] directly after [NewTracker] would take a superfluous second snapshot.
 func NewTracker(t assert.TestingTB, db gsql.Handle) (*Tracker, Snapshot) {
-	panic("TODO")
+	ctx := t.Context()
+	t.Helper()
+
+	topo, err := newTopology(ctx, db)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	s, err := newSnapshot(ctx, db, topo)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	return &Tracker{t, db, s, topo}, s
 }
 
 // DBChanges produces a diff of the current database contents against the state at the last Tracker call,
@@ -36,3 +54,38 @@ func (t *Tracker) DBChanges() Snapshot {
 func (t *Tracker) DBContent() Snapshot {
 	panic("TODO")
 }
+
+// sqlLiteral implements [sql.Scanner] by storing a representation of the captured value as an SQL literal.
+// For time.Time, the UNIX timestamp is stored instead.
+type sqlLiteral string
+
+// Scan implements the [sql.Scanner] interface.
+func (l *sqlLiteral) Scan(src any) error {
+	switch src := src.(type) {
+	case int64:
+		*l = sqlLiteral(strconv.FormatInt(src, 10))
+	case float64:
+		*l = sqlLiteral(fmt.Sprintf("%g", src))
+	case bool:
+		if src {
+			*l = "TRUE"
+		} else {
+			*l = "FALSE"
+		}
+	case []byte:
+		*l = makeSQLStringLiteral(string(src))
+	case string:
+		*l = makeSQLStringLiteral(src)
+	case time.Time:
+		*l = sqlLiteral(strconv.FormatInt(src.Unix(), 10))
+	case nil:
+		*l = "NULL"
+	default:
+		return fmt.Errorf("sqlLiteral.Scan(): do not know how to serialize type %T", src)
+	}
+	return nil
+}
+
+func makeSQLStringLiteral(in string) sqlLiteral {
+	return sqlLiteral("'" + strings.ReplaceAll(in, "'", "''") + "'")
+}
diff --git a/pgtest/snapshot.go b/pgtest/snapshot.go
index f27ecec..ed0752c 100644
--- a/pgtest/snapshot.go
+++ b/pgtest/snapshot.go
@@ -4,9 +4,11 @@
 package pgtest
 
 import (
+	"context"
 	"fmt"
 
 	"go.xyrillian.de/gg/assert"
+	"go.xyrillian.de/gg/gsql"
 )
 
 // Snapshot contains a set of SQL statements.
@@ -15,6 +17,10 @@ type Snapshot struct {
 	t assert.TestingTB
 }
 
+func newSnapshot(ctx context.Context, db gsql.Handle, topo topology) (Snapshot, error) {
+	panic("TODO")
+}
+
 // AssertEmpty is a shorthand for AssertEqual("").
 func (s Snapshot) AssertEmpty() {
 	s.t.Helper()
diff --git a/pgtest/topology.go b/pgtest/topology.go
new file mode 100644
index 0000000..05e462c
--- /dev/null
+++ b/pgtest/topology.go
@@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: 2026 Stefan Majewsky <majewsky@gmx.net>
+// SPDX-License-Identifier: Apache-2.0
+
+package pgtest
+
+import (
+	"context"
+	"fmt"
+	"slices"
+
+	"go.xyrillian.de/gg/errext"
+	"go.xyrillian.de/gg/gsql"
+)
+
+type topology struct {
+	TableInfoByName map[string]tableInfo
+}
+
+type tableInfo struct {
+	Columns []columnInfo
+}
+
+type columnInfo struct {
+	Name         string
+	DefaultValue sqlLiteral
+	IsPrimaryKey bool
+}
+
+const (
+	topologyGetColumnsQuery = `
+		SELECT table_name, column_name, column_default
+		  FROM information_schema.columns
+		 WHERE table_schema = 'public'
+		 ORDER BY table_name, ordinal_position
+	`
+	topologyGetObviousPrimaryKeysQuery = `
+		SELECT DISTINCT table_name, column_name
+		  FROM information_schema.key_column_usage
+		 WHERE table_schema = 'public' AND position_in_unique_constraint IS NULL AND constraint_name = table_name || '_pkey'
+		 ORDER BY 1, 2
+	`
+	topologyGetPossiblePrimaryKeysQuery = `
+		SELECT DISTINCT table_name, column_name
+		  FROM information_schema.key_column_usage
+		 WHERE table_schema = 'public' AND position_in_unique_constraint IS NULL
+		 ORDER BY 1, 2
+	`
+)
+
+func newTopology(ctx context.Context, db gsql.Handle) (topology, error) {
+	result := topology{
+		TableInfoByName: make(map[string]tableInfo),
+	}
+
+	// enumerate tables and their columns
+	columnInfosByTableName, err := topologyGetColumnInfo(ctx, db)
+	if err != nil {
+		return topology{}, fmt.Errorf("while querying information_schema.columns: %w", err)
+	}
+	for tableName, columnInfos := range columnInfosByTableName {
+		result.TableInfoByName[tableName] = tableInfo{Columns: columnInfos}
+	}
+
+	// find obvious primary keys (columns that are included in UNIQUE constraints with the name `${TABLE}_pkey`)
+	obviousPKColumnsByTableName, err := topologyGetColumnsMatching(ctx, db, topologyGetObviousPrimaryKeysQuery)
+	if err != nil {
+		return topology{}, fmt.Errorf("while querying information_schema.key_column_usage for obvious primary keys: %w", err)
+	}
+	for tableName, columnNames := range obviousPKColumnsByTableName {
+		for idx, col := range result.TableInfoByName[tableName].Columns {
+			if slices.Contains(columnNames, col.Name) {
+				col.IsPrimaryKey = true
+				result.TableInfoByName[tableName].Columns[idx] = col
+			}
+		}
+	}
+
+	// as a fallback, find possible primary keys (columns that are included in any UNIQUE constraint)
+	possiblePKColumnsByTableName, err := topologyGetColumnsMatching(ctx, db, topologyGetPossiblePrimaryKeysQuery)
+	if err != nil {
+		return topology{}, fmt.Errorf("while querying information_schema.key_column_usage for possible primary keys: %w", err)
+	}
+	for tableName, columnNames := range possiblePKColumnsByTableName {
+		if len(obviousPKColumnsByTableName[tableName]) > 0 {
+			continue
+		}
+		for idx, col := range result.TableInfoByName[tableName].Columns {
+			if slices.Contains(columnNames, col.Name) {
+				col.IsPrimaryKey = true
+				result.TableInfoByName[tableName].Columns[idx] = col
+			}
+		}
+	}
+
+	return result, nil
+}
+
+func topologyGetColumnInfo(ctx context.Context, db gsql.Handle) (map[string][]columnInfo, error) {
+	rows, err := db.GSQLQuery(ctx, topologyGetColumnsQuery, nil)
+	if err != nil {
+		return nil, err
+	}
+	result := make(map[string][]columnInfo)
+	for rows.Next() {
+		var (
+			tableName string
+			col       columnInfo
+		)
+		err := rows.Scan(&tableName, &col.Name, &col.DefaultValue)
+		if err != nil {
+			return nil, errext.WithCleanup(err, "rows.Close", rows.Close())
+		}
+		result[tableName] = append(result[tableName], col)
+	}
+	return result, errext.WithCleanup(nil, "rows.Err", rows.Err())
+}
+
+func topologyGetColumnsMatching(ctx context.Context, db gsql.Handle, query string) (map[string][]string, error) {
+	rows, err := db.GSQLQuery(ctx, query, nil)
+	if err != nil {
+		return nil, err
+	}
+	result := make(map[string][]string)
+	for rows.Next() {
+		var tableName, columnName string
+		err := rows.Scan(&tableName, &columnName)
+		if err != nil {
+			return nil, errext.WithCleanup(err, "rows.Close", rows.Close())
+		}
+		result[tableName] = append(result[tableName], columnName)
+	}
+	return result, errext.WithCleanup(nil, "rows.Err", rows.Err())
+}