From 1cd9419020859ad6ca7cd8f47fb2b4a100bccce9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 22:33:05 +0000 Subject: [PATCH] core: speed up the analysis catalog with prepared statements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The core catalog backs every lookup with a query against its in-memory SQLite database, and handed the raw *sql.DB to the generated queries. Go's database/sql prepares and discards a statement per call, so SQLite re-parsed and re-planned the same dozen queries thousands of times per run: profiling analyzer.Prepare put 65% of samples inside sqlite3Parser and 85% under database/sql dispatch, with the analyzer itself barely on the profile. - Prepare each catalog statement once and rebind it, for both the seed path and the analyzer's lookups. - Install dialect seeds in a single transaction, and drop journalling for a database that is rebuilt from scratch on every run. - Pin the pool to one connection. Each ":memory:" connection is its own empty database, so a second one would have seen a catalog with no tables in it — a latent bug under any concurrent use. On the analyzer side, resolving a column no longer allocates a candidate slice per reference, the scope holds the catalog's column list instead of copying it into a scope-local type, a target's field list is flattened once instead of three times, and parameters are tracked by value. Analyzing 500 GoogleSQL queries drops from 181ms to 82ms end to end, and 100 queries from 58ms to 34ms, with byte-identical output on GoogleSQL and ClickHouse. Microbenchmarks (internal/core/analyzer/bench_test.go): CatalogNew 6001532 ns/op (catalog + dialect seed) CatalogNew 2863319 ns/op Prepare 1579238 ns/op (5 statements, warm catalog) Prepare 398200 ns/op Lookup volume is unchanged at 13 catalog queries per statement, and analysis still allocates ~400 times per statement — nearly all of it database/sql scanning a row per lookup. Memoizing lookups removes both, but needs cache invalidation and a read-only contract on the slices the catalog returns; that is not worth it for an in-memory database, so it is left out. Behaviour is covered by the existing end-to-end tests: analyze_basic, analyze_select and analyze_dml exercise both dialects through TestReplay. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Fp7UGoiFnUwK9ZkDBoGy7u --- internal/core/analyzer/analyzer.go | 9 +- internal/core/analyzer/bench_test.go | 125 ++++++++++++++++++++ internal/core/analyzer/dml.go | 25 ++-- internal/core/analyzer/expr.go | 12 +- internal/core/analyzer/projection.go | 48 ++++---- internal/core/analyzer/scope.go | 68 ++++------- internal/core/attribute.go | 1 + internal/core/catalog.go | 68 +++++++++-- internal/core/db.go | 171 +++++++++++++++++++++++++++ internal/core/operator.go | 2 + internal/core/proc.go | 4 + 11 files changed, 434 insertions(+), 99 deletions(-) create mode 100644 internal/core/analyzer/bench_test.go create mode 100644 internal/core/db.go diff --git a/internal/core/analyzer/analyzer.go b/internal/core/analyzer/analyzer.go index 533557e417..9b03b445cd 100644 --- a/internal/core/analyzer/analyzer.go +++ b/internal/core/analyzer/analyzer.go @@ -13,7 +13,7 @@ func Prepare(cat *core.Catalog, stmt ast.Node) (core.PrepareResult, error) { } a := &analyzer{ cat: cat, - params: map[int]*core.Parameter{}, + params: map[int]core.Parameter{}, } switch s := stmt.(type) { case *ast.SelectStmt: @@ -46,7 +46,7 @@ type analyzer struct { cat *core.Catalog scope *scope columns []core.Column - params map[int]*core.Parameter + params map[int]core.Parameter command core.Command } @@ -58,7 +58,7 @@ func (a *analyzer) result() core.PrepareResult { } } -func orderedParams(m map[int]*core.Parameter) []core.Parameter { +func orderedParams(m map[int]core.Parameter) []core.Parameter { if len(m) == 0 { return nil } @@ -71,7 +71,7 @@ func orderedParams(m map[int]*core.Parameter) []core.Parameter { out := make([]core.Parameter, 0, len(m)) for i := 1; i <= maxN; i++ { if p, ok := m[i]; ok { - out = append(out, *p) + out = append(out, p) } } return out @@ -112,6 +112,7 @@ func (a *analyzer) analyzeSelect(s *ast.SelectStmt) error { if targets == nil { return fmt.Errorf("select: empty target list") } + a.columns = make([]core.Column, 0, len(targets)) for _, t := range targets { rt, ok := t.(*ast.ResTarget) if !ok { diff --git a/internal/core/analyzer/bench_test.go b/internal/core/analyzer/bench_test.go new file mode 100644 index 0000000000..fa4a67dad4 --- /dev/null +++ b/internal/core/analyzer/bench_test.go @@ -0,0 +1,125 @@ +package analyzer_test + +import ( + "strings" + "testing" + + "github.com/sqlc-dev/sqlc/internal/config" + "github.com/sqlc-dev/sqlc/internal/core" + "github.com/sqlc-dev/sqlc/internal/core/analyzer" + coreschema "github.com/sqlc-dev/sqlc/internal/core/schema" + "github.com/sqlc-dev/sqlc/internal/engine/googlesql" + "github.com/sqlc-dev/sqlc/internal/sql/ast" + "github.com/sqlc-dev/sqlc/internal/sql/rewrite" + "github.com/sqlc-dev/sqlc/internal/sql/validate" +) + +const benchSchema = ` +CREATE TABLE users ( + id INT64 NOT NULL, + name STRING NOT NULL, + bio STRING, +) PRIMARY KEY (id); + +CREATE TABLE posts ( + id INT64 NOT NULL, + user_id INT64 NOT NULL, + title STRING(255), + created TIMESTAMP NOT NULL, +) PRIMARY KEY (id); +` + +const benchQueries = ` +SELECT * FROM users; +SELECT COUNT(*) AS total FROM users; +SELECT u.name, p.* FROM users u JOIN posts p ON p.user_id = u.id WHERE u.name = @name; +SELECT id, name, bio FROM users WHERE id = @id AND name = @name; +SELECT p.title, p.created FROM posts p WHERE p.user_id = @uid AND p.title = @t; +` + +func parseAll(t testing.TB, src string) []ast.Node { + p := googlesql.NewParser() + stmts, err := p.Parse(strings.NewReader(src)) + if err != nil { + t.Fatal(err) + } + out := make([]ast.Node, 0, len(stmts)) + for i := range stmts { + out = append(out, stmts[i].Raw) + } + return out +} + +// parseQueries mirrors the compiler pipeline: parse, then rewrite named +// parameters into ParamRef nodes before handing the statement to the analyzer. +func parseQueries(t testing.TB, src string) []ast.Node { + out := make([]ast.Node, 0) + for _, n := range parseAll(t, src) { + raw, ok := n.(*ast.RawStmt) + if !ok { + t.Fatalf("not a raw statement: %T", n) + } + numbers, dollar, err := validate.ParamRef(raw) + if err != nil { + t.Fatal(err) + } + rewritten, _, _ := rewrite.NamedParameters(config.EngineGoogleSQL, raw, numbers, dollar) + out = append(out, rewritten) + } + return out +} + +func newBenchCatalog(t testing.TB) (*core.Catalog, []ast.Node) { + cat, err := core.New(googlesql.Dialect()) + if err != nil { + t.Fatal(err) + } + for _, n := range parseAll(t, benchSchema) { + if err := coreschema.Apply(cat, n); err != nil { + t.Fatal(err) + } + } + return cat, parseQueries(t, benchQueries) +} + +// BenchmarkCatalogNew measures per-compile catalog setup: open the SQLite +// catalog, install the schema, and run the dialect seed. +func BenchmarkCatalogNew(b *testing.B) { + for b.Loop() { + cat, err := core.New(googlesql.Dialect()) + if err != nil { + b.Fatal(err) + } + cat.Close() + } +} + +// BenchmarkApplySchema measures loading user DDL into the catalog. +func BenchmarkApplySchema(b *testing.B) { + stmts := parseAll(b, benchSchema) + for b.Loop() { + cat, err := core.New(googlesql.Dialect()) + if err != nil { + b.Fatal(err) + } + for _, n := range stmts { + if err := coreschema.Apply(cat, n); err != nil { + b.Fatal(err) + } + } + cat.Close() + } +} + +// BenchmarkPrepare measures steady-state query analysis against a warm catalog. +func BenchmarkPrepare(b *testing.B) { + cat, queries := newBenchCatalog(b) + defer cat.Close() + for b.Loop() { + for _, q := range queries { + if _, err := analyzer.Prepare(cat, q); err != nil { + b.Fatal(err) + } + } + } +} diff --git a/internal/core/analyzer/dml.go b/internal/core/analyzer/dml.go index 7ea4285e7e..2f92afec8a 100644 --- a/internal/core/analyzer/dml.go +++ b/internal/core/analyzer/dml.go @@ -3,6 +3,7 @@ package analyzer import ( "fmt" + "github.com/sqlc-dev/sqlc/internal/core" "github.com/sqlc-dev/sqlc/internal/sql/ast" ) @@ -89,12 +90,12 @@ func (a *analyzer) relationScope(relations, extra *ast.List) (*scope, error) { return sc, nil } -func insertTargets(rel scopeRel, cols *ast.List) ([]scopeCol, error) { +func insertTargets(rel scopeRel, cols *ast.List) ([]core.ClassColumn, error) { items := listItems(cols) if len(items) == 0 { return rel.cols, nil } - out := make([]scopeCol, 0, len(items)) + out := make([]core.ClassColumn, 0, len(items)) for _, item := range items { rt, ok := item.(*ast.ResTarget) if !ok || rt.Name == nil { @@ -109,7 +110,7 @@ func insertTargets(rel scopeRel, cols *ast.List) ([]scopeCol, error) { return out, nil } -func (a *analyzer) bindInsertValues(n ast.Node, rel scopeRel, targets []scopeCol) error { +func (a *analyzer) bindInsertValues(n ast.Node, rel scopeRel, targets []core.ClassColumn) error { if n == nil { return nil } @@ -126,7 +127,7 @@ func (a *analyzer) bindInsertValues(n ast.Node, rel scopeRel, targets []scopeCol continue } for i, v := range values.Items { - var target *scopeCol + var target *core.ClassColumn if i < len(targets) { target = &targets[i] } @@ -138,7 +139,7 @@ func (a *analyzer) bindInsertValues(n ast.Node, rel scopeRel, targets []scopeCol return nil } -func (a *analyzer) bindValue(rel scopeRel, target *scopeCol, v ast.Node) error { +func (a *analyzer) bindValue(rel scopeRel, target *core.ClassColumn, v ast.Node) error { if target != nil { switch value := v.(type) { case *ast.ParamRef: @@ -165,21 +166,21 @@ func (a *analyzer) projectReturning(l *ast.List) error { return nil } -func findColumn(rel scopeRel, name string) (scopeCol, bool) { +func findColumn(rel scopeRel, name string) (core.ClassColumn, bool) { for _, col := range rel.cols { - if col.name == name { + if col.Name == name { return col, true } } - return scopeCol{}, false + return core.ClassColumn{}, false } -func columnType(rel scopeRel, col scopeCol) exprType { +func columnType(rel scopeRel, col core.ClassColumn) exprType { return exprType{ - typeOID: col.typeOID, - nullable: !col.notNull, + typeOID: col.TypeOID, + nullable: !col.NotNull, sourceClassOID: rel.classOID, - sourceAttributeOID: col.attOID, + sourceAttributeOID: col.AttOID, sourceTableAlias: rel.alias, } } diff --git a/internal/core/analyzer/expr.go b/internal/core/analyzer/expr.go index 037e2f34ad..a030dd55de 100644 --- a/internal/core/analyzer/expr.go +++ b/internal/core/analyzer/expr.go @@ -113,10 +113,10 @@ func (a *analyzer) typeColumnRef(c *ast.ColumnRef) (exprType, error) { return exprType{}, fmt.Errorf("unknown column %q", column) } return exprType{ - typeOID: col.typeOID, - nullable: !col.notNull, + typeOID: col.TypeOID, + nullable: !col.NotNull, sourceClassOID: rel.classOID, - sourceAttributeOID: col.attOID, + sourceAttributeOID: col.AttOID, sourceTableAlias: rel.alias, }, nil } @@ -141,7 +141,7 @@ func flattenFields(fields *ast.List) []string { func (a *analyzer) typeParamRef(p *ast.ParamRef) (exprType, error) { cur, ok := a.params[p.Number] if !ok { - cur = &core.Parameter{Number: p.Number} + cur = core.Parameter{Number: p.Number} a.params[p.Number] = cur } return exprType{typeOID: cur.TypeOID, nullable: !cur.NotNull}, nil @@ -150,8 +150,7 @@ func (a *analyzer) typeParamRef(p *ast.ParamRef) (exprType, error) { func (a *analyzer) inferParam(number int, t exprType) { cur, ok := a.params[number] if !ok { - cur = &core.Parameter{Number: number} - a.params[number] = cur + cur = core.Parameter{Number: number} } if cur.TypeOID == 0 && t.typeOID != 0 { cur.TypeOID = t.typeOID @@ -171,6 +170,7 @@ func (a *analyzer) inferParam(number int, t exprType) { } } } + a.params[number] = cur } func (a *analyzer) typeAExpr(e *ast.A_Expr) (exprType, error) { diff --git a/internal/core/analyzer/projection.go b/internal/core/analyzer/projection.go index 3ed4034a76..049ca16ddc 100644 --- a/internal/core/analyzer/projection.go +++ b/internal/core/analyzer/projection.go @@ -1,14 +1,20 @@ package analyzer import ( + "slices" + "github.com/sqlc-dev/sqlc/internal/core" "github.com/sqlc-dev/sqlc/internal/sql/ast" ) func (a *analyzer) projectTarget(rt *ast.ResTarget) error { + // A column reference's field list is flattened once here and threaded + // through the star check, the star expansion and the output name. + var fields []string if cr, ok := rt.Val.(*ast.ColumnRef); ok { - if isStarRef(cr) { - a.emitStar(cr) + fields = flattenFields(cr.Fields) + if isStar(fields) { + a.emitStar(fields) return nil } } @@ -18,7 +24,7 @@ func (a *analyzer) projectTarget(rt *ast.ResTarget) error { return err } col := core.Column{ - Name: targetName(rt), + Name: targetName(rt, fields), TypeOID: t.typeOID, NotNull: !t.nullable, SourceClassOID: t.sourceClassOID, @@ -56,15 +62,14 @@ func (a *analyzer) decorateSource(col *core.Column, attOID int64, tableAlias str col.IsAutoIncrement = ad.AutoIncrement } -func targetName(rt *ast.ResTarget) string { +// targetName picks the output name for a target. fields is the already +// flattened field list when rt.Val is a column reference, and nil otherwise. +func targetName(rt *ast.ResTarget, fields []string) string { if rt.Name != nil && *rt.Name != "" { return *rt.Name } - if cr, ok := rt.Val.(*ast.ColumnRef); ok { - parts := flattenFields(cr.Fields) - if len(parts) > 0 { - return parts[len(parts)-1] - } + if len(fields) > 0 { + return fields[len(fields)-1] } if fc, ok := rt.Val.(*ast.FuncCall); ok { if name := funcCallName(fc); name != "" { @@ -74,33 +79,32 @@ func targetName(rt *ast.ResTarget) string { return "?column?" } -func isStarRef(c *ast.ColumnRef) bool { - parts := flattenFields(c.Fields) - return len(parts) > 0 && parts[len(parts)-1] == "*" +func isStar(fields []string) bool { + return len(fields) > 0 && fields[len(fields)-1] == "*" } -func (a *analyzer) emitStar(cr *ast.ColumnRef) { - parts := flattenFields(cr.Fields) +func (a *analyzer) emitStar(fields []string) { relName := "" - if len(parts) > 1 { - relName = parts[0] + if len(fields) > 1 { + relName = fields[0] } for _, rel := range a.scope.rels { if relName != "" && rel.alias != relName { continue } + a.columns = slices.Grow(a.columns, len(rel.cols)) for _, c := range rel.cols { col := core.Column{ - Name: c.name, - TypeOID: c.typeOID, - NotNull: c.notNull, + Name: c.Name, + TypeOID: c.TypeOID, + NotNull: c.NotNull, SourceClassOID: rel.classOID, - SourceAttributeOID: c.attOID, + SourceAttributeOID: c.AttOID, } - if name, err := a.cat.TypeName(c.typeOID); err == nil { + if name, err := a.cat.TypeName(c.TypeOID); err == nil { col.DataType = name } - a.decorateSource(&col, c.attOID, rel.alias) + a.decorateSource(&col, c.AttOID, rel.alias) a.columns = append(a.columns, col) } } diff --git a/internal/core/analyzer/scope.go b/internal/core/analyzer/scope.go index 772d6e2dc5..72c6647ef0 100644 --- a/internal/core/analyzer/scope.go +++ b/internal/core/analyzer/scope.go @@ -14,19 +14,15 @@ type scope struct { type scopeRel struct { alias string classOID int64 - cols []scopeCol -} - -type scopeCol struct { - name string - attOID int64 - typeOID int64 - notNull bool + // cols is the catalog's column list, held as-is rather than copied into a + // scope-local column type. + cols []core.ClassColumn } func (a *analyzer) buildScope(from *ast.List) (*scope, error) { - sc := &scope{} - for _, item := range listItems(from) { + items := listItems(from) + sc := &scope{rels: make([]scopeRel, 0, len(items))} + for _, item := range items { if err := a.appendFromItem(sc, item); err != nil { return nil, err } @@ -81,7 +77,7 @@ func (a *analyzer) bindRangeVar(rv *ast.RangeVar) (scopeRel, error) { rel.alias = *rv.Alias.Aliasname } - cols, err := a.classColumns(classOID) + cols, err := a.cat.ClassColumns(classOID) if err != nil { return scopeRel{}, err } @@ -89,48 +85,28 @@ func (a *analyzer) bindRangeVar(rv *ast.RangeVar) (scopeRel, error) { return rel, nil } -func (a *analyzer) classColumns(classOID int64) ([]scopeCol, error) { - cols, err := a.cat.ClassColumns(classOID) - if err != nil { - return nil, err - } - out := make([]scopeCol, 0, len(cols)) - for _, c := range cols { - out = append(out, scopeCol{ - name: c.Name, - attOID: c.AttOID, - typeOID: c.TypeOID, - notNull: c.NotNull, - }) - } - return out, nil -} - -func (s *scope) resolveColumn(relation, column string) (rel scopeRel, col scopeCol, ok bool, err error) { - var matches []struct { - rel scopeRel - col scopeCol - } +// resolveColumn finds the single column named column, optionally qualified by +// relation. It reports an error when more than one relation in scope offers +// that name. +func (s *scope) resolveColumn(relation, column string) (rel scopeRel, col core.ClassColumn, ok bool, err error) { + found := 0 for _, r := range s.rels { if relation != "" && r.alias != relation { continue } for _, c := range r.cols { - if c.name == column { - matches = append(matches, struct { - rel scopeRel - col scopeCol - }{r, c}) + if c.Name != column { + continue } + found++ + if found > 1 { + return scopeRel{}, core.ClassColumn{}, false, fmt.Errorf("ambiguous column reference %q", column) + } + rel, col = r, c } } - if len(matches) == 0 { - return scopeRel{}, scopeCol{}, false, nil - } - if len(matches) > 1 { - return scopeRel{}, scopeCol{}, false, fmt.Errorf("ambiguous column reference %q", column) + if found == 0 { + return scopeRel{}, core.ClassColumn{}, false, nil } - return matches[0].rel, matches[0].col, true, nil + return rel, col, true, nil } - -var _ = core.Column{} diff --git a/internal/core/attribute.go b/internal/core/attribute.go index 599479c26c..d8450404f6 100644 --- a/internal/core/attribute.go +++ b/internal/core/attribute.go @@ -154,6 +154,7 @@ type ClassColumn struct { NotNull bool } +// ClassColumns returns a relation's columns in ordinal order. func (c *Catalog) ClassColumns(classOID int64) ([]ClassColumn, error) { rows, err := c.q.ClassAttributes(context.Background(), classOID) if err != nil { diff --git a/internal/core/catalog.go b/internal/core/catalog.go index e87af5b7d6..a7c4683558 100644 --- a/internal/core/catalog.go +++ b/internal/core/catalog.go @@ -1,6 +1,7 @@ package core import ( + "context" "database/sql" "fmt" @@ -13,8 +14,9 @@ import ( //go:generate go run github.com/sqlc-dev/sqlc/cmd/sqlc generate type Catalog struct { - db *sql.DB - q *catalogdb.Queries + db *sql.DB + stmts *stmtCache + q *catalogdb.Queries } type Option func(*Catalog) error @@ -28,26 +30,74 @@ func New(opts ...Option) (*Catalog, error) { if err != nil { return nil, fmt.Errorf("core: open catalog: %w", err) } + // Every ":memory:" connection is its own empty database, so the pool has to + // be pinned to a single connection that is never retired — otherwise a + // second connection would see a catalog with no tables in it. + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(1) + db.SetConnMaxIdleTime(0) + db.SetConnMaxLifetime(0) + + // The catalog is scratch state rebuilt on every run and never read back + // from disk, so durability buys nothing and costs a journal write per + // statement. + for _, pragma := range []string{ + "PRAGMA journal_mode = OFF", + "PRAGMA synchronous = OFF", + } { + if _, err := db.Exec(pragma); err != nil { + db.Close() + return nil, fmt.Errorf("core: %s: %w", pragma, err) + } + } + + // catalogdef.Schema is a batch of DDL statements, so it goes to the + // connection directly rather than through the prepared-statement cache. if _, err := db.Exec(catalogdef.Schema); err != nil { db.Close() return nil, fmt.Errorf("core: init schema: %w", err) } - c := &Catalog{db: db, q: catalogdb.New(db)} - if err := c.bootstrap(); err != nil { + + stmts := newStmtCache(db) + c := &Catalog{db: db, stmts: stmts, q: catalogdb.New(stmts)} + + // Run the seeds inside one transaction rather than paying a commit for each + // of the hundreds of rows a dialect registers. + ctx := context.Background() + if err := stmts.begin(ctx); err != nil { db.Close() - return nil, fmt.Errorf("core: bootstrap: %w", err) + return nil, fmt.Errorf("core: begin setup: %w", err) + } + if err := c.setup(opts); err != nil { + stmts.rollback() + db.Close() + return nil, err + } + if err := stmts.commit(); err != nil { + db.Close() + return nil, fmt.Errorf("core: commit setup: %w", err) + } + return c, nil +} + +func (c *Catalog) setup(opts []Option) error { + if err := c.bootstrap(); err != nil { + return fmt.Errorf("core: bootstrap: %w", err) } for i, opt := range opts { if err := opt(c); err != nil { - db.Close() - return nil, fmt.Errorf("core: option %d: %w", i, err) + return fmt.Errorf("core: option %d: %w", i, err) } } - return c, nil + return nil } func (c *Catalog) Close() error { - return c.db.Close() + err := c.stmts.Close() + if cerr := c.db.Close(); err == nil { + err = cerr + } + return err } func (c *Catalog) DB() *sql.DB { diff --git a/internal/core/db.go b/internal/core/db.go new file mode 100644 index 0000000000..1f58522ee9 --- /dev/null +++ b/internal/core/db.go @@ -0,0 +1,171 @@ +package core + +import ( + "context" + "database/sql" + "sync" +) + +// stmtCache is a catalogdb.DBTX that keeps one prepared statement per distinct +// query string. +// +// The catalog runs the same dozen or so statements over and over — seeding a +// dialect is hundreds of inserts through a handful of queries, and analysis +// issues one lookup per table, column, type, operator and function a statement +// touches. Handing the raw *sql.DB to catalogdb means SQLite re-parses and +// re-plans that SQL text on every call, which dominated both paths. Preparing +// once and rebinding is the whole win. +// +// Statements live for the lifetime of the catalog and are released by Close. +type stmtCache struct { + db *sql.DB + + mu sync.RWMutex + stmts map[string]*sql.Stmt + + // While a transaction is open every statement must be prepared on the + // transaction's own connection: the catalog pins the pool to a single + // connection, so reaching for the *sql.DB mid-transaction would block + // forever waiting for a connection the transaction already holds. These + // statements are owned by the transaction and die with it. + tx *sql.Tx + txStmts map[string]*sql.Stmt +} + +func newStmtCache(db *sql.DB) *stmtCache { + return &stmtCache{db: db, stmts: make(map[string]*sql.Stmt)} +} + +func (c *stmtCache) prepared(ctx context.Context, query string) (*sql.Stmt, error) { + c.mu.RLock() + tx, cache := c.tx, c.stmts + if tx != nil { + cache = c.txStmts + } + stmt, ok := cache[query] + c.mu.RUnlock() + if ok { + return stmt, nil + } + + if tx != nil { + stmt, err := tx.PrepareContext(ctx, query) + if err != nil { + return nil, err + } + c.mu.Lock() + defer c.mu.Unlock() + // Bail out if the transaction ended while we were preparing; the + // statement is dead either way. + if c.tx != tx { + return stmt, nil + } + if existing, ok := c.txStmts[query]; ok { + stmt.Close() + return existing, nil + } + c.txStmts[query] = stmt + return stmt, nil + } + + stmt, err := c.db.PrepareContext(ctx, query) + if err != nil { + return nil, err + } + + c.mu.Lock() + defer c.mu.Unlock() + // Another goroutine may have prepared the same query while we were not + // holding the lock; keep the winner so each query has a single statement. + if existing, ok := c.stmts[query]; ok { + stmt.Close() + return existing, nil + } + c.stmts[query] = stmt + return stmt, nil +} + +func (c *stmtCache) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) { + stmt, err := c.prepared(ctx, query) + if err != nil { + return nil, err + } + return stmt.ExecContext(ctx, args...) +} + +func (c *stmtCache) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) { + stmt, err := c.prepared(ctx, query) + if err != nil { + return nil, err + } + return stmt.QueryContext(ctx, args...) +} + +func (c *stmtCache) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row { + stmt, err := c.prepared(ctx, query) + if err != nil { + // *sql.Row carries an error but has no exported constructor, so fall + // back to the unprepared path and let it report the failure on Scan. + return c.db.QueryRowContext(ctx, query, args...) + } + return stmt.QueryRowContext(ctx, args...) +} + +// PrepareContext returns a statement owned by the caller; it is deliberately +// not cached, since the caller is responsible for closing it. +func (c *stmtCache) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) { + return c.db.PrepareContext(ctx, query) +} + +// begin opens a transaction that every subsequent statement runs inside, until +// commit or rollback. Batching writes this way saves SQLite a commit per +// statement, which is most of the cost of installing a dialect seed. +func (c *stmtCache) begin(ctx context.Context) error { + tx, err := c.db.BeginTx(ctx, nil) + if err != nil { + return err + } + c.mu.Lock() + c.tx = tx + c.txStmts = make(map[string]*sql.Stmt) + c.mu.Unlock() + return nil +} + +// endTx detaches the current transaction and returns it, or nil if none is +// open. The transaction's statements are released along with it. +func (c *stmtCache) endTx() *sql.Tx { + c.mu.Lock() + defer c.mu.Unlock() + tx := c.tx + c.tx, c.txStmts = nil, nil + return tx +} + +func (c *stmtCache) commit() error { + if tx := c.endTx(); tx != nil { + return tx.Commit() + } + return nil +} + +func (c *stmtCache) rollback() { + if tx := c.endTx(); tx != nil { + tx.Rollback() + } +} + +func (c *stmtCache) Close() error { + c.rollback() + + c.mu.Lock() + defer c.mu.Unlock() + var firstErr error + for query, stmt := range c.stmts { + if err := stmt.Close(); err != nil && firstErr == nil { + firstErr = err + } + delete(c.stmts, query) + } + return firstErr +} diff --git a/internal/core/operator.go b/internal/core/operator.go index 7eae634aeb..784ea892a6 100644 --- a/internal/core/operator.go +++ b/internal/core/operator.go @@ -42,6 +42,8 @@ type OperatorOverload struct { ProcOID int64 } +// FindOperators returns the overloads of name matching the given operand types, +// where a zero OID means "any". func (c *Catalog) FindOperators(name string, leftTypeOID, rightTypeOID int64) ([]OperatorOverload, error) { rows, err := c.q.FindOperators(context.Background(), catalogdb.FindOperatorsParams{ Name: name, diff --git a/internal/core/proc.go b/internal/core/proc.go index a726c9f459..baee74e132 100644 --- a/internal/core/proc.go +++ b/internal/core/proc.go @@ -80,6 +80,8 @@ type ProcOverload struct { ArgTypes []int64 } +// FindProcs returns the overloads of name, optionally restricted to the given +// namespaces. func (c *Catalog) FindProcs(name string, namespaceOIDs []int64) ([]ProcOverload, error) { ctx := context.Background() lname := strings.ToLower(name) @@ -90,6 +92,7 @@ func (c *Catalog) FindProcs(name string, namespaceOIDs []int64) ([]ProcOverload, if err != nil { return nil, fmt.Errorf("find procs %q: %w", name, err) } + out = make([]ProcOverload, 0, len(rows)) for _, r := range rows { out = append(out, ProcOverload{ OID: r.Oid, @@ -111,6 +114,7 @@ func (c *Catalog) FindProcs(name string, namespaceOIDs []int64) ([]ProcOverload, if err != nil { return nil, fmt.Errorf("find procs %q: %w", name, err) } + out = make([]ProcOverload, 0, len(rows)) for _, r := range rows { out = append(out, ProcOverload{ OID: r.Oid,