add e2e test for thrift resultRowLimit

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -402,11 +402,6 @@ def __init__(
 
         self.connection: Connection = connection
 
-        if not connection.session.use_sea and row_limit is not None:
-            logger.warning(
-                "Row limit is only supported for SEA protocol. Ignoring row_limit."
-            )
-
         self.rowcount: int = -1  # Return -1 as this is not supported
         self.buffer_size_bytes: int = result_buffer_size_bytes
         self.active_result_set: Union[ResultSet, None] = None
@@ -802,6 +797,7 @@ def execute(
             parameters=prepared_params,
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
+            row_limit=self.row_limit,
         )
 
         if self.active_result_set and self.active_result_set.is_staging_operation:
@@ -858,6 +854,7 @@ def execute_async(
             parameters=prepared_params,
             async_op=True,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
+            row_limit=self.row_limit,
         )
 
         return self

src/databricks/sql/session.py

@@ -76,9 +76,7 @@ def __init__(
             tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
         )
 
-        self.use_sea = kwargs.get("use_sea", False)
         self.backend = self._create_backend(
-            self.use_sea,
             server_hostname,
             http_path,
             all_headers,
@@ -91,7 +89,6 @@ def __init__(
 
     def _create_backend(
         self,
-        use_sea: bool,
         server_hostname: str,
         http_path: str,
         all_headers: List[Tuple[str, str]],
@@ -100,6 +97,8 @@ def _create_backend(
         kwargs: dict,
     ) -> DatabricksClient:
         """Create and return the appropriate backend client."""
+        use_sea = kwargs.get("use_sea", False)
+
         databricks_client_class: Type[DatabricksClient]
         if use_sea:
             logger.debug("Creating SEA backend client")

tests/e2e/test_driver.py

@@ -113,10 +113,12 @@ def connection(self, extra_params=()):
             conn.close()
 
     @contextmanager
-    def cursor(self, extra_params=()):
+    def cursor(self, extra_params=(), extra_cursor_params=()):
         with self.connection(extra_params) as conn:
             cursor = conn.cursor(
-                arraysize=self.arraysize, buffer_size_bytes=self.buffer_size_bytes
+                arraysize=self.arraysize,
+                buffer_size_bytes=self.buffer_size_bytes,
+                **extra_cursor_params,
             )
             try:
                 yield cursor
@@ -945,6 +947,60 @@ def test_result_set_close(self):
             finally:
                 cursor.close()
 
+    def test_row_limit_with_larger_result(self):
+        """Test that row_limit properly constrains results when query would return more rows"""
+        row_limit = 1000
+        with self.cursor(extra_cursor_params={"row_limit": row_limit}) as cursor:
+            # Execute a query that returns more than row_limit rows
+            cursor.execute("SELECT * FROM range(2000)")
+            rows = cursor.fetchall()
+
+            # Check if the number of rows is limited to row_limit
+            assert len(rows) == row_limit, f"Expected {row_limit} rows, got {len(rows)}"
+
+    def test_row_limit_with_smaller_result(self):
+        """Test that row_limit doesn't affect results when query returns fewer rows than limit"""
+        row_limit = 100
+        expected_rows = 50
+        with self.cursor(extra_cursor_params={"row_limit": row_limit}) as cursor:
+            # Execute a query that returns fewer than row_limit rows
+            cursor.execute(f"SELECT * FROM range({expected_rows})")
+            rows = cursor.fetchall()
+
+            # Check if all rows are returned (not limited by row_limit)
+            assert (
+                len(rows) == expected_rows
+            ), f"Expected {expected_rows} rows, got {len(rows)}"
+
+    @skipUnless(pysql_supports_arrow(), "arrow test needs arrow support")
+    def test_row_limit_with_arrow_larger_result(self):
+        """Test that row_limit properly constrains arrow results when query would return more rows"""
+        row_limit = 800
+        with self.cursor(extra_cursor_params={"row_limit": row_limit}) as cursor:
+            # Execute a query that returns more than row_limit rows
+            cursor.execute("SELECT * FROM range(1500)")
+            arrow_table = cursor.fetchall_arrow()
+
+            # Check if the number of rows in the arrow table is limited to row_limit
+            assert (
+                arrow_table.num_rows == row_limit
+            ), f"Expected {row_limit} rows, got {arrow_table.num_rows}"
+
+    @skipUnless(pysql_supports_arrow(), "arrow test needs arrow support")
+    def test_row_limit_with_arrow_smaller_result(self):
+        """Test that row_limit doesn't affect arrow results when query returns fewer rows than limit"""
+        row_limit = 200
+        expected_rows = 100
+        with self.cursor(extra_cursor_params={"row_limit": row_limit}) as cursor:
+            # Execute a query that returns fewer than row_limit rows
+            cursor.execute(f"SELECT * FROM range({expected_rows})")
+            arrow_table = cursor.fetchall_arrow()
+
+            # Check if all rows are returned (not limited by row_limit)
+            assert (
+                arrow_table.num_rows == expected_rows
+            ), f"Expected {expected_rows} rows, got {arrow_table.num_rows}"
+
 
 # use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
 # the 429/503 subsuites separate since they execute under different circumstances.