From da7a40309de6ca8063d6dcf6678de96a463344e6 Mon Sep 17 00:00:00 2001 From: Jim Fulton Date: Tue, 17 Aug 2021 14:43:09 -0400 Subject: [PATCH] feat: Add geography support (#228) --- docs/alembic.rst | 2 +- docs/geography.rst | 87 +++++++ docs/index.rst | 2 + docs/reference.rst | 12 + docs/samples | 1 + noxfile.py | 53 ++-- owlbot.py | 62 +++-- samples/__init__.py | 0 samples/pytest.ini | 0 samples/snippets/__init__.py | 20 ++ samples/snippets/conftest.py | 48 ++++ samples/snippets/geography.py | 66 +++++ samples/snippets/noxfile.py | 263 ++++++++++++++++++++ samples/snippets/requirements-test.txt | 12 + samples/snippets/requirements.txt | 62 +++++ samples/snippets/test_geography.py | 27 ++ setup.py | 5 + sqlalchemy_bigquery/__init__.py | 10 +- sqlalchemy_bigquery/base.py | 14 +- sqlalchemy_bigquery/geography.py | 242 ++++++++++++++++++ tests/system/conftest.py | 8 + tests/system/test_alembic.py | 8 +- tests/system/test_geography.py | 304 +++++++++++++++++++++++ tests/system/test_sqlalchemy_bigquery.py | 11 +- tests/unit/conftest.py | 11 + tests/unit/test_geography.py | 179 +++++++++++++ 26 files changed, 1439 insertions(+), 70 deletions(-) create mode 100644 docs/geography.rst create mode 100644 docs/reference.rst create mode 120000 docs/samples create mode 100644 samples/__init__.py create mode 100644 samples/pytest.ini create mode 100644 samples/snippets/__init__.py create mode 100644 samples/snippets/conftest.py create mode 100644 samples/snippets/geography.py create mode 100644 samples/snippets/noxfile.py create mode 100644 samples/snippets/requirements-test.txt create mode 100644 samples/snippets/requirements.txt create mode 100644 samples/snippets/test_geography.py create mode 100644 sqlalchemy_bigquery/geography.py create mode 100644 tests/system/test_geography.py create mode 100644 tests/unit/test_geography.py diff --git a/docs/alembic.rst b/docs/alembic.rst index 2f1e03ad..e83953a0 100644 --- a/docs/alembic.rst +++ b/docs/alembic.rst @@ -1,5 +1,5 @@ Alembic support ---------------- +^^^^^^^^^^^^^^^ `Alembic `_ is a lightweight database migration tool for usage with the SQLAlchemy Database Toolkit for diff --git a/docs/geography.rst b/docs/geography.rst new file mode 100644 index 00000000..aef79749 --- /dev/null +++ b/docs/geography.rst @@ -0,0 +1,87 @@ +Working with Geographic data +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +BigQuery provides a `GEOGRAPHY data type +`_ +for `working with geographic data +`_, including: + +- Points, +- Linestrings, +- Polygons, and +- Collections of points, linestrings, and polygons. + +Geographic data uses the `WGS84 +`_ coordinate system. + +To define a geography column, use the `GEOGRAPHY` data type imported +from the `sqlalchemy_bigquery` module: + +.. literalinclude:: samples/snippets/geography.py + :language: python + :dedent: 4 + :start-after: [START bigquery_sqlalchemy_create_table_with_geography] + :end-before: [END bigquery_sqlalchemy_create_table_with_geography] + +BigQuery has a variety of `SQL geographic functions +`_ +for working with geographic data. Among these are functions for +converting between SQL geometry objects and `standard text (WKT) and +binary (WKB) representations +`_. + +Geography data is typically represented in Python as text strings in +WKT format or as `WKB` objects, which contain binary data in WKB +format. Querying geographic data returns `WKB` objects and `WKB` +objects may be used in queries. When +calling spatial functions that expect geographic arguments, text +arguments are automatically coerced to geography. + +Inserting data +~~~~~~~~~~~~~~ + +When inserting geography data, you can pass WKT strings, `WKT` objects, +or `WKB` objects: + +.. literalinclude:: samples/snippets/geography.py + :language: python + :dedent: 4 + :start-after: [START bigquery_sqlalchemy_insert_geography] + :end-before: [END bigquery_sqlalchemy_insert_geography] + +Note that in the `lake3` example, we got a `WKB` object by creating a +`WKT` object and getting its `wkb` property. Normally, we'd get `WKB` +objects as results of previous queries. + +Queries +~~~~~~~ + +When performing spacial queries, and geography objects are expected, +you can to pass `WKB` or `WKT` objects: + +.. literalinclude:: samples/snippets/geography.py + :language: python + :dedent: 4 + :start-after: [START bigquery_sqlalchemy_query_geography_wkb] + :end-before: [END bigquery_sqlalchemy_query_geography_wkb] + +In this example, we passed the `geog` attribute of `lake2`, which is a WKB object. + +Or you can pass strings in WKT format: + +.. literalinclude:: samples/snippets/geography.py + :language: python + :dedent: 4 + :start-after: [START bigquery_sqlalchemy_query_geography_text] + :end-before: [END bigquery_sqlalchemy_query_geography_text] + +Installing geography support +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To get geography support, you need to install `sqlalchemy-bigquery` +with the `geography` extra, or separately install `GeoAlchemy2` and +`shapely`. + +.. code-block:: console + + pip install 'sqlalchemy-bigquery[geography]' diff --git a/docs/index.rst b/docs/index.rst index eef073c6..4fe42891 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,7 +3,9 @@ :maxdepth: 2 README + geography alembic + reference Changelog --------- diff --git a/docs/reference.rst b/docs/reference.rst new file mode 100644 index 00000000..9f8cabef --- /dev/null +++ b/docs/reference.rst @@ -0,0 +1,12 @@ +API Reference +^^^^^^^^^^^^^ + +Geography +~~~~~~~~~ + +.. autoclass:: sqlalchemy_bigquery.geography.GEOGRAPHY + :exclude-members: bind_expression, ElementType, bind_processor + +.. automodule:: sqlalchemy_bigquery.geography + :members: WKB, WKT + :exclude-members: GEOGRAPHY diff --git a/docs/samples b/docs/samples new file mode 120000 index 00000000..e804737e --- /dev/null +++ b/docs/samples @@ -0,0 +1 @@ +../samples \ No newline at end of file diff --git a/noxfile.py b/noxfile.py index 7c2097ab..8404d6d2 100644 --- a/noxfile.py +++ b/noxfile.py @@ -82,22 +82,6 @@ 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 do not 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. @@ -114,8 +98,13 @@ def default(session): constraints_path, ) - install_alembic_for_python_38(session, constraints_path) - session.install("-e", ".", "-c", constraints_path) + if session.python == "3.8": + extras = "[alembic]" + elif session.python == "3.9": + extras = "[geography]" + else: + extras = "" + session.install("-e", f".{extras}", "-c", constraints_path) # Run py.test against the unit tests. session.run( @@ -167,8 +156,13 @@ 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) + if session.python == "3.8": + extras = "[alembic]" + elif session.python == "3.9": + extras = "[geography]" + else: + extras = "" + session.install("-e", f".{extras}", "-c", constraints_path) # Run py.test against the system tests. if system_test_exists: @@ -216,7 +210,13 @@ def compliance(session): "-c", constraints_path, ) - session.install("-e", ".", "-c", constraints_path) + if session.python == "3.8": + extras = "[alembic]" + elif session.python == "3.9": + extras = "[geography]" + else: + extras = "" + session.install("-e", f".{extras}", "-c", constraints_path) session.run( "py.test", @@ -251,7 +251,9 @@ def docs(session): """Build the docs for this library.""" session.install("-e", ".") - session.install("sphinx==4.0.1", "alabaster", "recommonmark") + session.install( + "sphinx==4.0.1", "alabaster", "geoalchemy2", "shapely", "recommonmark" + ) shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True) session.run( @@ -274,7 +276,12 @@ def docfx(session): session.install("-e", ".") session.install( - "sphinx==4.0.1", "alabaster", "recommonmark", "gcp-sphinx-docfx-yaml" + "sphinx==4.0.1", + "alabaster", + "geoalchemy2", + "shapely", + "recommonmark", + "gcp-sphinx-docfx-yaml", ) shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True) diff --git a/owlbot.py b/owlbot.py index e63929a0..03cc243e 100644 --- a/owlbot.py +++ b/owlbot.py @@ -18,7 +18,7 @@ import synthtool as s from synthtool import gcp - +from synthtool.languages import python REPO_ROOT = pathlib.Path(__file__).parent.absolute() @@ -27,10 +27,19 @@ # ---------------------------------------------------------------------------- # Add templated files # ---------------------------------------------------------------------------- +extras = [] +extras_by_python = { + "3.8": ["alembic"], + "3.9": ["geography"], +} templated_files = common.py_library( unit_test_python_versions=["3.6", "3.7", "3.8", "3.9"], system_test_python_versions=["3.8", "3.9"], - cov_level=100 + cov_level=100, + unit_test_extras=extras, + unit_test_extras_by_python=extras_by_python, + system_test_extras=extras, + system_test_extras_by_python=extras_by_python, ) s.move(templated_files, excludes=[ # sqlalchemy-bigquery was originally licensed MIT @@ -77,37 +86,6 @@ def place_before(path, text, *before_text, escape=None): "nox.options.stop_on_first_error = True", ) -install_alembic_for_python_38 = ''' -def install_alembic_for_python_38(session, constraints_path): - """ - install alembic for Python 3.8 unit and system tests - - We do not 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) - - -''' - -place_before( - "noxfile.py", - "def default", - install_alembic_for_python_38, - ) - -place_before( - "noxfile.py", - ' session.install("-e", ".", ', - " install_alembic_for_python_38(session, constraints_path)", - escape='(') - old_sessions = ''' "unit", "system", @@ -125,6 +103,9 @@ def install_alembic_for_python_38(session, constraints_path): s.replace( ["noxfile.py"], old_sessions, new_sessions) +# Maybe we can get rid of this when we don't need pytest-rerunfailures, +# which we won't need when BQ retries itself: +# https://github.com/googleapis/python-bigquery/pull/837 compliance = ''' @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) def compliance(session): @@ -153,7 +134,13 @@ def compliance(session): "-c", constraints_path, ) - session.install("-e", ".", "-c", constraints_path) + if session.python == "3.8": + extras = "[alembic]" + elif session.python == "3.9": + extras = "[geography]" + else: + extras = "" + session.install("-e", f".{extras}", "-c", constraints_path) session.run( "py.test", @@ -180,6 +167,7 @@ def compliance(session): escape="()", ) +s.replace(["noxfile.py"], '"alabaster"', '"alabaster", "geoalchemy2", "shapely"') @@ -201,6 +189,12 @@ def compliance(session): """ ) +# ---------------------------------------------------------------------------- +# Samples templates +# ---------------------------------------------------------------------------- + +python.py_samples(skip_readmes=True) + # ---------------------------------------------------------------------------- # Final cleanup # ---------------------------------------------------------------------------- diff --git a/samples/__init__.py b/samples/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/samples/pytest.ini b/samples/pytest.ini new file mode 100644 index 00000000..e69de29b diff --git a/samples/snippets/__init__.py b/samples/snippets/__init__.py new file mode 100644 index 00000000..fa3a9cd6 --- /dev/null +++ b/samples/snippets/__init__.py @@ -0,0 +1,20 @@ +# 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. + +__version__ = "1.0.0-a1" diff --git a/samples/snippets/conftest.py b/samples/snippets/conftest.py new file mode 100644 index 00000000..dc78bc4e --- /dev/null +++ b/samples/snippets/conftest.py @@ -0,0 +1,48 @@ +# 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. +""" +SQLAlchemy dialect for Google BigQuery +""" + +from google.cloud import bigquery +import pytest +import sqlalchemy +import test_utils.prefixer + +prefixer = test_utils.prefixer.Prefixer("python-bigquery-sqlalchemy", "tests/system") + + +@pytest.fixture(scope="session") +def client(): + return bigquery.Client() + + +@pytest.fixture(scope="session") +def dataset_id(client: bigquery.Client): + project_id = client.project + dataset_id = prefixer.create_prefix() + dataset = bigquery.Dataset(f"{project_id}.{dataset_id}") + dataset = client.create_dataset(dataset) + yield dataset_id + client.delete_dataset(dataset_id, delete_contents=True) + + +@pytest.fixture(scope="session") +def engine(dataset_id): + return sqlalchemy.create_engine(f"bigquery:///{dataset_id}") diff --git a/samples/snippets/geography.py b/samples/snippets/geography.py new file mode 100644 index 00000000..d6adc115 --- /dev/null +++ b/samples/snippets/geography.py @@ -0,0 +1,66 @@ +# 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. + + +def example(engine): + # [START bigquery_sqlalchemy_create_table_with_geography] + from sqlalchemy.ext.declarative import declarative_base + from sqlalchemy import Column, String + from sqlalchemy_bigquery import GEOGRAPHY + + Base = declarative_base() + + class Lake(Base): + __tablename__ = "lakes" + + name = Column(String, primary_key=True) + geog = Column(GEOGRAPHY) + + # [END bigquery_sqlalchemy_create_table_with_geography] + Lake.__table__.create(engine) + + # [START bigquery_sqlalchemy_insert_geography] + from sqlalchemy.orm import sessionmaker + from sqlalchemy_bigquery import WKT + + Session = sessionmaker(bind=engine) + session = Session() + + lake = Lake(name="Majeur", geog="POLYGON((0 0,1 0,1 1,0 1,0 0))") + lake2 = Lake(name="Garde", geog=WKT("POLYGON((1 0,3 0,3 2,1 2,1 0))")) + b = WKT("POLYGON((3 0,6 0,6 3,3 3,3 0))").wkb + lake3 = Lake(name="Orta", geog=b) + + session.add_all((lake, lake2, lake3)) + session.commit() + # [END bigquery_sqlalchemy_insert_geography] + + # [START bigquery_sqlalchemy_query_geography_wkb] + from sqlalchemy import func + + lakes_touching_lake2 = list( + session.query(Lake).filter(func.ST_Touches(Lake.geog, lake2.geog)) + ) + # [END bigquery_sqlalchemy_query_geography_wkb] + # [START bigquery_sqlalchemy_query_geography_text] + lakes_containing = list( + session.query(Lake).filter(func.ST_Contains(Lake.geog, "POINT(4 1)")) + ) + # [END bigquery_sqlalchemy_query_geography_text] + return lakes_touching_lake2, lakes_containing diff --git a/samples/snippets/noxfile.py b/samples/snippets/noxfile.py new file mode 100644 index 00000000..9fc7f178 --- /dev/null +++ b/samples/snippets/noxfile.py @@ -0,0 +1,263 @@ +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import print_function + +import os +from pathlib import Path +import sys +from typing import Callable, Dict, List, Optional + +import nox + + +# WARNING - WARNING - WARNING - WARNING - WARNING +# WARNING - WARNING - WARNING - WARNING - WARNING +# DO NOT EDIT THIS FILE EVER! +# WARNING - WARNING - WARNING - WARNING - WARNING +# WARNING - WARNING - WARNING - WARNING - WARNING + +BLACK_VERSION = "black==19.10b0" + +# Copy `noxfile_config.py` to your directory and modify it instead. + +# `TEST_CONFIG` dict is a configuration hook that allows users to +# modify the test configurations. The values here should be in sync +# with `noxfile_config.py`. Users will copy `noxfile_config.py` into +# their directory and modify it. + +TEST_CONFIG = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["2.7"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": False, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} + + +try: + # Ensure we can import noxfile_config in the project's directory. + sys.path.append(".") + from noxfile_config import TEST_CONFIG_OVERRIDE +except ImportError as e: + print("No user noxfile_config found: detail: {}".format(e)) + TEST_CONFIG_OVERRIDE = {} + +# Update the TEST_CONFIG with the user supplied values. +TEST_CONFIG.update(TEST_CONFIG_OVERRIDE) + + +def get_pytest_env_vars() -> Dict[str, str]: + """Returns a dict for pytest invocation.""" + ret = {} + + # Override the GCLOUD_PROJECT and the alias. + env_key = TEST_CONFIG["gcloud_project_env"] + # This should error out if not set. + ret["GOOGLE_CLOUD_PROJECT"] = os.environ[env_key] + + # Apply user supplied envs. + ret.update(TEST_CONFIG["envs"]) + return ret + + +# DO NOT EDIT - automatically generated. +# All versions used to tested samples. +ALL_VERSIONS = ["2.7", "3.6", "3.7", "3.8", "3.9"] + +# Any default versions that should be ignored. +IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] + +TESTED_VERSIONS = sorted([v for v in ALL_VERSIONS if v not in IGNORED_VERSIONS]) + +INSTALL_LIBRARY_FROM_SOURCE = bool(os.environ.get("INSTALL_LIBRARY_FROM_SOURCE", False)) +# +# Style Checks +# + + +def _determine_local_import_names(start_dir: str) -> List[str]: + """Determines all import names that should be considered "local". + + This is used when running the linter to insure that import order is + properly checked. + """ + file_ext_pairs = [os.path.splitext(path) for path in os.listdir(start_dir)] + return [ + basename + for basename, extension in file_ext_pairs + if extension == ".py" + or os.path.isdir(os.path.join(start_dir, basename)) + and basename not in ("__pycache__") + ] + + +# Linting with flake8. +# +# We ignore the following rules: +# E203: whitespace before ‘:’ +# E266: too many leading ‘#’ for block comment +# E501: line too long +# I202: Additional newline in a section of imports +# +# We also need to specify the rules which are ignored by default: +# ['E226', 'W504', 'E126', 'E123', 'W503', 'E24', 'E704', 'E121'] +FLAKE8_COMMON_ARGS = [ + "--show-source", + "--builtin=gettext", + "--max-complexity=20", + "--import-order-style=google", + "--exclude=.nox,.cache,env,lib,generated_pb2,*_pb2.py,*_pb2_grpc.py", + "--ignore=E121,E123,E126,E203,E226,E24,E266,E501,E704,W503,W504,I202", + "--max-line-length=88", +] + + +@nox.session +def lint(session: nox.sessions.Session) -> None: + if not TEST_CONFIG["enforce_type_hints"]: + session.install("flake8", "flake8-import-order") + else: + session.install("flake8", "flake8-import-order", "flake8-annotations") + + local_names = _determine_local_import_names(".") + args = FLAKE8_COMMON_ARGS + [ + "--application-import-names", + ",".join(local_names), + ".", + ] + session.run("flake8", *args) + + +# +# Black +# + + +@nox.session +def blacken(session: nox.sessions.Session) -> None: + session.install(BLACK_VERSION) + python_files = [path for path in os.listdir(".") if path.endswith(".py")] + + session.run("black", *python_files) + + +# +# Sample Tests +# + + +PYTEST_COMMON_ARGS = ["--junitxml=sponge_log.xml"] + + +def _session_tests( + session: nox.sessions.Session, post_install: Callable = None +) -> None: + if TEST_CONFIG["pip_version_override"]: + pip_version = TEST_CONFIG["pip_version_override"] + session.install(f"pip=={pip_version}") + """Runs py.test for a particular project.""" + if os.path.exists("requirements.txt"): + if os.path.exists("constraints.txt"): + session.install("-r", "requirements.txt", "-c", "constraints.txt") + else: + session.install("-r", "requirements.txt") + + if os.path.exists("requirements-test.txt"): + if os.path.exists("constraints-test.txt"): + session.install("-r", "requirements-test.txt", "-c", "constraints-test.txt") + else: + session.install("-r", "requirements-test.txt") + + if INSTALL_LIBRARY_FROM_SOURCE: + session.install("-e", _get_repo_root()) + + if post_install: + post_install(session) + + session.run( + "pytest", + *(PYTEST_COMMON_ARGS + session.posargs), + # Pytest will return 5 when no tests are collected. This can happen + # on travis where slow and flaky tests are excluded. + # See http://doc.pytest.org/en/latest/_modules/_pytest/main.html + success_codes=[0, 5], + env=get_pytest_env_vars(), + ) + + +@nox.session(python=ALL_VERSIONS) +def py(session: nox.sessions.Session) -> None: + """Runs py.test for a sample using the specified version of Python.""" + if session.python in TESTED_VERSIONS: + _session_tests(session) + else: + session.skip( + "SKIPPED: {} tests are disabled for this sample.".format(session.python) + ) + + +# +# Readmegen +# + + +def _get_repo_root() -> Optional[str]: + """ Returns the root folder of the project. """ + # Get root of this repository. Assume we don't have directories nested deeper than 10 items. + p = Path(os.getcwd()) + for i in range(10): + if p is None: + break + if Path(p / ".git").exists(): + return str(p) + # .git is not available in repos cloned via Cloud Build + # setup.py is always in the library's root, so use that instead + # https://github.com/googleapis/synthtool/issues/792 + if Path(p / "setup.py").exists(): + return str(p) + p = p.parent + raise Exception("Unable to detect repository root.") + + +GENERATED_READMES = sorted([x for x in Path(".").rglob("*.rst.in")]) + + +@nox.session +@nox.parametrize("path", GENERATED_READMES) +def readmegen(session: nox.sessions.Session, path: str) -> None: + """(Re-)generates the readme for a sample.""" + session.install("jinja2", "pyyaml") + dir_ = os.path.dirname(path) + + if os.path.exists(os.path.join(dir_, "requirements.txt")): + session.install("-r", os.path.join(dir_, "requirements.txt")) + + in_file = os.path.join(dir_, "README.rst.in") + session.run( + "python", _get_repo_root() + "/scripts/readme-gen/readme_gen.py", in_file + ) diff --git a/samples/snippets/requirements-test.txt b/samples/snippets/requirements-test.txt new file mode 100644 index 00000000..ceda9ac3 --- /dev/null +++ b/samples/snippets/requirements-test.txt @@ -0,0 +1,12 @@ +attrs==21.2.0 +google-cloud-testutils==1.0.0 +importlib-metadata==4.6.1 +iniconfig==1.1.1 +packaging==21.0 +pluggy==0.13.1 +py==1.10.0 +pyparsing==2.4.7 +pytest==6.2.4 +toml==0.10.2 +typing-extensions==3.10.0.0 +zipp==3.5.0 diff --git a/samples/snippets/requirements.txt b/samples/snippets/requirements.txt new file mode 100644 index 00000000..621dd3ba --- /dev/null +++ b/samples/snippets/requirements.txt @@ -0,0 +1,62 @@ +aiocontextvars==0.2.2 +attrs==21.2.0 +cachetools==4.2.2 +certifi==2021.5.30 +cffi==1.14.6 +charset-normalizer==2.0.3 +click==8.0.1 +click-plugins==1.1.1 +cligj==0.7.2 +contextvars==2.4 +dataclasses==0.6 +Deprecated==1.2.12 +Fiona==1.8.20 +future==0.18.2 +GeoAlchemy2==0.9.3 +geopandas==0.9.0 +google-api-core==1.31.1 +google-auth==1.34.0 +google-cloud-bigquery==2.23.2 +google-cloud-bigquery-storage==2.6.2 +google-cloud-core==1.7.2 +google-crc32c==1.1.2 +google-resumable-media==1.3.2 +googleapis-common-protos==1.53.0 +greenlet==1.1.0 +grpcio==1.39.0 +idna==3.2 +immutables==0.15 +importlib-metadata==4.6.1 +libcst==0.3.19 +munch==2.5.0 +mypy-extensions==0.4.3 +numpy==1.19.5 +opentelemetry-api==1.4.1 +opentelemetry-instrumentation==0.23b2 +opentelemetry-sdk==1.4.1 +opentelemetry-semantic-conventions==0.23b2 +packaging==21.0 +pandas==1.1.5 +proto-plus==1.19.0 +protobuf==3.17.3 +pyarrow==5.0.0 +pyasn1==0.4.8 +pyasn1-modules==0.2.8 +pybigquery==0.10.0 +pycparser==2.20 +pyparsing==2.4.7 +pyproj==3.0.1 +python-dateutil==2.8.2 +pytz==2021.1 +PyYAML==5.4.1 +requests==2.26.0 +rsa==4.7.2 +Shapely==1.7.1 +six==1.16.0 +SQLAlchemy==1.4.22 +tqdm==4.61.2 +typing-extensions==3.10.0.0 +typing-inspect==0.7.1 +urllib3==1.26.6 +wrapt==1.12.1 +zipp==3.5.0 diff --git a/samples/snippets/test_geography.py b/samples/snippets/test_geography.py new file mode 100644 index 00000000..7a570b81 --- /dev/null +++ b/samples/snippets/test_geography.py @@ -0,0 +1,27 @@ +# 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. + + +def test_geography(engine): + from . import geography + + lakes_touching_lake2, lakes_containing = geography.example(engine) + + assert sorted(lake.name for lake in lakes_touching_lake2) == ["Majeur", "Orta"] + assert [lake.name for lake in lakes_containing] == ["Orta"] diff --git a/setup.py b/setup.py index 437c0df0..378e3a23 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import io +import itertools import os import re from setuptools import setup @@ -44,6 +45,9 @@ def readme(): return f.read() +extras = dict(geography=["GeoAlchemy2", "shapely"], alembic=["alembic"],) +extras["all"] = set(itertools.chain.from_iterable(extras.values())) + setup( name=name, version=version, @@ -79,6 +83,7 @@ def readme(): "sqlalchemy>=1.2.0,<1.5.0dev", "future", ], + extras_require=extras, python_requires=">=3.6, <3.10", tests_require=["pytz"], entry_points={ diff --git a/sqlalchemy_bigquery/__init__.py b/sqlalchemy_bigquery/__init__.py index fb08e576..e3dd3f2d 100644 --- a/sqlalchemy_bigquery/__init__.py +++ b/sqlalchemy_bigquery/__init__.py @@ -20,7 +20,7 @@ SQLAlchemy dialect for Google BigQuery """ -from .version import __version__ +from .version import __version__ # noqa from .base import BigQueryDialect from .base import ( @@ -42,7 +42,6 @@ ) __all__ = [ - "__version__", "BigQueryDialect", "STRING", "BOOL", @@ -61,6 +60,13 @@ "BIGNUMERIC", ] +try: + from .geography import GEOGRAPHY, WKB, WKT # noqa +except ImportError: + pass +else: + __all__.extend(["GEOGRAPHY", "WKB", "WKT"]) + try: import pybigquery # noqa except ImportError: diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index db7336f6..e6403796 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -55,6 +55,11 @@ from .parse_url import parse_url from sqlalchemy_bigquery import _helpers +try: + from .geography import GEOGRAPHY +except ImportError: + pass + FIELD_ILLEGAL_CHARACTERS = re.compile(r"[^\w]+") @@ -126,6 +131,7 @@ def format_label(self, label, name=None): "BIGNUMERIC": types.Numeric, } +# By convention, dialect-provided types are spelled with all upper case. STRING = _type_map["STRING"] BOOL = _type_map["BOOL"] BOOLEAN = _type_map["BOOLEAN"] @@ -142,6 +148,11 @@ def format_label(self, label, name=None): NUMERIC = _type_map["NUMERIC"] BIGNUMERIC = _type_map["NUMERIC"] +try: + _type_map["GEOGRAPHY"] = GEOGRAPHY +except NameError: + pass + class BigQueryExecutionContext(DefaultExecutionContext): def create_cursor(self): @@ -429,6 +440,7 @@ def visit_bindparam( if bq_type[-1] == ">" and bq_type.startswith("ARRAY<"): # Values get arrayified at a lower level. bq_type = bq_type[6:-1] + bq_type = self.__remove_type_parameter(bq_type) assert_(param != "%s", f"Unexpected param: {param}") @@ -479,7 +491,7 @@ def visit_BINARY(self, type_, **kw): return f"BYTES({type_.length})" return "BYTES" - visit_VARBINARY = visit_BINARY + visit_VARBINARY = visit_BLOB = visit_BINARY def visit_NUMERIC(self, type_, **kw): if (type_.precision is not None) and isinstance( diff --git a/sqlalchemy_bigquery/geography.py b/sqlalchemy_bigquery/geography.py new file mode 100644 index 00000000..9a10c236 --- /dev/null +++ b/sqlalchemy_bigquery/geography.py @@ -0,0 +1,242 @@ +# 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 geoalchemy2 +from geoalchemy2.shape import to_shape +import geoalchemy2.functions +from shapely import wkb, wkt +import sqlalchemy.ext.compiler +from sqlalchemy.sql.elements import BindParameter + +SRID = 4326 # WGS84, https://spatialreference.org/ref/epsg/wgs-84/ + + +class WKB(geoalchemy2.WKBElement): + """ + Well-Known-Binary data wrapper. + + WKB objects hold geographic data in a binary format known as + "Well-Known Binary", + . + These objects are returned from queries and can be used in insert + and queries. + + The WKB class is a subclass of the Geoalchemy2 WKBElement class + customized for working with BigQuery. + """ + + geom_from_extended_version = "ST_GeogFromWKB" + + def __init__(self, data): + super().__init__(data, SRID, True) + + @property + def wkt(self): + """ + Return the WKB object as a WKT object. + """ + return WKT(to_shape(self).wkt) + + +class WKT(geoalchemy2.WKTElement): + """ + Well-Known-Text data wrapper. + + WKT objects hold geographic data in a text format known as + "Well-Known Text", + . + + You generally shouldn't need to create WKT objects directly, as + text arguments to geographic functions and inserts to GEOGRAPHY + columns are automatically coerced to geographic data. + + The WKT class is a subclass of the Geoalchemy2 WKTElement class + customized for working with BigQuery. + """ + + geom_from_extended_version = "ST_GeogFromText" + + def __init__(self, data): + super().__init__(data, SRID, True) + + @property + def wkb(self): + """ + Return the WKT object as a WKB object. + """ + return WKB(wkb.dumps(wkt.loads(self.data))) + + +class GEOGRAPHY(geoalchemy2.Geography): + """ + GEOGRAPHY type + + Use this to define BigQuery GEOGRAPHY columns:: + + class Lake(Base): + __tablename__ = 'lakes' + + name = Column(String) + geog = column(GEOGRAPHY) + + + """ + + def __init__(self): + super().__init__( + geometry_type=None, spatial_index=False, srid=SRID, + ) + self.extended = True + + # Un-inherit the bind function that adds an ST_GeogFromText. + # It's unnecessary and causes BigQuery to error. + # + # Some things to note about this: + # + # 1. bind_expression can't always know the value. When multiple + # rows are being inserted, the values may be different in each + # row. As a consequence, we have to treat all the values as WKT. + # + # 2. This applies equally to explicitly converting with + # st_geogfromtext, or implicitly with the geography parameter + # conversion. + # + # 3. We handle different types using bind_processor, below. + # + bind_expression = sqlalchemy.sql.type_api.TypeEngine.bind_expression + + def bind_processor(self, dialect): + """ + SqlAlchemy plugin that controls how values are converted to parameters + + When we bind values, we always bind as text. We have to do + this because when we decide how to bind, we don't always know + what the values will be. + + This is not a user-facing method. + """ + + def process(bindvalue): + if isinstance(bindvalue, WKT): + return bindvalue.data + elif isinstance(bindvalue, WKB): + return bindvalue.wkt.data + else: + return bindvalue + + return process + + @staticmethod + def ElementType(data, srid=SRID, extended=True): + """ + Plugin for the Geoalchemy2 framework for constructing WKB objects. + + The framework wants a callable, which it assumes is a class + (this the name), for constructing a geographic element. + + We don't want `WKB` to accept extra arguments that it checks + and ignores, so we do that in this wrapper. + + This is not a user-facing method. + """ + if srid != SRID: + raise AssertionError("Bad srid", srid) + if not extended: + raise AssertionError("Extended must be True.") + return WKB(data) + + +@sqlalchemy.ext.compiler.compiles(geoalchemy2.functions.GenericFunction, "bigquery") +def _fixup_st_arguments(element, compiler, **kw): + """ + Compiler-plugin for the BigQuery that overrides how geographic functions are handled + + Geographic function (ST_...) get turned into + `geoalchemy2.functions.GenericFunction` objects by + Geoalchemy2. The code here overrides how they're handeled. + + We want arguments passed to have the GEOGRAPHY type associated + with them, when appropriate, where "when appropriate" is + determined by the `function documentation + `_.. + + This is not a user-facing function. + """ + argument_types = _argument_types.get(element.name.lower()) + if argument_types: + for argument_type, argument in zip(argument_types, element.clauses.clauses): + if isinstance(argument, BindParameter) and ( + argument.type is not argument_type + or not isinstance(argument.type, argument_type) + ): + argument.type = argument_type() + + return compiler.visit_function(element, **kw) + + +_argument_types = dict( + st_area=(GEOGRAPHY,), + st_asbinary=(GEOGRAPHY,), + st_asgeojson=(GEOGRAPHY,), + st_astext=(GEOGRAPHY,), + st_boundary=(GEOGRAPHY,), + st_centroid=(GEOGRAPHY,), + st_centroid_agg=(GEOGRAPHY,), + st_closestpoint=(GEOGRAPHY, GEOGRAPHY,), + st_clusterdbscan=(GEOGRAPHY,), + st_contains=(GEOGRAPHY, GEOGRAPHY,), + st_convexhull=(GEOGRAPHY,), + st_coveredby=(GEOGRAPHY, GEOGRAPHY,), + st_covers=(GEOGRAPHY, GEOGRAPHY,), + st_difference=(GEOGRAPHY, GEOGRAPHY,), + st_dimension=(GEOGRAPHY,), + st_disjoint=(GEOGRAPHY, GEOGRAPHY,), + st_distance=(GEOGRAPHY, GEOGRAPHY,), + st_dump=(GEOGRAPHY,), + st_dwithin=(GEOGRAPHY, GEOGRAPHY,), + st_endpoint=(GEOGRAPHY,), + st_equals=(GEOGRAPHY, GEOGRAPHY,), + st_exteriorring=(GEOGRAPHY,), + st_geohash=(GEOGRAPHY,), + st_intersection=(GEOGRAPHY, GEOGRAPHY,), + st_intersects=(GEOGRAPHY, GEOGRAPHY,), + st_intersectsbox=(GEOGRAPHY,), + st_iscollection=(GEOGRAPHY,), + st_isempty=(GEOGRAPHY,), + st_length=(GEOGRAPHY,), + st_makeline=(GEOGRAPHY, GEOGRAPHY,), + st_makepolygon=(GEOGRAPHY, GEOGRAPHY,), + st_makepolygonoriented=(GEOGRAPHY,), + st_maxdistance=(GEOGRAPHY, GEOGRAPHY,), + st_npoints=(GEOGRAPHY,), + st_numpoints=(GEOGRAPHY,), + st_perimeter=(GEOGRAPHY,), + st_pointn=(GEOGRAPHY,), + st_simplify=(GEOGRAPHY,), + st_snaptogrid=(GEOGRAPHY,), + st_startpoint=(GEOGRAPHY,), + st_touches=(GEOGRAPHY, GEOGRAPHY,), + st_union=(GEOGRAPHY, GEOGRAPHY,), + st_union_agg=(GEOGRAPHY,), + st_within=(GEOGRAPHY, GEOGRAPHY,), + st_x=(GEOGRAPHY,), + st_y=(GEOGRAPHY,), +) + +__all__ = ["GEOGRAPHY", "WKB", "WKT"] diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 3b3bda8e..077f06a8 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -122,6 +122,14 @@ def bigquery_regional_dataset(bigquery_client, bigquery_schema): bigquery_client.delete_dataset(dataset_id, delete_contents=True) +@pytest.fixture(autouse=True) +def cleanup_extra_tables(bigquery_client, bigquery_dataset): + common = "sample", "sample_one_row", "sample_view", "sample_dml_empty" + for table in bigquery_client.list_tables(bigquery_dataset): + if table.table_id not in common: + bigquery_client.delete_table(table) + + @pytest.fixture(scope="session", autouse=True) def cleanup_datasets(bigquery_client: bigquery.Client): for dataset in bigquery_client.list_datasets(): diff --git a/tests/system/test_alembic.py b/tests/system/test_alembic.py index db9ceb4f..81c686d1 100644 --- a/tests/system/test_alembic.py +++ b/tests/system/test_alembic.py @@ -22,13 +22,10 @@ import pytest from sqlalchemy import Column, DateTime, Integer, String -try: - import alembic # noqa -except ImportError: - alembic = None - import google.api_core.exceptions +alembic = pytest.importorskip("alembic") + @pytest.fixture def alembic_table(bigquery_dataset, bigquery_client): @@ -62,7 +59,6 @@ def get_table(table_name, data="table"): 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. diff --git a/tests/system/test_geography.py b/tests/system/test_geography.py new file mode 100644 index 00000000..18bcc7d4 --- /dev/null +++ b/tests/system/test_geography.py @@ -0,0 +1,304 @@ +# 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 pytest + +geoalchemy2 = pytest.importorskip("geoalchemy2") + + +def test_geoalchemy2_core(bigquery_dataset): + """Make sure GeoAlchemy 2 Core Tutorial works as adapted to only having geography + + https://geoalchemy-2.readthedocs.io/en/latest/core_tutorial.html + + Note: + + - Bigquery doesn't have ST_BUFFER + """ + + # Connect to the DB + + from sqlalchemy import create_engine + + engine = create_engine(f"bigquery:///{bigquery_dataset}") + + # Create the Table + + from sqlalchemy import Table, Column, String, MetaData + from sqlalchemy_bigquery import GEOGRAPHY + + metadata = MetaData() + lake_table = Table( + "lake", metadata, Column("name", String), Column("geog", GEOGRAPHY) + ) + + lake_table.create(engine) + + # Insertions + + conn = engine.connect() + + conn.execute( + lake_table.insert().values( + name="Majeur", geog="POLYGON((0 0,1 0,1 1,0 1,0 0))", + ) + ) + + conn.execute( + lake_table.insert(), + [ + {"name": "Garde", "geog": "POLYGON((1 0,3 0,3 2,1 2,1 0))"}, + {"name": "Orta", "geog": "POLYGON((3 0,6 0,6 3,3 3,3 0))"}, + ], + ) + + # Selections + + from sqlalchemy.sql import select + + assert sorted( + (r.name, r.geog.desc[:4]) for r in conn.execute(select([lake_table])) + ) == [("Garde", "0103"), ("Majeur", "0103"), ("Orta", "0103")] + + # Spatial query + + from sqlalchemy import func + + [[result]] = conn.execute( + select([lake_table.c.name], func.ST_Contains(lake_table.c.geog, "POINT(4 1)")) + ) + assert result == "Orta" + + assert sorted( + (r.name, int(r.area)) + for r in conn.execute( + select([lake_table.c.name, lake_table.c.geog.ST_AREA().label("area")]) + ) + ) == [("Garde", 49452374328), ("Majeur", 12364036567), ("Orta", 111253664228)] + + # Extra: Make sure we can save a retrieved value back: + + [[geog]] = conn.execute(select([lake_table.c.geog], lake_table.c.name == "Garde")) + conn.execute(lake_table.insert().values(name="test", geog=geog)) + assert ( + int( + list( + conn.execute( + select([lake_table.c.geog.st_area()], lake_table.c.name == "test") + ) + )[0][0] + ) + == 49452374328 + ) + + # and, while we're at it, that we can insert WKTs, although we + # normally wouldn't want to. + from sqlalchemy_bigquery import WKT + + conn.execute( + lake_table.insert().values( + name="test2", geog=WKT("POLYGON((1 0,3 0,3 2,1 2,1 0))"), + ) + ) + assert ( + int( + list( + conn.execute( + select([lake_table.c.geog.st_area()], lake_table.c.name == "test2") + ) + )[0][0] + ) + == 49452374328 + ) + + +def test_geoalchemy2_orm(bigquery_dataset): + """Make sure GeoAlchemy 2 ORM Tutorial works as adapted to only having geometry + + https://geoalchemy-2.readthedocs.io/en/latest/orm_tutorial.html + """ + + # Connect to the DB + + from sqlalchemy import create_engine + + engine = create_engine(f"bigquery:///{bigquery_dataset}") + + # Declare a Mapping + + from sqlalchemy.ext.declarative import declarative_base + from sqlalchemy import Column, Integer, String + from sqlalchemy_bigquery import GEOGRAPHY + + Base = declarative_base() + + class Lake(Base): + __tablename__ = "lake" + # The ORM insists on an id, but bigquery doesn't auto-assign + # ids, so we'll have to provide them below. + id = Column(Integer, primary_key=True) + name = Column(String) + geog = Column(GEOGRAPHY) + + # Create the Table in the Database + + Lake.__table__.create(engine) + + # Create an Instance of the Mapped Class + + lake = Lake(id=1, name="Majeur", geog="POLYGON((0 0,1 0,1 1,0 1,0 0))") + + # Create a Session + + from sqlalchemy.orm import sessionmaker + + Session = sessionmaker(bind=engine) + + session = Session() + + # Add New Objects + + session.add(lake) + session.commit() + + our_lake = session.query(Lake).filter_by(name="Majeur").first() + assert our_lake.name == "Majeur" + + from geoalchemy2 import WKBElement + + assert isinstance(our_lake.geog, WKBElement) + + session.add_all( + [ + Lake(id=2, name="Garde", geog="POLYGON((1 0,3 0,3 2,1 2,1 0))"), + Lake(id=3, name="Orta", geog="POLYGON((3 0,6 0,6 3,3 3,3 0))"), + ] + ) + + session.commit() + + # Query + + query = session.query(Lake).order_by(Lake.name) + assert [lake.name for lake in query] == ["Garde", "Majeur", "Orta"] + + assert [lake.name for lake in session.query(Lake).order_by(Lake.name).all()] == [ + "Garde", + "Majeur", + "Orta", + ] + + # Make Spatial Queries + + from sqlalchemy import func + + query = session.query(Lake).filter(func.ST_Contains(Lake.geog, "POINT(4 1)")) + + assert [lake.name for lake in query] == ["Orta"] + + query = ( + session.query(Lake) + .filter(Lake.geog.ST_Intersects("LINESTRING(2 1,4 1)")) + .order_by(Lake.name) + ) + assert [lake.name for lake in query] == ["Garde", "Orta"] + + lake = session.query(Lake).filter_by(name="Garde").one() + assert session.scalar(lake.geog.ST_Intersects("LINESTRING(2 1,4 1)")) + + # Use Other Spatial Functions + query = session.query(Lake.name, func.ST_Area(Lake.geog).label("area")).order_by( + Lake.name + ) + assert [(name, int(area)) for name, area in query] == [ + ("Garde", 49452374328), + ("Majeur", 12364036567), + ("Orta", 111253664228), + ] + + query = session.query(Lake.name, Lake.geog.ST_Area().label("area")).order_by( + Lake.name + ) + assert [(name, int(area)) for name, area in query] == [ + ("Garde", 49452374328), + ("Majeur", 12364036567), + ("Orta", 111253664228), + ] + + +def test_geoalchemy2_orm_w_relationship(bigquery_dataset): + from sqlalchemy import create_engine + + engine = create_engine(f"bigquery:///{bigquery_dataset}") + + from sqlalchemy import Column, Integer, String + from sqlalchemy_bigquery import GEOGRAPHY + + from sqlalchemy.ext.declarative import declarative_base + + Base = declarative_base() + + class Treasure(Base): + __tablename__ = "treasure" + id = Column(Integer, primary_key=True) + geog = Column(GEOGRAPHY) + + Treasure.__table__.create(engine) + + from sqlalchemy.orm import relationship, backref + + class Lake(Base): + __tablename__ = "lake" + id = Column(Integer, primary_key=True) + name = Column(String) + geog = Column(GEOGRAPHY) + treasures = relationship( + "Treasure", + primaryjoin="func.ST_Contains(foreign(Lake.geog), Treasure.geog).as_comparison(1, 2)", + backref=backref("lake_rel", uselist=False), + viewonly=True, + uselist=True, + ) + + Lake.__table__.create(engine) + + from sqlalchemy.orm import sessionmaker + + Session = sessionmaker(bind=engine) + + session = Session() + + session.add_all( + [ + Treasure(id=21, geog="Point(1.5 1)"), + Treasure(id=22, geog="Point(2.5 1.5)"), + Treasure(id=31, geog="Point(4.5 2)"), + Treasure(id=42, geog="Point(5.5 1.5)"), + Lake(id=2, name="Garde", geog="POLYGON((1 0,3 0,3 2,1 2,1 0))"), + Lake(id=3, name="Orta", geog="POLYGON((3 0,6 0,6 3,3 3,3 0))"), + ] + ) + + session.commit() + + lakes = session.query(Lake).order_by(Lake.name).all() + assert [(lake.id, sorted(t.id for t in lake.treasures)) for lake in lakes] == [ + (2, [21, 22]), + (3, [31, 42]), + ] diff --git a/tests/system/test_sqlalchemy_bigquery.py b/tests/system/test_sqlalchemy_bigquery.py index 0fe878b2..06024368 100644 --- a/tests/system/test_sqlalchemy_bigquery.py +++ b/tests/system/test_sqlalchemy_bigquery.py @@ -234,7 +234,7 @@ def test_engine_with_dataset(engine_using_test_dataset, bigquery_dataset): table_one_row = Table( "sample_one_row", MetaData(bind=engine_using_test_dataset), autoload=True ) - rows = table_one_row.select().execute().fetchall() + rows = table_one_row.select(use_labels=True).execute().fetchall() assert list(rows[0]) == ONE_ROW_CONTENTS_EXPANDED table_one_row = Table( @@ -242,7 +242,7 @@ def test_engine_with_dataset(engine_using_test_dataset, bigquery_dataset): MetaData(bind=engine_using_test_dataset), autoload=True, ) - rows = table_one_row.select().execute().fetchall() + rows = table_one_row.select(use_labels=True).execute().fetchall() # verify that we are pulling from the specifically-named dataset, # instead of pulling from the default dataset of the engine (which # does not have this table at all) @@ -279,7 +279,12 @@ def test_reflect_select(table, table_using_test_dataset): assert isinstance(table.c["nested_record.record.name"].type, types.String) assert isinstance(table.c.array.type, types.ARRAY) - rows = table.select().execute().fetchall() + # Force unique column labels using `use_labels` below to deal + # with BQ sometimes complaining about duplicate column names + # when a destination table is specified, even though no + # destination table is specified. When this test was written, + # `use_labels` was forced by the dialect. + rows = table.select(use_labels=True).execute().fetchall() assert len(rows) == 1000 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 1c78b12d..e5de882d 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -71,6 +71,17 @@ def ex(sql, *args, **kw): conn.close() +@pytest.fixture() +def last_query(faux_conn): + def last_query(sql, params=None, offset=1): + actual_sql, actual_params = faux_conn.test_data["execute"][-offset] + assert actual_sql == sql + if params is not None: + assert actual_params == params + + return last_query + + @pytest.fixture() def metadata(): return sqlalchemy.MetaData() diff --git a/tests/unit/test_geography.py b/tests/unit/test_geography.py new file mode 100644 index 00000000..3ee2cce6 --- /dev/null +++ b/tests/unit/test_geography.py @@ -0,0 +1,179 @@ +# 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 pytest + +from conftest import setup_table + +geoalchemy2 = pytest.importorskip("geoalchemy2") + + +def test_geoalchemy2_core(faux_conn, last_query): + """Make sure GeoAlchemy 2 Core Tutorial works as adapted to only having geometry + """ + conn = faux_conn + + # Create the Table + + from sqlalchemy import Column, String + from sqlalchemy_bigquery import GEOGRAPHY + + lake_table = setup_table( + conn, "lake", Column("name", String), Column("geog", GEOGRAPHY) + ) + + # Insertions + + conn.execute( + lake_table.insert().values( + name="Majeur", geog="POLYGON((0 0,1 0,1 1,0 1,0 0))", + ) + ) + + last_query( + "INSERT INTO `lake` (`name`, `geog`)" + " VALUES (%(name:STRING)s, %(geog:geography)s)", + ({"geog": "POLYGON((0 0,1 0,1 1,0 1,0 0))", "name": "Majeur"}), + ) + + conn.execute( + lake_table.insert(), + [ + {"name": "Garde", "geog": "POLYGON((1 0,3 0,3 2,1 2,1 0))"}, + {"name": "Orta", "geog": "POLYGON((3 0,6 0,6 3,3 3,3 0))"}, + ], + ) + last_query( + "INSERT INTO `lake` (`name`, `geog`)" + " VALUES (%(name:STRING)s, %(geog:geography)s)", + {"name": "Garde", "geog": "POLYGON((1 0,3 0,3 2,1 2,1 0))"}, + offset=2, + ) + last_query( + "INSERT INTO `lake` (`name`, `geog`)" + " VALUES (%(name:STRING)s, %(geog:geography)s)", + {"name": "Orta", "geog": "POLYGON((3 0,6 0,6 3,3 3,3 0))"}, + ) + + # Selections + + from sqlalchemy.sql import select + + try: + conn.execute(select([lake_table])) + except Exception: + pass # sqlite had no special functions :) + last_query( + "SELECT `lake`.`name`, ST_AsBinary(`lake`.`geog`) AS `geog` \n" "FROM `lake`" + ) + + # Spatial query + + from sqlalchemy import func + + try: + conn.execute( + select( + [lake_table.c.name], func.ST_Contains(lake_table.c.geog, "POINT(4 1)") + ) + ) + except Exception: + pass # sqlite had no special functions :) + last_query( + "SELECT `lake`.`name` \n" + "FROM `lake` \n" + "WHERE ST_Contains(`lake`.`geog`, %(ST_Contains_1:geography)s)", + {"ST_Contains_1": "POINT(4 1)"}, + ) + + try: + conn.execute( + select([lake_table.c.name, lake_table.c.geog.ST_AREA().label("area")]) + ) + except Exception: + pass # sqlite had no special functions :) + last_query("SELECT `lake`.`name`, ST_Area(`lake`.`geog`) AS `area` \nFROM `lake`") + + # Extra: Make sure we can save a retrieved value back: + + from sqlalchemy_bigquery import WKB, WKT + + geog = WKT("point(0 0)").wkb + assert isinstance(geog, WKB) + assert geog.data == ( + b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00" + b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + ) + conn.execute(lake_table.insert().values(name="test", geog=geog)) + last_query( + "INSERT INTO `lake` (`name`, `geog`)" + " VALUES (%(name:STRING)s, %(geog:geography)s)", + {"name": "test", "geog": "POINT (0 0)"}, + ) + + # and, while we're at it, that we can insert WKTs, although we + # normally wouldn't want to. + + conn.execute( + lake_table.insert().values( + name="test2", geog=WKT("POLYGON((1 0,3 0,3 2,1 2,1 0))"), + ) + ) + last_query( + "INSERT INTO `lake` (`name`, `geog`)" + " VALUES (%(name:STRING)s, %(geog:geography)s)", + {"name": "test2", "geog": "POLYGON((1 0,3 0,3 2,1 2,1 0))"}, + ) + + +def test_GEOGRAPHY_ElementType_bad_srid(): + from sqlalchemy_bigquery import GEOGRAPHY + + with pytest.raises(AssertionError, match="Bad srid"): + GEOGRAPHY.ElementType("data", srid=-1) + + +def test_GEOGRAPHY_ElementType_bad_extended(): + from sqlalchemy_bigquery import GEOGRAPHY + + with pytest.raises(AssertionError, match="Extended must be True."): + GEOGRAPHY.ElementType("data", extended=False) + + +def test_GEOGRAPHY_ElementType(): + from sqlalchemy_bigquery import GEOGRAPHY, WKB + + data = GEOGRAPHY.ElementType("data") + assert isinstance(data, WKB) + assert (data.data, data.srid, data.extended) == ("data", 4326, True) + + +def test_calling_st_functions_that_dont_take_geographies(faux_conn, last_query): + from sqlalchemy import select, func + + try: + faux_conn.execute(select([func.ST_GEOGFROMTEXT("point(0 0)")])) + except Exception: + pass # sqlite had no special functions :) + + last_query( + "SELECT ST_AsBinary(ST_GeogFromText(%(ST_GeogFromText_2:STRING)s))" + " AS `ST_GeogFromText_1`", + dict(ST_GeogFromText_2="point(0 0)"), + )