feat: support UPDATE + JOIN in BigQuery dialect (#1083)

Linchin · gcf-owl-bot[bot] · web-flow · commit d766d21053f7 · 2024-06-13T09:53:24.000-07:00
* feat: support UPDATE JOIN * vendor code from sqlalchemy * remove code and add comments * add vendored folder to path in MANIFEST.in * remove extra code * include third_party in pytest * update MANIFEST.in * add init file * add init file * add pyproject.toml * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * update noxfile and owlbot * lint * fix owlbot * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * undo changes for testing * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * add comments to pyproject.toml --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -16,7 +16,8 @@
 
 # Generated by synthtool. DO NOT EDIT!
 include README.rst LICENSE
-recursive-include google *.json *.proto py.typed
+recursive-include third_party/sqlalchemy_bigquery_vendored *
+recursive-include sqlalchemy_bigquery *.json *.proto py.typed
 recursive-include tests *
 global-exclude *.py[co]
 global-exclude __pycache__
diff --git a/noxfile.py b/noxfile.py
@@ -31,7 +31,14 @@
 FLAKE8_VERSION = "flake8==6.1.0"
 BLACK_VERSION = "black[jupyter]==23.7.0"
 ISORT_VERSION = "isort==5.11.0"
-LINT_PATHS = ["docs", "sqlalchemy_bigquery", "tests", "noxfile.py", "setup.py"]
+LINT_PATHS = [
+    "third_party",
+    "docs",
+    "sqlalchemy_bigquery",
+    "tests",
+    "noxfile.py",
+    "setup.py",
+]
 
 DEFAULT_PYTHON_VERSION = "3.8"
 
diff --git a/owlbot.py b/owlbot.py
@@ -15,6 +15,7 @@
 """This script is used to synthesize generated parts of this library."""
 
 import pathlib
+import re
 
 import synthtool as s
 from synthtool import gcp
@@ -76,13 +77,24 @@
     "import re\nimport shutil",
 )
 
+s.replace(
+    ["noxfile.py"],
+    "LINT_PATHS = \[",
+    "LINT_PATHS = [\"third_party\", "
+)
 
 s.replace(
     ["noxfile.py"],
     "--cov=google",
     "--cov=sqlalchemy_bigquery",
 )
 
