diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index 879bcb40..c36ca1b1 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -1388,18 +1388,36 @@ def __init__(self, *args, **kwargs): pass else: from alembic.ddl import impl - from alembic.ddl.base import ColumnType, format_type, alter_table, alter_column + from alembic.ddl.base import ( + ColumnName, + ColumnType, + format_column_name, + format_type, + alter_table, + alter_column, + ) class SqlalchemyBigqueryImpl(impl.DefaultImpl): __dialect__ = "bigquery" + @compiles(ColumnName, "bigquery") + def visit_column_name(element: ColumnName, compiler: DDLCompiler, **kw) -> str: + """Replaces the visit_column_name() function in alembic/alembic/ddl/base.py. + See https://github.com/googleapis/python-bigquery-sqlalchemy/issues/1097""" + + return "%s RENAME COLUMN %s TO %s" % ( + alter_table(compiler, element.table_name, element.schema), + format_column_name(compiler, element.column_name), + format_column_name(compiler, element.newname), + ) + @compiles(ColumnType, "bigquery") def visit_column_type(element: ColumnType, compiler: DDLCompiler, **kw) -> str: """Replaces the visit_column_type() function in alembic/alembic/ddl/base.py. The alembic version ends in TYPE , but bigquery requires this syntax: SET DATA TYPE """ - return "%s %s %s" % ( # pragma: NO COVER + return "%s %s %s" % ( alter_table(compiler, element.table_name, element.schema), alter_column(compiler, element.column_name), "SET DATA TYPE %s" % format_type(compiler, element.type_), diff --git a/tests/system/test_alembic.py b/tests/system/test_alembic.py index 30308c68..0838be8f 100644 --- a/tests/system/test_alembic.py +++ b/tests/system/test_alembic.py @@ -66,6 +66,10 @@ def test_alembic_scenario(alembic_table): table mods within a short time. """ from alembic import op + from sqlalchemy_bigquery.base import SqlalchemyBigqueryImpl + + # Register the BigQuery Implementation to test the customized methods. + op.implementation_for(SqlalchemyBigqueryImpl) assert alembic_table("account") is None @@ -100,9 +104,11 @@ def test_alembic_scenario(alembic_table): "account", Column("last_transaction_date", DateTime, comment="when updated") ) + # Tests customized visit_column_name() + op.alter_column("account", "name", new_column_name="friendly_name") assert alembic_table("account", "schema") == [ SchemaField("id", "INTEGER", "REQUIRED"), - SchemaField("name", "STRING(50)", "REQUIRED", description="The name"), + SchemaField("friendly_name", "STRING(50)", "REQUIRED", description="The name"), SchemaField("description", "STRING(200)"), SchemaField("last_transaction_date", "DATETIME", description="when updated"), ] @@ -131,7 +137,7 @@ def test_alembic_scenario(alembic_table): assert alembic_table("account") is None assert alembic_table("accounts", "schema") == [ SchemaField("id", "INTEGER", "REQUIRED"), - SchemaField("name", "STRING(50)", "REQUIRED", description="The name"), + SchemaField("friendly_name", "STRING(50)", "REQUIRED", description="The name"), SchemaField("description", "STRING(200)"), SchemaField("last_transaction_date", "DATETIME", description="when updated"), ] @@ -162,6 +168,7 @@ def test_alembic_scenario(alembic_table): # if allowed by BigQuery's type coercion rules op.create_table("identifiers", Column("id", Integer)) + # Tests customized visit_column_type() op.alter_column("identifiers", "id", type_=Numeric) assert alembic_table("identifiers", "schema") == [SchemaField("id", "NUMERIC")] diff --git a/tests/unit/test_alembic.py b/tests/unit/test_alembic.py new file mode 100644 index 00000000..9db66933 --- /dev/null +++ b/tests/unit/test_alembic.py @@ -0,0 +1,37 @@ +# Copyright (c) 2021 The sqlalchemy-bigquery Authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +from alembic.ddl.base import ColumnName, ColumnType +from sqlalchemy import Integer + + +def test_alter_table_rename_column_bigquery(ddl_compiler): + from sqlalchemy_bigquery.base import visit_column_name + + column = ColumnName(name="t", column_name="a", newname="b") + res = visit_column_name(element=column, compiler=ddl_compiler) + assert res == "ALTER TABLE `t` RENAME COLUMN `a` TO `b`" + + +def test_alter_column_new_type_bigquery(ddl_compiler): + from sqlalchemy_bigquery.base import visit_column_type + + column = ColumnType(name="t", column_name="a", type_=Integer()) + res = visit_column_type(element=column, compiler=ddl_compiler) + assert res == "ALTER TABLE `t` ALTER COLUMN `a` SET DATA TYPE INT64"