postgresql: port the extension catalog to seed files

internal/core/catalog.go

@@ -22,6 +22,12 @@ type Catalog struct {
 	// dialectOID is the dialect this catalog was seeded with. A catalog is
 	// built for one dialect, so dialect-wide lookups need no other input.
 	dialectOID int64
+
+	// loadExtension applies a named extension's seed, installed by the
+	// dialect's seed for the engines that ship extension data. It runs at most
+	// once per extension name: a schema is free to say CREATE EXTENSION twice.
+	loadExtension func(name string) error
+	extensions    map[string]bool
 }
 
 type Option func(*Catalog) error
@@ -149,3 +155,30 @@ func (c *Catalog) bootstrap() error {
 	_, err := c.CreateNamespace("public")
 	return err
 }
+
+// SeededDialectOID returns the dialect this catalog was seeded with, which is
+// what anything registered after seeding — an extension's types and functions
+// — records as its dialect.
+func (c *Catalog) SeededDialectOID() int64 {
+	return c.dialectOID
+}
+
+// SetExtensionLoader installs the function CREATE EXTENSION calls to apply an
+// extension's seed. A dialect without extension data leaves it unset.
+func (c *Catalog) SetExtensionLoader(fn func(name string) error) {
+	c.loadExtension = fn
+}
+
+// LoadExtension applies the named extension's seed to the catalog. An
+// extension already applied, or a dialect with no loader, contributes
+// nothing; an extension the loader does not know is the loader's to ignore.
+func (c *Catalog) LoadExtension(name string) error {
+	if c.loadExtension == nil || c.extensions[name] {
+		return nil
+	}
+	if c.extensions == nil {
+		c.extensions = map[string]bool{}
+	}
+	c.extensions[name] = true
+	return c.loadExtension(name)
+}

internal/core/catalogdef/query.sql

@@ -57,6 +57,11 @@ LIMIT 1;
 -- name: TypeNameByOID :one
 SELECT name FROM sql_type WHERE oid = ?;
 
+-- name: TypeOIDsInCategory :many
+SELECT oid FROM sql_type
+WHERE dialect_oid = ? AND category = ?
+ORDER BY oid;
+
 -- name: LookupType :one
 SELECT oid, name, category, typtype, preferred
 FROM sql_type

internal/core/dialect.go

