corrections to tests and addition of 2.0 conditionals

Author: chalmerlowe

packages/sqlalchemy-bigquery/tests/unit/conftest.py

@@ -26,6 +26,7 @@
 import sqlalchemy
 
 from sqlalchemy_bigquery.base import BigQueryDDLCompiler, BigQueryDialect
+from sqlalchemy.sql.compiler import Compiled
 
 from . import fauxdbi
 
@@ -69,6 +70,18 @@ def ex(sql, *args, **kw):
 
                 ex("create table comments" " (key string primary key, comment string)")
 
+                # Modernize faux_conn for SQLAlchemy 2.0+ by allowing it to execute
+                # Compiled objects (common in this test suite)
+                if sqlalchemy_version >= packaging.version.parse("2.0"):
+                    original_execute = conn.execute
+
+                    def execute(statement, *args, **kw):
+                        if isinstance(statement, Compiled):
+                            statement = sqlalchemy.text(str(statement))
+                        return original_execute(statement, *args, **kw)
+
+                    conn.execute = execute
+
                 yield conn
                 conn.close()
 

packages/sqlalchemy-bigquery/tests/unit/test__struct.py

@@ -16,13 +16,19 @@
 # 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 packaging.version
 import pytest
 
 import sqlalchemy
 
 
+sqlalchemy_version = packaging.version.parse(sqlalchemy.__version__)
+SQL_2_0 = sqlalchemy_version >= packaging.version.parse("2.0")
+
+
 def _test_struct():
     from sqlalchemy_bigquery import STRUCT
 
@@ -90,9 +96,24 @@ def test_struct_traversal_project(faux_conn, expr, sql):
 @pytest.mark.parametrize(
     "expr,sql",
     [
-        (_col()["name"] == "x", "(`t`.`person`.name) = %(param_1:STRING)s"),
-        (_col()["Name"] == "x", "(`t`.`person`.Name) = %(param_1:STRING)s"),
-        (_col().NAME == "x", "(`t`.`person`.NAME) = %(param_1:STRING)s"),
+        (
+            _col()["name"] == "bob",
+            "(`t`.`person`.name) = %(param_1:STRING)s"
+            if not SQL_2_0
+            else "`t`.`person`.name = %(param_1:STRING)s",
+        ),
+        (
+            _col()["Name"] == "bob",
+            "(`t`.`person`.Name) = %(param_1:STRING)s"
+            if not SQL_2_0
+            else "`t`.`person`.Name = %(param_1:STRING)s",
+        ),
+        (
+            _col()["NAME"] == "bob",
+            "(`t`.`person`.NAME) = %(param_1:STRING)s"
+            if not SQL_2_0
+            else "`t`.`person`.NAME = %(param_1:STRING)s",
+        ),
         (
             _col().children[0] == dict(name="foo", bdate=datetime.date(2020, 1, 1)),
             "(`t`.`person`.children)[OFFSET(%(param_1:INT64)s)]"
@@ -106,11 +127,17 @@ def test_struct_traversal_project(faux_conn, expr, sql):
         (
             _col().children[0]["bdate"] == datetime.date(2021, 8, 30),
             "(((`t`.`person`.children)[OFFSET(%(param_1:INT64)s)]).bdate)"
+            " = %(param_2:DATE)s"
+            if not SQL_2_0
+            else "((`t`.`person`.children)[OFFSET(%(param_1:INT64)s)]).bdate"
             " = %(param_2:DATE)s",
         ),
         (
             _col().children[0].bdate == datetime.date(2021, 8, 30),
             "(((`t`.`person`.children)[OFFSET(%(param_1:INT64)s)]).bdate)"
+            " = %(param_2:DATE)s"
+            if not SQL_2_0
+            else "((`t`.`person`.children)[OFFSET(%(param_1:INT64)s)]).bdate"
             " = %(param_2:DATE)s",
         ),
     ],

packages/sqlalchemy-bigquery/tests/unit/test_compliance.py

@@ -56,7 +56,9 @@ def test_distinct_selectable_in_unions(faux_conn):
     s2 = select(table).where(table.c.id == 3).distinct()
 
     u1 = union(s1, s2).limit(2)
-    assert_result(faux_conn, u1.order_by(u1.c.id), [(2, 2, 3), (3, 3, 4)])
+    assert_result(
+        faux_conn, u1.order_by(u1.selected_columns.id), [(2, 2, 3), (3, 3, 4)]
+    )
 
 
 def test_limit_offset_aliased_selectable_in_unions(faux_conn):
@@ -79,7 +81,9 @@ def test_limit_offset_aliased_selectable_in_unions(faux_conn):
     )
 
     u1 = union(s1, s2).limit(2)
-    assert_result(faux_conn, u1.order_by(u1.c.id), [(2, 2, 3), (3, 3, 4)])
+    assert_result(
+        faux_conn, u1.order_by(u1.selected_columns.id), [(2, 2, 3), (3, 3, 4)]
+    )
 
 
 def test_percent_sign_round_trip(faux_conn, metadata):