init hybrid

Author: varun-edachali-dbx

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

@@ -125,6 +125,8 @@ def __init__(
 
         super().__init__(ssl_options=ssl_options, **kwargs)
 
+        self.use_hybrid_disposition = kwargs.get("use_hybrid_disposition", False)
+
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
 
@@ -449,7 +451,11 @@ def execute_command(
             ResultFormat.ARROW_STREAM if use_cloud_fetch else ResultFormat.JSON_ARRAY
         ).value
         disposition = (
-            ResultDisposition.EXTERNAL_LINKS
+            (
+                ResultDisposition.HYBRID
+                if self.use_hybrid_disposition
+                else ResultDisposition.EXTERNAL_LINKS
+            )
             if use_cloud_fetch
             else ResultDisposition.INLINE
         ).value
@@ -630,7 +636,7 @@ def get_execution_result(
             arraysize=cursor.arraysize,
         )
 
-    def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
+    def get_chunk_links(self, statement_id: str, chunk_index: int) -> List[ExternalLink]:
         """
         Get links for chunks starting from the specified index.
         Args:
@@ -647,17 +653,7 @@ def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
         response = GetChunksResponse.from_dict(response_data)
 
         links = response.external_links
-        link = next((l for l in links if l.chunk_index == chunk_index), None)
-        if not link:
-            raise ServerOperationError(
-                f"No link found for chunk index {chunk_index}",
-                {
-                    "operation-id": statement_id,
-                    "diagnostic-info": None,
-                },
-            )
-
-        return link
+        return links
 
     # == Metadata Operations ==
 

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

@@ -28,7 +28,7 @@ class ResultFormat(Enum):
 class ResultDisposition(Enum):
     """Enum for result disposition values."""
 
-    # TODO: add support for hybrid disposition
+    HYBRID = "INLINE_OR_EXTERNAL_LINKS"
     EXTERNAL_LINKS = "EXTERNAL_LINKS"
     INLINE = "INLINE"
 

src/databricks/sql/client.py

@@ -87,6 +87,10 @@ def __init__(
         Connect to a Databricks SQL endpoint or a Databricks cluster.
 
         Parameters:
+            :param use_sea: `bool`, optional (default is False)
+                Use the SEA backend instead of the Thrift backend.
+            :param use_hybrid_disposition: `bool`, optional (default is False)
+                Use the hybrid disposition instead of the inline disposition.
             :param server_hostname: Databricks instance host name.
             :param http_path: Http path either to a DBSQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef)
                 or to a DBR interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)