Add a type override function for domain types.

Author: Boris Smidt

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