Add a type override function for domain types.

README.md

@@ -76,3 +76,42 @@ class Status(str, enum.Enum):
     OPEN = "op!en"
     CLOSED = "clo@sed"
 ```
+
+### Map domains (and other unknown types) to Python types
+
+Option: `domain_overrides`
+
+sqlc does not pass `CREATE DOMAIN` definitions (their base type or `CHECK`
+constraints) to code generation plugins, so columns using a domain are emitted
+as `Any` and a `unknown PostgreSQL type` warning is logged. The
+`domain_overrides` option lets you map a PostgreSQL type name to a
+fully-qualified Python type. The required `import` is added automatically,
+including for nested modules.
+
+```yaml
+options:
+  package: authors
+  domain_overrides:
+    job_status: my.module.JobStatus
+    positive_int: decimal.Decimal
+```
+
+Given a domain `job_status` used by a `status` column, this generates:
+
+```py
+import decimal
+
+import my.module
+
+
+@dataclasses.dataclass()
+class Job:
+    id: int
+    status: my.module.JobStatus
+    priority: Optional[decimal.Decimal]
+```
+
+The key is matched against the column's data type, its bare type name, and its
+schema-qualified name (e.g. `public.job_status`), so you can key the override
+however is most convenient. This also works for any other type sqlc reports as
+unknown, not just domains.

internal/config.go

@@ -10,4 +10,10 @@ type Config struct {
 	EmitStrEnum                 bool     `json:"emit_str_enum"`
 	QueryParameterLimit         *int32   `json:"query_parameter_limit"`
 	InflectionExcludeTableNames []string `json:"inflection_exclude_table_names"`
+
+	// DomainOverrides maps a PostgreSQL type name (typically a DOMAIN, whose
+	// definition sqlc does not pass to plugins) to a fully-qualified Python
+	// type. For example {"job_status": "my.module.JobStatus"} emits a
+	// "import my.module" and annotates the column as "my.module.JobStatus".
+	DomainOverrides map[string]string `json:"domain_overrides"`
 }

internal/endtoend/testdata/domain_overrides/db/models.py

@@ -0,0 +1,15 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.31.1
+import dataclasses
+import decimal
+from typing import Optional
+
+import my.module
+
+
+@dataclasses.dataclass()
+class Job:
+    id: int
+    status: my.module.JobStatus
+    priority: Optional[decimal.Decimal]

internal/endtoend/testdata/domain_overrides/db/query.py

@@ -0,0 +1,72 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.31.1
+# source: query.sql
+from typing import AsyncIterator, Iterator, Optional
+
+import my.module
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+
+from db import models
+
+
+GET_JOB = """-- name: get_job \\:one
+SELECT id, status, priority FROM jobs
+WHERE id = :p1 LIMIT 1
+"""
+
+
+LIST_JOBS_BY_STATUS = """-- name: list_jobs_by_status \\:many
+SELECT id, status, priority FROM jobs
+WHERE status = :p1
+ORDER BY priority
+"""
+
+
+class Querier:
+    def __init__(self, conn: sqlalchemy.engine.Connection):
+        self._conn = conn
+
+    def get_job(self, *, id: int) -> Optional[models.Job]:
+        row = self._conn.execute(sqlalchemy.text(GET_JOB), {"p1": id}).first()
+        if row is None:
+            return None
+        return models.Job(
+            id=row[0],
+            status=row[1],
+            priority=row[2],
+        )
+
+    def list_jobs_by_status(self, *, status: my.module.JobStatus) -> Iterator[models.Job]:
+        result = self._conn.execute(sqlalchemy.text(LIST_JOBS_BY_STATUS), {"p1": status})
+        for row in result:
+            yield models.Job(
+                id=row[0],
+                status=row[1],
+                priority=row[2],
+            )
+
+
+class AsyncQuerier:
+    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
+        self._conn = conn
+
+    async def get_job(self, *, id: int) -> Optional[models.Job]:
+        row = (await self._conn.execute(sqlalchemy.text(GET_JOB), {"p1": id})).first()
+        if row is None:
+            return None
+        return models.Job(
+            id=row[0],
+            status=row[1],
+            priority=row[2],
+        )
+
+    async def list_jobs_by_status(self, *, status: my.module.JobStatus) -> AsyncIterator[models.Job]:
+        rows = (await self._conn.execute(sqlalchemy.text(LIST_JOBS_BY_STATUS), {"p1": status})).all()
+        for row in rows:
+            yield models.Job(
+                id=row[0],
+                status=row[1],
+                priority=row[2],
+            )

internal/endtoend/testdata/domain_overrides/query.sql

@@ -0,0 +1,8 @@
+-- name: GetJob :one
+SELECT * FROM jobs
+WHERE id = $1 LIMIT 1;
+
+-- name: ListJobsByStatus :many
+SELECT * FROM jobs
+WHERE status = $1
+ORDER BY priority;

internal/endtoend/testdata/domain_overrides/schema.sql

@@ -0,0 +1,19 @@
+CREATE DOMAIN job_status AS text
+CHECK (
+    VALUE IN (
+                 'QUEUED',
+                 'PENDING',
+                 'RUNNING',
+                 'COMPLETED',
+                 'FAILED'
+        )) NOT NULL;
+
+CREATE DOMAIN positive_int AS integer
+CHECK (VALUE > 0);
+
+
+CREATE TABLE jobs (
+          id       BIGSERIAL PRIMARY KEY,
+          status   job_status NOT NULL,
+          priority positive_int
+);

internal/endtoend/testdata/domain_overrides/sqlc.yaml

@@ -0,0 +1,20 @@
+version: "2"
+plugins:
+  - name: py
+    wasm:
+      url: file://../../../../bin/sqlc-gen-python.wasm
+      sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
+sql:
+  - schema: schema.sql
+    queries: query.sql
+    engine: postgresql
+    codegen:
+      - plugin: py
+        out: db
+        options:
+          package: db
+          emit_sync_querier: true
+          emit_async_querier: true
+          domain_overrides:
+            job_status: my.module.JobStatus
+            positive_int: decimal.Decimal

internal/endtoend/testdata/emit_pydantic_models/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/emit_str_enum/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
   - name: py
     wasm:
       url: file://../../../../bin/sqlc-gen-python.wasm
-      sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+      sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
   - schema: schema.sql
     queries: query.sql

internal/endtoend/testdata/exec_result/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/exec_rows/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/inflection_exclude_table_names/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/query_parameter_limit_two/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/query_parameter_limit_undefined/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/query_parameter_limit_zero/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/query_parameter_no_limit/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "00c7c16380c4593d7a86b82e2f650c5655c179cdb0d63b1513d6987ec9be0f46"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/gen.go

@@ -180,19 +180,19 @@ func (q Query) ArgDictNode() *pyast.Node {
 	}
 }
 
-func makePyType(req *plugin.GenerateRequest, col *plugin.Column) pyType {
-	typ := pyInnerType(req, col)
+func makePyType(conf Config, req *plugin.GenerateRequest, col *plugin.Column) pyType {
+	typ := pyInnerType(conf, req, col)
 	return pyType{
 		InnerType: typ,
 		IsArray:   col.IsArray,
 		IsNull:    !col.NotNull,
 	}
 }
 
-func pyInnerType(req *plugin.GenerateRequest, col *plugin.Column) string {
+func pyInnerType(conf Config, req *plugin.GenerateRequest, col *plugin.Column) string {
 	switch req.Settings.Engine {
 	case "postgresql":
-		return postgresType(req, col)
+		return postgresType(conf, req, col)
 	default:
 		log.Println("unsupported engine type")
 		return "Any"
@@ -285,7 +285,7 @@ func buildModels(conf Config, req *plugin.GenerateRequest) []Struct {
 				Comment: table.Comment,
 			}
 			for _, column := range table.Columns {
-				typ := makePyType(req, column) // TODO: This used to call compiler.ConvertColumn?
+				typ := makePyType(conf, req, column) // TODO: This used to call compiler.ConvertColumn?
 				typ.InnerType = strings.TrimPrefix(typ.InnerType, "models.")
 				s.Fields = append(s.Fields, Field{
 					Name:    column.Name,
@@ -321,7 +321,7 @@ type pyColumn struct {
 	*plugin.Column
 }
 
-func columnsToStruct(req *plugin.GenerateRequest, name string, columns []pyColumn) *Struct {
+func columnsToStruct(conf Config, req *plugin.GenerateRequest, name string, columns []pyColumn) *Struct {
 	gs := Struct{
 		Name: name,
 	}
@@ -344,7 +344,7 @@ func columnsToStruct(req *plugin.GenerateRequest, name string, columns []pyColum
 		}
 		gs.Fields = append(gs.Fields, Field{
 			Name: fieldName,
-			Type: makePyType(req, c.Column),
+			Type: makePyType(conf, req, c.Column),
 		})
 		seen[colName]++
 	}
@@ -406,14 +406,14 @@ func buildQueries(conf Config, req *plugin.GenerateRequest, structs []Struct) ([
 			gq.Args = []QueryValue{{
 				Emit:   true,
 				Name:   "arg",
-				Struct: columnsToStruct(req, query.Name+"Params", cols),
+				Struct: columnsToStruct(conf, req, query.Name+"Params", cols),
 			}}
 		} else {
 			args := make([]QueryValue, 0, len(query.Params))
 			for _, p := range query.Params {
 				args = append(args, QueryValue{
 					Name: paramName(p),
-					Typ:  makePyType(req, p.Column),
+					Typ:  makePyType(conf, req, p.Column),
 				})
 			}
 			gq.Args = args
@@ -423,7 +423,7 @@ func buildQueries(conf Config, req *plugin.GenerateRequest, structs []Struct) ([
 			c := query.Columns[0]
 			gq.Ret = QueryValue{
 				Name: columnName(c, 0),
-				Typ:  makePyType(req, c),
+				Typ:  makePyType(conf, req, c),
 			}
 		} else if len(query.Columns) > 1 {
 			var gs *Struct
@@ -438,7 +438,7 @@ func buildQueries(conf Config, req *plugin.GenerateRequest, structs []Struct) ([
 				for i, f := range s.Fields {
 					c := query.Columns[i]
 					// HACK: models do not have "models." on their types, so trim that so we can find matches
-					trimmedPyType := makePyType(req, c)
+					trimmedPyType := makePyType(conf, req, c)
 					trimmedPyType.InnerType = strings.TrimPrefix(trimmedPyType.InnerType, "models.")
 					sameName := f.Name == columnName(c, i)
 					sameType := f.Type == trimmedPyType
@@ -461,7 +461,7 @@ func buildQueries(conf Config, req *plugin.GenerateRequest, structs []Struct) ([
 						Column: c,
 					})
 				}
-				gs = columnsToStruct(req, query.Name+"Row", columns)
+				gs = columnsToStruct(conf, req, query.Name+"Row", columns)
 				emit = true
 			}
 			gq.Ret = QueryValue{
@@ -1038,6 +1038,25 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 					),
 				)
 				f.Returns = subscriptNode("AsyncIterator", q.Ret.Annotation())
+			case ":stream":
+				stream := connMethodNode("stream", q.ConstantName, q.ArgDictNode())
+				f.Body = append(f.Body,
+					assignNode("result", poet.Await(stream)),
+					poet.Node(
+						&pyast.AsyncFor{
+							Target: poet.Name("row"),
+							Iter:   poet.Name("result"),
+							Body: []*pyast.Node{
+								poet.Expr(
+									poet.Yield(
+										q.Ret.RowNode("row"),
+									),
+								),
+							},
+						},
+					),
+				)
+				f.Returns = subscriptNode("AsyncIterator", q.Ret.Annotation())
 			case ":exec":
 				f.Body = append(f.Body, poet.Await(exec))
 				f.Returns = poet.Constant(nil)