+s.replace(
+    ["noxfile.py"],
+    """os.path.join("tests", "unit"),""",
+    """os.path.join("tests", "unit"),
+        os.path.join("third_party", "sqlalchemy_bigquery_vendored"),""",
+)
 
 s.replace(
     ["noxfile.py"],
@@ -284,6 +296,15 @@ def system_noextras(session):
 """,
 )
 
+
+# Make sure build includes all necessary files.
+s.replace(
+    ["MANIFEST.in"],
+    re.escape("recursive-include google"),
+    """recursive-include third_party/sqlalchemy_bigquery_vendored *
+recursive-include sqlalchemy_bigquery""",
+)
+
 # ----------------------------------------------------------------------------
 # Samples templates
 # ----------------------------------------------------------------------------
diff --git a/pyproject.toml b/pyproject.toml
@@ -0,0 +1,5 @@
+# Added so third_party folder is included when running `pip install -e .`
+# See PR #1083 for more detail
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
diff --git a/setup.py b/setup.py
@@ -22,6 +22,7 @@
 import itertools
 import os
 import re
+import setuptools
 from setuptools import setup
 
 # Package metadata.
@@ -67,6 +68,16 @@ def readme():
 
 extras["all"] = set(itertools.chain.from_iterable(extras.values()))
 
+packages = [
+    package
+    for package in setuptools.find_namespace_packages()
+    if package.startswith("sqlalchemy_bigquery")
+] + [
+    package
+    for package in setuptools.find_namespace_packages("third_party")
+    if package.startswith("sqlalchemy_bigquery_vendored")
+]
+
 setup(
     name=name,
     version=version,
@@ -75,7 +86,11 @@ def readme():
     long_description_content_type="text/x-rst",
     author="The Sqlalchemy-Bigquery Authors",
     author_email="googleapis-packages@google.com",
-    packages=["sqlalchemy_bigquery"],
+    package_dir={
+        "sqlalchemy-bigquery": "sqlalchemy_bigquery",
+        "sqlalchemy_bigquery_vendored": "third_party/sqlalchemy_bigquery_vendored",
+    },
+    packages=packages,
     url="https://github.com/googleapis/python-bigquery-sqlalchemy",
     keywords=["bigquery", "sqlalchemy"],
     classifiers=[
diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py
@@ -60,6 +60,7 @@
 
 from .parse_url import parse_url
 from . import _helpers, _struct, _types
+import sqlalchemy_bigquery_vendored.sqlalchemy.postgresql.base as vendored_postgresql
 
 # Illegal characters is intended to be all characters that are not explicitly
 # allowed as part of the flexible column names.
@@ -189,7 +190,7 @@ def pre_exec(self):
         )
 
 
-class BigQueryCompiler(_struct.SQLCompiler, SQLCompiler):
+class BigQueryCompiler(_struct.SQLCompiler, vendored_postgresql.PGCompiler):
     compound_keywords = SQLCompiler.compound_keywords.copy()
     compound_keywords[selectable.CompoundSelect.UNION] = "UNION DISTINCT"
     compound_keywords[selectable.CompoundSelect.UNION_ALL] = "UNION ALL"
diff --git a/tests/unit/test_compiler.py b/tests/unit/test_compiler.py
@@ -161,6 +161,35 @@ def prepare_implicit_join_base_query(
     return q
 
 
+# Test vendored method update_from_clause()
+# from sqlalchemy_bigquery_vendored.sqlalchemy.postgresql.base.PGCompiler
+def test_update_from_clause(faux_conn, metadata):
+    table1 = setup_table(
+        faux_conn,
+        "table1",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.String),
+        sqlalchemy.Column("bar", sqlalchemy.Integer),
+    )
+    table2 = setup_table(
+        faux_conn,
+        "table2",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.String),
+        sqlalchemy.Column("bar", sqlalchemy.Integer),
+    )
+
+    stmt = (
+        sqlalchemy.update(table1)
+        .where(table1.c.foo == table2.c.foo)
+        .where(table2.c.bar == 1)
+        .values(bar=2)
+    )
+    expected_sql = "UPDATE `table1` SET `bar`=%(bar:INT64)s FROM `table2` WHERE `table1`.`foo` = `table2`.`foo` AND `table2`.`bar` = %(bar_1:INT64)s"
+    found_sql = stmt.compile(faux_conn).string
+    assert found_sql == expected_sql
+
+
 @sqlalchemy_before_2_0
 def test_no_implicit_join_asterix_for_inner_unnest_before_2_0(faux_conn, metadata):
     # See: https://github.com/googleapis/python-bigquery-sqlalchemy/issues/368
diff --git a/third_party/__init__.py b/third_party/__init__.py
diff --git a/third_party/sqlalchemy_bigquery_vendored/__init__.py b/third_party/sqlalchemy_bigquery_vendored/__init__.py
diff --git a/third_party/sqlalchemy_bigquery_vendored/py.typed b/third_party/sqlalchemy_bigquery_vendored/py.typed
diff --git a/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/AUTHORS b/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/AUTHORS
@@ -0,0 +1,30 @@
+SQLAlchemy was created by Michael Bayer.
+
+Major contributing authors include:
+
+- Mike Bayer
+- Jason Kirtland
+- Michael Trier
+- Diana Clarke
+- Gaetan de Menten
+- Lele Gaifax
+- Jonathan Ellis
+- Gord Thompson
+- Federico Caselli
+- Philip Jenvey
+- Rick Morrison
+- Chris Withers
+- Ants Aasma
+- Sheila Allen
+- Paul Johnston
+- Tony Locke
+- Hajime Nakagami
+- Vraj Mohan
+- Robert Leftwich
+- Taavi Burns
+- Jonathan Vanasco
+- Jeff Widman
+- Scott Dugas
+- Dobes Vandermeer
+- Ville Skytta
+- Rodrigo Menezes
diff --git a/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/LICENSE b/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/LICENSE
@@ -0,0 +1,19 @@
+Copyright 2005-2024 SQLAlchemy authors and contributors <see AUTHORS file>.
+
+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.
diff --git a/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/__init__.py b/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/__init__.py
@@ -0,0 +1,6 @@
+# __init__.py
+# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
diff --git a/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/postgresql/__init__.py b/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/postgresql/__init__.py
diff --git a/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/postgresql/base.py b/third_party/sqlalchemy_bigquery_vendored/sqlalchemy/postgresql/base.py
@@ -0,0 +1,19 @@
+# dialects/postgresql/base.py
+# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: ignore-errors
+
+from sqlalchemy.sql import compiler
+
+
+class PGCompiler(compiler.SQLCompiler):
+    def update_from_clause(
+        self, update_stmt, from_table, extra_froms, from_hints, **kw
+    ):
+        kw["asfrom"] = True
+        return "FROM " + ", ".join(
+            t._compiler_dispatch(self, fromhints=from_hints, **kw) for t in extra_froms
+        )
