Merge branch 'exec-resp-norm' into cloudfetch-sea

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -3,11 +3,9 @@
 
 Implementations of this class are responsible for:
 - Managing connections to Databricks SQL services
-- Handling authentication
 - Executing SQL queries and commands
 - Retrieving query results
 - Fetching metadata about catalogs, schemas, tables, and columns
-- Managing error handling and retries
 """
 
 from abc import ABC, abstractmethod
@@ -178,29 +176,28 @@ def get_columns(
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
     ) -> "ResultSet":
-        pass
-
-    # == Properties ==
-    @property
-    @abstractmethod
-    def staging_allowed_local_path(self) -> Union[None, str, List[str]]:
         """
-        Gets the allowed local paths for staging operations.
+        Retrieves a list of columns, optionally filtered by catalog, schema, table, and column name patterns.
 
-        Returns:
-            Union[None, str, List[str]]: The allowed local paths for staging operations,
-            or None if staging is not allowed
-        """
-        pass
+        This method fetches metadata about columns available in the specified table,
+        or all tables if not specified.
 
-    @property
-    @abstractmethod
-    def ssl_options(self) -> SSLOptions:
-        """
-        Gets the SSL options for this client.
+        Args:
+            session_id: The session identifier
+            max_rows: Maximum number of rows to fetch in a single batch
+            max_bytes: Maximum number of bytes to fetch in a single batch
+            cursor: The cursor object that will handle the results
+            catalog_name: Optional catalog name pattern to filter by
+            schema_name: Optional schema name pattern to filter by
+            table_name: Optional table name pattern to filter by
+            column_name: Optional column name pattern to filter by
 
         Returns:
-            SSLOptions: The SSL configuration options
+            ResultSet: An object containing the column metadata
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error retrieving the columns
         """
         pass
 

src/databricks/sql/backend/databricks_client.py

@@ -69,7 +69,6 @@ def __init__(
         http_headers: List[Tuple[str, str]],
         auth_provider,
         ssl_options: SSLOptions,
-        staging_allowed_local_path: Union[None, str, List[str]] = None,
         **kwargs,
     ):
         """
