introduce row_limit

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -85,6 +85,7 @@ def execute_command(
         parameters: List,
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
+        row_limit: Optional[int] = None,
     ) -> Union["ResultSet", None]:
         """
         Executes a SQL command or query within the specified session.
@@ -103,6 +104,7 @@ def execute_command(
             parameters: List of parameters to bind to the query
             async_op: Whether to execute the command asynchronously
             enforce_embedded_schema_correctness: Whether to enforce schema correctness
+            row_limit: Maximum number of rows to fetch overall. Only supported for SEA protocol.
 
         Returns:
             If async_op is False, returns a ResultSet object containing the

src/databricks/sql/backend/sea/backend.py

@@ -405,6 +405,7 @@ def execute_command(
         parameters: List[Dict[str, Any]],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
+        row_limit: Optional[int] = None,
     ) -> Union["ResultSet", None]:
         """
         Execute a SQL command using the SEA backend.
@@ -462,7 +463,7 @@ def execute_command(
             format=format,
             wait_timeout=(WaitTimeout.ASYNC if async_op else WaitTimeout.SYNC).value,
             on_wait_timeout="CONTINUE",
-            row_limit=max_rows,
+            row_limit=row_limit,
             parameters=sea_parameters if sea_parameters else None,
             result_compression=result_compression,
         )

src/databricks/sql/client.py

@@ -341,6 +341,7 @@ def cursor(
         self,
         arraysize: int = DEFAULT_ARRAY_SIZE,
         buffer_size_bytes: int = DEFAULT_RESULT_BUFFER_SIZE_BYTES,
+        row_limit: int = None,
     ) -> "Cursor":
         """
         Return a new Cursor object using the connection.
@@ -355,6 +356,7 @@ def cursor(
             self.session.backend,
             arraysize=arraysize,
             result_buffer_size_bytes=buffer_size_bytes,
+            row_limit=row_limit,
         )
         self._cursors.append(cursor)
         return cursor
@@ -388,6 +390,7 @@ def __init__(
         backend: DatabricksClient,
         result_buffer_size_bytes: int = DEFAULT_RESULT_BUFFER_SIZE_BYTES,
         arraysize: int = DEFAULT_ARRAY_SIZE,
+        row_limit: Optional[int] = None,
     ) -> None:
         """
         These objects represent a database cursor, which is used to manage the context of a fetch
@@ -397,11 +400,17 @@ def __init__(
         visible by other cursors or connections.
         """
 
+        if not self.connection.session.use_sea and row_limit is not None:
+            logger.warning(
+                "Row limit is only supported for SEA protocol. Ignoring row_limit."
+            )
+
         self.connection = connection
         self.rowcount = -1  # Return -1 as this is not supported
         self.buffer_size_bytes = result_buffer_size_bytes
         self.active_result_set: Union[ResultSet, None] = None
         self.arraysize = arraysize
+        self.row_limit = row_limit
         # Note that Cursor closed => active result set closed, but not vice versa
         self.open = True
         self.executing_command_id = None

src/databricks/sql/session.py

@@ -97,10 +97,10 @@ def _create_backend(
         kwargs: dict,
     ) -> DatabricksClient:
         """Create and return the appropriate backend client."""
-        use_sea = kwargs.get("use_sea", False)
+        self.use_sea = kwargs.get("use_sea", False)
 
         databricks_client_class: Type[DatabricksClient]
-        if use_sea:
+        if self.use_sea:
             logger.debug("Creating SEA backend client")
             databricks_client_class = SeaDatabricksClient
         else: