Merge pull request #70 from obendidi/sqlalchemy_v2

sqlalchemy v2 (fix tests)

Author: chrisguidry

prefect_sqlalchemy/database.py

@@ -16,6 +16,7 @@
 else:
     from pydantic import AnyUrl, Field, SecretStr
 
+from sqlalchemy import __version__ as SQLALCHEMY_VERSION
 from sqlalchemy.engine import Connection, Engine, create_engine
 from sqlalchemy.engine.cursor import CursorResult
 from sqlalchemy.engine.url import URL, make_url
@@ -67,6 +68,8 @@ async def _execute(
     if async_supported:
         result = await result
         await connection.commit()
+    elif SQLALCHEMY_VERSION.startswith("2."):
+        connection.commit()
     return result
 
 
@@ -497,6 +500,8 @@ async def _async_sync_execute(
         if self._driver_is_async:
             result_set = await result_set
             await connection.commit()  # very important
+        elif SQLALCHEMY_VERSION.startswith("2."):
+            connection.commit()
         return result_set
 
     @asynccontextmanager

requirements.txt

@@ -1,2 +1,2 @@
 prefect>=2.13.5
-sqlalchemy>=1.4.31,<2
+sqlalchemy>=1.4.31,<3

tests/test_database.py

@@ -4,6 +4,7 @@
 import cloudpickle
 import pytest
 from prefect import flow, task
+from sqlalchemy import __version__ as SQLALCHEMY_VERSION
 from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
@@ -88,6 +89,9 @@ def execute(self, query, params):
         )
         return cursor_result
 
+    def commit(self):
+        pass
+
 
 @pytest.fixture()
 def sqlalchemy_credentials_async():
@@ -340,13 +344,21 @@ async def managed_connector_with_data(self, connector_with_data, request):
         assert connector_with_data._engine is None
 
     @pytest.mark.parametrize("begin", [True, False])
-    def test_get_connection(self, begin, managed_connector_with_data):
+    async def test_get_connection(self, begin, managed_connector_with_data):
         connection = managed_connector_with_data.get_connection(begin=begin)
         if begin:
             engine_type = (
                 AsyncEngine if managed_connector_with_data._driver_is_async else Engine
             )
-            assert isinstance(connection, engine_type._trans_ctx)
+
+            if SQLALCHEMY_VERSION.startswith("1."):
+                assert isinstance(connection, engine_type._trans_ctx)
+            elif managed_connector_with_data._driver_is_async:
+                async with connection as conn:
+                    assert isinstance(conn, engine_type._connection_cls)
+            else:
+                with connection as conn:
+                    assert isinstance(conn, engine_type._connection_cls)
         else:
             engine_type = (
                 AsyncConnection
@@ -356,15 +368,22 @@ def test_get_connection(self, begin, managed_connector_with_data):
             assert isinstance(connection, engine_type)
 
     @pytest.mark.parametrize("begin", [True, False])
-    def test_get_client(self, begin, managed_connector_with_data):
+    async def test_get_client(self, begin, managed_connector_with_data):
         connection = managed_connector_with_data.get_client(
             client_type="connection", begin=begin
         )
         if begin:
             engine_type = (
                 AsyncEngine if managed_connector_with_data._driver_is_async else Engine
             )
-            assert isinstance(connection, engine_type._trans_ctx)
+            if SQLALCHEMY_VERSION.startswith("1."):
+                assert isinstance(connection, engine_type._trans_ctx)
+            elif managed_connector_with_data._driver_is_async:
+                async with connection as conn:
+                    assert isinstance(conn, engine_type._connection_cls)
+            else:
+                with connection as conn:
+                    assert isinstance(conn, engine_type._connection_cls)
         else:
             engine_type = (
                 AsyncConnection