introduce Result Set interface and concrete thrift implementation

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

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -5,6 +5,11 @@
 from databricks.sql.utils import ExecuteResponse
 from databricks.sql.types import SSLOptions
 
+# Forward reference for type hints
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from databricks.sql.result_set import ResultSet
+
 
 class DatabricksClient(ABC):
     # == Connection and Session Management ==
@@ -35,7 +40,7 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
-    ) -> Any:
+    ) -> "ResultSet":  # Changed return type to ResultSet
         pass
 
     @abstractmethod

src/databricks/sql/backend/thrift_backend.py

@@ -42,6 +42,7 @@
 )
 from databricks.sql.types import SSLOptions
 from databricks.sql.backend.databricks_client import DatabricksClient
+from databricks.sql.result_set import ThriftResultSet
 
 logger = logging.getLogger(__name__)
 
@@ -833,7 +834,7 @@ def get_execution_result(self, op_handle, cursor):
             ssl_options=self._ssl_options,
         )
 
-        return ExecuteResponse(
+        execute_response = ExecuteResponse(
             arrow_queue=queue,
             status=resp.status,
             has_been_closed_server_side=False,
@@ -844,6 +845,15 @@ def get_execution_result(self, op_handle, cursor):
             description=description,
             arrow_schema_bytes=schema_bytes,
         )
+        
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch
+        )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
         if initial_operation_status_resp:
@@ -938,8 +948,18 @@ def execute_command(
 
         if async_op:
             self._handle_execute_response_async(resp, cursor)
+            return None
         else:
-            return self._handle_execute_response(resp, cursor)
+            execute_response = self._handle_execute_response(resp, cursor)
+            
+            return ThriftResultSet(
+                connection=cursor.connection,
+                execute_response=execute_response,
+                thrift_client=self,
+                buffer_size_bytes=max_bytes,
+                arraysize=max_rows,
+                use_cloud_fetch=use_cloud_fetch
+            )
 
     def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
         assert session_handle is not None
@@ -951,7 +971,17 @@ def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
             ),
         )
         resp = self.make_request(self._client.GetCatalogs, req)
-        return self._handle_execute_response(resp, cursor)
+        
+        execute_response = self._handle_execute_response(resp, cursor)
+        
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch
+        )
 
     def get_schemas(
         self,
@@ -973,7 +1003,17 @@ def get_schemas(
             schemaName=schema_name,
         )
         resp = self.make_request(self._client.GetSchemas, req)
-        return self._handle_execute_response(resp, cursor)
+        
+        execute_response = self._handle_execute_response(resp, cursor)
+        
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch
+        )
 
     def get_tables(
         self,
@@ -999,7 +1039,17 @@ def get_tables(
             tableTypes=table_types,
         )
         resp = self.make_request(self._client.GetTables, req)
-        return self._handle_execute_response(resp, cursor)
+        
+        execute_response = self._handle_execute_response(resp, cursor)
+        
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch
+        )
 
     def get_columns(
         self,
@@ -1025,7 +1075,17 @@ def get_columns(
             columnName=column_name,
         )
         resp = self.make_request(self._client.GetColumns, req)
-        return self._handle_execute_response(resp, cursor)
+        
+        execute_response = self._handle_execute_response(resp, cursor)
+        
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch
+        )
 
     def _handle_execute_response(self, resp, cursor):
         cursor.active_op_handle = resp.operationHandle