introduce req resp models, update example tester script

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -6,23 +6,26 @@
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
+
 def test_sea_query_execution():
     """
     Test executing a query using the SEA backend.
-    
+
     This function connects to a Databricks SQL endpoint using the SEA backend,
     executes a simple query, 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.")
+        logger.error(
+            "Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
+        )
         sys.exit(1)
-    
+
     try:
         # Create connection with SEA backend
         connection = Connection(
@@ -32,41 +35,46 @@ def test_sea_query_execution():
             catalog=catalog,
             schema="default",
             use_sea=True,
-            user_agent_entry="SEA-Test-Client"
+            user_agent_entry="SEA-Test-Client",
+            use_cloud_fetch=False,
+        )
+
+        logger.info(
+            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
-        
-        logger.info(f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}")
         logger.info(f"backend type: {type(connection.session.backend)}")
-        
+
         # Create a cursor and execute a simple query
-        cursor = connection.cursor()
-        
+        cursor = connection.cursor(buffer_size_bytes=0)
+
         logger.info("Executing query: SELECT 1 as test_value")
         cursor.execute("SELECT 1 as test_value")
-        
+
         # We don't fetch results yet since we haven't implemented the fetch functionality
         logger.info("Query executed successfully")
-        
+
         # Close cursor and connection
         cursor.close()
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
     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 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
@@ -76,17 +84,19 @@ def test_sea_session():
     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(
@@ -95,29 +105,33 @@ 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
+            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"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
-    test_sea_query_execution()
+    test_sea_query_execution()

src/databricks/sql/backend/models/__init__.py

@@ -10,7 +10,7 @@
     ExternalLink,
     ResultData,
     ColumnInfo,
-    ResultManifest
+    ResultManifest,
 )
 
 from databricks.sql.backend.models.requests import (
@@ -20,35 +20,33 @@
     CancelStatementRequest,
     CloseStatementRequest,
     CreateSessionRequest,
-    DeleteSessionRequest
+    DeleteSessionRequest,
 )
 
 from databricks.sql.backend.models.responses import (
     ExecuteStatementResponse,
     GetStatementResponse,
-    CreateSessionResponse
+    CreateSessionResponse,
 )
 
 __all__ = [
     # Base models
-    'ServiceError',
-    'StatementStatus',
-    'ExternalLink',
-    'ResultData',
-    'ColumnInfo',
-    'ResultManifest',
-    
+    "ServiceError",
+    "StatementStatus",
+    "ExternalLink",
+    "ResultData",
+    "ColumnInfo",
+    "ResultManifest",
     # Request models
-    'StatementParameter',
-    'ExecuteStatementRequest',
-    'GetStatementRequest',
-    'CancelStatementRequest',
-    'CloseStatementRequest',
-    'CreateSessionRequest',
-    'DeleteSessionRequest',
-    
+    "StatementParameter",
+    "ExecuteStatementRequest",
+    "GetStatementRequest",
+    "CancelStatementRequest",
+    "CloseStatementRequest",
+    "CreateSessionRequest",
+    "DeleteSessionRequest",
     # Response models
-    'ExecuteStatementResponse',
-    'GetStatementResponse',
-    'CreateSessionResponse'
-]
+    "ExecuteStatementResponse",
+    "GetStatementResponse",
+    "CreateSessionResponse",
+]

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

@@ -11,13 +11,15 @@
 @dataclass
 class ServiceError:
     """Error information returned by the SEA API."""
+
     message: str
     error_code: Optional[str] = None
 
 
 @dataclass
 class StatementStatus:
     """Status information for a statement execution."""
+
     state: str
     error: Optional[ServiceError] = None
     sql_state: Optional[str] = None
@@ -26,6 +28,7 @@ class StatementStatus:
 @dataclass
 class ExternalLink:
     """External link information for result data."""
+
     external_link: str
     expiration: str
     chunk_index: int
@@ -34,13 +37,15 @@ class ExternalLink:
 @dataclass
 class ResultData:
     """Result data from a statement execution."""
+
     data: Optional[List[List[Any]]] = None
     external_links: Optional[List[ExternalLink]] = None
 
 
 @dataclass
 class ColumnInfo:
     """Information about a column in the result set."""
+
     name: str
     type_name: str
     type_text: str
@@ -53,8 +58,9 @@ class ColumnInfo:
 @dataclass
 class ResultManifest:
     """Manifest information for a result set."""
+
     schema: List[ColumnInfo]
     total_row_count: int
     total_byte_count: int
     truncated: bool = False
-    chunk_count: Optional[int] = None
+    chunk_count: Optional[int] = None