SNOW-629084 follow sqlalchemy 2.0 migration guide to verify the requirements are met (#317)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -9,6 +9,9 @@ Source code is also available at:
 
 # Release Notes
 
+- v1.4.1(Unreleased)
+  - snowflake-sqlalchemy is now SQLAlchemy 2.0 compatible.
+
 - v1.4.0(July 20, 2022)
 
   - Added support for `regexp_match`, `regexp_replace` in `sqlalchemy.sql.expression.ColumnOperators`.

setup.cfg

@@ -3,7 +3,7 @@ universal = 1
 
 [metadata]
 name = snowflake-sqlalchemy
-version = 1.4.0
+version = 1.4.1
 description = Snowflake SQLAlchemy Dialect
 long_description = file: DESCRIPTION.md
 long_description_content_type = text/markdown

src/snowflake/sqlalchemy/snowdialect.py

@@ -329,15 +329,15 @@ def _get_schema_primary_keys(self, connection, schema, **kw):
         ans = {}
         key_sequence_order_map = defaultdict(list)
         for row in result:
-            table_name = self.normalize_name(row["table_name"])
-            key_sequence_order_map[table_name].append(row["key_sequence"])
+            table_name = self.normalize_name(row._mapping["table_name"])
+            key_sequence_order_map[table_name].append(row._mapping["key_sequence"])
             if table_name not in ans:
                 ans[table_name] = {
                     "constrained_columns": [],
-                    "name": self.normalize_name(row["constraint_name"]),
+                    "name": self.normalize_name(row._mapping["constraint_name"]),
                 }
             ans[table_name]["constrained_columns"].append(
-                self.normalize_name(row["column_name"])
+                self.normalize_name(row._mapping["column_name"])
             )
 
         for k, v in ans.items():
@@ -368,16 +368,16 @@ def _get_schema_unique_constraints(self, connection, schema, **kw):
         )
         unique_constraints = {}
         for row in result:
-            name = self.normalize_name(row["constraint_name"])
+            name = self.normalize_name(row._mapping["constraint_name"])
             if name not in unique_constraints:
                 unique_constraints[name] = {
-                    "column_names": [self.normalize_name(row["column_name"])],
+                    "column_names": [self.normalize_name(row._mapping["column_name"])],
                     "name": name,
-                    "table_name": self.normalize_name(row["table_name"]),
+                    "table_name": self.normalize_name(row._mapping["table_name"]),
                 }
             else:
                 unique_constraints[name]["column_names"].append(
-                    self.normalize_name(row["column_name"])
+                    self.normalize_name(row._mapping["column_name"])
                 )
 
         ans = defaultdict(list)
@@ -409,12 +409,14 @@ def _get_schema_foreign_keys(self, connection, schema, **kw):
         foreign_key_map = {}
         key_sequence_order_map = defaultdict(list)
         for row in result:
-            name = self.normalize_name(row["fk_name"])
-            key_sequence_order_map[name].append(row["key_sequence"])
+            name = self.normalize_name(row._mapping["fk_name"])
+            key_sequence_order_map[name].append(row._mapping["key_sequence"])
             if name not in foreign_key_map:
-                referred_schema = self.normalize_name(row["pk_schema_name"])
+                referred_schema = self.normalize_name(row._mapping["pk_schema_name"])
                 foreign_key_map[name] = {
-                    "constrained_columns": [self.normalize_name(row["fk_column_name"])],
+                    "constrained_columns": [
+                        self.normalize_name(row._mapping["fk_column_name"])
+                    ],
                     # referred schema should be None in context where it doesn't need to be specified
                     # https://docs.sqlalchemy.org/en/14/core/reflection.html#reflection-schema-qualified-interaction
                     "referred_schema": (
@@ -423,23 +425,31 @@ def _get_schema_foreign_keys(self, connection, schema, **kw):
                         not in (self.default_schema_name, current_schema)
                         else None
                     ),
-                    "referred_table": self.normalize_name(row["pk_table_name"]),
-                    "referred_columns": [self.normalize_name(row["pk_column_name"])],
+                    "referred_table": self.normalize_name(
+                        row._mapping["pk_table_name"]
+                    ),
+                    "referred_columns": [
+                        self.normalize_name(row._mapping["pk_column_name"])
+                    ],
                     "name": name,
-                    "table_name": self.normalize_name(row["fk_table_name"]),
+                    "table_name": self.normalize_name(row._mapping["fk_table_name"]),
                 }
                 options = {}
