feat: add overrides config option

Author: oliverlambson

README.md

@@ -76,3 +76,47 @@ class Status(str, enum.Enum):
     OPEN = "op!en"
     CLOSED = "clo@sed"
 ```
+
+### Override Column Types
+
+Option: `overrides`
+
+You can override the SQL to Python type mapping for specific columns using the `overrides` option. This is useful for columns with JSON data or other custom types.
+
+Example configuration:
+
+```yaml
+options:
+  package: authors
+  emit_pydantic_models: true
+  overrides:
+    - column: "some_table.payload"
+      py_import: "my_lib.models"
+      py_type: "Payload"
+```
+
+This will:
+1. Override the column `payload` in `some_table` to use the type `Payload`
+2. Add an import for `my_lib.models` to the models file
+
+Example output:
+
+```python
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.28.0
+
+import datetime
+import pydantic
+from typing import Any
+
+import my_lib.models
+
+
+class SomeTable(pydantic.BaseModel):
+    id: int
+    created_at: datetime.datetime
+    payload: my_lib.models.Payload
+```
+
+This is similar to the [overrides functionality in the Go version of sqlc](https://docs.sqlc.dev/en/stable/howto/overrides.html#overriding-types).

internal/config.go

@@ -1,13 +1,20 @@
 package python
 
+type OverrideColumn struct {
+	Column    string `json:"column"`
+	PyType    string `json:"py_type"`
+	PyImport  string `json:"py_import"`
+}
+
 type Config struct {
-	EmitExactTableNames         bool     `json:"emit_exact_table_names"`
-	EmitSyncQuerier             bool     `json:"emit_sync_querier"`
-	EmitAsyncQuerier            bool     `json:"emit_async_querier"`
-	Package                     string   `json:"package"`
-	Out                         string   `json:"out"`
-	EmitPydanticModels          bool     `json:"emit_pydantic_models"`
-	EmitStrEnum                 bool     `json:"emit_str_enum"`
-	QueryParameterLimit         *int32   `json:"query_parameter_limit"`
-	InflectionExcludeTableNames []string `json:"inflection_exclude_table_names"`
+	EmitExactTableNames         bool             `json:"emit_exact_table_names"`
+	EmitSyncQuerier             bool             `json:"emit_sync_querier"`
+	EmitAsyncQuerier            bool             `json:"emit_async_querier"`
+	Package                     string           `json:"package"`
+	Out                         string           `json:"out"`
+	EmitPydanticModels          bool             `json:"emit_pydantic_models"`
+	EmitStrEnum                 bool             `json:"emit_str_enum"`
+	QueryParameterLimit         *int32           `json:"query_parameter_limit"`
+	InflectionExcludeTableNames []string         `json:"inflection_exclude_table_names"`
+	Overrides                   []OverrideColumn `json:"overrides"`
 }

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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 sql:
   - schema: schema.sql
     queries: query.sql

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

@@ -0,0 +1,13 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.28.0
+import datetime
+import pydantic
+
+import my_lib.models
+
+
+class Book(pydantic.BaseModel):
+    id: int
+    created_at: datetime.datetime
+    payload: my_lib.models.Payload

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

@@ -0,0 +1,98 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.28.0
+# source: query.sql
+from typing import AsyncIterator, Iterator, Optional
+
+import my_lib.models
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+
+from db import models
+
+
+CREATE_BOOK = """-- name: create_book \\:one
+INSERT INTO books (payload)
+VALUES (:p1)
+RETURNING id, created_at, payload
+"""
+
+
+GET_BOOK = """-- name: get_book \\:one
+SELECT id, created_at, payload FROM books
+WHERE id = :p1 LIMIT 1
+"""
+
+
+LIST_BOOKS = """-- name: list_books \\:many
+SELECT id, created_at, payload FROM books
+ORDER BY id
+"""
+
+
+class Querier:
+    def __init__(self, conn: sqlalchemy.engine.Connection):
+        self._conn = conn
+
+    def create_book(self, *, payload: my_lib.models.Payload) -> Optional[models.Book]:
+        row = self._conn.execute(sqlalchemy.text(CREATE_BOOK), {"p1": payload}).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            created_at=row[1],
+            payload=row[2],
+        )
+
+    def get_book(self, *, id: int) -> Optional[models.Book]:
+        row = self._conn.execute(sqlalchemy.text(GET_BOOK), {"p1": id}).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            created_at=row[1],
+            payload=row[2],
+        )
+
+    def list_books(self) -> Iterator[models.Book]:
+        result = self._conn.execute(sqlalchemy.text(LIST_BOOKS))
+        for row in result:
+            yield models.Book(
+                id=row[0],
+                created_at=row[1],
+                payload=row[2],
+            )
+
+
+class AsyncQuerier:
+    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
+        self._conn = conn
+
+    async def create_book(self, *, payload: my_lib.models.Payload) -> Optional[models.Book]:
+        row = (await self._conn.execute(sqlalchemy.text(CREATE_BOOK), {"p1": payload})).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            created_at=row[1],
+            payload=row[2],
+        )
+
+    async def get_book(self, *, id: int) -> Optional[models.Book]:
+        row = (await self._conn.execute(sqlalchemy.text(GET_BOOK), {"p1": id})).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            created_at=row[1],
+            payload=row[2],
+        )
+
+    async def list_books(self) -> AsyncIterator[models.Book]:
+        result = await self._conn.stream(sqlalchemy.text(LIST_BOOKS))
+        async for row in result:
+            yield models.Book(
+                id=row[0],
+                created_at=row[1],
+                payload=row[2],
+            )

internal/endtoend/testdata/emit_type_overrides/my_lib/__init__.py

@@ -0,0 +1,7 @@
+from datetime import date
+
+from pydantic import BaseModel
+
+class Payload(BaseModel):
+    name: str
+    release_date: date

internal/endtoend/testdata/emit_type_overrides/my_lib/models.py

@@ -0,0 +1,12 @@
+-- name: GetBook :one
+SELECT * FROM books
+WHERE id = $1 LIMIT 1;
+
+-- name: ListBooks :many
+SELECT * FROM books
+ORDER BY id;
+
+-- name: CreateBook :one
+INSERT INTO books (payload)
+VALUES (sqlc.arg(payload))
+RETURNING *;

internal/endtoend/testdata/emit_type_overrides/query.sql

@@ -0,0 +1,5 @@
+CREATE TABLE books (
+    id SERIAL PRIMARY KEY,
+    created_at TIMESTAMPTZ NOT NULL,
+    payload JSONB NOT NULL
+);

internal/endtoend/testdata/emit_type_overrides/schema.sql

@@ -0,0 +1,22 @@
+version: "2"
+plugins:
+  - name: py
+    wasm:
+      url: file://../../../../bin/sqlc-gen-python.wasm
+      sha256: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
+sql:
+  - schema: schema.sql
+    queries: query.sql
+    engine: postgresql
+    codegen:
+      - plugin: py
+        out: db
+        options:
+          package: db
+          emit_pydantic_models: true
+          emit_sync_querier: true
+          emit_async_querier: true
+          overrides:
+            - column: "books.payload"
+              py_import: "my_lib.models"
+              py_type: "Payload"

internal/endtoend/testdata/emit_type_overrides/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+    sha256: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 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: "31717935ced1923fdaea102da0d345776c173df7fa6668120be8f7900a7fe938"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/query_parameter_no_limit/sqlc.yaml

@@ -181,6 +181,40 @@ func (q Query) ArgDictNode() *pyast.Node {
 }
 
 func makePyType(req *plugin.GenerateRequest, col *plugin.Column) pyType {
+	// Parse the configuration
+	var conf Config
+	if len(req.PluginOptions) > 0 {
+		if err := json.Unmarshal(req.PluginOptions, &conf); err != nil {
+			log.Printf("failed to parse plugin options: %s", err)
+		}
+	}
+
+	// Check for overrides
+	if len(conf.Overrides) > 0 && col.Table != nil {
+		tableName := col.Table.Name
+		if col.Table.Schema != "" && col.Table.Schema != req.Catalog.DefaultSchema {
+			tableName = col.Table.Schema + "." + tableName
+		}
+
+		// Look for a matching override
+		for _, override := range conf.Overrides {
+			overrideKey := tableName + "." + col.Name
+			if override.Column == overrideKey {
+				// Found a match, use the override
+				typeStr := override.PyType
+				if override.PyImport != "" && !strings.Contains(typeStr, ".") {
+					typeStr = override.PyImport + "." + override.PyType
+				}
+				return pyType{
+					InnerType: typeStr,
+					IsArray:   col.IsArray,
+					IsNull:    !col.NotNull,
+				}
+			}
+		}
+	}
+	
+	// No override found, use the standard type mapping
 	typ := pyInnerType(req, col)
 	return pyType{
 		InnerType: typ,