googleapis · jimfulton · May 25, 2021 · May 20, 2021 · May 24, 2021 · May 24, 2021
diff --git a/docs/alembic.rst b/docs/alembic.rst
@@ -0,0 +1,51 @@
+Alembic support
+---------------
+
+`Alembic <https://alembic.sqlalchemy.org>`_ is a lightweight database
+migration tool for usage with the SQLAlchemy Database Toolkit for
+Python.  It can use this BigQuery SQLAlchemy support to manage
+BigQuery shemas.
+
+Some features, like management of constrains and indexes, aren't
+supported because `BigQuery doesn't support them
+<https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language>`_.
+
+Supported operations:
+
+`add_column(table_name, column, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.add_column>`_
+
+`alter_column(table_name, column_name, nullable=None, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.alter_column>`_
+
+`bulk_insert(table, rows, multiinsert=True)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.bulk_insert>`_
+
+`create_table(table_name, *columns, **kw)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.create_table>`_
+
+`create_table_comment(table_name, comment, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.create_table_comment>`_
+
+`drop_column(table_name, column_name, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.drop_column>`_
+
+`drop_table(table_name, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.drop_table>`_
+
+`drop_table_comment(table_name, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.drop_table_comment>`_
+
+`execute(sqltext, execution_options=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.execute>`_
+
+`rename_table(old_table_name, new_table_name, schema=None)
+<https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.rename_table>`_
+
+Note that some of the operations above have limited capability, again
+do to `BigQuery limitations
+<https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language>`_.
+
+The `execute` operation allows access to BigQuery-specific
+`data-definition-language
+<https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language>`_.
diff --git a/docs/index.rst b/docs/index.rst
@@ -2,6 +2,8 @@
 
 .. include:: multiprocessing.rst
 
+.. include:: alembic.rst
+
 API Reference
 -------------
 .. toctree::

diff --git a/noxfile.py b/noxfile.py
@@ -82,14 +82,30 @@ def lint_setup_py(session):
     session.run("python", "setup.py", "check", "--restructuredtext", "--strict")
 
 
+def install_alembic_for_python_38(session, constraints_path):
+    """
+    install alembic for Python 3.8 unit and system tests
+
+    We don't require alembic and most tests should run without it, however
+
+    - We run some unit tests (Python 3.8) to cover the alembic
+      registration that happens when alembic is installed.
+
+    - We have a system test that demonstrates working with alembic and
+      proves that the things we think should work do work. :)
+    """
+    if session.python == "3.8":
+        session.install("alembic", "-c", constraints_path)
+
+
 def default(session):
     # Install all test dependencies, then install this package in-place.
 
     constraints_path = str(
         CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt"
     )
     session.install("mock", "pytest", "pytest-cov", "-c", constraints_path)
-
+    install_alembic_for_python_38(session, constraints_path)
     session.install("-e", ".", "-c", constraints_path)
 
     # Run py.test against the unit tests.
@@ -142,6 +158,7 @@ def system(session):
     # Install all test dependencies, then install this package into the
     # virtualenv's dist-packages.
     session.install("mock", "pytest", "google-cloud-testutils", "-c", constraints_path)
+    install_alembic_for_python_38(session, constraints_path)
     session.install("-e", ".", "-c", constraints_path)
 
     # Run py.test against the system tests.

diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py
@@ -966,3 +966,14 @@ def get_view_definition(self, connection, view_name, schema=None, **kw):
             view_name = f"{self.dataset_id}.{view_name}"
         view = client.get_table(view_name)
         return view.view_query
+
+
+try:
+    import alembic  # noqa
+except ImportError:
+    pass
+else:
+    from alembic.ddl import impl
+
+    class PyBigQueryImpl(impl.DefaultImpl):
+        __dialect__ = "bigquery"
diff --git a/tests/system/conftest.py b/tests/system/conftest.py
@@ -1,8 +1,21 @@
-# Copyright 2021 The PyBigQuery Authors
+# Copyright (c) 2021 The PyBigQuery Authors
 #
-# Use of this source code is governed by an MIT-style
-# license that can be found in the LICENSE file or at
-# https://opensource.org/licenses/MIT.
+# 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.
 
 import datetime
 import pathlib
@@ -55,7 +68,7 @@ def bigquery_schema(bigquery_client: bigquery.Client):
     return bigquery_client.schema_from_json(DATA_DIR / "schema.json")
 
 
-@pytest.fixture(scope="session", autouse=True)
+@pytest.fixture(scope="session")
 def bigquery_dataset(
     bigquery_client: bigquery.Client, bigquery_schema: List[bigquery.SchemaField]
 ):
