Update tests and docs regarding availability of the argument query_mo…

…difier (#1914)

Although we `transform` and `transform_file` supported the
`query_modifier` argument, we didn't have tests nor documentation
describing this. This PR adds tests and docs for these use cases.

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=<my-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=<my-query-tag>])``.

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=<my-query-tag>])``.

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=<my-query-tag>])``.

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")