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) + } +}