remove un-necessary changes

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -1,115 +1,66 @@
-"""
-Main script to run all SEA connector tests.
-
-This script runs all the individual test modules and displays
-a summary of test results with visual indicators.
-"""
 import os
 import sys
 import logging
-import subprocess
-from typing import List, Tuple
+from databricks.sql.client import Connection
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-TEST_MODULES = [
-    "test_sea_session",
-    "test_sea_sync_query",
-    "test_sea_async_query",
-    "test_sea_metadata",
-]
-
-
-def run_test_module(module_name: str) -> bool:
-    """Run a test module and return success status."""
-    module_path = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)), "tests", f"{module_name}.py"
-    )
-
-    # Simply run the module as a script - each module handles its own test execution
-    result = subprocess.run(
-        [sys.executable, module_path], capture_output=True, text=True
-    )
-
-    # Log the output from the test module
-    if result.stdout:
-        for line in result.stdout.strip().split("\n"):
-            logger.info(line)
-
-    if result.stderr:
-        for line in result.stderr.strip().split("\n"):
-            logger.error(line)
-
-    return result.returncode == 0
-
-
-def run_tests() -> List[Tuple[str, bool]]:
-    """Run all tests and return results."""
-    results = []
-
-    for module_name in TEST_MODULES:
-        try:
-            logger.info(f"\n{'=' * 50}")
-            logger.info(f"Running test: {module_name}")
-            logger.info(f"{'-' * 50}")
-
-            success = run_test_module(module_name)
-            results.append((module_name, success))
-
-            status = "✅ PASSED" if success else "❌ FAILED"
-            logger.info(f"Test {module_name}: {status}")
-
-        except Exception as e:
-            logger.error(f"Error loading or running test {module_name}: {str(e)}")
-            import traceback
-
-            logger.error(traceback.format_exc())
-            results.append((module_name, False))
-
-    return results
-
-
-def print_summary(results: List[Tuple[str, bool]]) -> None:
-    """Print a summary of test results."""
-    logger.info(f"\n{'=' * 50}")
-    logger.info("TEST SUMMARY")
-    logger.info(f"{'-' * 50}")
-
-    passed = sum(1 for _, success in results if success)
-    total = len(results)
-
-    for module_name, success in results:
-        status = "✅ PASSED" if success else "❌ FAILED"
-        logger.info(f"{status} - {module_name}")
-
-    logger.info(f"{'-' * 50}")
-    logger.info(f"Total: {total} | Passed: {passed} | Failed: {total - passed}")
-    logger.info(f"{'=' * 50}")
-
-
-if __name__ == "__main__":
-    # Check if required environment variables are set
-    required_vars = [
-        "DATABRICKS_SERVER_HOSTNAME",
-        "DATABRICKS_HTTP_PATH",
-        "DATABRICKS_TOKEN",
-    ]
-    missing_vars = [var for var in required_vars if not os.environ.get(var)]
-
-    if missing_vars:
-        logger.error(
-            f"Missing required environment variables: {', '.join(missing_vars)}"
+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.")
+        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(
+            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" # add custom user agent
         )
-        logger.error("Please set these variables before running the tests.")
+        
+        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")
 
-    # Run all tests
-    results = run_tests()
-
-    # Print summary
-    print_summary(results)
-
-    # Exit with appropriate status code
-    all_passed = all(success for _, success in results)
-    sys.exit(0 if all_passed else 1)
+if __name__ == "__main__":
+    test_sea_session()

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

@@ -42,12 +42,29 @@ class ExternalLink:
     http_headers: Optional[Dict[str, str]] = None
 
 
+@dataclass
+class ChunkInfo:
+    """Information about a chunk in the result set."""
+
+    chunk_index: int
+    byte_count: int
+    row_offset: int
+    row_count: int
+
+
 @dataclass
 class ResultData:
     """Result data from a statement execution."""
 
     data: Optional[List[List[Any]]] = None
     external_links: Optional[List[ExternalLink]] = None
+    byte_count: Optional[int] = None
+    chunk_index: Optional[int] = None
+    next_chunk_index: Optional[int] = None
+    next_chunk_internal_link: Optional[str] = None
+    row_count: Optional[int] = None
+    row_offset: Optional[int] = None
+    attachment: Optional[bytes] = None
 
 
 @dataclass
@@ -73,5 +90,6 @@ class ResultManifest:
     total_byte_count: int
     total_chunk_count: int
     truncated: bool = False
-    chunks: Optional[List[Dict[str, Any]]] = None
+    chunks: Optional[List[ChunkInfo]] = None
     result_compression: Optional[str] = None
+    is_volume_operation: Optional[bool] = None

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

@@ -10,7 +10,7 @@
 
 @dataclass
 class StatementParameter:
-    """Parameter for a SQL statement."""
+    """Representation of a parameter for a SQL statement."""
 
     name: str
     value: Optional[str] = None
@@ -19,7 +19,7 @@ class StatementParameter:
 
 @dataclass
 class ExecuteStatementRequest:
-    """Request to execute a SQL statement."""
+    """Representation of a request to execute a SQL statement."""
 
     session_id: str
     statement: str
@@ -65,7 +65,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class GetStatementRequest:
-    """Request to get information about a statement."""
+    """Representation of a request to get information about a statement."""
 
     statement_id: str
 
@@ -76,7 +76,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CancelStatementRequest:
-    """Request to cancel a statement."""
+    """Representation of a request to cancel a statement."""
 
     statement_id: str
 
@@ -87,7 +87,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CloseStatementRequest:
-    """Request to close a statement."""
+    """Representation of a request to close a statement."""
 
     statement_id: str
 
@@ -98,7 +98,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CreateSessionRequest:
-    """Request to create a new session."""
+    """Representation of a request to create a new session."""
 
     warehouse_id: str
     session_confs: Optional[Dict[str, str]] = None
@@ -123,7 +123,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class DeleteSessionRequest:
-    """Request to delete a session."""
+    """Representation of a request to delete a session."""
 
     warehouse_id: str
     session_id: str

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

@@ -14,6 +14,7 @@
     ResultData,
     ServiceError,
     ExternalLink,
+    ChunkInfo,
 )
 
 
