Skip to content

Commit dbf0498

Browse files
jaogoypre-commit-ci[bot]agronholm
authored
Add support for rendering dialect kwargs and info, and introduce keep-dialect-types option (#438)
* [Add support for rendering dialect kwargs and info; introduce keep-dialect-types option ### Summary - Render `Table`/`Column` dialect-specific kwargs and `info` in generated code. - Add `keep-dialect-types` generator option to preserve dialect-specific column types instead of adapting to generic SQLAlchemy types. ### Motivation - Some dialect features (e.g., engine/storage params, partitioning, materialized view metadata) are carried in SQLAlchemy `dialect_kwargs` or `info` but weren’t emitted by sqlacodegen. - For custom dialects with their own Column types, the current type “adaptation” step collapses them into generic SQLAlchemy types, losing fidelity. ### What’s changed - Render dialect kwargs and info - Table (Core): include table-level `dialect_kwargs` and `info` in `Table(...)` kwargs. - Column (Core/ORM): include column-level `dialect_kwargs` and `info` in `Column(...)` / `mapped_column(...)`. - Declarative: include table-level `dialect_kwargs` and `info` in `__table_args__` dict. - New option: keep-dialect-types - Generator option: `keep_dialect_types`. - CLI flag: `--options=keep_dialect_types`. - Behavior: gates only the type adaptation step in `fix_column_types()`; Boolean/Enum inference and PostgreSQL sequence handling remain intact. ### Usage examples - Core (Table) ```python t_orders = Table( 'orders', metadata, Column('id', INTEGER, primary_key=True, starrocks_aggr_type='SUM'), schema='public', comment='orders table', info={'table_kind': 'view'} ) ``` - Declarative (`__table_args__`) ```python class Orders(Base): __tablename__ = 'orders' __table_args__ = ( { 'schema': 'public', 'comment': 'orders table', 'info': {'table_kind': 'view'}, 'starrocks_properties': {'replication_num': '1'}, } ) id: Mapped[INTEGER] = mapped_column(primary_key=True, starrocks_aggr_type='SUM') ``` - Preserve dialect types ```bash sqlacodegen postgresql://... --options=keep_dialect_types ``` When the flag is set, dialect-specific types (e.g., `starrocks.datatype.INTEGER`) are preserved and imported from their original module rather than adapted to generic SQLAlchemy types. ### Backward compatibility - Default behavior is unchanged: - If `keep_dialect_types` is not set, type adaptation continues as before. - Output remains stable except that existing `info`/`dialect_kwargs` now render when present (additive enhancement). - No breaking API changes. Signed-off-by: jaogoy <jaogoy@gmail.com> * Add support for rendering dialect kwargs and info; introduce option ### Summary - Render `Table`/`Column` dialect-specific kwargs and `info` in generated code. - Add `include-dialect-options` and `keep-dialect-types` option to render dialect options and preserve dialect-specific column types instead of adapting to generic SQLAlchemy types. ### Motivation - Some dialect features (e.g., engine/storage params, partitioning, materialized view metadata) are carried in SQLAlchemy `dialect_kwargs` or `info` but weren’t emitted by sqlacodegen. - For custom dialects with their own Column types, the current type “adaptation” step collapses them into generic SQLAlchemy types, losing fidelity. ### What’s changed - Render dialect kwargs and info - Table (Core): include table-level `dialect_kwargs` and `info` in `Table(...)` kwargs. - Column (Core/ORM): include column-level `dialect_kwargs` and `info` in `Column(...)` / `mapped_column(...)`. - Declarative: include table-level `dialect_kwargs` and `info` in `__table_args__` dict. - New option: include-dialect-options, keep-dialect-types - Behavior: gates only the type adaptation step in `fix_column_types()`; Boolean/Enum inference and PostgreSQL sequence handling remain intact. ### Usage examples - Core (Table) ```python t_orders = Table( 'orders', metadata, Column('id', INTEGER, primary_key=True, starrocks_aggr_type='SUM'), schema='public', comment='orders table', info={'table_kind': 'view'} ) ``` - Declarative (`__table_args__`) ```python class Orders(Base): __tablename__ = 'orders' __table_args__ = ( { 'schema': 'public', 'comment': 'orders table', 'info': {'table_kind': 'view'}, 'starrocks_properties': {'replication_num': '1'}, } ) id: Mapped[INTEGER] = mapped_column(primary_key=True, starrocks_aggr_type='SUM') ``` - Preserve dialect types ```bash sqlacodegen postgresql://... --options=include-dialect-options,keep-dialect-types ``` When the `include-dialect-options` flag is set, it will render the dialect options (e.g. `starrocks_properties`) and `info` dict (which will store the `table_kind`, to indicate it's a table or a view/mv, and the `definition` for a view/mv). When the `keep-dialect-types` flag is set, dialect-specific types (e.g., `starrocks.datatype.INTEGER`) are preserved and imported from their original module rather than adapted to generic SQLAlchemy types. ### Backward compatibility - Default behavior is unchanged: - If `include-dialect-options` is not set, type dialect options and info will not be rendered as before. - If `keep-dialect-types` is not set, type adaptation continues as before. - No breaking API changes. Signed-off-by: jaogoy <jaogoy@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add test cases for dialect options and readme Signed-off-by: jaogoy <jaogoy@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix pre-commit format Signed-off-by: jaogoy <jaogoy@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix format for readme and change Signed-off-by: jaogoy <jaogoy@gmail.com> * fix format for change Signed-off-by: jaogoy <jaogoy@gmail.com> * Wrapped the changelog lines * remove redundent options, and refine format Signed-off-by: jaogoy <jaogoy@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * use mock dialect in test only remove importing StarRocksDialect in test Signed-off-by: jaogoy <jaogoy@gmail.com> * change option name to be underscores about include_dialect_options and keep_dialect_types Signed-off-by: jaogoy <jaogoy@gmail.com> * fix option names in README and CHANGES fix `include_dialect_options` and `keep_dialect_types` Signed-off-by: jaogoy <jaogoy@gmail.com> --------- Signed-off-by: jaogoy <jaogoy@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>
1 parent d7a6024 commit dbf0498

5 files changed

Lines changed: 341 additions & 7 deletions

File tree

‎CHANGES.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ Version history
77
- Fix Postgres ``DOMAIN`` adaptation regression introduced in SQLAlchemy 2.0.42 (PR by @sheinbergon)
88
- Support disabling special naming logic for single column many-to-one and one-to-one relationships
99
(PR by @Henkhogan, revised by @sheinbergon)
10+
- Add ``include_dialect_options`` option to render ``Table`` and ``Column``
11+
dialect-specific kwargs and ``info`` in generated code. (PR by @jaogoy)
12+
- Add ``keep_dialect_types`` option to preserve dialect-specific column types instead of
13+
adapting to generic SQLAlchemy types. (PR by @jaogoy)
1014

1115
**3.1.1**
1216

‎README.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,8 @@ values must be delimited by commas, e.g. ``--options noconstraints,nobidi``):
105105
* ``noindexes``: ignore indexes
106106
* ``noidsuffix``: prevent the special naming logic for single column many-to-one
107107
and one-to-one relationships (see `Relationship naming logic`_ for details)
108+
* ``include_dialect_options``: render a table' dialect options, such as ``starrocks_partition`` for StarRocks' specific options.
109+
* ``keep_dialect_types``: preserve dialect-specific column types instead of adapting to generic SQLAlchemy types.
108110

109111
* ``declarative``
110112

‎src/sqlacodegen/generators.py‎

Lines changed: 77 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,13 @@ def generate(self) -> str:
119119

120120
@dataclass(eq=False)
121121
class TablesGenerator(CodeGenerator):
122-
valid_options: ClassVar[set[str]] = {"noindexes", "noconstraints", "nocomments"}
122+
valid_options: ClassVar[set[str]] = {
123+
"noindexes",
124+
"noconstraints",
125+
"nocomments",
126+
"include_dialect_options",
127+
"keep_dialect_types",
128+
}
123129
stdlib_module_names: ClassVar[set[str]] = get_stdlib_module_names()
124130

125131
def __init__(
@@ -135,6 +141,13 @@ def __init__(
135141
self.imports: dict[str, set[str]] = defaultdict(set)
136142
self.module_imports: set[str] = set()
137143

144+
# Render SchemaItem.info and dialect kwargs (Table/Column) into output
145+
self.include_dialect_options_and_info: bool = (
146+
"include_dialect_options" in self.options
147+
)
148+
# Keep dialect-specific types instead of adapting to generic SQLAlchemy types
149+
self.keep_dialect_types: bool = "keep_dialect_types" in self.options
150+
138151
@property
139152
def views_supported(self) -> bool:
140153
return True
@@ -393,6 +406,10 @@ def render_table(self, table: Table) -> str:
393406
if table_comment:
394407
kwargs["comment"] = repr(table.comment)
395408

409+
# add info + dialect kwargs for callable context (opt-in)
410+
if self.include_dialect_options_and_info:
411+
self._add_dialect_kwargs_and_info(table, kwargs, values_for_dict=False)
412+
396413
return render_callable("Table", *args, kwargs=kwargs, indentation=" ")
397414

398415
def render_index(self, index: Index) -> str:
@@ -498,6 +515,10 @@ def render_column(
498515
if comment:
499516
kwargs["comment"] = repr(comment)
500517

518+
# add column info + dialect kwargs for callable context (opt-in)
519+
if self.include_dialect_options_and_info:
520+
self._add_dialect_kwargs_and_info(column, kwargs, values_for_dict=False)
521+
501522
return self.render_column_callable(is_table, *args, **kwargs)
502523

503524
def render_column_callable(self, is_table: bool, *args: Any, **kwargs: Any) -> str:
@@ -615,6 +636,51 @@ def add_fk_options(*opts: Any) -> None:
615636

616637
return render_callable(constraint.__class__.__name__, *args, kwargs=kwargs)
617638

639+
def _add_dialect_kwargs_and_info(
640+
self, obj: Any, target_kwargs: dict[str, object], *, values_for_dict: bool
641+
) -> None:
642+
"""
643+
Merge SchemaItem-like object's .info and .dialect_kwargs into target_kwargs.
644+
- values_for_dict=True: keep raw values so pretty-printer emits repr() (for __table_args__ dict)
645+
- values_for_dict=False: set values to repr() strings (for callable kwargs)
646+
"""
647+
info_dict = getattr(obj, "info", None)
648+
if info_dict:
649+
target_kwargs["info"] = info_dict if values_for_dict else repr(info_dict)
650+
651+
dialect_keys: list[str]
652+
try:
653+
dialect_keys = sorted(getattr(obj, "dialect_kwargs"))
654+
except Exception:
655+
return
656+
657+
dialect_kwargs = getattr(obj, "dialect_kwargs", {})
658+
for key in dialect_keys:
659+
try:
660+
value = dialect_kwargs[key]
661+
except Exception:
662+
continue
663+
664+
# Render values:
665+
# - callable context (values_for_dict=False): produce a string expression.
666+
# primitives use repr(value); custom objects stringify then repr().
667+
# - dict context (values_for_dict=True): pass raw primitives / str;
668+
# custom objects become str(value) so pformat quotes them.
669+
if values_for_dict:
670+
if isinstance(value, type(None) | bool | int | float):
671+
target_kwargs[key] = value
672+
elif isinstance(value, str | dict | list):
673+
target_kwargs[key] = value
674+
else:
675+
target_kwargs[key] = str(value)
676+
else:
677+
if isinstance(
678+
value, type(None) | bool | int | float | str | dict | list
679+
):
680+
target_kwargs[key] = repr(value)
681+
else:
682+
target_kwargs[key] = repr(str(value))
683+
618684
def should_ignore_table(self, table: Table) -> bool:
619685
# Support for Alembic and sqlalchemy-migrate -- never expose the schema version
620686
# tables
@@ -680,10 +746,11 @@ def fix_column_types(self, table: Table) -> None:
680746
continue
681747

682748
for column in table.c:
683-
try:
684-
column.type = self.get_adapted_type(column.type)
685-
except CompileError:
686-
pass
749+
if not self.keep_dialect_types:
750+
try:
751+
column.type = self.get_adapted_type(column.type)
752+
except CompileError:
753+
continue
687754

688755
# PostgreSQL specific fix: detect sequences from server_default
689756
if column.server_default and self.bind.dialect.name == "postgresql":
@@ -1193,7 +1260,7 @@ def render_class_variables(self, model: ModelClass) -> str:
11931260

11941261
def render_table_args(self, table: Table) -> str:
11951262
args: list[str] = []
1196-
kwargs: dict[str, str] = {}
1263+
kwargs: dict[str, object] = {}
11971264

11981265
# Render constraints
11991266
for constraint in sorted(table.constraints, key=get_constraint_sort_key):
@@ -1219,6 +1286,10 @@ def render_table_args(self, table: Table) -> str:
12191286
if table.comment:
12201287
kwargs["comment"] = table.comment
12211288

1289+
# add info + dialect kwargs for dict context (__table_args__) (opt-in)
1290+
if self.include_dialect_options_and_info:
1291+
self._add_dialect_kwargs_and_info(table, kwargs, values_for_dict=True)
1292+
12221293
if kwargs:
12231294
formatted_kwargs = pformat(kwargs)
12241295
if not args:

‎tests/test_generator_declarative.py‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,180 @@ class SimpleItems(Base):
111111
)
112112

113113

114+
@pytest.mark.parametrize("generator", [["include_dialect_options"]], indirect=True)
115+
def test_include_dialect_options_and_info_table_and_column(
116+
generator: CodeGenerator,
117+
) -> None:
118+
from .test_generator_tables import _PartitionInfo
119+
120+
Table(
121+
"t_opts",
122+
generator.metadata,
123+
Column("id", INTEGER, primary_key=True, starrocks_is_agg_key=True),
124+
Column("name", VARCHAR, starrocks_agg_type="REPLACE"),
125+
starrocks_aggregate_key="id",
126+
starrocks_partition_by=_PartitionInfo("RANGE(id)"),
127+
starrocks_security="DEFINER",
128+
starrocks_PROPERTIES={"replication_num": "3", "storage_medium": "SSD"},
129+
info={
130+
"table_kind": "MATERIALIZED VIEW",
131+
"definition": "SELECT id, name FROM t_opts_base_table",
132+
},
133+
)
134+
135+
validate_code(
136+
generator.generate(),
137+
"""\
138+
from typing import Optional
139+
140+
from sqlalchemy import Integer, String
141+
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
142+
143+
class Base(DeclarativeBase):
144+
pass
145+
146+
147+
class TOpts(Base):
148+
__tablename__ = 't_opts'
149+
__table_args__ = {'info': {'definition': 'SELECT id, name FROM t_opts_base_table',
150+
'table_kind': 'MATERIALIZED VIEW'},
151+
'starrocks_PROPERTIES': {'replication_num': '3', 'storage_medium': 'SSD'},
152+
'starrocks_aggregate_key': 'id',
153+
'starrocks_partition_by': 'RANGE(id)',
154+
'starrocks_security': 'DEFINER'}
155+
156+
id: Mapped[int] = mapped_column(Integer, primary_key=True, starrocks_is_agg_key=True)
157+
name: Mapped[Optional[str]] = mapped_column(String, starrocks_agg_type='REPLACE')
158+
""",
159+
)
160+
161+
162+
@pytest.mark.parametrize("generator", [["include_dialect_options"]], indirect=True)
163+
def test_include_dialect_options_and_info_with_hyphen(generator: CodeGenerator) -> None:
164+
Table(
165+
"t_opts2",
166+
generator.metadata,
167+
Column("id", INTEGER, primary_key=True),
168+
mysql_engine="InnoDB",
169+
info={"table_kind": "View"},
170+
)
171+
172+
validate_code(
173+
generator.generate(),
174+
"""\
175+
from sqlalchemy import Integer
176+
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
177+
178+
class Base(DeclarativeBase):
179+
pass
180+
181+
182+
class TOpts2(Base):
183+
__tablename__ = 't_opts2'
184+
__table_args__ = {'info': {'table_kind': 'View'}, 'mysql_engine': 'InnoDB'}
185+
186+
id: Mapped[int] = mapped_column(Integer, primary_key=True)
187+
""",
188+
)
189+
190+
191+
def test_include_dialect_options_not_enabled_skips(generator: CodeGenerator) -> None:
192+
from .test_generator_tables import _PartitionInfo
193+
194+
Table(
195+
"t_plain",
196+
generator.metadata,
197+
Column(
198+
"id",
199+
INTEGER,
200+
primary_key=True,
201+
info={"abc": True},
202+
starrocks_is_agg_key=True,
203+
),
204+
starrocks_engine="OLAP",
205+
starrocks_partition_by=_PartitionInfo("RANGE(id)"),
206+
)
207+
208+
validate_code(
209+
generator.generate(),
210+
"""\
211+
from sqlalchemy import Integer
212+
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
213+
214+
class Base(DeclarativeBase):
215+
pass
216+
217+
218+
class TPlain(Base):
219+
__tablename__ = 't_plain'
220+
221+
id: Mapped[int] = mapped_column(Integer, primary_key=True)
222+
""",
223+
)
224+
225+
226+
def test_keep_dialect_types_adapts_mysql_integer_default(
227+
generator: CodeGenerator,
228+
) -> None:
229+
from sqlalchemy.dialects.mysql import INTEGER as MYSQL_INTEGER
230+
231+
Table(
232+
"num",
233+
generator.metadata,
234+
Column("id", INTEGER, primary_key=True),
235+
Column("val", MYSQL_INTEGER(), nullable=False),
236+
)
237+
238+
validate_code(
239+
generator.generate(),
240+
"""\
241+
from sqlalchemy import Integer
242+
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
243+
244+
class Base(DeclarativeBase):
245+
pass
246+
247+
248+
class Num(Base):
249+
__tablename__ = 'num'
250+
251+
id: Mapped[int] = mapped_column(Integer, primary_key=True)
252+
val: Mapped[int] = mapped_column(Integer, nullable=False)
253+
""",
254+
)
255+
256+
257+
@pytest.mark.parametrize("generator", [["keep_dialect_types"]], indirect=True)
258+
def test_keep_dialect_types_keeps_mysql_integer(generator: CodeGenerator) -> None:
259+
from sqlalchemy.dialects.mysql import INTEGER as MYSQL_INTEGER
260+
261+
Table(
262+
"num2",
263+
generator.metadata,
264+
Column("id", INTEGER, primary_key=True),
265+
Column("val", MYSQL_INTEGER(), nullable=False),
266+
)
267+
268+
validate_code(
269+
generator.generate(),
270+
"""\
271+
from sqlalchemy import INTEGER
272+
from sqlalchemy.dialects.mysql import INTEGER
273+
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
274+
275+
class Base(DeclarativeBase):
276+
pass
277+
278+
279+
class Num2(Base):
280+
__tablename__ = 'num2'
281+
282+
id: Mapped[int] = mapped_column(INTEGER, primary_key=True)
283+
val: Mapped[int] = mapped_column(INTEGER, nullable=False)
284+
""",
285+
)
286+
287+
114288
def test_onetomany(generator: CodeGenerator) -> None:
115289
Table(
116290
"simple_items",

0 commit comments

Comments
 (0)