add client side filtering for session confs, add note on warehouses over endoints

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

Author: varun-edachali-dbx

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

@@ -1,6 +1,6 @@
 import logging
 import re
-from typing import Dict, Tuple, List, Optional, TYPE_CHECKING
+from typing import Dict, Tuple, List, Optional, TYPE_CHECKING, Set
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
@@ -9,6 +9,9 @@
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
 from databricks.sql.exc import ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
+from databricks.sql.backend.sea.utils.constants import (
+    ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
+)
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 
@@ -21,6 +24,34 @@
 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.
@@ -111,7 +142,8 @@ def _extract_warehouse_id(self, http_path: str) -> str:
         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"/path/to/endpoints/{{warehouse_id}}."
+            f"Note: SEA only works for warehouses."
         )
         logger.error(error_message)
         raise ValueError(error_message)
@@ -152,6 +184,8 @@ def open_session(
             schema,
         )
 
+        session_configuration = _filter_session_configuration(session_configuration)
+
         request_data = CreateSessionRequest(
             warehouse_id=self.warehouse_id,
             session_confs=session_configuration,
@@ -205,6 +239,29 @@ def close_session(self, session_id: SessionId) -> None:
             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 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())
+
     # == Not Implemented Operations ==
     # These methods will be implemented in future iterations
 

src/databricks/sql/backend/sea/utils/constants.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",
+}

tests/unit/test_sea_backend.py

@@ -101,8 +101,12 @@ def test_open_session_with_all_parameters(self, sea_client, mock_http_client):
         # Set up mock response
         mock_http_client._make_request.return_value = {"session_id": "test-session-456"}
 
-        # Call the method with all parameters
-        session_config = {"ANSI_MODE": "FALSE", "STATEMENT_TIMEOUT": "3600"}
+        # Call the method with all parameters, including both supported and unsupported configurations
+        session_config = {
+            "ANSI_MODE": "FALSE",  # Supported parameter
+            "STATEMENT_TIMEOUT": "3600",  # Supported parameter
+            "unsupported_param": "value",  # Unsupported parameter
+        }
         catalog = "test_catalog"
         schema = "test_schema"
 
@@ -113,10 +117,14 @@ def test_open_session_with_all_parameters(self, sea_client, mock_http_client):
         assert session_id.backend_type == BackendType.SEA
         assert session_id.guid == "test-session-456"
 
-        # Verify the HTTP request
+        # Verify the HTTP request - only supported parameters should be included
+        # and keys should be in lowercase
         expected_data = {
             "warehouse_id": "abc123",
-            "session_confs": session_config,
+            "session_confs": {
+                "ansi_mode": "FALSE",
+                "statement_timeout": "3600",
+            },
             "catalog": catalog,
             "schema": schema,
         }
@@ -166,3 +174,32 @@ def test_close_session_invalid_id_type(self, sea_client):
             sea_client.close_session(session_id)
 
         assert "Not a valid SEA session ID" in str(excinfo.value)
+
+    def test_session_configuration_helpers(self):
+        """Test the session configuration helper methods."""
+        # Test getting default value for a supported parameter
+        default_value = SeaDatabricksClient.get_default_session_configuration_value(
+            "ANSI_MODE"
+        )
+        assert default_value == "true"
+
+        # Test getting default value for an unsupported parameter
+        default_value = SeaDatabricksClient.get_default_session_configuration_value(
+            "UNSUPPORTED_PARAM"
+        )
+        assert default_value is None
+
+        # Test getting the list of allowed configurations
+        allowed_configs = SeaDatabricksClient.get_allowed_session_configurations()
+
+        expected_keys = {
+            "ANSI_MODE",
+            "ENABLE_PHOTON",
+            "LEGACY_TIME_PARSER_POLICY",
+            "MAX_FILE_PARTITION_BYTES",
+            "READ_ONLY_EXTERNAL_METASTORE",
+            "STATEMENT_TIMEOUT",
+            "TIMEZONE",
+            "USE_CACHED_RESULT",
+        }
+        assert set(allowed_configs) == expected_keys