diff --git a/README.md b/README.md index c9f2531..4446c1c 100644 --- a/README.md +++ b/README.md @@ -76,3 +76,9 @@ class Status(str, enum.Enum): OPEN = "op!en" CLOSED = "clo@sed" ``` + +### Omit unused table and enum structs + +Option: `omit_unused_structs` + +If `true`, sqlc won't generate table and enum structs that aren't used in queries for a given package. diff --git a/internal/config.go b/internal/config.go index 1a8a565..cde8561 100644 --- a/internal/config.go +++ b/internal/config.go @@ -8,6 +8,7 @@ type Config struct { Out string `json:"out"` EmitPydanticModels bool `json:"emit_pydantic_models"` EmitStrEnum bool `json:"emit_str_enum"` + OmitUnusedStructs bool `json:"omit_unused_structs,omitempty"` QueryParameterLimit *int32 `json:"query_parameter_limit"` InflectionExcludeTableNames []string `json:"inflection_exclude_table_names"` } diff --git a/internal/endtoend/testdata/omit_unused_structs/db/models.py b/internal/endtoend/testdata/omit_unused_structs/db/models.py new file mode 100644 index 0000000..e6f76b9 --- /dev/null +++ b/internal/endtoend/testdata/omit_unused_structs/db/models.py @@ -0,0 +1,66 @@ +# Code generated by sqlc. DO NOT EDIT. +# versions: +# sqlc v1.28.0 +import dataclasses +import enum +from typing import List, Optional + + +class ArrayEnum(enum.StrEnum): + O = "o" + P = "p" + + +class QueryParamEnumTableEnum(enum.StrEnum): + G = "g" + H = "h" + + +class QueryParamStructEnumTableEnum(enum.StrEnum): + I = "i" + J = "j" + + +class QueryReturnEnumTableEnum(enum.StrEnum): + K = "k" + L = "l" + + +class QueryReturnFullTableEnum(enum.StrEnum): + E = "e" + F = "f" + + +class QueryReturnStructEnumTableEnum(enum.StrEnum): + K = "k" + L = "l" + + +class QuerySqlcEmbedEnum(enum.StrEnum): + M = "m" + N = "n" + + +@dataclasses.dataclass() +class ArrayEnumTable: + id: int + value: Optional[List[ArrayEnum]] + + +@dataclasses.dataclass() +class QueryParamEnumTable: + id: int + other: QueryParamEnumTableEnum + value: Optional[QueryParamEnumTableEnum] + + +@dataclasses.dataclass() +class QueryReturnFullTable: + id: int + value: Optional[QueryReturnFullTableEnum] + + +@dataclasses.dataclass() +class QuerySqlcEmbedTable: + id: int + value: Optional[QuerySqlcEmbedEnum] diff --git a/internal/endtoend/testdata/omit_unused_structs/db/query.py b/internal/endtoend/testdata/omit_unused_structs/db/query.py new file mode 100644 index 0000000..6f4dca3 --- /dev/null +++ b/internal/endtoend/testdata/omit_unused_structs/db/query.py @@ -0,0 +1,168 @@ +# Code generated by sqlc. DO NOT EDIT. +# versions: +# sqlc v1.28.0 +# source: query.sql +import dataclasses +from typing import Any, AsyncIterator, Iterator, Optional + +import sqlalchemy +import sqlalchemy.ext.asyncio + +from db import models + + +QUERY_ENUM_ARRAY_TABLE = """-- name: query_enum_array_table \\:many +SELECT id, value FROM array_enum_table +""" + + +QUERY_PARAM_ENUM_TABLE = """-- name: query_param_enum_table \\:one +SELECT id, other, value FROM query_param_enum_table WHERE value = :p1 +""" + + +QUERY_PARAM_STRUCT_ENUM_TABLE = """-- name: query_param_struct_enum_table \\:one +SELECT id FROM query_param_struct_enum_table WHERE id = :p1 AND value = :p2 +""" + + +QUERY_RETURN_ENUM_TABLE = """-- name: query_return_enum_table \\:one +SELECT value FROM query_return_enum_table WHERE id = :p1 +""" + + +QUERY_RETURN_FULL_TABLE = """-- name: query_return_full_table \\:many +SELECT id, value FROM query_return_full_table +""" + + +QUERY_RETURN_STRUCT_ENUM_TABLE = """-- name: query_return_struct_enum_table \\:one +SELECT value, another FROM query_return_struct_enum_table WHERE id = :p1 +""" + + +@dataclasses.dataclass() +class query_return_struct_enum_tableRow: + value: Optional[models.QueryReturnStructEnumTableEnum] + another: Optional[int] + + +QUERY_SQLC_EMBED_TABLE = """-- name: query_sqlc_embed_table \\:one +SELECT query_sqlc_embed_table.id, query_sqlc_embed_table.value FROM query_sqlc_embed_table WHERE id = :p1 +""" + + +class Querier: + def __init__(self, conn: sqlalchemy.engine.Connection): + self._conn = conn + + def query_enum_array_table(self) -> Iterator[models.ArrayEnumTable]: + result = self._conn.execute(sqlalchemy.text(QUERY_ENUM_ARRAY_TABLE)) + for row in result: + yield models.ArrayEnumTable( + id=row[0], + value=row[1], + ) + + def query_param_enum_table(self, *, value: Optional[models.QueryParamEnumTableEnum]) -> Optional[models.QueryParamEnumTable]: + row = self._conn.execute(sqlalchemy.text(QUERY_PARAM_ENUM_TABLE), {"p1": value}).first() + if row is None: + return None + return models.QueryParamEnumTable( + id=row[0], + other=row[1], + value=row[2], + ) + + def query_param_struct_enum_table(self, *, id: int, value: Optional[models.QueryParamStructEnumTableEnum]) -> Optional[int]: + row = self._conn.execute(sqlalchemy.text(QUERY_PARAM_STRUCT_ENUM_TABLE), {"p1": id, "p2": value}).first() + if row is None: + return None + return row[0] + + def query_return_enum_table(self, *, id: int) -> Optional[Optional[models.QueryReturnEnumTableEnum]]: + row = self._conn.execute(sqlalchemy.text(QUERY_RETURN_ENUM_TABLE), {"p1": id}).first() + if row is None: + return None + return row[0] + + def query_return_full_table(self) -> Iterator[models.QueryReturnFullTable]: + result = self._conn.execute(sqlalchemy.text(QUERY_RETURN_FULL_TABLE)) + for row in result: + yield models.QueryReturnFullTable( + id=row[0], + value=row[1], + ) + + def query_return_struct_enum_table(self, *, id: int) -> Optional[query_return_struct_enum_tableRow]: + row = self._conn.execute(sqlalchemy.text(QUERY_RETURN_STRUCT_ENUM_TABLE), {"p1": id}).first() + if row is None: + return None + return query_return_struct_enum_tableRow( + value=row[0], + another=row[1], + ) + + def query_sqlc_embed_table(self, *, id: int) -> Optional[Optional[Any]]: + row = self._conn.execute(sqlalchemy.text(QUERY_SQLC_EMBED_TABLE), {"p1": id}).first() + if row is None: + return None + return row[0] + + +class AsyncQuerier: + def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection): + self._conn = conn + + async def query_enum_array_table(self) -> AsyncIterator[models.ArrayEnumTable]: + result = await self._conn.stream(sqlalchemy.text(QUERY_ENUM_ARRAY_TABLE)) + async for row in result: + yield models.ArrayEnumTable( + id=row[0], + value=row[1], + ) + + async def query_param_enum_table(self, *, value: Optional[models.QueryParamEnumTableEnum]) -> Optional[models.QueryParamEnumTable]: + row = (await self._conn.execute(sqlalchemy.text(QUERY_PARAM_ENUM_TABLE), {"p1": value})).first() + if row is None: + return None + return models.QueryParamEnumTable( + id=row[0], + other=row[1], + value=row[2], + ) + + async def query_param_struct_enum_table(self, *, id: int, value: Optional[models.QueryParamStructEnumTableEnum]) -> Optional[int]: + row = (await self._conn.execute(sqlalchemy.text(QUERY_PARAM_STRUCT_ENUM_TABLE), {"p1": id, "p2": value})).first() + if row is None: + return None + return row[0] + + async def query_return_enum_table(self, *, id: int) -> Optional[Optional[models.QueryReturnEnumTableEnum]]: + row = (await self._conn.execute(sqlalchemy.text(QUERY_RETURN_ENUM_TABLE), {"p1": id})).first() + if row is None: + return None + return row[0] + + async def query_return_full_table(self) -> AsyncIterator[models.QueryReturnFullTable]: + result = await self._conn.stream(sqlalchemy.text(QUERY_RETURN_FULL_TABLE)) + async for row in result: + yield models.QueryReturnFullTable( + id=row[0], + value=row[1], + ) + + async def query_return_struct_enum_table(self, *, id: int) -> Optional[query_return_struct_enum_tableRow]: + row = (await self._conn.execute(sqlalchemy.text(QUERY_RETURN_STRUCT_ENUM_TABLE), {"p1": id})).first() + if row is None: + return None + return query_return_struct_enum_tableRow( + value=row[0], + another=row[1], + ) + + async def query_sqlc_embed_table(self, *, id: int) -> Optional[Optional[Any]]: + row = (await self._conn.execute(sqlalchemy.text(QUERY_SQLC_EMBED_TABLE), {"p1": id})).first() + if row is None: + return None + return row[0] diff --git a/internal/endtoend/testdata/omit_unused_structs/query.sql b/internal/endtoend/testdata/omit_unused_structs/query.sql new file mode 100644 index 0000000..60104c8 --- /dev/null +++ b/internal/endtoend/testdata/omit_unused_structs/query.sql @@ -0,0 +1,20 @@ +-- name: query_return_full_table :many +SELECT * FROM query_return_full_table; + +-- name: query_param_enum_table :one +SELECT * FROM query_param_enum_table WHERE value = $1; + +-- name: query_param_struct_enum_table :one +SELECT id FROM query_param_struct_enum_table WHERE id = $1 AND value = $2; + +-- name: query_return_enum_table :one +SELECT value FROM query_return_enum_table WHERE id = $1; + +-- name: query_return_struct_enum_table :one +SELECT value, another FROM query_return_struct_enum_table WHERE id = $1; + +-- name: query_sqlc_embed_table :one +SELECT sqlc.embed(query_sqlc_embed_table) FROM query_sqlc_embed_table WHERE id = $1; + +-- name: query_enum_array_table :many +SELECT * FROM array_enum_table; diff --git a/internal/endtoend/testdata/omit_unused_structs/schema.sql b/internal/endtoend/testdata/omit_unused_structs/schema.sql new file mode 100644 index 0000000..4ce5921 --- /dev/null +++ b/internal/endtoend/testdata/omit_unused_structs/schema.sql @@ -0,0 +1,69 @@ +CREATE TYPE unused_enum AS ENUM ( + 'a', 'b' +); + +CREATE TYPE unused_table_enum AS ENUM ( + 'c', 'd' +); +CREATE TABLE unused_table ( + id INTEGER PRIMARY KEY, + value unused_table_enum +); + +CREATE TYPE query_return_full_table_enum AS ENUM ( + 'e', 'f' +); +CREATE TABLE query_return_full_table ( + id INTEGER PRIMARY KEY, + value query_return_full_table_enum +); + +CREATE TYPE query_param_enum_table_enum AS ENUM ( + 'g', 'h' +); +CREATE TABLE query_param_enum_table ( + id INTEGER PRIMARY KEY, + other query_param_enum_table_enum NOT NULL, + value query_param_enum_table_enum +); + +CREATE TYPE query_param_struct_enum_table_enum AS ENUM ( + 'i', 'j' +); +CREATE TABLE query_param_struct_enum_table ( + id INTEGER PRIMARY KEY, + value query_param_struct_enum_table_enum +); + +CREATE TYPE query_return_enum_table_enum AS ENUM ( + 'k', 'l' +); +CREATE TABLE query_return_enum_table ( + id INTEGER PRIMARY KEY, + value query_return_enum_table_enum +); + +CREATE TYPE query_return_struct_enum_table_enum AS ENUM ( + 'k', 'l' +); +CREATE TABLE query_return_struct_enum_table ( + id INTEGER PRIMARY KEY, + value query_return_struct_enum_table_enum, + another INTEGER +); + +CREATE TYPE query_sqlc_embed_enum AS ENUM ( + 'm', 'n' +); +CREATE TABLE query_sqlc_embed_table ( + id INTEGER PRIMARY KEY, + value query_sqlc_embed_enum +); + +CREATE TYPE array_enum AS ENUM ( + 'o', 'p' +); +CREATE TABLE array_enum_table ( + id INTEGER PRIMARY KEY, + value array_enum[] +); diff --git a/internal/endtoend/testdata/omit_unused_structs/sqlc.yaml b/internal/endtoend/testdata/omit_unused_structs/sqlc.yaml new file mode 100644 index 0000000..cc4ec11 --- /dev/null +++ b/internal/endtoend/testdata/omit_unused_structs/sqlc.yaml @@ -0,0 +1,19 @@ +version: "2" +plugins: + - name: py + wasm: + url: file://../../../../bin/sqlc-gen-python.wasm + sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" +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 + emit_str_enum: true + omit_unused_structs: true diff --git a/internal/gen.go b/internal/gen.go index 6e50fae..982426a 100644 --- a/internal/gen.go +++ b/internal/gen.go @@ -49,10 +49,25 @@ func (t pyType) Annotation() *pyast.Node { return ann } +// typeNames returns InnerType plus the models.-stripped form. +// Optional/List are flags on pyType, not wrappers in the name. +func (t pyType) typeNames() []string { + if t.InnerType == "" { + return nil + } + names := []string{t.InnerType} + if bare := strings.TrimPrefix(t.InnerType, "models."); bare != t.InnerType { + names = append(names, bare) + } + return names +} + type Field struct { Name string Type pyType Comment string + // Embed is the table model behind sqlc.embed; used only by omit_unused_structs. + Embed *Struct } type Struct struct { @@ -67,6 +82,8 @@ type QueryValue struct { Name string Struct *Struct Typ pyType + // Embed is the table model behind a single-column sqlc.embed return. + Embed *Struct } func (v QueryValue) Annotation() *pyast.Node { @@ -321,7 +338,24 @@ type pyColumn struct { *plugin.Column } -func columnsToStruct(req *plugin.GenerateRequest, name string, columns []pyColumn) *Struct { +func findEmbedModel(structs []Struct, embed *plugin.Identifier, defaultSchema string) *Struct { + if embed == nil { + return nil + } + embedSchema := defaultSchema + if embed.Schema != "" { + embedSchema = embed.Schema + } + for i := range structs { + s := &structs[i] + if embed.Name == s.Table.Name && embedSchema == s.Table.Schema { + return s + } + } + return nil +} + +func columnsToStruct(req *plugin.GenerateRequest, name string, columns []pyColumn, structs []Struct) *Struct { gs := Struct{ Name: name, } @@ -343,8 +377,9 @@ func columnsToStruct(req *plugin.GenerateRequest, name string, columns []pyColum fieldName = fmt.Sprintf("%s_%d", fieldName, suffix) } gs.Fields = append(gs.Fields, Field{ - Name: fieldName, - Type: makePyType(req, c.Column), + Name: fieldName, + Type: makePyType(req, c.Column), + Embed: findEmbedModel(structs, c.EmbedTable, req.Catalog.DefaultSchema), }) seen[colName]++ } @@ -406,7 +441,7 @@ 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(req, query.Name+"Params", cols, structs), }} } else { args := make([]QueryValue, 0, len(query.Params)) @@ -422,8 +457,9 @@ func buildQueries(conf Config, req *plugin.GenerateRequest, structs []Struct) ([ if len(query.Columns) == 1 { c := query.Columns[0] gq.Ret = QueryValue{ - Name: columnName(c, 0), - Typ: makePyType(req, c), + Name: columnName(c, 0), + Typ: makePyType(req, c), + Embed: findEmbedModel(structs, c.EmbedTable, req.Catalog.DefaultSchema), } } else if len(query.Columns) > 1 { var gs *Struct @@ -461,7 +497,7 @@ func buildQueries(conf Config, req *plugin.GenerateRequest, structs []Struct) ([ Column: c, }) } - gs = columnsToStruct(req, query.Name+"Row", columns) + gs = columnsToStruct(req, query.Name+"Row", columns, structs) emit = true } gq.Ret = QueryValue{ @@ -1081,6 +1117,72 @@ func HashComment(s string) string { return "# " + strings.ReplaceAll(s, "\n", "\n# ") } +func filterUnusedStructs(enums []Enum, models []Struct, queries []Query) ([]Enum, []Struct) { + keepTypes := make(map[string]struct{}) + + keep := func(t string) { + keepTypes[t] = struct{}{} + if bare := strings.TrimPrefix(t, "models."); bare != t { + keepTypes[bare] = struct{}{} + } + } + + keepPyType := func(t pyType) { + for _, name := range t.typeNames() { + keep(name) + } + } + + var keepStruct func(*Struct) + keepStruct = func(s *Struct) { + if s == nil { + return + } + keep(s.Name) + for _, f := range s.Fields { + keepPyType(f.Type) + keepStruct(f.Embed) + } + } + + keepValue := func(v QueryValue) { + if v.isEmpty() { + return + } + if v.IsStruct() { + keepStruct(v.Struct) + } else { + keepPyType(v.Typ) + } + keepStruct(v.Embed) + } + + for _, query := range queries { + for _, arg := range query.Args { + keepValue(arg) + } + keepValue(query.Ret) + } + + keepEnums := make([]Enum, 0, len(enums)) + for _, enum := range enums { + _, keepBare := keepTypes[enum.Name] + _, keepQualified := keepTypes["models."+enum.Name] + if keepBare || keepQualified { + keepEnums = append(keepEnums, enum) + } + } + + keepStructs := make([]Struct, 0, len(models)) + for _, st := range models { + if _, ok := keepTypes[st.Name]; ok { + keepStructs = append(keepStructs, st) + } + } + + return keepEnums, keepStructs +} + func Generate(_ context.Context, req *plugin.GenerateRequest) (*plugin.GenerateResponse, error) { var conf Config if len(req.PluginOptions) > 0 { @@ -1096,6 +1198,10 @@ func Generate(_ context.Context, req *plugin.GenerateRequest) (*plugin.GenerateR return nil, err } + if conf.OmitUnusedStructs { + enums, models = filterUnusedStructs(enums, models, queries) + } + i := &importer{ Models: models, Queries: queries,