fix merge artifacts

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

Author: varun-edachali-dbx

src/databricks/sql/backend/filters.py

@@ -20,6 +20,7 @@
 from databricks.sql.utils import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.types import ExecuteResponse, CommandId
 from databricks.sql.backend.sea.models.base import ResultData
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
 
 if TYPE_CHECKING:
     from databricks.sql.result_set import ResultSet, SeaResultSet
@@ -69,19 +70,21 @@ def _filter_sea_result_set(
             status=result_set.status,
             description=result_set.description,
             has_more_rows=result_set._has_more_rows,
-            results_queue=JsonQueue(filtered_rows),
             has_been_closed_server_side=result_set.has_been_closed_server_side,
             lz4_compressed=False,
             is_staging_operation=False,
         )
 
-        return SeaResultSet(
+        # Create a new SeaResultSet with the filtered data
+        filtered_result_set = SeaResultSet(
             connection=result_set.connection,
             execute_response=execute_response,
-            sea_client=result_set.backend,
+            sea_client=cast(SeaDatabricksClient, result_set.backend),
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
         )
+        
+        return filtered_result_set
 
     @staticmethod
     def filter_by_column_values(

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

@@ -1,7 +1,10 @@
 import logging
 import uuid
 import time
-from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING
+import re
+from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING, Set
+
+from databricks.sql.backend.sea.utils.constants import ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
@@ -43,6 +46,32 @@
 logger = logging.getLogger(__name__)
 
 
+def _filter_session_configuration(
+    session_configuration: Optional[Dict[str, str]]
+) -> Optional[Dict[str, str]]:
+    if not session_configuration:
+        return None
+
+    filtered_session_configuration = {}
+    ignored_configs: Set[str] = set()
+
+    for key, value in session_configuration.items():
+        if key.upper() in ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP:
+            filtered_session_configuration[key.lower()] = value
+        else:
+            ignored_configs.add(key)
+
+    if ignored_configs:
+        logger.warning(
+            "Some session configurations were ignored because they are not supported: %s",
+            ignored_configs,
+        )
+        logger.warning(
+            "Supported session configurations are: %s",
+            list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys()),
+        )
+
+    return filtered_session_configuration
 class SeaDatabricksClient(DatabricksClient):
     """
     Statement Execution API (SEA) implementation of the DatabricksClient interface.
@@ -60,7 +89,7 @@ class SeaDatabricksClient(DatabricksClient):
     CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
     CHUNKS_PATH_WITH_ID = STATEMENT_PATH + "/{}/result/chunks"
     CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}"
-
+    
     def __init__(
         self,
         server_hostname: str,
@@ -92,6 +121,7 @@ def __init__(
         )
 
         self._max_download_threads = kwargs.get("max_download_threads", 10)
+        self.ssl_options = ssl_options
 
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
@@ -232,9 +262,45 @@ def close_session(self, session_id: SessionId) -> None:
         self.http_client._make_request(
             method="DELETE",
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
-            params=request.to_dict(),
+            data=request_data.to_dict(),
         )
 
+    @staticmethod
+    def get_default_session_configuration_value(name: str) -> Optional[str]:
+        """
+        Get the default value for a session configuration parameter.
+
+        Args:
+            name: The name of the session configuration parameter
+
+        Returns:
+            The default value if the parameter is supported, None otherwise
+        """
+        return ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.get(name.upper())
+
+    @staticmethod
+    def is_session_configuration_parameter_supported(name: str) -> bool:
+        """
+        Check if a session configuration parameter is supported.
+
+        Args:
+            name: The name of the session configuration parameter
+
+        Returns:
+            True if the parameter is supported, False otherwise
+        """
+        return name.upper() in ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP
+
+    @staticmethod
+    def get_allowed_session_configurations() -> List[str]:
+        """
+        Get the list of allowed session configuration parameters.
+
+        Returns:
+            List of allowed session configuration parameter names
+        """
+        return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
+
     def fetch_chunk_links(
         self, statement_id: str, chunk_index: int
     ) -> List["ExternalLink"]:
@@ -491,10 +557,11 @@ def _results_message_to_execute_response(self, sea_response, command_id):
             status=state,
             description=description,
             has_more_rows=False,
-            results_queue=results_queue,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=False,
+            arrow_schema_bytes=schema_bytes,
+            result_format=manifest_data.get("format"),
         )
 
     def execute_command(

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

@@ -39,8 +39,12 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
                 error_code=error_data.get("error_code"),
             )
 
+        state = CommandState.from_sea_state(status_data.get("state", ""))
+        if state is None: 
+            raise ValueError(f"Invalid state: {status_data.get('state', '')}")
+
         status = StatementStatus(
-            state=CommandState.from_sea_state(status_data.get("state", "")),
+            state=state,
             error=error,
             sql_state=status_data.get("sql_state"),
         )
@@ -119,8 +123,12 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
                 error_code=error_data.get("error_code"),
             )
 
+        state = CommandState.from_sea_state(status_data.get("state", ""))
+        if state is None: 
+            raise ValueError(f"Invalid state: {status_data.get('state', '')}")
+
         status = StatementStatus(
-            state=CommandState.from_sea_state(status_data.get("state", "")),
+            state=state,
             error=error,
             sql_state=status_data.get("sql_state"),
         )

src/databricks/sql/cloud_fetch_queue.py

@@ -6,7 +6,7 @@
 """
 
 from abc import ABC
-from typing import Any, List, Optional, Union, TYPE_CHECKING
+from typing import Any, List, Optional, Tuple, Union, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from databricks.sql.backend.sea.backend import SeaDatabricksClient
@@ -137,7 +137,7 @@ def __init__(
         max_download_threads: int,
         ssl_options: SSLOptions,
         lz4_compressed: bool = True,
-        description: Optional[List[List[Any]]] = None,
+        description: Optional[List[Tuple[Any, ...]]] = None,
     ):
         """
         Initialize the base CloudFetchQueue.
@@ -227,7 +227,7 @@ def __init__(
         statement_id: str,
         total_chunk_count: int,
         lz4_compressed: bool = False,
-        description: Optional[List[List[Any]]] = None,
+        description: Optional[List[Tuple[Any, ...]]] = None,
     ):
         """
         Initialize the SEA CloudFetchQueue.
@@ -496,7 +496,7 @@ def __init__(
         start_row_offset: int = 0,
         result_links: Optional[List[TSparkArrowResultLink]] = None,
         lz4_compressed: bool = True,
-        description: Optional[List[List[Any]]] = None,
+        description: Optional[List[Tuple[Any, ...]]] = None,
     ):
         """
         Initialize the Thrift CloudFetchQueue.

src/databricks/sql/result_set.py

@@ -6,7 +6,9 @@
 import pandas
 
 from databricks.sql.backend.sea.backend import SeaDatabricksClient
+from databricks.sql.backend.sea.models.base import ExternalLink, ResultData, ResultManifest
 from databricks.sql.cloud_fetch_queue import SeaCloudFetchQueue
+from databricks.sql.utils import SeaResultSetQueueFactory
 
 try:
     import pyarrow
@@ -465,6 +467,77 @@ def __init__(
             if execute_response.command_id
             else None
         )
+        
+        # Get the response data from the SEA backend
+        response_data = sea_client.http_client._make_request(
+            method="GET",
+            path=sea_client.STATEMENT_PATH_WITH_ID.format(self.statement_id),
+            data={"statement_id": self.statement_id},
+        )
+        
+        # Build the results queue
+        results_queue = None
+        
+        # Extract data from the response
+        result_data = response_data.get("result", {})
+        manifest_data = response_data.get("manifest", {})
+        
+        if result_data:
+            # Convert external links
+            external_links = None
+            if "external_links" in result_data:
+                external_links = []
+                for link_data in result_data["external_links"]:
+                    external_links.append(
+                        ExternalLink(
+                            external_link=link_data.get("external_link", ""),
+                            expiration=link_data.get("expiration", ""),
+                            chunk_index=link_data.get("chunk_index", 0),
+                            byte_count=link_data.get("byte_count", 0),
+                            row_count=link_data.get("row_count", 0),
+                            row_offset=link_data.get("row_offset", 0),
+                            next_chunk_index=link_data.get("next_chunk_index"),
+                            next_chunk_internal_link=link_data.get("next_chunk_internal_link"),
+                            http_headers=link_data.get("http_headers", {}),
+                        )
+                    )
+
+            # Create the result data object
+            result_data_obj = ResultData(
+                data=result_data.get("data_array"), external_links=external_links
+            )
+
+            # Create the manifest object
+            manifest_obj = ResultManifest(
+                format=manifest_data.get("format", ""),
+                schema=manifest_data.get("schema", {}),
+                total_row_count=manifest_data.get("total_row_count", 0),
+                total_byte_count=manifest_data.get("total_byte_count", 0),
+                total_chunk_count=manifest_data.get("total_chunk_count", 0),
+                truncated=manifest_data.get("truncated", False),
+                chunks=manifest_data.get("chunks"),
+                result_compression=manifest_data.get("result_compression"),
+            )
+
+            # Build the queue based on the response data
+            from typing import cast, List
+            
+            # Convert description to the expected format
+            desc = None
+            if execute_response.description:
+                desc = cast(List[Tuple[Any, ...]], execute_response.description)
+            
+            results_queue = SeaResultSetQueueFactory.build_queue(
+                result_data_obj,
+                manifest_obj,
+                str(self.statement_id),
+                description=desc,
+                schema_bytes=execute_response.arrow_schema_bytes if execute_response.arrow_schema_bytes else None,
+                max_download_threads=sea_client.max_download_threads,
+                ssl_options=sea_client.ssl_options,
+                sea_client=sea_client,
+                lz4_compressed=execute_response.lz4_compressed,
+            )
 
         # Call parent constructor with common attributes
         super().__init__(
@@ -476,14 +549,12 @@ def __init__(
             status=execute_response.status,
             has_been_closed_server_side=execute_response.has_been_closed_server_side,
             has_more_rows=execute_response.has_more_rows,
-            results_queue=execute_response.results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
         )
 
         # Initialize queue for result data if not provided
-        if not self.results:
-            self.results = JsonQueue([])
+        self.results = results_queue or JsonQueue([])
 
     def _convert_to_row_objects(self, rows):
         """

src/databricks/sql/utils.py

@@ -63,7 +63,7 @@ def build_queue(
         max_download_threads: Optional[int] = None,
         ssl_options: Optional[SSLOptions] = None,
         lz4_compressed: bool = True,
-        description: Optional[List[Tuple]] = None,
+        description: Optional[List[Tuple[Any, ...]]] = None,
     ) -> ResultSetQueue:
         """
         Factory method to build a result set queue for Thrift backend.
@@ -131,7 +131,7 @@ def build_queue(
         sea_result_data: ResultData,
         manifest: ResultManifest,
         statement_id: str,
-        description: Optional[List[List[Any]]] = None,
+        description: Optional[List[Tuple[Any, ...]]] = None,
         schema_bytes: Optional[bytes] = None,
         max_download_threads: Optional[int] = None,
         ssl_options: Optional[SSLOptions] = None,