Repository: astronomer/astro-sdk

Issue:
Set a configurable limit to run_raw_sql rows output
**Context**
By default, `aql.run_raw_sql` does not limit the number of rows it can return.

[Comment on commit](https://github.com/astronomer/astro-sdk/commit/5b93df870dac5787dc1ec4ff11f238b9c360ebb0#r83305025)

**Version**
* Astro: 1.0.2

**Acceptance criteria**

* [x] Have an Airflow config "astro_sdk.run_raw_sql_limit_rows" which allows users to limit the amount of rows returned when running `run_raw_sql`
* [x] Set the default value not to be limited
* [x] Raise a warning if unlimited, explaining the risks to the user



Discussion:
This will need a breaking change and hence for now, we will add a warning.

Resolving patch:
diff --git a/python-sdk/src/astro/settings.py b/python-sdk/src/astro/settings.py
--- a/python-sdk/src/astro/settings.py
+++ b/python-sdk/src/astro/settings.py
@@ -31,3 +31,9 @@
 LOAD_TABLE_AUTODETECT_ROWS_COUNT = conf.getint(
     section="astro_sdk", key="load_table_autodetect_rows_count", fallback=1000
 )
+
+
+#: Reduce responses sizes returned by aql.run_raw_sql to avoid trashing the Airflow DB if the BaseXCom is used.
+RAW_SQL_MAX_RESPONSE_SIZE = conf.getint(
+    section="astro_sdk", key="run_raw_sql_response_size", fallback=-1
+)
diff --git a/python-sdk/src/astro/sql/operators/base_decorator.py b/python-sdk/src/astro/sql/operators/base_decorator.py
--- a/python-sdk/src/astro/sql/operators/base_decorator.py
+++ b/python-sdk/src/astro/sql/operators/base_decorator.py
@@ -27,7 +27,8 @@ def __init__(
         handler: Function | None = None,
         database: str | None = None,
         schema: str | None = None,
-        response_limit: int = 0,
+        response_limit: int = -1,
+        response_size: int = -1,
         sql: str = "",
         **kwargs: Any,
     ):
@@ -42,6 +43,8 @@ def __init__(
         self.database = self.op_kwargs.pop("database", database)
         self.schema = self.op_kwargs.pop("schema", schema)
         self.response_limit = self.op_kwargs.pop("response_limit", response_limit)
+        self.response_size = self.op_kwargs.pop("response_size", response_size)
+
         self.op_args: dict[str, Table | pd.DataFrame] = {}
 
         # We purposely do NOT render upstream_tasks otherwise we could have a case where a user
diff --git a/python-sdk/src/astro/sql/operators/raw_sql.py b/python-sdk/src/astro/sql/operators/raw_sql.py
--- a/python-sdk/src/astro/sql/operators/raw_sql.py
+++ b/python-sdk/src/astro/sql/operators/raw_sql.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from collections.abc import Iterable, Mapping
 from typing import Any, Callable
 
@@ -9,6 +10,7 @@
     from airflow.decorators.base import task_decorator_factory
     from airflow.decorators import _TaskDecorator as TaskDecorator
 
+from astro import settings
 from astro.exceptions import IllegalLoadToDatabaseException
 from astro.sql.operators.base_decorator import BaseSQLDecoratedOperator
 
@@ -28,11 +30,22 @@ def execute(self, context: dict) -> Any:
         result = self.database_impl.run_sql(
             sql_statement=self.sql, parameters=self.parameters
         )
+        if self.response_size == -1 and not settings.IS_CUSTOM_XCOM_BACKEND:
+            logging.warning(
+                "Using `run_raw_sql` without `response_size` can result in excessive amount of data being recorded "
+                "to the Airflow metadata database, leading to issues to the orchestration of tasks. It is possible to "
+                "avoid this problem by either setting `response_size` to a small integer or by using a custom XCom "
+                "backend."
+            )
+
         if self.handler:
             response = self.handler(result)
-            if self.response_limit and len(response) > self.response_limit:
+            if self.response_limit >= 0 and len(response) > self.response_limit:
                 raise IllegalLoadToDatabaseException()  # pragma: no cover
-            return response
+            if self.response_size >= 0:
+                return response[: self.response_size]
+            else:
+                return response
         else:
             return None
 
@@ -44,6 +57,7 @@ def run_raw_sql(
     database: str | None = None,
     schema: str | None = None,
     handler: Callable | None = None,
+    response_size: int = settings.RAW_SQL_MAX_RESPONSE_SIZE,
     **kwargs: Any,
 ) -> TaskDecorator:
     """
@@ -67,20 +81,26 @@ def my_sql_statement(table1: Table) -> Table:
     result of a provided ``handler`` function or it will not return anything at all.
 
 
-    :param python_callable: This parameter is filled in automatically when you use the transform function as a decorator
-        This is where the python function gets passed to the wrapping function
-    :param conn_id: Connection ID for the database you want to connect to. If you do not pass in a value for this object
-        we can infer the connection ID from the first table passed into the python_callable function.
-        (required if there are no table arguments)
+    :param python_callable: This parameter is filled in automatically when you use the transform function as a
+        decorator. This is where the python function gets passed to the wrapping function
+    :param conn_id: Connection ID for the database you want to connect to.
+        If you do not pass in a value for this object we can infer the connection ID from the first table
+        passed into the python_callable function. (required if there are no table arguments)
     :param parameters: parameters to pass into the SQL query
     :param database: Database within the SQL instance you want to access. If left blank we will default to the
-        table.metatadata.database in the first Table passed to the function (required if there are no table arguments)
+        table.metatadata.database in the first Table passed to the function
+        (required if there are no table arguments)
     :param schema: Schema within the SQL instance you want to access. If left blank we will default to the
-        table.metatadata.schema in the first Table passed to the function (required if there are no table arguments)
+        table.metatadata.schema in the first Table passed to the function
+        (required if there are no table arguments)
     :param handler: Handler function to process the result of the SQL query. For more information please consult
         https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Result
+    :param response_size: Used to trim the responses returned to avoid trashing the Airflow DB.
+        The default value is -1, which means the response is not changed. Otherwise, if the response is a list,
+        returns up to the desired amount of items. If the response is a string, trims it to the desired size.
     :param kwargs:
-    :return: By default returns None unless there is a handler function, in which case returns the result of the handler
+    :return: By default returns None unless there is a handler function,
+        in which case returns the result of the handler
     """
 
     kwargs.update(
@@ -90,6 +110,7 @@ def my_sql_statement(table1: Table) -> Table:
             "database": database,
             "schema": schema,
             "handler": handler,
+            "response_size": response_size,
         }
     )
     return task_decorator_factory(


Test patch:
diff --git a/python-sdk/tests/sql/operators/test_raw_sql.py b/python-sdk/tests/sql/operators/test_raw_sql.py
new file mode 100644
--- /dev/null
+++ b/python-sdk/tests/sql/operators/test_raw_sql.py
@@ -0,0 +1,82 @@
+import logging
+import pathlib
+
+import pytest
+from airflow.decorators import task
+
+from astro import sql as aql
+from astro.constants import Database
+from astro.files import File
+from tests.sql.operators import utils as test_utils
+
+CWD = pathlib.Path(__file__).parent
+DATA_FILEPATH = pathlib.Path(CWD.parent.parent, "data/sample.csv")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    "database_table_fixture",
+    [{"database": Database.SQLITE, "file": File(path=str(DATA_FILEPATH))}],
+    indirect=True,
+    ids=["sqlite"],
+)
+def test_run_raw_sql_without_limit(caplog, sample_dag, database_table_fixture):
+    _, test_table = database_table_fixture
+    caplog.set_level(logging.WARNING)
+
+    @aql.run_raw_sql
+    def raw_sql_query(input_table):
+        return "SELECT * from {{input_table}}"
+
+    @task
+    def assert_num_rows(results):
+        assert len(results) == 3
+        assert results == [
+            (1, "First"),
+            (2, "Second"),
+            (3, "Third with unicode पांचाल"),
+        ]
+
+    with sample_dag:
+        results = raw_sql_query(
+            input_table=test_table,
+            handler=lambda cursor: cursor.fetchall(),
+        )
+        assert_num_rows(results)
+
+    test_utils.run_dag(sample_dag)
+
+    expected_warning = (
+        "excessive amount of data being recorded to the Airflow metadata database"
+    )
+    assert expected_warning in caplog.text
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    "database_table_fixture",
+    [{"database": Database.SQLITE, "file": File(path=str(DATA_FILEPATH))}],
+    indirect=True,
+    ids=["sqlite"],
+)
+def test_run_raw_sql_with_limit(sample_dag, database_table_fixture):
+    _, test_table = database_table_fixture
+
+    @aql.run_raw_sql
+    def raw_sql_query(input_table):
+        return "SELECT * from {{input_table}}"
+
+    @task
+    def assert_num_rows(results):
+        assert len(results) == 1
+        assert results == [(1, "First")]
+
+    with sample_dag:
+        results = raw_sql_query(
+            input_table=test_table,
+            response_size=1,
+            handler=lambda cursor: cursor.fetchall(),
+        )
+        assert_num_rows(results)
+
+    test_utils.run_dag(sample_dag)