Add an option to load SQL queries from a file for SQLQueryDataSet

kedro/extras/datasets/pandas/sql_dataset.py

@@ -29,13 +29,20 @@
 
 import copy
 import re
+from pathlib import PurePosixPath
 from typing import Any, Dict, Optional
 
+import fsspec
 import pandas as pd
 from sqlalchemy import create_engine
 from sqlalchemy.exc import NoSuchModuleError
 
-from kedro.io.core import AbstractDataSet, DataSetError
+from kedro.io.core import (
+    AbstractDataSet,
+    DataSetError,
+    get_filepath_str,
+    get_protocol_and_path,
+)
 
 __all__ = ["SQLTableDataSet", "SQLQueryDataSet"]
 
@@ -278,13 +285,19 @@ class SQLQueryDataSet(AbstractDataSet):
 
     """
 
-    def __init__(
-        self, sql: str, credentials: Dict[str, Any], load_args: Dict[str, Any] = None
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        sql: str = None,
+        filepath: str = None,
+        credentials: Dict[str, Any] = None,
+        load_args: Dict[str, Any] = None,
+        fs_args: Dict[str, Any] = None,
     ) -> None:
         """Creates a new ``SQLQueryDataSet``.
 
         Args:
             sql: The sql query statement.
+            filepath: A path to a file with a sql query statement
             credentials: A dictionary with a ``SQLAlchemy`` connection string.
                 Users are supposed to provide the connection string 'con'
                 through credentials. It overwrites `con` parameter in
@@ -297,14 +310,23 @@ def __init__(
                 https://pandas.pydata.org/pandas-docs/stable/generated/pandas.read_sql_query.html
                 To find all supported connection string formats, see here:
                 https://docs.sqlalchemy.org/en/13/core/engines.html#database-urls
+            fs_args: Extra arguments to pass into underlying filesystem class constructor
+                (e.g. `{"project": "my-project"}` for ``GCSFileSystem``), as well as
+                to pass to the filesystem's `open` method through nested keys
+                `open_args_load` and `open_args_save`.
+                Here you can find all available arguments for `open`:
+                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
+                and to `w` when saving.
 
         Raises:
             DataSetError: When either ``sql`` or ``con`` parameters is emtpy.
         """
 
-        if not sql:
+        if not (sql or filepath):
             raise DataSetError(
-                "`sql` argument cannot be empty. Please provide a sql query"
+                "`sql` and `filepath` arguments cannot both be empty."
+                "Please provide a sql query or path to a sql query file."
             )
 
         if not (credentials and "con" in credentials and credentials["con"]):
@@ -321,7 +343,23 @@ def __init__(
             else default_load_args
         )
 
-        self._load_args["sql"] = sql
+        # load sql query from file
+        if not sql:
+            # filesystem for loading sql file
+            _fs_args = copy.deepcopy(fs_args) or {}
+            _fs_open_args_load = _fs_args.pop("open_args_load", {})
+            _fs_credentials = _fs_args.pop("credentials", {})
+            protocol, path = get_protocol_and_path(str(filepath))
+
+            self._protocol = protocol
+            self._fs = fsspec.filesystem(self._protocol, **_fs_credentials, **_fs_args)
+
+            _fs_open_args_load.setdefault("mode", "r")
+            self._fs_open_args_load = _fs_open_args_load
+
+            self._load_args["filepath"] = path
+        else:
+            self._load_args["sql"] = sql
         self._load_args["con"] = credentials["con"]
 
     def _describe(self) -> Dict[str, Any]:
@@ -331,8 +369,17 @@ def _describe(self) -> Dict[str, Any]:
         return dict(sql=self._load_args["sql"], load_args=load_args)
 
     def _load(self) -> pd.DataFrame:
+        load_args = self._load_args.copy()
+
+        if "sql" not in load_args:
+            filepath = load_args.pop("filepath")
+            load_path = get_filepath_str(PurePosixPath(filepath), self._protocol)
+
+            with self._fs.open(load_path, **self._fs_open_args_load) as fs_file:
+                load_args["sql"] = fs_file.read()
+
         try:
-            return pd.read_sql_query(**self._load_args)
+            return pd.read_sql_query(**load_args)
         except ImportError as import_error:
             raise _get_missing_module_error(import_error) from import_error
         except NoSuchModuleError as exc:

tests/extras/datasets/pandas/test_sql_dataset.py

@@ -51,6 +51,14 @@ def dummy_dataframe():
     return pd.DataFrame({"col1": [1, 2], "col2": [4, 5], "col3": [5, 6]})
 
 
+@pytest.fixture
+def sql_file(tmp_path):
+    file = tmp_path / "test.sql"
+    with file.open("w") as f:
+        f.write(SQL_QUERY)
+    return f.as_posix()
+
+
 @pytest.fixture(params=[{}])
 def table_data_set(request):
     kwargs = dict(table_name=TABLE_NAME, credentials=dict(con=CONNECTION))
@@ -65,6 +73,13 @@ def query_data_set(request):
     return SQLQueryDataSet(**kwargs)
 
 
+@pytest.fixture(params=[{}])
+def query_file_data_set(request, sql_file):
+    kwargs = dict(filepath=sql_file, credentials=dict(con=CONNECTION))
+    kwargs.update(request.param)
+    return SQLQueryDataSet(**kwargs)
+
+
 class TestSQLTableDataSetLoad:
     @staticmethod
     def _assert_pd_called_once():
@@ -244,10 +259,13 @@ def _assert_pd_called_once():
         _callable.assert_called_once_with(sql=SQL_QUERY, con=CONNECTION)
 
     def test_empty_query_error(self):
-        """Check the error when instantiating with empty query"""
-        pattern = r"`sql` argument cannot be empty\. Please provide a sql query"
+        """Check the error when instantiating with empty query or file"""
+        pattern = (
+            r"`sql` and `filepath` arguments cannot both be empty\."
+            r"Please provide a sql query or path to a sql query file\."
+        )
         with pytest.raises(DataSetError, match=pattern):
-            SQLQueryDataSet(sql="", credentials=dict(con=CONNECTION))
+            SQLQueryDataSet(sql="", filepath="", credentials=dict(con=CONNECTION))
 
     def test_empty_con_error(self):
         """Check the error when instantiating with empty connection string"""
@@ -264,6 +282,12 @@ def test_load(self, mocker, query_data_set):
         query_data_set.load()
         self._assert_pd_called_once()
 
+    def test_load_query_file(self, mocker, query_file_data_set):
+        """Test `load` method with a query file"""
+        mocker.patch("pandas.read_sql_query")
+        query_file_data_set.load()
+        self._assert_pd_called_once()
+
     def test_load_driver_missing(self, mocker, query_data_set):
         """Test that if an unknown module/driver is encountered by SQLAlchemy
         then the error should contain the original error message"""