Merge branch 'main' into execution-sea

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -969,7 +969,13 @@ def execute_command(
         parameters=[],
         async_op=False,
         enforce_embedded_schema_correctness=False,
+<<<<<<< HEAD
         row_limit: Optional[int] = None,
+||||||| 576eafc
+    ):
+        assert session_handle is not None
+=======
+>>>>>>> main
     ) -> Union["ResultSet", None]:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
@@ -1018,10 +1024,17 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
+<<<<<<< HEAD
             execute_response, is_direct_results = self._handle_execute_response(
                 resp, cursor
             )
+||||||| 576eafc
+            return self._handle_execute_response(resp, cursor)
+=======
+            execute_response = self._handle_execute_response(resp, cursor)
+>>>>>>> main
 
+<<<<<<< HEAD
             t_row_set = None
             if resp.directResults and resp.directResults.resultSet:
                 t_row_set = resp.directResults.resultSet.results
@@ -1045,6 +1058,26 @@ def get_catalogs(
         max_rows: int,
         max_bytes: int,
         cursor: "Cursor",
+||||||| 576eafc
+    def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
+        assert session_handle is not None
+=======
+            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_id: SessionId,
+        max_rows: int,
+        max_bytes: int,
+        cursor: Cursor,
+>>>>>>> main
     ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
@@ -1058,6 +1091,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
+<<<<<<< HEAD
         execute_response, is_direct_results = self._handle_execute_response(
             resp, cursor
         )
@@ -1077,6 +1111,19 @@ def get_catalogs(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+||||||| 576eafc
+        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,
+>>>>>>> main
         )
 
     def get_schemas(
@@ -1088,8 +1135,14 @@ def get_schemas(
         catalog_name=None,
         schema_name=None,
     ) -> "ResultSet":
+<<<<<<< HEAD
         from databricks.sql.result_set import ThriftResultSet
 
+||||||| 576eafc
+    ):
+        assert session_handle is not None
+=======
+>>>>>>> main
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1104,6 +1157,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
+<<<<<<< HEAD
         execute_response, is_direct_results = self._handle_execute_response(
             resp, cursor
         )
@@ -1123,6 +1177,19 @@ def get_schemas(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+||||||| 576eafc
+        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,
+>>>>>>> main
         )
 
     def get_tables(
@@ -1136,8 +1203,14 @@ def get_tables(
         table_name=None,
         table_types=None,
     ) -> "ResultSet":
+<<<<<<< HEAD
         from databricks.sql.result_set import ThriftResultSet
 
+||||||| 576eafc
+    ):
+        assert session_handle is not None
+=======
+>>>>>>> main
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1154,6 +1227,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
+<<<<<<< HEAD
         execute_response, is_direct_results = self._handle_execute_response(
             resp, cursor
         )
@@ -1173,6 +1247,19 @@ def get_tables(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+||||||| 576eafc
+        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,
+>>>>>>> main
         )
 
     def get_columns(
@@ -1186,8 +1273,14 @@ def get_columns(
         table_name=None,
         column_name=None,
     ) -> "ResultSet":
+<<<<<<< HEAD
         from databricks.sql.result_set import ThriftResultSet
 
+||||||| 576eafc
+    ):
+        assert session_handle is not None
+=======
+>>>>>>> main
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1204,6 +1297,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
+<<<<<<< HEAD
         execute_response, is_direct_results = self._handle_execute_response(
             resp, cursor
         )
@@ -1223,6 +1317,19 @@ def get_columns(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+||||||| 576eafc
+        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,
+>>>>>>> main
         )
 
     def _handle_execute_response(self, resp, cursor):