diff --git a/docs/alembic.rst b/docs/alembic.rst new file mode 100644 index 00000000..2f1e03ad --- /dev/null +++ b/docs/alembic.rst @@ -0,0 +1,51 @@ +Alembic support +--------------- + +`Alembic `_ 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 +`_. + +Supported operations: + +`add_column(table_name, column, schema=None) +`_ + +`alter_column(table_name, column_name, nullable=None, schema=None) +`_ + +`bulk_insert(table, rows, multiinsert=True) +`_ + +`create_table(table_name, *columns, **kw) +`_ + +`create_table_comment(table_name, comment, schema=None) +`_ + +`drop_column(table_name, column_name, schema=None) +`_ + +`drop_table(table_name, schema=None) +`_ + +`drop_table_comment(table_name, schema=None) +`_ + +`execute(sqltext, execution_options=None) +`_ + +`rename_table(old_table_name, new_table_name, schema=None) +`_ + +Note that some of the operations above have limited capability, again +do to `BigQuery limitations +`_. + +The `execute` operation allows access to BigQuery-specific +`data-definition-language +`_. diff --git a/docs/index.rst b/docs/index.rst index 5de1b9aa..0e255b3e 100644 --- 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 index f4d7b3b6..fd0f0101 100644 --- a/noxfile.py +++ b/noxfile.py @@ -82,6 +82,22 @@ 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. @@ -89,7 +105,7 @@ def default(session): 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 index fb00b553..60b8aab0 100644 --- 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 index 646842a8..bd7e25a4 100644 --- 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 new file mode 100644 index 00000000..2ea7af18 --- /dev/null +++ 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")