Merge pull request #49 from PrefectHQ/fix_fetch

Fix fetch flow

Author: ahuang11

CHANGELOG.md

@@ -19,6 +19,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Security
 
+## 0.2.4
+
+Released on February 15th, 2022.
+
+### Fixed
+
+- Initializing `_unique_results` and `_exit_stack` so `fetch_*` is able to run - [#49](https://github.com/PrefectHQ/prefect-sqlalchemy/pull/49)
+
 ## 0.2.3
 
 Released on February 10th, 2022.

prefect_sqlalchemy/database.py

@@ -301,6 +301,12 @@ def block_initialization(self):
         except ValueError:
             self._driver_is_async = False
 
+        if self._unique_results is None:
+            self._unique_results = {}
+
+        if self._exit_stack is None:
+            self._start_exit_stack()
+
     def _start_exit_stack(self):
         """
         Starts an AsyncExitStack or ExitStack depending on whether driver is async.
@@ -378,10 +384,6 @@ def sqlalchemy_credentials_flow():
 
         if self._engine is None:
             self._engine = engine
-        if self._exit_stack is None:
-            self._start_exit_stack()
-        if self._unique_results is None:
-            self._unique_results = {}
 
         return engine
 
@@ -938,3 +940,9 @@ def __getstate__(self):
     def __setstate__(self, data: dict):
         """Upon loading back, restart the engine and results."""
         self.__dict__.update(data)
+
+        if self._unique_results is None:
+            self._unique_results = {}
+
+        if self._exit_stack is None:
+            self._start_exit_stack()

tests/test_database.py

@@ -3,7 +3,7 @@
 
 import cloudpickle
 import pytest
-from prefect import flow
+from prefect import flow, task
 from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
@@ -261,22 +261,30 @@ async def test_connector_init(self):
         credentials_url = SqlAlchemyConnector(connection_info=connection_url)
         assert credentials_components._rendered_url == credentials_url._rendered_url
 
-    def test_delay_start(self, caplog):
+    @pytest.mark.parametrize("method", ["fetch_all", "execute"])
+    def test_delay_start(self, caplog, method):
         with SqlAlchemyConnector(
             connection_info=ConnectionComponents(
                 driver=SyncDriver.SQLITE_PYSQLITE,
                 database=":memory:",
             ),
         ) as connector:
+            assert connector._unique_results == {}
+            assert isinstance(connector._exit_stack, ExitStack)
             connector.reset_connections()
-            assert caplog.records[0].msg == "There were no connections to reset."
+            assert (
+                caplog.records[0].msg == "Reset opened connections and their results."
+            )
             assert connector._engine is None
-            assert connector._unique_results is None
-            assert connector._exit_stack is None
-            connector.execute("SELECT 1")
-            assert isinstance(connector._engine, Engine)
             assert connector._unique_results == {}
             assert isinstance(connector._exit_stack, ExitStack)
+            getattr(connector, method)("SELECT 1")
+            assert isinstance(connector._engine, Engine)
+            if method == "execute":
+                assert connector._unique_results == {}
+            else:
+                assert len(connector._unique_results) == 1
+            assert isinstance(connector._exit_stack, ExitStack)
 
     @pytest.fixture(params=[SyncDriver.SQLITE_PYSQLITE, AsyncDriver.SQLITE_AIOSQLITE])
     async def connector_with_data(self, tmp_path, request):
@@ -547,3 +555,52 @@ def test_sync_compatible_reset_connections(self, tmp_path):
         assert len(conn._unique_results) == 1
         conn.reset_connections()
         assert len(conn._unique_results) == 0
+
+    def test_flow_without_initialized_engine(self, tmp_path):
+        @task
+        def setup_table(block_name: str) -> None:
+            with SqlAlchemyConnector.load(block_name) as connector:
+                connector.execute(
+                    "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);"  # noqa
+                )
+                connector.execute(
+                    "INSERT INTO customers (name, address) VALUES (:name, :address);",
+                    parameters={"name": "Marvin", "address": "Highway 42"},
+                )
+                connector.execute_many(
+                    "INSERT INTO customers (name, address) VALUES (:name, :address);",
+                    seq_of_parameters=[
+                        {"name": "Ford", "address": "Highway 42"},
+                        {"name": "Unknown", "address": "Highway 42"},
+                    ],
+                )
+
+        @task
+        def fetch_data(block_name: str) -> list:
+            all_rows = []
+            with SqlAlchemyConnector.load(block_name) as connector:
+                while True:
+                    # Repeated fetch* calls using the same operation will
+                    # skip re-executing and instead return the next set of results
+                    new_rows = connector.fetch_many("SELECT * FROM customers", size=2)
+                    if len(new_rows) == 0:
+                        break
+                    all_rows.append(new_rows)
+            return all_rows
+
+        @flow
+        def sqlalchemy_flow(block_name: str) -> list:
+            SqlAlchemyConnector(
+                connection_info=ConnectionComponents(
+                    driver=SyncDriver.SQLITE_PYSQLITE,
+                    database=str(tmp_path / "test.db"),
+                )
+            ).save(block_name)
+            setup_table(block_name)
+            all_rows = fetch_data(block_name)
+            return all_rows
+
+        assert sqlalchemy_flow("connector") == [
+            [("Marvin", "Highway 42"), ("Ford", "Highway 42")],
+            [("Unknown", "Highway 42")],
+        ]