-                if self.normalize_name(row["delete_rule"]) != "NO ACTION":
-                    options["ondelete"] = self.normalize_name(row["delete_rule"])
-                if self.normalize_name(row["update_rule"]) != "NO ACTION":
-                    options["onupdate"] = self.normalize_name(row["update_rule"])
+                if self.normalize_name(row._mapping["delete_rule"]) != "NO ACTION":
+                    options["ondelete"] = self.normalize_name(
+                        row._mapping["delete_rule"]
+                    )
+                if self.normalize_name(row._mapping["update_rule"]) != "NO ACTION":
+                    options["onupdate"] = self.normalize_name(
+                        row._mapping["update_rule"]
+                    )
                 foreign_key_map[name]["options"] = options
             else:
                 foreign_key_map[name]["constrained_columns"].append(
-                    self.normalize_name(row["fk_column_name"])
+                    self.normalize_name(row._mapping["fk_column_name"])
                 )
                 foreign_key_map[name]["referred_columns"].append(
-                    self.normalize_name(row["pk_column_name"])
+                    self.normalize_name(row._mapping["pk_column_name"])
                 )
 
         ans = {}
@@ -847,7 +857,11 @@ def get_table_comment(self, connection, table_name, schema=None, **kw):
             # the "table" being reflected is actually a view
             result = self._get_view_comment(connection, table_name, schema)
 
-        return {"text": result["comment"] if result and result["comment"] else None}
+        return {
+            "text": result._mapping["comment"]
+            if result and result._mapping["comment"]
+            else None
+        }
 
 
 @sa_vnt.listens_for(Table, "before_create")

tests/conftest.py

@@ -6,6 +6,7 @@
 import sys
 import time
 import uuid
+from functools import partial
 from logging import getLogger
 
 import pytest
@@ -27,6 +28,18 @@
 else:
     TEST_SCHEMA = "sqlalchemy_tests_" + str(uuid.uuid4()).replace("-", "_")
 
+create_engine_with_future_flag = create_engine
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--run_v20_sqlalchemy",
+        help="Use only 2.0 SQLAlchemy APIs, any legacy features (< 2.0) will not be supported."
+        "Turning on this option will set future flag to True on Engine and Session objects according to"
+        "the migration guide: https://docs.sqlalchemy.org/en/14/changelog/migration_20.html",
+        action="store_true",
+    )
+
 
 @pytest.fixture(scope="session")
 def on_travis():
@@ -139,7 +152,7 @@ def get_engine(user=None, password=None, account=None, schema=None):
 
     from sqlalchemy.pool import NullPool
 
-    engine = create_engine(
+    engine = create_engine_with_future_flag(
         URL(
             user=ret["user"],
             password=ret["password"],
@@ -203,6 +216,19 @@ def sql_compiler():
     ).replace("\n", "")
 
 
+@pytest.fixture(scope="session")
+def run_v20_sqlalchemy(pytestconfig):
+    return pytestconfig.option.run_v20_sqlalchemy
+
+
+def pytest_sessionstart(session):
+    # patch the create_engine with future flag
+    global create_engine_with_future_flag
+    create_engine_with_future_flag = partial(
+        create_engine, future=session.config.option.run_v20_sqlalchemy
+    )
+
+
 def running_on_public_ci() -> bool:
     """Whether or not tests are currently running on one of our public CIs."""
     return os.getenv("GITHUB_ACTIONS") == "true"

tests/test_compiler.py

@@ -96,6 +96,6 @@ def test_quoted_name_label(engine_testaccount):
 
     for t in test_cases:
         col = column("colname").label(t["label"])
-        sel_from_tbl = select([col]).group_by(col).select_from(table("abc"))
+        sel_from_tbl = select(col).group_by(col).select_from(table("abc"))
         compiled_result = sel_from_tbl.compile()
         assert str(compiled_result) == t["output"]

tests/test_copy.py

@@ -64,7 +64,7 @@ def test_copy_into_location(engine_testaccount, sql_compiler):
         "(KMS_KEY_ID='1234abcd-12ab-34cd-56ef-1234567890ab' TYPE='AWS_SSE_KMS')"
     )
     copy_stmt_2 = CopyIntoStorage(
-        from_=select([food_items]).where(food_items.c.id == 1),  # Test sub-query
+        from_=select(food_items).where(food_items.c.id == 1),  # Test sub-query
         into=AWSBucket.from_uri("s3://backup")
         .credentials(aws_role="some_iam_role")
         .encryption_aws_sse_s3(),