From 657a7aa052ee28ce99d30800b0d4d7424ba9990a Mon Sep 17 00:00:00 2001 From: Tianzhou Date: Thu, 15 Oct 2026 22:50:35 +0800 Subject: [PATCH] feat: log the inspector's query call tree under --debug (#285, #644, #654) Every catalog query the inspector issues is now attributed to the BuildIR step that issued it. With --debug, the finished tree is logged when BuildIR returns: one line per step with its query count, summed query time and wall time, and each step with more than one query broken down by query shape (sqlc query name, or the collapsed leading SQL). Steps inside a concurrent group are listed heaviest first. On a schema with grants, the per-row role and constraint-column lookups behind #644 and #654 show up by name with their counts, so slow plans over high-latency connections can be diagnosed from the CLI without server-side statement logging - the original ask in #285. Co-Authored-By: Claude Fable 5.1 --- ir/inspector.go | 108 ++++++++++++++-------- ir/trace.go | 226 +++++++++++++++++++++++++++++++++++++++++++++++ ir/trace_test.go | 170 +++++++++++++++++++++++++++++++++++ 3 files changed, 467 insertions(+), 37 deletions(-) create mode 100644 ir/trace.go create mode 100644 ir/trace_test.go diff --git a/ir/inspector.go b/ir/inspector.go index da590775..c7e0a4bd 100644 --- a/ir/inspector.go +++ b/ir/inspector.go @@ -10,6 +10,7 @@ import ( "strings" "sync" + "github.com/pgplex/pgschema/internal/logger" "github.com/pgplex/pgschema/ir/queries" "golang.org/x/sync/errgroup" ) @@ -35,9 +36,13 @@ const ( // Inspector builds IR from database queries type Inspector struct { - db *sql.DB + // db is the connection wrapped in tracingDB, so every query is attributed + // to the BuildIR step that issued it (see trace.go). + db queries.DBTX queries *queries.Queries ignoreConfig *IgnoreConfig + // trace is the query call tree of the most recent BuildIR call. + trace *traceSpan // managedSchema is the logical schema pgschema is planning/applying, // which can differ from targetSchema (the schema this particular // BuildIR call actually introspects): when introspecting pgschema's own @@ -74,9 +79,10 @@ type Inspector struct { // NewInspector creates a new schema inspector with optional ignore configuration func NewInspector(db *sql.DB, ignoreConfig *IgnoreConfig) *Inspector { + traced := &tracingDB{inner: db} return &Inspector{ - db: db, - queries: queries.New(db), + db: traced, + queries: queries.New(traced), ignoreConfig: ignoreConfig, } } @@ -94,68 +100,81 @@ func (i *Inspector) BuildIR(ctx context.Context, targetSchema string) (*IR, erro i.managedSchema = targetSchema } + // Every query below is attributed to a named step of this call tree. With + // --debug the tree is logged when BuildIR returns, so the step responsible + // for most round trips over a slow connection is visible (see trace.go). + root := newTraceSpan("BuildIR schema=" + targetSchema) + i.trace = root + ctx = context.WithValue(ctx, traceSpanKey{}, root) + defer func() { + root.end() + if logger.IsDebug() { + root.log(logger.Get()) + } + }() + schema := NewIR() // Sequential prerequisites - if err := i.buildMetadata(ctx, schema); err != nil { + if err := traced(ctx, "metadata", func(ctx context.Context) error { return i.buildMetadata(ctx, schema) }); err != nil { return nil, fmt.Errorf("failed to build metadata: %w", err) } - if err := i.validateSchemaExists(ctx, targetSchema); err != nil { + if err := traced(ctx, "validate schema", func(ctx context.Context) error { return i.validateSchemaExists(ctx, targetSchema) }); err != nil { return nil, err } - if err := i.buildSchemas(ctx, schema, targetSchema); err != nil { + if err := traced(ctx, "schemas", func(ctx context.Context) error { return i.buildSchemas(ctx, schema, targetSchema) }); err != nil { return nil, fmt.Errorf("failed to build schemas: %w", err) } - if err := i.populateExtensionSchemas(ctx); err != nil { + if err := traced(ctx, "extension schemas", func(ctx context.Context) error { return i.populateExtensionSchemas(ctx) }); err != nil { return nil, fmt.Errorf("failed to query extension schemas: %w", err) } - if err := i.buildTables(ctx, schema, targetSchema); err != nil { + if err := traced(ctx, "tables", func(ctx context.Context) error { return i.buildTables(ctx, schema, targetSchema) }); err != nil { return nil, fmt.Errorf("failed to build tables: %w", err) } // Concurrent Group 1: Table Details group1 := queryGroup{ name: "table details", - funcs: []func(context.Context, *IR, string) error{ - i.buildColumns, - i.buildConstraints, - i.buildPartitions, + steps: []buildStep{ + {"columns", i.buildColumns}, + {"constraints", i.buildConstraints}, + {"partitions", i.buildPartitions}, }, } // Concurrent Group 2: Independent Objects group2 := queryGroup{ name: "independent objects", - funcs: []func(context.Context, *IR, string) error{ - i.buildSequences, - i.buildFunctions, - i.buildProcedures, - i.buildAggregates, - i.buildTypes, - i.buildDefaultPrivileges, - i.buildPrivileges, - i.buildColumnPrivileges, + steps: []buildStep{ + {"sequences", i.buildSequences}, + {"functions", i.buildFunctions}, + {"procedures", i.buildProcedures}, + {"aggregates", i.buildAggregates}, + {"types", i.buildTypes}, + {"default privileges", i.buildDefaultPrivileges}, + {"privileges", i.buildPrivileges}, + {"column privileges", i.buildColumnPrivileges}, }, } // Concurrent Group 3: View and table-dependent objects (views must load first) group3 := queryGroup{ name: "views", - funcs: []func(context.Context, *IR, string) error{ - i.buildViews, + steps: []buildStep{ + {"views", i.buildViews}, }, } // Group 4: Objects that depend on both tables AND views (must run after views are loaded) group4 := queryGroup{ name: "table-and-view-dependent objects", - funcs: []func(context.Context, *IR, string) error{ - i.buildTriggers, - i.buildRLSPolicies, + steps: []buildStep{ + {"triggers", i.buildTriggers}, + {"rls policies", i.buildRLSPolicies}, }, } @@ -175,12 +194,12 @@ func (i *Inspector) BuildIR(ctx context.Context, targetSchema string) (*IR, erro return nil, err } - if err := i.buildPartitionParentColumns(ctx, schema, targetSchema); err != nil { + if err := traced(ctx, "partition parent columns", func(ctx context.Context) error { return i.buildPartitionParentColumns(ctx, schema, targetSchema) }); err != nil { return nil, fmt.Errorf("failed to build partition parent columns: %w", err) } // Build function dependencies after functions are loaded - if err := i.buildFunctionDependencies(ctx, schema, targetSchema); err != nil { + if err := traced(ctx, "function dependencies", func(ctx context.Context) error { return i.buildFunctionDependencies(ctx, schema, targetSchema) }); err != nil { return nil, err } @@ -195,7 +214,7 @@ func (i *Inspector) BuildIR(ctx context.Context, targetSchema string) (*IR, erro } // Build indexes after views are loaded (indexes can reference materialized views) - if err := i.buildIndexes(ctx, schema, targetSchema); err != nil { + if err := traced(ctx, "indexes", func(ctx context.Context) error { return i.buildIndexes(ctx, schema, targetSchema) }); err != nil { return nil, fmt.Errorf("failed to build indexes: %w", err) } @@ -205,25 +224,40 @@ func (i *Inspector) BuildIR(ctx context.Context, targetSchema string) (*IR, erro return schema, nil } -// queryGroup represents a group of queries that can be executed concurrently +// buildStep is one named introspection step of BuildIR. The name labels the +// step's node in the query call tree. +type buildStep struct { + name string + fn func(context.Context, *IR, string) error +} + +// queryGroup represents a group of steps that can be executed concurrently type queryGroup struct { name string - funcs []func(context.Context, *IR, string) error + steps []buildStep } -// executeConcurrentGroup executes a group of functions concurrently +// executeConcurrentGroup executes a group's steps concurrently, each in its +// own trace span under a span for the group func (i *Inspector) executeConcurrentGroup(ctx context.Context, schema *IR, targetSchema string, group queryGroup) error { + ctx, groupSpan := startSpan(ctx, group.name+" (concurrent)") + groupSpan.concurrent = true + defer groupSpan.end() + var wg sync.WaitGroup - errChan := make(chan error, len(group.funcs)) + errChan := make(chan error, len(group.steps)) - for _, fn := range group.funcs { + for _, step := range group.steps { wg.Add(1) - go func(f func(context.Context, *IR, string) error) { + go func(step buildStep) { defer wg.Done() - if err := f(ctx, schema, targetSchema); err != nil { + err := traced(ctx, step.name, func(ctx context.Context) error { + return step.fn(ctx, schema, targetSchema) + }) + if err != nil { errChan <- err } - }(fn) + }(step) } wg.Wait() diff --git a/ir/trace.go b/ir/trace.go new file mode 100644 index 00000000..56508290 --- /dev/null +++ b/ir/trace.go @@ -0,0 +1,226 @@ +package ir + +import ( + "context" + "database/sql" + "log/slog" + "sort" + "strings" + "sync" + "time" + + "github.com/pgplex/pgschema/ir/queries" +) + +// Query tracing for the inspector. +// +// Every catalog query the inspector issues goes through tracingDB, which +// attributes it to the traceSpan carried by the calling context: one span per +// BuildIR step ("tables", "constraints", "privileges", ...), nested under its +// concurrent group and the BuildIR root. With --debug the finished tree is +// logged when BuildIR returns, one line per node with its query count and +// time, and each step with more than one query is broken down by query shape. +// Over a high-latency connection that shows directly which step - and which +// per-row lookup inside it - is paying for the round trips (issues #285, +// #644, #654), without needing server-side statement logging. + +type traceSpanKey struct{} + +// traceSpan is one node of the inspector's query call tree. +type traceSpan struct { + name string + start time.Time + wall time.Duration + // concurrent marks a span whose children ran in parallel. Their order of + // completion carries no meaning, so they are logged heaviest first. + concurrent bool + + mu sync.Mutex + children []*traceSpan + queries int // queries issued directly in this span + queryTime time.Duration // time spent in those queries + shapes map[string]*queryShape +} + +// queryShape aggregates the queries of one span that share a shape key. +type queryShape struct { + key string + count int + time time.Duration +} + +func newTraceSpan(name string) *traceSpan { + return &traceSpan{name: name, start: time.Now(), shapes: map[string]*queryShape{}} +} + +// startSpan opens a child span under the span carried by ctx, if any, and +// returns a context carrying the child. +func startSpan(ctx context.Context, name string) (context.Context, *traceSpan) { + child := newTraceSpan(name) + if parent := spanFromContext(ctx); parent != nil { + parent.mu.Lock() + parent.children = append(parent.children, child) + parent.mu.Unlock() + } + return context.WithValue(ctx, traceSpanKey{}, child), child +} + +func spanFromContext(ctx context.Context) *traceSpan { + s, _ := ctx.Value(traceSpanKey{}).(*traceSpan) + return s +} + +// traced runs fn inside a new child span named name. +func traced(ctx context.Context, name string, fn func(context.Context) error) error { + ctx, span := startSpan(ctx, name) + defer span.end() + return fn(ctx) +} + +func (s *traceSpan) end() { s.wall = time.Since(s.start) } + +func (s *traceSpan) record(sqlText string, d time.Duration) { + key := queryShapeKey(sqlText) + s.mu.Lock() + defer s.mu.Unlock() + s.queries++ + s.queryTime += d + shape := s.shapes[key] + if shape == nil { + shape = &queryShape{key: key} + s.shapes[key] = shape + } + shape.count++ + shape.time += d +} + +// totals returns the query count and query time of this span including all +// descendants. +func (s *traceSpan) totals() (int, time.Duration) { + s.mu.Lock() + n, d := s.queries, s.queryTime + children := append([]*traceSpan(nil), s.children...) + s.mu.Unlock() + for _, c := range children { + cn, cd := c.totals() + n += cn + d += cd + } + return n, d +} + +// log writes the span tree at debug level, one line per node, indented by +// depth. Every line starts with "inspector:" so the tree can be grepped out +// of the rest of the debug output. +func (s *traceSpan) log(log *slog.Logger) { + s.logNode(log, 0) +} + +func (s *traceSpan) logNode(log *slog.Logger, depth int) { + n, d := s.totals() + attrs := []any{"queries", n, "query_time", d.Round(time.Millisecond)} + if s.wall > 0 { + attrs = append(attrs, "wall", s.wall.Round(time.Millisecond)) + } + log.Debug("inspector: "+strings.Repeat(" ", depth)+s.name, attrs...) + + s.mu.Lock() + children := append([]*traceSpan(nil), s.children...) + shapes := make([]*queryShape, 0, len(s.shapes)) + for _, sh := range s.shapes { + shapes = append(shapes, sh) + } + own := s.queries + s.mu.Unlock() + + if s.concurrent { + sort.SliceStable(children, func(a, b int) bool { + na, _ := children[a].totals() + nb, _ := children[b].totals() + return na > nb + }) + } + for _, c := range children { + c.logNode(log, depth+1) + } + + // A step that issues a single query needs no breakdown. One that issues + // several is split by shape, so a per-row lookup appears with its count + // next to the batched query it accompanies. + if own > 1 { + sort.Slice(shapes, func(a, b int) bool { + if shapes[a].count != shapes[b].count { + return shapes[a].count > shapes[b].count + } + return shapes[a].key < shapes[b].key + }) + for _, sh := range shapes { + log.Debug("inspector: "+strings.Repeat(" ", depth+1)+sh.key, + "queries", sh.count, "query_time", sh.time.Round(time.Millisecond)) + } + } +} + +// queryShapeKey reduces a SQL text to a short label: the sqlc query name when +// the text starts with sqlc's "-- name:" header, otherwise the leading part of +// the statement with whitespace collapsed. +func queryShapeKey(sqlText string) string { + const header = "-- name: " + if strings.HasPrefix(sqlText, header) { + rest := sqlText[len(header):] + if i := strings.IndexAny(rest, " \n"); i > 0 { + return rest[:i] + } + } + collapsed := strings.Join(strings.Fields(sqlText), " ") + const maxLen = 70 + if len(collapsed) > maxLen { + return collapsed[:maxLen] + "..." + } + return collapsed +} + +// tracingDB wraps the inspector's connection and records every call against +// the span in the calling context. The recorded time covers the call itself - +// sending the statement and receiving the first response - not the later +// reading of result rows, so for the inspector's single-row and small-result +// queries it is close to one network round trip. +type tracingDB struct { + inner queries.DBTX +} + +var _ queries.DBTX = (*tracingDB)(nil) + +func (t *tracingDB) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { + start := time.Now() + res, err := t.inner.ExecContext(ctx, query, args...) + t.record(ctx, query, start) + return res, err +} + +func (t *tracingDB) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) { + start := time.Now() + stmt, err := t.inner.PrepareContext(ctx, query) + t.record(ctx, query, start) + return stmt, err +} + +func (t *tracingDB) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { + start := time.Now() + rows, err := t.inner.QueryContext(ctx, query, args...) + t.record(ctx, query, start) + return rows, err +} + +func (t *tracingDB) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row { + start := time.Now() + row := t.inner.QueryRowContext(ctx, query, args...) + t.record(ctx, query, start) + return row +} + +func (t *tracingDB) record(ctx context.Context, query string, start time.Time) { + if s := spanFromContext(ctx); s != nil { + s.record(query, time.Since(start)) + } +} diff --git a/ir/trace_test.go b/ir/trace_test.go new file mode 100644 index 00000000..a7d8c1f1 --- /dev/null +++ b/ir/trace_test.go @@ -0,0 +1,170 @@ +package ir + +import ( + "bytes" + "context" + "database/sql" + "log/slog" + "strings" + "testing" + "time" +) + +// fakeDBTX is a connection stand-in that answers every call immediately. +type fakeDBTX struct{} + +func (fakeDBTX) ExecContext(context.Context, string, ...any) (sql.Result, error) { + return nil, nil +} +func (fakeDBTX) PrepareContext(context.Context, string) (*sql.Stmt, error) { return nil, nil } +func (fakeDBTX) QueryContext(context.Context, string, ...any) (*sql.Rows, error) { + return nil, nil +} +func (fakeDBTX) QueryRowContext(context.Context, string, ...any) *sql.Row { return nil } + +func TestQueryShapeKey(t *testing.T) { + cases := map[string]string{ + "-- name: GetPrivilegesForSchema :many\nWITH acl_data AS (...)": "GetPrivilegesForSchema", + "SELECT rolname FROM pg_roles WHERE oid = $1": "SELECT rolname FROM pg_roles WHERE oid = $1", + "SELECT version()\n": "SELECT version()", + "SELECT " + strings.Repeat("x", 100): "SELECT " + strings.Repeat("x", 63) + "...", + } + for in, want := range cases { + if got := queryShapeKey(in); got != want { + t.Errorf("queryShapeKey(%q) = %q, want %q", in, got, want) + } + } +} + +func TestTracingDBAttributesQueriesToSpans(t *testing.T) { + db := &tracingDB{inner: fakeDBTX{}} + root := newTraceSpan("root") + ctx := context.WithValue(context.Background(), traceSpanKey{}, root) + + // A query outside any span is simply not recorded. + db.QueryRowContext(context.Background(), "SELECT 1") + + // One query directly in the root. + db.QueryRowContext(ctx, "SELECT version()") + + // A step with one batched query and three per-row lookups. + err := traced(ctx, "privileges", func(ctx context.Context) error { + _, _ = db.QueryContext(ctx, "-- name: GetPrivilegesForSchema :many\nSELECT 1") + for range 3 { + db.QueryRowContext(ctx, "SELECT rolname FROM pg_roles WHERE oid = $1", 42) + } + return nil + }) + if err != nil { + t.Fatal(err) + } + + // A concurrent group: a light step finishes first, a heavier one second. + groupCtx, group := startSpan(ctx, "group") + group.concurrent = true + err = traced(groupCtx, "partitions", func(ctx context.Context) error { + _, _ = db.QueryContext(ctx, "-- name: GetPartitionedTablesForSchema :many\nSELECT 1") + return nil + }) + if err != nil { + t.Fatal(err) + } + err = traced(groupCtx, "columns", func(ctx context.Context) error { + _, _ = db.QueryContext(ctx, "-- name: GetColumnsForSchema :many\nSELECT 1") + _, _ = db.ExecContext(ctx, "SET x = 1") + _, _ = db.PrepareContext(ctx, "SELECT 2") + return nil + }) + if err != nil { + t.Fatal(err) + } + group.end() + root.end() + + if root.queries != 1 { + t.Errorf("root own queries = %d, want 1", root.queries) + } + if n, _ := root.totals(); n != 9 { + t.Errorf("root total queries = %d, want 9", n) + } + if len(root.children) != 2 { + t.Fatalf("root children = %d, want 2", len(root.children)) + } + + priv := root.children[0] + if priv.name != "privileges" || priv.queries != 4 { + t.Errorf("privileges span = %q with %d queries, want 4", priv.name, priv.queries) + } + if sh := priv.shapes["SELECT rolname FROM pg_roles WHERE oid = $1"]; sh == nil || sh.count != 3 { + t.Errorf("per-row lookup shape not aggregated: %+v", sh) + } + if sh := priv.shapes["GetPrivilegesForSchema"]; sh == nil || sh.count != 1 { + t.Errorf("sqlc query shape not keyed by name: %+v", sh) + } + + if root.children[1] != group || group.queries != 0 || len(group.children) != 2 { + t.Errorf("group span not nested as expected: own=%d children=%d", group.queries, len(group.children)) + } + if n, _ := group.totals(); n != 4 { + t.Errorf("group total queries = %d, want 4", n) + } + if root.wall <= 0 { + t.Errorf("root wall time not recorded") + } + + // The logged tree: one line per node, indented, with the per-shape + // breakdown only under steps that issued more than one query. + var buf bytes.Buffer + logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{ + Level: slog.LevelDebug, + ReplaceAttr: func(_ []string, a slog.Attr) slog.Attr { + if a.Key == slog.TimeKey { + return slog.Attr{} + } + return a + }, + })) + root.log(logger) + out := buf.String() + for _, want := range []string{ + `msg="inspector: root" queries=9`, + `msg="inspector: privileges" queries=4`, + `msg="inspector: SELECT rolname FROM pg_roles WHERE oid = $1" queries=3`, + `msg="inspector: GetPrivilegesForSchema" queries=1`, + `msg="inspector: group" queries=4`, + `msg="inspector: columns" queries=3`, + `msg="inspector: partitions" queries=1`, + } { + if !strings.Contains(out, want) { + t.Errorf("log output missing %q:\n%s", want, out) + } + } + // Children of a concurrent group are logged heaviest first, regardless of + // which finished first. + if strings.Index(out, `"inspector: columns"`) > strings.Index(out, `"inspector: partitions"`) { + t.Errorf("concurrent group children not sorted by query count:\n%s", out) + } + // The root issued one query, so it gets no shape line of its own. + if strings.Contains(out, `"inspector: SELECT version()"`) { + t.Errorf("single-query span should not be broken down by shape:\n%s", out) + } +} + +func TestTraceSpanConcurrentRecording(t *testing.T) { + span := newTraceSpan("s") + done := make(chan struct{}) + for range 8 { + go func() { + for range 100 { + span.record("SELECT 1", time.Microsecond) + } + done <- struct{}{} + }() + } + for range 8 { + <-done + } + if span.queries != 800 || span.shapes["SELECT 1"].count != 800 { + t.Errorf("concurrent recording lost updates: queries=%d", span.queries) + } +}