Repository: kedro-org/kedro

Issue:
Adding execution_options to kedro.extras.datasets.pandas.SQLQueryDataSet
## Description
Is it possible to add execution_options to kedro.extras.datasets.pandas.SQLQueryDataSet?

For example, I would like to add stream_results=True to the connection string.
```python
  engine = create_engine(
      "postgresql://postgres:pass@localhost/example"
  )
  conn = engine.connect().execution_options(stream_results=True)
```
## Context
This article https://pythonspeed.com/articles/pandas-sql-chunking/ displays that its helpful to read large amounts of data using pandas with stream_results. I would like to use this method with SQLQueryDataSet.


Discussion:
Hi @jacobweiss2305 thank you for raising this, and thank you for including the article link for context - that was definitely an interesting read! This makes sense to me, I'll discuss it with the team but tagging @deepyaman too in the meantime, who I think might be interested in this too?
For visibility, this was also asked on SO: https://stackoverflow.com/questions/67182358/adding-stream-results-true-execution-options-to-kedro-extras-datasets-pandas-s/67182708#67182708
@mzjp2, per your SO response, would I need to create a custom dataset using Kedro IO?
Hi @jacobweiss2305, are you still interested in having this feature? And if so, would you be able to make a contribution PR for it? 
Yep I can make a pr
Hello! I see that this issue received no contribution so far. I am preparing a PR, will submit soon.

Resolving patch:
diff --git a/kedro/extras/datasets/pandas/sql_dataset.py b/kedro/extras/datasets/pandas/sql_dataset.py
--- a/kedro/extras/datasets/pandas/sql_dataset.py
+++ b/kedro/extras/datasets/pandas/sql_dataset.py
@@ -286,7 +286,19 @@ class SQLQueryDataSet(AbstractDataSet[None, pd.DataFrame]):
         >>>   type: pandas.SQLQueryDataSet
         >>>   sql: "select shuttle, shuttle_id from spaceflights.shuttles;"
         >>>   credentials: db_credentials
-        >>>   layer: raw
+
+    Advanced example using the `stream_results` and `chunksize` option to reduce memory usage
+
+    .. code-block:: yaml
+
+        >>> shuttle_id_dataset:
+        >>>   type: pandas.SQLQueryDataSet
+        >>>   sql: "select shuttle, shuttle_id from spaceflights.shuttles;"
+        >>>   credentials: db_credentials
+        >>>   execution_options:
+        >>>     stream_results: true
+        >>>   load_args:
+        >>>     chunksize: 1000
 
     Sample database credentials entry in ``credentials.yml``:
 
