Simplify default rows limit return result

providers/common/sql/docs/operators.rst

@@ -288,7 +288,7 @@ Parameters
 ----------
 * ``datasource_configs`` (list[DataSourceConfig], required): List of datasource configurations
 * ``queries`` (list[str], required): List of SQL queries to run on the data
-* ``max_rows_check`` (int, optional): Maximum number of rows to check for each query. Default is 100. If any query returns more than this number of rows, it will be skipped in the results returned by the operator. This is to prevent returning too many rows in the results which can cause xcom rendering issues in Airflow UI.
+* ``max_rows_check`` (int, optional): Maximum number of rows returned for each query. Default is 100. If a query returns more rows, the engine logs a warning and returns only the first ``max_rows_check`` rows. This prevents returning too many rows, which can cause xcom rendering issues in the Airflow UI.
 * ``engine`` (DataFusionEngine, optional): Query engine to use. Default is "datafusion". Currently, only "datafusion" is supported.
 * ``result_output_format`` (str, optional): Output format for the results. Default is ``tabulate``. Supported formats are ``tabulate``, ``json``.
 

providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py

@@ -101,11 +101,22 @@ def _register_data_source_format(self, datasource_config: DataSourceConfig):
             datasource_config.table_name,
         )
 
-    def execute_query(self, query: str) -> dict[str, list[Any]]:
+    def execute_query(self, query: str, max_rows: int | None = None) -> dict[str, list[Any]]:
         """Execute a query and return the result as a dictionary."""
         try:
             self.log.info("Executing query: %s", query)
             df = self.session_context.sql(query)
+
+            if max_rows is not None:
+                result = df.limit(max_rows + 1).to_pydict()
+                if result and len(next(iter(result.values()))) > max_rows:
+                    self.log.warning(
+                        "Query returned more than %s rows. Returning first %s rows.",
+                        max_rows,
+                        max_rows,
+                    )
+                    return {column: values[:max_rows] for column, values in result.items()}
+                return result
             return df.to_pydict()
         except Exception as e:
             raise QueryExecutionException(f"Error while executing query: {e}")

providers/common/sql/src/airflow/providers/common/sql/operators/analytics.py

@@ -38,7 +38,8 @@ class AnalyticsOperator(BaseOperator):
 
     :param datasource_configs: List of datasource configurations to register.
     :param queries: List of SQL queries to execute.
