fix: alembic column rename (#1141)

* fix: alembic column rename

* fix system test

* correct system test to test visit_column_name() and visit_column_type()

* add unit tests

* add unit tests

---------

Co-authored-by: Jacob Hayes <[email protected]>

Author: Linchin

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 <element type>, but bigquery requires this syntax:
         SET DATA TYPE <element 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_),

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")]
 

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"