@@ -96,7 +109,7 @@ def bigquery_empty_table(
     return table_id
 
 
-@pytest.fixture(scope="session", autouse=True)
+@pytest.fixture(scope="session")
 def bigquery_regional_dataset(bigquery_client, bigquery_schema):
     project_id = bigquery_client.project
     dataset_id = f"test_pybigquery_location_{temp_suffix()}"

diff --git a/tests/system/test_alembic.py b/tests/system/test_alembic.py
@@ -0,0 +1,170 @@
+# Copyright (c) 2021 The PyBigQuery 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.
+
+import contextlib
+
+import pytest
+from sqlalchemy import Column, DateTime, Integer, String
+
+try:
+    import alembic  # noqa
+except ImportError:
+    alembic = None
+
+import google.api_core.exceptions
+
+
+@pytest.fixture
+def alembic_table(bigquery_dataset, bigquery_client):
+    import sqlalchemy
+    import alembic.migration
+    import alembic.operations
+
+    def get_table(table_name, data="table"):
+        try:
+            table_id = f"{bigquery_dataset}.{table_name}"
+            if data == "rows":
+                return [dict(r) for r in bigquery_client.list_rows(table_id)]
+            else:
+                table = bigquery_client.get_table(table_id)
+                if data == "table":
+                    return table
+                elif data == "schema":
+                    return [
+                        repr(s).replace(", (), None)", ")").replace(", None)", ")")
+                        for s in table.schema
+                    ]
+                else:
+                    raise ValueError(data)
+        except google.api_core.exceptions.NotFound:
+            return None
+
+    engine = sqlalchemy.create_engine(f"bigquery:///{bigquery_dataset}")
+    with contextlib.closing(engine.connect()) as conn:
+        migration_context = alembic.migration.MigrationContext.configure(conn, {})
+        with alembic.operations.Operations.context(migration_context):
+            yield get_table
+
+
+@pytest.mark.skipif(alembic is None, reason="Alembic isn't installed.")
+def test_alembic_scenario(alembic_table):
+    """
+    Exercise all of the operations we support.
+
+    It's a little awkward because we have to avoid doing too many
+    operations on the same table to avoid tripping over limits on
+    table mods within a short time.
+    """
+    from alembic import op
+
+    assert alembic_table("account") is None
+
+    account = op.create_table(
+        "account",
+        Column("id", Integer, nullable=False),
+        Column("name", String(50), nullable=False, comment="The name"),
+        Column("description", String(200)),
+    )
+    assert alembic_table("account", "schema") == [
+        "SchemaField('id', 'INTEGER', 'REQUIRED')",
+        "SchemaField('name', 'STRING(50)', 'REQUIRED', 'The name')",
+        "SchemaField('description', 'STRING(200)', 'NULLABLE')",
+    ]
+
+    op.bulk_insert(
+        account,
+        [
+            dict(id=1, name="home", description="the home account"),
+            dict(id=2, name="operations", description="the ops account"),
+            dict(id=3, name="savings", description=None),
+        ],
+    )
+
+    assert alembic_table("account", "rows") == [
+        {"description": "the home account", "id": 1, "name": "home"},
+        {"description": "the ops account", "id": 2, "name": "operations"},
+        {"description": None, "id": 3, "name": "savings"},
+    ]
+
+    op.add_column(
+        "account", Column("last_transaction_date", DateTime, comment="when updated")
+    )
+
+    assert alembic_table("account", "schema") == [
+        "SchemaField('id', 'INTEGER', 'REQUIRED')",
+        "SchemaField('name', 'STRING(50)', 'REQUIRED', 'The name')",
+        "SchemaField('description', 'STRING(200)', 'NULLABLE')",
+        "SchemaField('last_transaction_date', 'DATETIME', 'NULLABLE', 'when updated')",
+    ]
+
+    op.create_table(
+        "account_w_comment",
+        Column("id", Integer, nullable=False),
+        Column("name", String(50), nullable=False, comment="The name"),
+        Column("description", String(200)),
+        comment="This table has comments",
+    )
+    assert alembic_table("account_w_comment").description == "This table has comments"
+    op.drop_table_comment("account_w_comment")
+    assert alembic_table("account_w_comment").description is None
+
+    op.drop_column("account_w_comment", "description")
+    assert alembic_table("account_w_comment", "schema") == [
+        "SchemaField('id', 'INTEGER', 'REQUIRED')",
+        "SchemaField('name', 'STRING(50)', 'REQUIRED', 'The name')",
+    ]
+
+    op.drop_table("account_w_comment")
+    assert alembic_table("account_w_comment") is None
+
+    op.rename_table("account", "accounts")
+    assert alembic_table("account") is None
+    assert alembic_table("accounts", "schema") == [
+        "SchemaField('id', 'INTEGER', 'REQUIRED')",
+        "SchemaField('name', 'STRING(50)', 'REQUIRED', 'The name')",
+        "SchemaField('description', 'STRING(200)', 'NULLABLE')",
+        "SchemaField('last_transaction_date', 'DATETIME', 'NULLABLE', 'when updated')",
+    ]
+    op.drop_table("accounts")
+    assert alembic_table("accounts") is None
+
+    op.execute(
+        """
+        create table transactions(
+            account INT64 NOT NULL,
+            transaction_time DATETIME NOT NULL,
+            amount NUMERIC(11, 2) NOT NULL
+            )
+        partition by DATE(transaction_time)
+        """
+    )
+
+    # The only thing we can alter about a column is we can make it
+    # nullable:
+    op.alter_column("transactions", "amount", True)
+    assert alembic_table("transactions", "schema") == [
+        "SchemaField('account', 'INTEGER', 'REQUIRED')",
+        "SchemaField('transaction_time', 'DATETIME', 'REQUIRED')",
+        "SchemaField('amount', 'NUMERIC(11, 2)', 'NULLABLE')",
+    ]
+
+    op.create_table_comment("transactions", "Transaction log")
+    assert alembic_table("transactions").description == "Transaction log"
+
+    op.drop_table("transactions")