From: Federico Caselli Date: Sat, 8 Feb 2025 11:57:31 +0000 (+0100) Subject: Render labels in autogenerate index X-Git-Tag: rel_1_15_0~5^2 X-Git-Url: http://git.ipfire.org/gitweb/gitweb.cgi?a=commitdiff_plain;h=762b2d2b11a24a37ce87820c265d42701683544b;p=thirdparty%2Fsqlalchemy%2Falembic.git Render labels in autogenerate index Index autogenerate will now render labels for expressions that use them. This is useful when applying operator classes in PostgreSQL that can be keyed on the label name. Fixes: #1603 Change-Id: I187a944a5021e643e264ee1ec97807f6573e5f2f --- diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py index 6ebfbf9f..7357c1d5 100644 --- a/alembic/autogenerate/render.py +++ b/alembic/autogenerate/render.py @@ -19,6 +19,7 @@ from sqlalchemy import schema as sa_schema from sqlalchemy import sql from sqlalchemy import types as sqltypes from sqlalchemy.sql.elements import conv +from sqlalchemy.sql.elements import Label from sqlalchemy.sql.elements import quoted_name from .. import util @@ -584,23 +585,28 @@ def _render_potential_expr( value: Any, autogen_context: AutogenContext, *, - wrap_in_text: bool = True, + wrap_in_element: bool = True, is_server_default: bool = False, is_index: bool = False, ) -> str: if isinstance(value, sql.ClauseElement): - if wrap_in_text: - template = "%(prefix)stext(%(sql)r)" + sql_text = autogen_context.migration_context.impl.render_ddl_sql_expr( + value, is_server_default=is_server_default, is_index=is_index + ) + if wrap_in_element: + prefix = _sqlalchemy_autogenerate_prefix(autogen_context) + element = "literal_column" if is_index else "text" + value_str = f"{prefix}{element}({sql_text!r})" + if ( + is_index + and isinstance(value, Label) + and type(value.name) is str + ): + return value_str + f".label({value.name!r})" + else: + return value_str else: - template = "%(sql)r" - - return template % { - "prefix": _sqlalchemy_autogenerate_prefix(autogen_context), - "sql": autogen_context.migration_context.impl.render_ddl_sql_expr( - value, is_server_default=is_server_default, is_index=is_index - ), - } - + return repr(sql_text) else: return repr(value) @@ -787,7 +793,7 @@ def _render_computed( computed: Computed, autogen_context: AutogenContext ) -> str: text = _render_potential_expr( - computed.sqltext, autogen_context, wrap_in_text=False + computed.sqltext, autogen_context, wrap_in_element=False ) kwargs = {} @@ -1101,7 +1107,7 @@ def _render_check_constraint( else "" ), "sqltext": _render_potential_expr( - constraint.sqltext, autogen_context, wrap_in_text=False + constraint.sqltext, autogen_context, wrap_in_element=False ), } diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py index 60aa1536..2623308f 100644 --- a/alembic/ddl/postgresql.py +++ b/alembic/ddl/postgresql.py @@ -846,5 +846,5 @@ def _render_potential_column( return render._render_potential_expr( value, autogen_context, - wrap_in_text=isinstance(value, (TextClause, FunctionElement)), + wrap_in_element=isinstance(value, (TextClause, FunctionElement)), ) diff --git a/alembic/op.pyi b/alembic/op.pyi index 92044469..d86bef46 100644 --- a/alembic/op.pyi +++ b/alembic/op.pyi @@ -27,7 +27,6 @@ if TYPE_CHECKING: from sqlalchemy.sql.elements import conv from sqlalchemy.sql.elements import TextClause from sqlalchemy.sql.expression import TableClause - from sqlalchemy.sql.functions import Function from sqlalchemy.sql.schema import Column from sqlalchemy.sql.schema import Computed from sqlalchemy.sql.schema import Identity @@ -650,7 +649,7 @@ def create_foreign_key( def create_index( index_name: Optional[str], table_name: str, - columns: Sequence[Union[str, TextClause, Function[Any]]], + columns: Sequence[Union[str, TextClause, ColumnElement[Any]]], *, schema: Optional[str] = None, unique: bool = False, diff --git a/alembic/operations/base.py b/alembic/operations/base.py index 9b52fa6f..456d1c75 100644 --- a/alembic/operations/base.py +++ b/alembic/operations/base.py @@ -43,7 +43,6 @@ if TYPE_CHECKING: from sqlalchemy.sql.expression import ColumnElement from sqlalchemy.sql.expression import TableClause from sqlalchemy.sql.expression import TextClause - from sqlalchemy.sql.functions import Function from sqlalchemy.sql.schema import Column from sqlalchemy.sql.schema import Computed from sqlalchemy.sql.schema import Identity @@ -1074,7 +1073,7 @@ class Operations(AbstractOperations): self, index_name: Optional[str], table_name: str, - columns: Sequence[Union[str, TextClause, Function[Any]]], + columns: Sequence[Union[str, TextClause, ColumnElement[Any]]], *, schema: Optional[str] = None, unique: bool = False, diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py index 60b856a8..bb4d825b 100644 --- a/alembic/operations/ops.py +++ b/alembic/operations/ops.py @@ -35,7 +35,6 @@ if TYPE_CHECKING: from sqlalchemy.sql.elements import conv from sqlalchemy.sql.elements import quoted_name from sqlalchemy.sql.elements import TextClause - from sqlalchemy.sql.functions import Function from sqlalchemy.sql.schema import CheckConstraint from sqlalchemy.sql.schema import Column from sqlalchemy.sql.schema import Computed @@ -933,7 +932,7 @@ class CreateIndexOp(MigrateOperation): operations: Operations, index_name: Optional[str], table_name: str, - columns: Sequence[Union[str, TextClause, Function[Any]]], + columns: Sequence[Union[str, TextClause, ColumnElement[Any]]], *, schema: Optional[str] = None, unique: bool = False, diff --git a/docs/build/unreleased/1603.rst b/docs/build/unreleased/1603.rst new file mode 100644 index 00000000..8393ecd4 --- /dev/null +++ b/docs/build/unreleased/1603.rst @@ -0,0 +1,7 @@ +.. change:: + :tags: usecase, autogenerate + :tickets: 1603 + + Index autogenerate will now render labels for expressions + that use them. This is useful when applying operator classes + in PostgreSQL that can be keyed on the label name. diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py index b5c4e572..f466da3c 100644 --- a/tests/test_autogen_render.py +++ b/tests/test_autogen_render.py @@ -94,6 +94,32 @@ class AutogenRenderTest(TestBase): "['active', 'code'], unique=False)", ) + def test_render_add_index_fn(self): + t = self.table(Column("other", String(100))) + idx = Index("test_fn_idx", t.c.code + t.c.other) + op_obj = ops.CreateIndexOp.from_index(idx) + eq_ignore_whitespace( + autogenerate.render_op_text(self.autogen_context, op_obj), + "op.create_index('test_fn_idx', 'test', " + "[sa.literal_column('code || other')], unique=False)", + ) + + def test_render_add_index_label(self): + t = self.table(Column("other", String(100))) + idx = Index( + "test_fn_idx", + (t.c.code + t.c.other).label("foo"), + t.c.id.label("bar"), + ) + op_obj = ops.CreateIndexOp.from_index(idx) + eq_ignore_whitespace( + autogenerate.render_op_text(self.autogen_context, op_obj), + "op.create_index('test_fn_idx', 'test', [" + "sa.literal_column('code || other').label('foo'), " + "sa.literal_column('id').label('bar')" + "], unique=False)", + ) + def test_render_add_index_if_not_exists(self): """ autogenerate.render._add_index @@ -170,7 +196,7 @@ class AutogenRenderTest(TestBase): eq_ignore_whitespace( autogenerate.render_op_text(self.autogen_context, op_obj), "op.create_index('test_active_code_idx', 'test', " - "['active', sa.text('lower(code)')], unique=False)", + "['active', sa.literal_column('lower(code)')], unique=False)", ) op_obj_rev = op_obj.reverse() eq_ignore_whitespace( @@ -186,7 +212,7 @@ class AutogenRenderTest(TestBase): eq_ignore_whitespace( autogenerate.render_op_text(self.autogen_context, op_obj), "op.create_index('test_lower_code_idx', 'test', " - "[sa.text('lower(code)')], unique=False)", + "[sa.literal_column('lower(code)')], unique=False)", ) op_obj_rev = op_obj.reverse() eq_ignore_whitespace( @@ -202,7 +228,7 @@ class AutogenRenderTest(TestBase): eq_ignore_whitespace( autogenerate.render_op_text(self.autogen_context, op_obj), "op.create_index('test_lower_code_idx', 'test', " - "[sa.text('CAST(code AS VARCHAR)')], unique=False)", + "[sa.literal_column('CAST(code AS VARCHAR)')], unique=False)", ) def test_render_add_index_desc(self): @@ -212,7 +238,7 @@ class AutogenRenderTest(TestBase): eq_ignore_whitespace( autogenerate.render_op_text(self.autogen_context, op_obj), "op.create_index('test_desc_code_idx', 'test', " - "[sa.text('code DESC')], unique=False)", + "[sa.literal_column('code DESC')], unique=False)", ) def test_drop_index(self): @@ -256,7 +282,7 @@ class AutogenRenderTest(TestBase): eq_ignore_whitespace( autogenerate.render_op_text(self.autogen_context, op_obj_rev), "op.create_index('test_active_code_idx', 'test', " - "['active', sa.text('lower(code)')], unique=False)", + "['active', sa.literal_column('lower(code)')], unique=False)", ) def test_drop_index_func(self): @@ -274,7 +300,7 @@ class AutogenRenderTest(TestBase): eq_ignore_whitespace( autogenerate.render_op_text(self.autogen_context, op_obj_rev), "op.create_index('test_lower_code_idx', 'test', " - "[sa.text('lower(code)')], unique=False)", + "[sa.literal_column('lower(code)')], unique=False)", ) @testing.emits_warning("Can't validate argument ") diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index e42ea9d3..9eec5f26 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -1330,7 +1330,7 @@ class PostgresqlAutogenRenderTest(TestBase): ops.CreateIndexOp.from_index(idx), ), "op.create_index('my_idx', 'tbl', " - "[sa.text(\"(c ->> 'foo')\")], unique=False)", + "[sa.literal_column(\"(c ->> 'foo')\")], unique=False)", ) @config.requirements.nulls_not_distinct_sa