Support type overrides in addition to column overrides

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
         go-version: '1.23.5'
     - uses: sqlc-dev/setup-sqlc@v4
       with:
-        sqlc-version: '1.29.0'
+        sqlc-version: '1.30.0'
     - run: make
     - run: make test
     - run: sqlc diff

README.md

@@ -99,3 +99,51 @@ 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 or database types 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"
+    - db_type: "jsonb"
+      py_import: "my_lib.models"
+      py_type: "Payload"
+```
+
+This will:
+1. Override the column `payload` in `some_table` to use the type `Payload`
+2. Override any column with the database type `jsonb` to use the type `Payload`
+3. Add an import for `my_lib.models` to the models file
+
+Example output:
+
+```python
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.30.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).

examples/src/authors/models.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import dataclasses
 from typing import Optional
 

examples/src/authors/query.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query.sql
 from typing import AsyncIterator, Iterator, Optional
 

examples/src/booktest/models.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import dataclasses
 import datetime
 import enum

examples/src/booktest/query.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query.sql
 import dataclasses
 import datetime

examples/src/jets/models.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import dataclasses
 
 

examples/src/jets/query-building.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query-building.sql
 from typing import AsyncIterator, Optional
 

examples/src/ondeck/city.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: city.sql
 from typing import AsyncIterator, Optional
 

examples/src/ondeck/models.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import dataclasses
 import datetime
 import enum

examples/src/ondeck/venue.py

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: venue.sql
 import dataclasses
 from typing import AsyncIterator, List, Optional

internal/config.go

@@ -1,7 +1,15 @@
 package python
 
+import (
+	"bytes"
+	json "encoding/json"
+	"fmt"
+	"strings"
+)
+
 type OverrideColumn struct {
 	Column   string `json:"column"`
+	DbType   string `json:"db_type"`
 	PyType   string `json:"py_type"`
 	PyImport string `json:"py_import"`
 }
@@ -19,3 +27,19 @@ type Config struct {
 	InflectionExcludeTableNames []string         `json:"inflection_exclude_table_names"`
 	Overrides                   []OverrideColumn `json:"overrides"`
 }
+
+func parseConfig(raw []byte) (Config, error) {
+	var conf Config
+	if len(raw) == 0 {
+		return conf, nil
+	}
+
+	dec := json.NewDecoder(bytes.NewReader(raw))
+	dec.DisallowUnknownFields()
+	if err := dec.Decode(&conf); err != nil {
+		msg := strings.TrimPrefix(err.Error(), "json: ")
+		return Config{}, fmt.Errorf("invalid plugin options: %s", msg)
+	}
+
+	return conf, nil
+}

internal/config_test.go

@@ -0,0 +1,32 @@
+package python
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestParseConfigDisallowUnknownFields(t *testing.T) {
+	_, err := parseConfig([]byte(`{"emit_sync_querier":true,"db_typ":"jsonb"}`))
+	if err == nil {
+		t.Fatal("expected unknown field error, got nil")
+	}
+	if !strings.Contains(err.Error(), "invalid plugin options") {
+		t.Fatalf("expected error to reference plugin options, got: %v", err)
+	}
+	if !strings.Contains(err.Error(), `unknown field "db_typ"`) {
+		t.Fatalf("expected unknown field in error, got: %v", err)
+	}
+}
+
+func TestParseConfigValid(t *testing.T) {
+	conf, err := parseConfig([]byte(`{"emit_sync_querier":true,"overrides":[{"db_type":"jsonb","py_type":"str"}]}`))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !conf.EmitSyncQuerier {
+		t.Fatal("expected emit_sync_querier to be true")
+	}
+	if len(conf.Overrides) != 1 || conf.Overrides[0].DbType != "jsonb" || conf.Overrides[0].PyType != "str" {
+		t.Fatal("unexpected parsed overrides")
+	}
+}

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import dataclasses
 
 

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query.sql
 from typing import Optional
 

internal/endtoend/testdata/dataclasses_with_reserved_keywords/sqlc.yaml

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

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import pydantic
 from typing import Optional
 

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query.sql
 from typing import AsyncIterator, Iterator, Optional
 

internal/endtoend/testdata/emit_pydantic_models/sqlc.yaml

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

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import pydantic
 
 

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query.sql
 from typing import Optional
 

internal/endtoend/testdata/emit_pydantic_models_with_reserved_keywords/sqlc.yaml

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

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 import dataclasses
 import enum
 from typing import Optional

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

@@ -1,6 +1,6 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
-#   sqlc v1.29.0
+#   sqlc v1.30.0
 # source: query.sql
 from typing import AsyncIterator, Iterator, Optional
 

internal/endtoend/testdata/emit_str_enum/sqlc.yaml

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

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

@@ -0,0 +1,24 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.30.0
+import enum
+import pydantic
+from typing import Optional
+
+
+class BookStatus(str, enum.Enum):
+    AVAILABLE = "available"
+    CHECKED_OUT = "checked_out"
+    OVERDUE = "overdue"
+
+
+class Book(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(
+        validate_by_alias=True,
+        validate_by_name=True,
+    )
+    id: int
+    title: str
+    status: Optional[BookStatus]
+    payload: my_lib.models.BookPayload
+    metadata: Optional[my_lib.models.JsonValue]