Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: Alembic support (#183)
  • Loading branch information
jimfulton committed May 25, 2021
1 parent 6d17084 commit 4d5a17c
Show file tree
Hide file tree
Showing 6 changed files with 271 additions and 7 deletions.
51 changes: 51 additions & 0 deletions 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>`_.
2 changes: 2 additions & 0 deletions docs/index.rst
Expand Up @@ -2,6 +2,8 @@

.. include:: multiprocessing.rst

.. include:: alembic.rst

API Reference
-------------
.. toctree::
Expand Down
19 changes: 18 additions & 1 deletion noxfile.py
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
11 changes: 11 additions & 0 deletions pybigquery/sqlalchemy_bigquery.py
Expand Up @@ -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"
25 changes: 19 additions & 6 deletions 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
Expand Down Expand Up @@ -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]
):
Expand Down Expand Up @@ -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()}"
Expand Down
170 changes: 170 additions & 0 deletions 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")

0 comments on commit 4d5a17c

Please sign in to comment.