@@ -82,25 +81,23 @@ def __init__(
             http_headers: List of HTTP headers to include in requests
             auth_provider: Authentication provider
             ssl_options: SSL configuration options
-            staging_allowed_local_path: Allowed local paths for staging operations
             **kwargs: Additional keyword arguments
         """
+
         logger.debug(
-            "SEADatabricksClient.__init__(server_hostname=%s, port=%s, http_path=%s)",
+            "SeaDatabricksClient.__init__(server_hostname=%s, port=%s, http_path=%s)",
             server_hostname,
             port,
             http_path,
         )
 
-        self._staging_allowed_local_path = staging_allowed_local_path
-        self._ssl_options = ssl_options
         self._max_download_threads = kwargs.get("max_download_threads", 10)
 
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
 
         # Initialize HTTP client
-        self.http_client = CustomHttpClient(
+        self.http_client = SeaHttpClient(
             server_hostname=server_hostname,
             port=port,
             http_path=http_path,
@@ -114,47 +111,38 @@ def _extract_warehouse_id(self, http_path: str) -> str:
         """
         Extract the warehouse ID from the HTTP path.
 
-        The warehouse ID is expected to be the last segment of the path when the
-        second-to-last segment is either 'warehouses' or 'endpoints'.
-
         Args:
             http_path: The HTTP path from which to extract the warehouse ID
 
         Returns:
             The extracted warehouse ID
 
         Raises:
-            Error: If the warehouse ID cannot be extracted from the path
+            ValueError: If the warehouse ID cannot be extracted from the path
         """
-        path_parts = http_path.strip("/").split("/")
-        warehouse_id = None
 
-        if len(path_parts) >= 3 and path_parts[-2] in ["warehouses", "endpoints"]:
-            warehouse_id = path_parts[-1]
+        warehouse_pattern = re.compile(r".*/warehouses/(.+)")
+        endpoint_pattern = re.compile(r".*/endpoints/(.+)")
+
+        for pattern in [warehouse_pattern, endpoint_pattern]:
+            match = pattern.match(http_path)
+            if not match:
+                continue
+            warehouse_id = match.group(1)
             logger.debug(
                 f"Extracted warehouse ID: {warehouse_id} from path: {http_path}"
             )
-
-        if not warehouse_id:
-            error_message = (
-                f"Could not extract warehouse ID from http_path: {http_path}. "
-                f"Expected format: /path/to/warehouses/{{warehouse_id}} or "
-                f"/path/to/endpoints/{{warehouse_id}}"
-            )
-            logger.error(error_message)
-            raise ValueError(error_message)
-
-        return warehouse_id
-
-    @property
-    def staging_allowed_local_path(self) -> Union[None, str, List[str]]:
-        """Get the allowed local paths for staging operations."""
-        return self._staging_allowed_local_path
-
-    @property
-    def ssl_options(self) -> SSLOptions:
-        """Get the SSL options for this client."""
-        return self._ssl_options
+            return warehouse_id
+
+        # If no match found, raise error
+        error_message = (
+            f"Could not extract warehouse ID from http_path: {http_path}. "
+            f"Expected format: /path/to/warehouses/{{warehouse_id}} or "
+            f"/path/to/endpoints/{{warehouse_id}}."
+            f"Note: SEA only works for warehouses."
+        )
+        logger.error(error_message)
+        raise ValueError(error_message)
 
     @property
     def max_download_threads(self) -> int:
@@ -171,7 +159,9 @@ def open_session(
         Opens a new session with the Databricks SQL service using SEA.
 
         Args:
-            session_configuration: Optional dictionary of configuration parameters for the session
+            session_configuration: Optional dictionary of configuration parameters for the session.
+                                   Only specific parameters are supported as documented at:
+                                   https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-parameters
             catalog: Optional catalog name to use as the initial catalog for the session
             schema: Optional schema name to use as the initial schema for the session
 
@@ -182,28 +172,29 @@ def open_session(
             Error: If the session configuration is invalid
             OperationalError: If there's an error establishing the session
         """
+
         logger.debug(
-            "SEADatabricksClient.open_session(session_configuration=%s, catalog=%s, schema=%s)",
+            "SeaDatabricksClient.open_session(session_configuration=%s, catalog=%s, schema=%s)",
             session_configuration,
             catalog,
             schema,
         )
 
-        request = CreateSessionRequest(
+        session_configuration = _filter_session_configuration(session_configuration)
+
+        request_data = CreateSessionRequest(
             warehouse_id=self.warehouse_id,
             session_confs=session_configuration,
             catalog=catalog,
             schema=schema,
         )
 
         response = self.http_client._make_request(
-            method="POST", path=self.SESSION_PATH, data=request.to_dict()
+            method="POST", path=self.SESSION_PATH, data=request_data.to_dict()
         )
 
-        # Parse the response
         session_response = CreateSessionResponse.from_dict(response)
         session_id = session_response.session_id
-
         if not session_id:
             raise ServerOperationError(
                 "Failed to create session: No session ID returned",
@@ -226,14 +217,16 @@ def close_session(self, session_id: SessionId) -> None:
             ValueError: If the session ID is invalid
             OperationalError: If there's an error closing the session
         """
-        logger.debug("SEADatabricksClient.close_session(session_id=%s)", session_id)
+
+        logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
 
         if session_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA session ID")
         sea_session_id = session_id.to_sea_session_id()
 
-        request = DeleteSessionRequest(
-            warehouse_id=self.warehouse_id, session_id=sea_session_id
+        request_data = DeleteSessionRequest(
+            warehouse_id=self.warehouse_id,
+            session_id=sea_session_id,
         )
 
         self.http_client._make_request(

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

@@ -129,13 +129,5 @@ class DeleteSessionRequest:
     session_id: str
 
     def to_dict(self) -> Dict[str, str]:
-        """
-        Convert the request to query parameters.
-
-        In the SEA API, only the warehouse_id is sent as a query parameter.
-        The session_id is included in the URL path.
-
-        Returns:
-            A dictionary containing the warehouse_id as a query parameter
-        """
+        """Convert the request to a dictionary for JSON serialization."""
         return {"warehouse_id": self.warehouse_id, "session_id": self.session_id}

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

@@ -0,0 +1,17 @@
+"""
+Constants for the Statement Execution API (SEA) backend.
+"""
+
+from typing import Dict
+
+# from https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-parameters
+ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP: Dict[str, str] = {
+    "ANSI_MODE": "true",
+    "ENABLE_PHOTON": "true",
+    "LEGACY_TIME_PARSER_POLICY": "Exception",
+    "MAX_FILE_PARTITION_BYTES": "128m",
+    "READ_ONLY_EXTERNAL_METASTORE": "false",
+    "STATEMENT_TIMEOUT": "0",
+    "TIMEZONE": "UTC",
+    "USE_CACHED_RESULT": "true",
+}

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

@@ -10,7 +10,7 @@
 logger = logging.getLogger(__name__)
 
 
-class CustomHttpClient:
+class SeaHttpClient:
     """
     HTTP client for Statement Execution API (SEA).
 
@@ -145,7 +145,7 @@ def _make_request(
             if response.content:
                 result = response.json()
                 # Log response content (but limit it for large responses)
-                content_str = json.dumps(result, indent=2, sort_keys=True)
+                content_str = json.dumps(result)
                 if len(content_str) > 1000:
                     logger.debug(
                         f"Response content (truncated): {content_str[:1000]}..."