-    :param max_rows_check: Maximum number of rows allowed in query results. Queries exceeding this will be skipped.
+    :param max_rows_check: Maximum number of rows returned per query. Queries exceeding this return
+        only the first N rows with a warning.
     :param engine: Optional DataFusion engine instance.
     :param result_output_format: List of output formats for results. Supported: 'tabulate', 'json'. Default is 'tabulate'.
     """
@@ -89,7 +90,7 @@ def execute(self, context: Context) -> str:
 
         # TODO make it parallel as there is no dependency between queries
         for query in self.queries:
-            result_dict = self._df_engine.execute_query(query)
+            result_dict = self._df_engine.execute_query(query, max_rows=self.max_rows_check)
             results.append({"query": query, "data": result_dict})
 
         match self.result_output_format:
@@ -100,36 +101,14 @@ def execute(self, context: Context) -> str:
             case _:
                 raise ValueError(f"Unsupported output format: {self.result_output_format}")
 
-    def _is_result_too_large(self, result_dict: dict[str, Any]) -> tuple[bool, int]:
-        """Check if a result exceeds the max_rows_check limit."""
-        if not result_dict:
-            return False, 0
-        num_rows = len(next(iter(result_dict.values())))
-        max_rows_exceeded = num_rows > self.max_rows_check
-        if max_rows_exceeded:
-            self.log.warning(
-                "Query returned %s rows, exceeding max_rows_check (%s). Skipping result output as large datasets are unsuitable for return.",
-                num_rows,
-                self.max_rows_check,
-            )
-        return max_rows_exceeded, num_rows
-
     def _build_tabulate_output(self, query_results: list[dict[str, Any]]) -> str:
         from tabulate import tabulate
 
         output_parts = []
         for item in query_results:
             query = item["query"]
             result_dict = item["data"]
-            too_large, row_count = self._is_result_too_large(result_dict)
-
-            if too_large:
-                output_parts.append(
-                    f"\n### Results: {query}\n\n"
-                    f"**Skipped**: {row_count} rows exceed max_rows_check ({self.max_rows_check})\n\n"
-                    f"{'-' * 40}\n"
-                )
-                continue
+            row_count = len(next(iter(result_dict.values()))) if result_dict else 0
 
             table_str = tabulate(
                 self._get_rows(result_dict, row_count),
@@ -151,18 +130,7 @@ def _build_json_output(self, query_results: list[dict[str, Any]]) -> str:
         for item in query_results:
             query = item["query"]
             result_dict = item["data"]
-            max_rows_exceeded, row_count = self._is_result_too_large(result_dict)
-
-            if max_rows_exceeded:
-                json_results.append(
-                    {
-                        "query": query,
-                        "status": "skipped_too_large",
-                        "row_count": row_count,
-                        "max_allowed": self.max_rows_check,
-                    }
-                )
-                continue
+            row_count = len(next(iter(result_dict.values()))) if result_dict else 0
 
             json_results.append(
                 {

providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py

@@ -141,15 +141,52 @@ def test_register_datasource_duplicate_table(self, mock_get_conn):
     def test_execute_query_success(self):
         engine = DataFusionEngine()
         engine.df_ctx = MagicMock(spec=SessionContext)
-        mock_df = MagicMock()
+        mock_df = MagicMock(spec=["limit", "to_pydict"])
         mock_df.to_pydict.return_value = {"col1": [1, 2]}
         engine.df_ctx.sql.return_value = mock_df
 
         result = engine.execute_query("SELECT * FROM test_table")
 
         engine.df_ctx.sql.assert_called_once_with("SELECT * FROM test_table")
+        mock_df.limit.assert_not_called()
         assert result == {"col1": [1, 2]}
 
+    def test_execute_query_with_max_rows(self):
+        engine = DataFusionEngine()
+        engine.df_ctx = MagicMock(spec=SessionContext)
+        mock_df = MagicMock(spec=["limit", "to_pydict"])
+        limited_df = MagicMock(spec=["to_pydict"])
+        limited_df.to_pydict.return_value = {"col1": [1, 2, 3]}
+        mock_df.limit.return_value = limited_df
+        engine.df_ctx.sql.return_value = mock_df
+
+        result = engine.execute_query("SELECT * FROM test_table", max_rows=3)
+
+        engine.df_ctx.sql.assert_called_once_with("SELECT * FROM test_table")
+        mock_df.limit.assert_called_once_with(4)
+        assert result == {"col1": [1, 2, 3]}
+
+    def test_execute_query_with_max_rows_logs_warning_when_exceeded(self):
+        engine = DataFusionEngine()
+        engine.df_ctx = MagicMock(spec=SessionContext)
+        mock_df = MagicMock(spec=["limit", "to_pydict"])
+        limited_df = MagicMock(spec=["to_pydict"])
+        limited_df.to_pydict.return_value = {"col1": [1, 2, 3, 4]}
+        mock_df.limit.return_value = limited_df
+        engine.df_ctx.sql.return_value = mock_df
+
+        with patch.object(engine.log, "warning") as mock_warning:
+            result = engine.execute_query("SELECT * FROM test_table", max_rows=3)
+
+        engine.df_ctx.sql.assert_called_once_with("SELECT * FROM test_table")
+        mock_df.limit.assert_called_once_with(4)
+        mock_warning.assert_called_once_with(
+            "Query returned more than %s rows. Returning first %s rows.",
+            3,
+            3,
+        )
+        assert result == {"col1": [1, 2, 3]}
+
     def test_execute_query_failure(self):
         engine = DataFusionEngine()
         engine.df_ctx = MagicMock(spec=SessionContext)

providers/common/sql/tests/unit/common/sql/operators/test_analytics.py

@@ -53,18 +53,19 @@ def test_execute_success(self, operator, mock_engine):
         result = operator.execute(context={})
 
         mock_engine.register_datasource.assert_called_once()
-        mock_engine.execute_query.assert_called_once_with("SELECT * FROM users_data")
+        mock_engine.execute_query.assert_called_once_with("SELECT * FROM users_data", max_rows=100)
         assert "col1" in result
         assert "col2" in result
 
-    def test_execute_max_rows_exceeded(self, operator, mock_engine):
+    def test_execute_uses_overridden_max_rows_limit(self, operator, mock_engine):
         operator.max_rows_check = 3
-        mock_engine.execute_query.return_value = {"col1": [1, 2, 3, 4]}
+        mock_engine.execute_query.return_value = {"col1": [1, 2, 3]}
 
         result = operator.execute(context={})
 
-        assert "Skipped" in result
-        assert "4 rows exceed max_rows_check (3)" in result
+        mock_engine.execute_query.assert_called_once_with("SELECT * FROM users_data", max_rows=3)
+        assert "Skipped" not in result
+        assert "col1" in result
 
     def test_json_output_format(self, mock_engine):
         datasource_config = DataSourceConfig(