diff --git a/python-sdk/docs/astro/sql/operators/raw_sql.rst b/python-sdk/docs/astro/sql/operators/raw_sql.rst index 9fef2f50e..54086c712 100644 --- a/python-sdk/docs/astro/sql/operators/raw_sql.rst +++ b/python-sdk/docs/astro/sql/operators/raw_sql.rst @@ -37,4 +37,4 @@ Parameters * **fail_on_empty** - Sometimes the handler function can raise an exception when the data is not returned by the database and we try to run ``fetchall()``. We can make sure that the handler function doesn't raise an exception by passing ``fail_on_empty==False``. The default value for this parameter is ``True``. -* **query_modifier** - The query_modifier allows you to create both pre_queries and post_queries across multiple statements. An example of where this would be useful is if you want to add query tags to a snowflake statement you can set ``session_modifier.pre_queries = ["ALTER SESSION SET QUERY_TAG=]``, which will ensure that any query run will contain this query tag. +* **query_modifier** - The ``query_modifier`` parameter allows you to define statements to run before and after the ``run_raw_sql`` main statement. To associate a Snowflake query tag, for instance, it is possible to use ``query_modifier=QueryModifier(pre_queries=["ALTER SESSION SET QUERY_TAG=])``. diff --git a/python-sdk/docs/astro/sql/operators/transform.rst b/python-sdk/docs/astro/sql/operators/transform.rst index 902d8297c..e59530172 100644 --- a/python-sdk/docs/astro/sql/operators/transform.rst +++ b/python-sdk/docs/astro/sql/operators/transform.rst @@ -5,7 +5,7 @@ ============================================================ When to use the ``transform`` operator -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-------------------------------------- The ``transform`` operator allows you to implement the **T** of an ELT system by running a SQL query. Each step of the transform pipeline creates a new table from the ``SELECT`` statement and enables tasks to pass those tables as if they were native Python objects. The ``transform`` operator treats values in the double brackets as Airflow jinja templates. You can find more details on templating at :ref:`templating`. @@ -45,3 +45,8 @@ Case 2: Passing a Pandas dataframe between tasks while completing data transform :end-before: [END transform_example_4] Please note that in case you want to pass SQL file in the transform decorator, use :ref:`transform_file_operator` + +Parameters +----------- + +* **query_modifier** - The ``query_modifier`` parameter allows you to define statements to run before and after the ``run_raw_sql`` main statement. To associate a Snowflake query tag, for instance, it is possible to use ``query_modifier=QueryModifier(pre_queries=["ALTER SESSION SET QUERY_TAG=])``. diff --git a/python-sdk/docs/astro/sql/operators/transform_file.rst b/python-sdk/docs/astro/sql/operators/transform_file.rst index 127d2dc32..203aa8687 100644 --- a/python-sdk/docs/astro/sql/operators/transform_file.rst +++ b/python-sdk/docs/astro/sql/operators/transform_file.rst @@ -5,7 +5,7 @@ ================================================================================= When to use the ``transform_file`` operator -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +------------------------------------------- The ``transform_file`` operator allows you to implement the **T** of an ELT system by running a SQL query from specified SQL file. Each step of the transform pipeline creates a new table from the ``SELECT`` statement and enables tasks to pass those tables as if they were native Python objects. The ``transform_file`` functions return a ``Table`` object that can be passed to future tasks. This table will be either an auto-generated temporary table, or will overwrite a table given in the ``output_table`` parameter. The ``transform_file`` operator treats values in the double brackets as Airflow jinja templates. You can find more details on templating at :ref:`templating`. @@ -14,3 +14,8 @@ The ``transform_file`` functions return a ``Table`` object that can be passed to :language: python :start-after: [START transform_file_example_1] :end-before: [END transform_file_example_1] + +Parameters +----------- + +* **query_modifier** - The ``query_modifier`` parameter allows you to define statements to run before and after the ``run_raw_sql`` main statement. To associate a Snowflake query tag, for instance, it is possible to use ``query_modifier=QueryModifier(pre_queries=["ALTER SESSION SET QUERY_TAG=])``. diff --git a/python-sdk/tests/sql/operators/test_transform.py b/python-sdk/tests/sql/operators/test_transform.py new file mode 100644 index 000000000..6e94b6fbc --- /dev/null +++ b/python-sdk/tests/sql/operators/test_transform.py @@ -0,0 +1,58 @@ +import tempfile +from unittest import mock + +from astro import sql as aql +from tests.sql.operators import utils as test_utils + + +class MockReturn: + _scalar = [] + + def scalar(self): + return self._scalar + + +@mock.patch("astro.databases.base.BaseDatabase.connection") +def test_transform_calls_with_query_tag(run_sql, sample_dag): + from astro.query_modifier import QueryModifier + + run_sql.execute.return_value = MockReturn() + + with sample_dag: + + @aql.transform( + conn_id="sqlite_default", + query_modifier=QueryModifier(pre_queries=["ALTER team_1", "ALTER team_2"]), + ) + def dummy_method(): + return "SELECT 1+1" + + dummy_method() + + test_utils.run_dag(sample_dag) + enriched_query = run_sql.method_calls[1].args[0].text + assert enriched_query.startswith("ALTER team_1;ALTER team_2;CREATE TABLE IF NOT EXISTS ") + assert enriched_query.endswith("AS SELECT 1+1") + + +@mock.patch("astro.databases.base.BaseDatabase.connection") +def test_transform_file_calls_with_query_tag(run_sql, sample_dag): + from astro.query_modifier import QueryModifier + + run_sql.execute.return_value = MockReturn() + + with tempfile.NamedTemporaryFile(suffix=".sql") as tmp_file: + tmp_file.write(b"SELECT 1+1") + tmp_file.flush() + + with sample_dag: + aql.transform_file( + file_path=tmp_file.name, + conn_id="sqlite_default", + query_modifier=QueryModifier(pre_queries=["ALTER team_1", "ALTER team_2"]), + ) + test_utils.run_dag(sample_dag) + + enriched_query = run_sql.method_calls[1].args[0].text + assert enriched_query.startswith("ALTER team_1;ALTER team_2;CREATE TABLE IF NOT EXISTS ") + assert enriched_query.endswith("AS SELECT 1+1")