@@ -326,6 +338,7 @@ def __init__(  # pylint: disable=too-many-arguments
         load_args: Dict[str, Any] = None,
         fs_args: Dict[str, Any] = None,
         filepath: str = None,
+        execution_options: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Creates a new ``SQLQueryDataSet``.
 
@@ -351,6 +364,12 @@ def __init__(  # pylint: disable=too-many-arguments
                 https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.open
                 All defaults are preserved, except `mode`, which is set to `r` when loading.
             filepath: A path to a file with a sql query statement.
+            execution_options: A dictionary with non-SQL advanced options for the connection to
+                be applied to the underlying engine. To find all supported execution
+                options, see here:
+                https://docs.sqlalchemy.org/en/12/core/connections.html#sqlalchemy.engine.Connection.execution_options
+                Note that this is not a standard argument supported by pandas API, but could be
+                useful for handling large datasets.
 
         Raises:
             DataSetError: When either ``sql`` or ``con`` parameters is empty.
@@ -395,6 +414,7 @@ def __init__(  # pylint: disable=too-many-arguments
             self._fs = fsspec.filesystem(self._protocol, **_fs_credentials, **_fs_args)
             self._filepath = path
         self._connection_str = credentials["con"]
+        self._execution_options = execution_options or {}
         self.create_connection(self._connection_str)
 
     @classmethod
@@ -421,11 +441,14 @@ def _describe(self) -> Dict[str, Any]:
             sql=str(load_args.pop("sql", None)),
             filepath=str(self._filepath),
             load_args=str(load_args),
+            execution_options=str(self._execution_options),
         )
 
     def _load(self) -> pd.DataFrame:
         load_args = copy.deepcopy(self._load_args)
-        engine = self.engines[self._connection_str]  # type: ignore
+        engine = self.engines[self._connection_str].execution_options(
+            **self._execution_options
+        )  # type: ignore
 
         if self._filepath:
             load_path = get_filepath_str(PurePosixPath(self._filepath), self._protocol)


Test patch:
diff --git a/tests/extras/datasets/pandas/test_sql_dataset.py b/tests/extras/datasets/pandas/test_sql_dataset.py
--- a/tests/extras/datasets/pandas/test_sql_dataset.py
+++ b/tests/extras/datasets/pandas/test_sql_dataset.py
@@ -1,5 +1,6 @@
 # pylint: disable=no-member
 from pathlib import PosixPath
+from unittest.mock import ANY
 
 import pandas as pd
 import pytest
@@ -11,6 +12,7 @@
 TABLE_NAME = "table_a"
 CONNECTION = "sqlite:///kedro.db"
 SQL_QUERY = "SELECT * FROM table_a"
+EXECUTION_OPTIONS = {"stream_results": True}
 FAKE_CONN_STR = "some_sql://scott:tiger@localhost/foo"
 ERROR_PREFIX = (
     r"A module\/driver is missing when connecting to your SQL server\.(.|\n)*"
@@ -267,21 +269,51 @@ def test_empty_con_error(self):
         with pytest.raises(DataSetError, match=pattern):
             SQLQueryDataSet(sql=SQL_QUERY, credentials=dict(con=""))
 
-    def test_load(self, mocker, query_data_set):
+    @pytest.mark.parametrize(
+        "query_data_set, has_execution_options",
+        [
+            ({"execution_options": EXECUTION_OPTIONS}, True),
+            ({"execution_options": {}}, False),
+            ({}, False),
+        ],
+        indirect=["query_data_set"],
+    )
+    def test_load(self, mocker, query_data_set, has_execution_options):
         """Test `load` method invocation"""
         mocker.patch("pandas.read_sql_query")
         query_data_set.load()
-        pd.read_sql_query.assert_called_once_with(
-            sql=SQL_QUERY, con=query_data_set.engines[CONNECTION]
-        )
 
-    def test_load_query_file(self, mocker, query_file_data_set):
+        # Check that data was loaded with the expected query, connection string and
+        # execution options:
+        pd.read_sql_query.assert_called_once_with(sql=SQL_QUERY, con=ANY)
+        con_arg = pd.read_sql_query.call_args_list[0][1]["con"]
+        assert str(con_arg.url) == CONNECTION
+        assert len(con_arg.get_execution_options()) == bool(has_execution_options)
+        if has_execution_options:
+            assert con_arg.get_execution_options() == EXECUTION_OPTIONS
+
+    @pytest.mark.parametrize(
+        "query_file_data_set, has_execution_options",
+        [
+            ({"execution_options": EXECUTION_OPTIONS}, True),
+            ({"execution_options": {}}, False),
+            ({}, False),
+        ],
+        indirect=["query_file_data_set"],
+    )
+    def test_load_query_file(self, mocker, query_file_data_set, has_execution_options):
         """Test `load` method with a query file"""
         mocker.patch("pandas.read_sql_query")
         query_file_data_set.load()
-        pd.read_sql_query.assert_called_once_with(
-            sql=SQL_QUERY, con=query_file_data_set.engines[CONNECTION]
-        )
+
+        # Check that data was loaded with the expected query, connection string and
+        # execution options:
+        pd.read_sql_query.assert_called_once_with(sql=SQL_QUERY, con=ANY)
+        con_arg = pd.read_sql_query.call_args_list[0][1]["con"]
+        assert str(con_arg.url) == CONNECTION
+        assert len(con_arg.get_execution_options()) == bool(has_execution_options)
+        if has_execution_options:
+            assert con_arg.get_execution_options() == EXECUTION_OPTIONS
 
     def test_load_driver_missing(self, mocker):
         """Test that if an unknown module/driver is encountered by SQLAlchemy
@@ -332,8 +364,8 @@ def test_str_representation_sql(self, query_data_set, sql_file):
         """Test the data set instance string representation"""
         str_repr = str(query_data_set)
         assert (
-            f"SQLQueryDataSet(filepath=None, load_args={{}}, sql={SQL_QUERY})"
-            in str_repr
+            "SQLQueryDataSet(execution_options={}, filepath=None, "
+            f"load_args={{}}, sql={SQL_QUERY})" in str_repr
         )
         assert CONNECTION not in str_repr
         assert sql_file not in str_repr
@@ -342,8 +374,8 @@ def test_str_representation_filepath(self, query_file_data_set, sql_file):
         """Test the data set instance string representation with filepath arg."""
         str_repr = str(query_file_data_set)
         assert (
-            f"SQLQueryDataSet(filepath={str(sql_file)}, load_args={{}}, sql=None)"
-            in str_repr
+            f"SQLQueryDataSet(execution_options={{}}, filepath={str(sql_file)}, "
+            "load_args={}, sql=None)" in str_repr
         )
         assert CONNECTION not in str_repr
         assert SQL_QUERY not in str_repr
@@ -358,8 +390,8 @@ def test_sql_and_filepath_args(self, sql_file):
             SQLQueryDataSet(sql=SQL_QUERY, filepath=sql_file)
 
     def test_create_connection_only_once(self, mocker):
-        """Test that two datasets that need to connect to the same db
-        (but different tables, for example) only create a connection once.
+        """Test that two datasets that need to connect to the same db (but different
+        tables and execution options, for example) only create a connection once.
         """
         mock_engine = mocker.patch(
             "kedro.extras.datasets.pandas.sql_dataset.create_engine"
@@ -367,8 +399,29 @@ def test_create_connection_only_once(self, mocker):
         first = SQLQueryDataSet(sql=SQL_QUERY, credentials=dict(con=CONNECTION))
         assert len(first.engines) == 1
 
+        # second engine has identical params to the first one
+        # => no new engine should be created
         second = SQLQueryDataSet(sql=SQL_QUERY, credentials=dict(con=CONNECTION))
-        assert len(second.engines) == 1
+        mock_engine.assert_called_once_with(CONNECTION)
+        assert second.engines == first.engines
         assert len(first.engines) == 1
 
-        mock_engine.assert_called_once_with(CONNECTION)
+        # third engine only differs by its query execution options
+        # => no new engine should be created
+        third = SQLQueryDataSet(
+            sql="a different query",
+            credentials=dict(con=CONNECTION),
+            execution_options=EXECUTION_OPTIONS,
+        )
+        assert mock_engine.call_count == 1
+        assert third.engines == first.engines
+        assert len(first.engines) == 1
+
+        # fourth engine has a different connection string
+        # => a new engine has to be created
+        fourth = SQLQueryDataSet(
+            sql=SQL_QUERY, credentials=dict(con="an other connection string")
+        )
+        assert mock_engine.call_count == 2
+        assert fourth.engines == first.engines
+        assert len(first.engines) == 2