@@ -54,6 +54,11 @@ func (c *Catalog) DialectFlag(dialectOID int64, key string) (string, error) {
 // same ones the seeded types have.
 const FlagComparisonOperators = "operators.comparison"
 
+// FlagCastCategories holds the categories whose types are all implicitly
+// castable to one another, as the dialect's seed declared them, so that a type
+// arriving after the seed — an extension's, say — can join its category.
+const FlagCastCategories = "casts.categories"
+
 // Constant kinds. A literal in a query has no declared type, so each dialect
 // names the type its literals take on.
 const (

internal/core/schema/schema.go

@@ -28,6 +28,11 @@ func Apply(cat *core.Catalog, n ast.Node) error {
 		return applyDropTable(cat, v)
 	case *ast.CreateEnumStmt:
 		return applyCreateEnum(cat, v)
+	case *ast.CreateExtensionStmt:
+		if v.Extname == nil {
+			return nil
+		}
+		return cat.LoadExtension(*v.Extname)
 	case *ast.CreateFunctionStmt:
 		return applyCreateFunction(cat, v)
 	case *ast.AlterTableStmt:

internal/core/seed/extension.go

@@ -0,0 +1,147 @@
+package seed
+
+import (
+	"fmt"
+	"io/fs"
+	"path"
+	"strings"
+
+	"github.com/sqlc-dev/sqlc/internal/core"
+)
+
+// applyExtension applies the named extension's directory to a catalog that
+// has already been seeded. Unlike the dialect's own seed, an extension lands
+// in a catalog full of types, so everything it names is resolved against what
+// is there before being created.
+func applyExtension(cat *core.Catalog, fsys fs.FS, name string) error {
+	dir := path.Join(ExtensionsDir, name)
+	if _, err := fs.Stat(fsys, dir); err != nil {
+		// An extension sqlc has no data for adds nothing, the way the legacy
+		// catalog has always treated one.
+		return nil
+	}
+	sub, err := fs.Sub(fsys, dir)
+	if err != nil {
+		return fmt.Errorf("seed: extension %s: %w", name, err)
+	}
+	e := &extension{cat: cat}
+	if err := stream(sub, TypesFile, e.addType); err != nil {
+		return fmt.Errorf("extension %s: %w", name, err)
+	}
+	if err := stream(sub, FunctionsFile, e.addFunction); err != nil {
+		return fmt.Errorf("extension %s: %w", name, err)
+	}
+	return nil
+}
+
+type extension struct {
+	cat *core.Catalog
+}
+
+// addType registers a type the extension defines. The type joins the
+// dialect's rules the way a type declared by a schema does: it gains the
+// comparison operators, and implicit casts to the types of its category when
+// the dialect's seed declared that category mutually castable.
+func (e *extension) addType(t Type) error {
+	for _, name := range append([]string{t.Name}, t.Aliases...) {
+		if _, err := e.cat.TypeOID(name); err == nil {
+			// The dialect, or an earlier extension, already ships this type.
+			continue
+		}
+		oid, err := e.cat.CreateUserType(name, t.Category)
+		if err != nil {
+			return fmt.Errorf("type %q: %w", name, err)
+		}
+		if err := e.categoryCasts(oid, t.Category); err != nil {
+			return fmt.Errorf("type %q: %w", name, err)
+		}
+	}
+	return nil
+}
+
+// categoryCasts makes a new type implicitly castable to and from the types
+// already in its category, under the same rule the dialect's seed applied to
+// the types it was born with.
+func (e *extension) categoryCasts(oid int64, category string) error {
+	categories, err := e.cat.DialectFlag(e.cat.SeededDialectOID(), core.FlagCastCategories)
+	if err != nil || categories == "" {
+		return err
+	}
+	if categories != "*" && !strings.Contains(categories, category) {
+		return nil
+	}
+	peers, err := e.cat.TypeOIDsInCategory(category)
+	if err != nil {
+		return err
+	}
+	for _, peer := range peers {
+		if peer == oid {
+			continue
+		}
+		for _, cast := range [][2]int64{{oid, peer}, {peer, oid}} {
+			if err := e.cat.CreateCast(core.CastSpec{
+				SourceTypeOID: cast[0],
+				TargetTypeOID: cast[1],
+				Context:       "i",
+				DialectOID:    e.cat.SeededDialectOID(),
+			}); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (e *extension) addFunction(fn Function) error {
+	returnOID, err := e.funcType(fn.Returns)
+	if err != nil {
+		return fmt.Errorf("function %q: %w", fn.Name, err)
+	}
+	args := make([]core.ProcArg, 0, len(fn.Args))
+	for _, a := range fn.Args {
+		// An out or table parameter is part of the result, not the call.
+		switch a.Mode {
+		case "o", "t":
+			continue
+		}
+		argOID, err := e.funcType(a.Type)
+		if err != nil {
+			return fmt.Errorf("function %q: %w", fn.Name, err)
+		}
+		args = append(args, core.ProcArg{
+			Name:       a.Name,
+			TypeOID:    argOID,
+			Mode:       a.Mode,
+			HasDefault: a.HasDefault,
+		})
+	}
+	_, err = e.cat.CreateProc(core.ProcSpec{
+		Name:           fn.Name,
+		DialectOID:     e.cat.SeededDialectOID(),
+		Kind:           fn.Kind,
+		ReturnTypeOID:  returnOID,
+		ReturnNullable: fn.Nullable,
+		Args:           args,
+	})
+	if err != nil {
+		return fmt.Errorf("function %q: %w", fn.Name, err)
+	}
+	return nil
+}
+
+// funcType resolves a type a function signature names, registering the ones
+// no seed bothered to list the way the dialect's own function list does.
+func (e *extension) funcType(name string) (int64, error) {
+	if name == "" {
+		return 0, nil
+	}
+	if oid, err := e.cat.TypeOID(name); err == nil {
+		return oid, nil
+	}
+	return e.cat.CreateTypeSpec(core.TypeSpec{
+		Name:       name,
+		Typtype:    "b",
+		Category:   "U",
+		DialectOID: e.cat.SeededDialectOID(),
+	})
+}

internal/core/seed/seed.go

@@ -15,6 +15,10 @@
 // The lists are JSONL — one record per line — and are applied as they are
 // read, so a dialect whose function list runs to thousands of entries is never
 // held in memory as a whole. Any of the lists may be left out.
+//
+// A dialect may also hold an extensions/ directory with one directory per
+// extension, each a smaller bundle of the same files, applied when a schema
+// says CREATE EXTENSION.
 package seed
 
 import (
@@ -42,6 +46,10 @@ const (
 	RelationsFile = "relations.jsonl"
 )
 
+// ExtensionsDir is the directory under a dialect holding one directory per
+// extension the dialect knows.
+const ExtensionsDir = "extensions"
+
 // Settings is dialect.json: what the dialect is called and the rules that
 // generate its operators and casts.
 type Settings struct {
@@ -139,14 +147,21 @@ type Arg struct {
 }
 
 // Dialect returns the catalog option that seeds the dialect described by the
-// JSONL files in dir.
+// JSONL files in dir. A dialect that ships extension data — a directory per
+// extension under extensions/ — also has CREATE EXTENSION wired up to it.
 func Dialect(fsys fs.FS, dir string) core.Option {
 	return core.WithSeed(func(cat *core.Catalog) error {
 		sub, err := fs.Sub(fsys, dir)
 		if err != nil {
 			return fmt.Errorf("seed: %s: %w", dir, err)
 		}
-		return apply(cat, sub)
+		if err := apply(cat, sub); err != nil {
+			return err
+		}
+		cat.SetExtensionLoader(func(name string) error {
+			return applyExtension(cat, sub, name)
+		})
+		return nil
 	})
 }
 
@@ -439,8 +454,15 @@ func (b *builder) consts() error {
 	// arrays, a SQLite column typed whatever the author felt like. Recording
 	// the comparison operators lets the catalog give those types the same ones.
 	if len(b.settings.Comparison) > 0 {
-		return b.cat.SetDialectFlag(b.dialectOID, core.FlagComparisonOperators,
-			strings.Join(b.settings.Comparison, ","))
+		if err := b.cat.SetDialectFlag(b.dialectOID, core.FlagComparisonOperators,
+			strings.Join(b.settings.Comparison, ",")); err != nil {
+			return err
+		}
+	}
+	// The cast categories are recorded for the same reason: a type an
+	// extension adds joins its category's casts.
+	if b.settings.CastCategories != "" {
+		return b.cat.SetDialectFlag(b.dialectOID, core.FlagCastCategories, b.settings.CastCategories)
 	}
 	return nil
 }

internal/core/types.go

@@ -128,6 +128,19 @@ func (c *Catalog) createComparisons(typeOID int64) error {
 	return nil
 }
 
+// TypeOIDsInCategory returns the types the catalog's dialect has in the named
+// category, in the order they were created.
+func (c *Catalog) TypeOIDsInCategory(category string) ([]int64, error) {
+	oids, err := c.q.TypeOIDsInCategory(context.Background(), catalogdb.TypeOIDsInCategoryParams{
+		DialectOid: nullInt64(c.dialectOID),
+		Category:   nullString(category),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("types in category %q: %w", category, err)
+	}
+	return oids, nil
+}
+
 func (c *Catalog) TypeOID(name string) (int64, error) {
 	oid, err := c.q.TypeOIDByName(context.Background(), strings.ToLower(name))
 	if err != nil {

internal/endtoend/testdata/analyze_extension/postgresql/exec.json

@@ -0,0 +1,5 @@
+{
+  "command": "analyze",
+  "args": ["--dialect", "postgresql", "--schema", "schema.sql", "query.sql"],
+  "contexts": ["base"]
+}

internal/endtoend/testdata/analyze_extension/postgresql/query.sql

@@ -0,0 +1,8 @@
+-- name: GetUserByEmail :one
+SELECT id, email FROM users WHERE email = $1;
+
+-- name: MatchingEmails :many
+SELECT id FROM users WHERE email = name;
+
+-- name: RankUsers :many
+SELECT id, similarity(name, $1) AS score FROM users;

internal/endtoend/testdata/analyze_extension/postgresql/schema.sql

@@ -0,0 +1,8 @@
+CREATE EXTENSION citext;
+CREATE EXTENSION pg_trgm;
+
+CREATE TABLE users (
+  id    BIGSERIAL PRIMARY KEY,
+  name  text   NOT NULL,
+  email citext NOT NULL
+);