@@ -43,15 +44,28 @@ def _parse_manifest(data: Dict[str, Any]) -> ResultManifest:
     """Parse manifest from response data."""
 
     manifest_data = data.get("manifest", {})
+    chunks = None
+    if "chunks" in manifest_data:
+        chunks = [
+            ChunkInfo(
+                chunk_index=chunk.get("chunk_index", 0),
+                byte_count=chunk.get("byte_count", 0),
+                row_offset=chunk.get("row_offset", 0),
+                row_count=chunk.get("row_count", 0),
+            )
+            for chunk in manifest_data.get("chunks", [])
+        ]
+
     return 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"),
+        chunks=chunks,
         result_compression=manifest_data.get("result_compression"),
+        is_volume_operation=manifest_data.get("is_volume_operation"),
     )
 
 
@@ -80,12 +94,19 @@ def _parse_result(data: Dict[str, Any]) -> ResultData:
     return ResultData(
         data=result_data.get("data_array"),
         external_links=external_links,
+        byte_count=result_data.get("byte_count"),
+        chunk_index=result_data.get("chunk_index"),
+        next_chunk_index=result_data.get("next_chunk_index"),
+        next_chunk_internal_link=result_data.get("next_chunk_internal_link"),
+        row_count=result_data.get("row_count"),
+        row_offset=result_data.get("row_offset"),
+        attachment=result_data.get("attachment"),
     )
 
 
 @dataclass
 class ExecuteStatementResponse:
-    """Response from executing a SQL statement."""
+    """Representation of the response from executing a SQL statement."""
 
     statement_id: str
     status: StatementStatus
@@ -105,7 +126,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
 
 @dataclass
 class GetStatementResponse:
-    """Response from getting information about a statement."""
+    """Representation of the response from getting information about a statement."""
 
     statement_id: str
     status: StatementStatus
@@ -125,7 +146,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
 
 @dataclass
 class CreateSessionResponse:
-    """Response from creating a new session."""
+    """Representation of the response from creating a new session."""
 
     session_id: str
 

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

@@ -1,7 +1,7 @@
 import json
 import logging
 import requests
-from typing import Callable, Dict, Any, Optional, Union, List, Tuple
+from typing import Callable, Dict, Any, Optional, List, Tuple
 from urllib.parse import urljoin
 
 from databricks.sql.auth.authenticators import AuthProvider