remove irrelevant changes

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -6,122 +6,34 @@
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-
-def test_sea_query_exec():
-    """
-    Test executing a query using the SEA backend with result compression.
-
-    This function connects to a Databricks SQL endpoint using the SEA backend,
-    executes a simple query with result compression enabled and disabled,
-    and verifies that execution completes successfully.
-    """
-    server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
-    http_path = os.environ.get("DATABRICKS_HTTP_PATH")
-    access_token = os.environ.get("DATABRICKS_TOKEN")
-    catalog = os.environ.get("DATABRICKS_CATALOG")
-
-    if not all([server_hostname, http_path, access_token]):
-        logger.error("Missing required environment variables.")
-        logger.error(
-            "Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
-        )
-        sys.exit(1)
-
-    try:
-        # Test with compression enabled
-        logger.info("Creating connection with LZ4 compression enabled")
-        connection = Connection(
-            server_hostname=server_hostname,
-            http_path=http_path,
-            access_token=access_token,
-            catalog=catalog,
-            schema="default",
-            use_sea=True,
-            user_agent_entry="SEA-Test-Client",
-            use_cloud_fetch=True,  # Enable cloud fetch to use compression
-            enable_query_result_lz4_compression=True,  # Enable LZ4 compression
-        )
-
-        logger.info(
-            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
-        )
-        logger.info(f"backend type: {type(connection.session.backend)}")
-
-        # Execute a simple query with compression enabled
-        cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
-        logger.info("Executing query with LZ4 compression: SELECT 1 as test_value")
-        cursor.execute("SELECT 1 as test_value")
-        logger.info("Query with compression executed successfully")
-        cursor.close()
-        connection.close()
-        logger.info("Successfully closed SEA session with compression enabled")
-
-        # Test with compression disabled
-        logger.info("Creating connection with LZ4 compression disabled")
-        connection = Connection(
-            server_hostname=server_hostname,
-            http_path=http_path,
-            access_token=access_token,
-            catalog=catalog,
-            schema="default",
-            use_sea=True,
-            user_agent_entry="SEA-Test-Client",
-            use_cloud_fetch=False,  # Enable cloud fetch
-            enable_query_result_lz4_compression=False,  # Disable LZ4 compression
-        )
-
-        logger.info(
-            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
-        )
-
-        # Execute a simple query with compression disabled
-        cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
-        logger.info("Executing query without compression: SELECT 1 as test_value")
-        cursor.execute("SELECT 1 as test_value")
-        logger.info("Query without compression executed successfully")
-        cursor.close()
-        connection.close()
-        logger.info("Successfully closed SEA session with compression disabled")
-
-    except Exception as e:
-        logger.error(f"Error during SEA query execution test: {str(e)}")
-        import traceback
-
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-    logger.info("SEA query execution test with compression completed successfully")
-
-
 def test_sea_session():
     """
     Test opening and closing a SEA session using the connector.
-
+    
     This function connects to a Databricks SQL endpoint using the SEA backend,
     opens a session, and then closes it.
-
+    
     Required environment variables:
     - DATABRICKS_SERVER_HOSTNAME: Databricks server hostname
     - DATABRICKS_HTTP_PATH: HTTP path for the SQL endpoint
     - DATABRICKS_TOKEN: Personal access token for authentication
     """
+
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
     access_token = os.environ.get("DATABRICKS_TOKEN")
     catalog = os.environ.get("DATABRICKS_CATALOG")
-
+    
     if not all([server_hostname, http_path, access_token]):
         logger.error("Missing required environment variables.")
-        logger.error(
-            "Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
-        )
+        logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.")
         sys.exit(1)
-
+    
     logger.info(f"Connecting to {server_hostname}")
     logger.info(f"HTTP Path: {http_path}")
     if catalog:
         logger.info(f"Using catalog: {catalog}")
-
+    
     try:
         logger.info("Creating connection with SEA backend...")
         connection = Connection(
@@ -130,33 +42,25 @@ def test_sea_session():
             access_token=access_token,
             catalog=catalog,
             schema="default",
-            use_sea=True,
-            user_agent_entry="SEA-Test-Client",  # add custom user agent
-        )
-
-        logger.info(
-            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
+            use_sea=True, 
+            user_agent_entry="SEA-Test-Client" # add custom user agent
         )
+        
+        logger.info(f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}")
         logger.info(f"backend type: {type(connection.session.backend)}")
-
+        
         # Close the connection
         logger.info("Closing the SEA session...")
         connection.close()
         logger.info("Successfully closed SEA session")
-
+        
     except Exception as e:
         logger.error(f"Error testing SEA session: {str(e)}")
         import traceback
-
         logger.error(traceback.format_exc())
         sys.exit(1)
-
+    
     logger.info("SEA session test completed successfully")
 
-
 if __name__ == "__main__":
-    # Test session management
     test_sea_session()
-
-    # Test query execution with compression
-    test_sea_query_exec()

src/databricks/sql/backend/databricks_client.py

@@ -16,6 +16,8 @@
 
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.types import SessionId, CommandId, CommandState
+from databricks.sql.utils import ExecuteResponse
+from databricks.sql.types import SSLOptions
 
 # Forward reference for type hints
 from typing import TYPE_CHECKING
@@ -86,6 +88,34 @@ def execute_command(
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
     ) -> Union["ResultSet", None]:
+        """
+        Executes a SQL command or query within the specified session.
+
+        This method sends a SQL command to the server for execution and handles
+        the response. It can operate in both synchronous and asynchronous modes.
+
+        Args:
+            operation: The SQL command or query to execute
+            session_id: The session identifier in which to execute the command
+            max_rows: Maximum number of rows to fetch in a single fetch batch
+            max_bytes: Maximum number of bytes to fetch in a single fetch batch
+            lz4_compression: Whether to use LZ4 compression for result data
+            cursor: The cursor object that will handle the results
+            use_cloud_fetch: Whether to use cloud fetch for retrieving large result sets
+            parameters: List of parameters to bind to the query
+            async_op: Whether to execute the command asynchronously
+            enforce_embedded_schema_correctness: Whether to enforce schema correctness
+
+        Returns:
+            If async_op is False, returns a ResultSet object containing the
+            query results and metadata. If async_op is True, returns None and the
+            results must be fetched later using get_execution_result().
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error executing the command
+            ServerOperationError: If the server encounters an error during execution
+        """
         pass
 
     @abstractmethod