remove model redundancies

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -45,7 +45,7 @@ def test_sea_query_execution():
         logger.info(f"backend type: {type(connection.session.backend)}")
 
         # Create a cursor and execute a simple query
-        cursor = connection.cursor(buffer_size_bytes=0)
+        cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
 
         logger.info("Executing query: SELECT 1 as test_value")
         cursor.execute("SELECT 1 as test_value")

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

@@ -21,6 +21,7 @@ class StatementParameter:
 class ExecuteStatementRequest:
     """Request to execute a SQL statement."""
 
+    # TODO: result_compression key
     warehouse_id: str
     statement: str
     session_id: str
@@ -75,36 +76,33 @@ def to_dict(self) -> Dict[str, Any]:
 class GetStatementRequest:
     """Request to get information about a statement."""
 
-    warehouse_id: str
     statement_id: str
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
-        return {"warehouse_id": self.warehouse_id, "statement_id": self.statement_id}
+        return {"statement_id": self.statement_id}
 
 
 @dataclass
 class CancelStatementRequest:
     """Request to cancel a statement."""
 
-    warehouse_id: str
     statement_id: str
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
-        return {"warehouse_id": self.warehouse_id, "statement_id": self.statement_id}
+        return {"statement_id": self.statement_id}
 
 
 @dataclass
 class CloseStatementRequest:
     """Request to close a statement."""
 
-    warehouse_id: str
     statement_id: str
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
-        return {"warehouse_id": self.warehouse_id, "statement_id": self.statement_id}
+        return {"statement_id": self.statement_id}
 
 
 @dataclass
@@ -139,7 +137,7 @@ class DeleteSessionRequest:
     warehouse_id: str
     session_id: str
 
-    def to_query_params(self) -> Dict[str, str]:
+    def to_dict(self) -> Dict[str, str]:
         """
         Convert the request to query parameters.
 
@@ -149,4 +147,4 @@ def to_query_params(self) -> Dict[str, str]:
         Returns:
             A dictionary containing the warehouse_id as a query parameter
         """
-        return {"warehouse_id": self.warehouse_id}
+        return {"warehouse_id": self.warehouse_id, "session_id": self.session_id}

src/databricks/sql/backend/sea_backend.py

@@ -222,7 +222,7 @@ def close_session(self, session_id: SessionId) -> None:
         self.http_client._make_request(
             method="DELETE",
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
-            data=request.to_query_params(),  # Use to_query_params to get only the warehouse_id
+            data=request.to_dict(),
         )
 
     def execute_command(
@@ -284,7 +284,7 @@ def execute_command(
             statement=operation,
             disposition=disposition,
             format=format,
-            wait_timeout="0s" if async_op else "30s",
+            wait_timeout="0s" if async_op else "10s",
             on_wait_timeout="CONTINUE",
             row_limit=max_rows if max_rows > 0 else None,
             byte_limit=max_bytes if max_bytes > 0 else None,
@@ -324,9 +324,7 @@ def execute_command(
             time.sleep(0.5)
 
             # Create the request model
-            get_request = GetStatementRequest(
-                warehouse_id=self.warehouse_id, statement_id=statement_id
-            )
+            get_request = GetStatementRequest(statement_id=statement_id)
 
             # Get the statement status
             poll_response_data = self.http_client._make_request(
@@ -368,9 +366,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         sea_statement_id = command_id.to_sea_statement_id()
 
         # Create the request model
-        request = CancelStatementRequest(
-            warehouse_id=self.warehouse_id, statement_id=sea_statement_id
-        )
+        request = CancelStatementRequest(statement_id=sea_statement_id)
 
         # Send the cancel request
         self.http_client._make_request(
@@ -395,9 +391,7 @@ def close_command(self, command_id: CommandId) -> None:
         sea_statement_id = command_id.to_sea_statement_id()
 
         # Create the request model
-        request = CloseStatementRequest(
-            warehouse_id=self.warehouse_id, statement_id=sea_statement_id
-        )
+        request = CloseStatementRequest(statement_id=sea_statement_id)
 
         # Send the close request - SEA uses DELETE for closing statements
         self.http_client._make_request(
@@ -425,9 +419,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         sea_statement_id = command_id.to_sea_statement_id()
 
         # Create the request model
-        request = GetStatementRequest(
-            warehouse_id=self.warehouse_id, statement_id=sea_statement_id
-        )
+        request = GetStatementRequest(statement_id=sea_statement_id)
 
         # Get the statement status
         response_data = self.http_client._make_request(
@@ -466,9 +458,7 @@ def get_execution_result(
         sea_statement_id = command_id.to_sea_statement_id()
 
         # Create the request model
-        request = GetStatementRequest(
-            warehouse_id=self.warehouse_id, statement_id=sea_statement_id
-        )
+        request = GetStatementRequest(statement_id=sea_statement_id)
 
         # Get the statement result
         response_data = self.http_client._make_request(

src/databricks/sql/backend/sea_result_set.py

@@ -5,6 +5,7 @@
 which handles the result data returned by the SEA API.
 """
 
+import json
 import logging
 from typing import Optional, List, Any, Dict, Tuple
 
@@ -60,9 +61,8 @@ def __init__(
         # Parse the manifest
         manifest_data = sea_response.get("manifest")
         if manifest_data:
-            schema_data = manifest_data.get("schema", [])
             columns = []
-            for col_data in schema_data:
+            for col_data in manifest_data.get("schema", {}).get("columns", []):
                 columns.append(
                     ColumnInfo(
                         name=col_data.get("name", ""),

tests/unit/test_sea_backend.py

@@ -5,10 +5,12 @@
 the Databricks SQL connector's SEA backend functionality.
 """
 
+import json
 import pytest
 from unittest.mock import patch, MagicMock, Mock
 
 from databricks.sql.backend.sea_backend import SeaDatabricksClient
+from databricks.sql.backend.sea_result_set import SeaResultSet
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
 from databricks.sql.types import SSLOptions
 from databricks.sql.auth.authenticators import AuthProvider
@@ -178,7 +180,7 @@ def test_close_session_valid_id(self, sea_client, mock_http_client):
         mock_http_client._make_request.assert_called_once_with(
             method="DELETE",
             path=sea_client.SESSION_PATH_WITH_ID.format("test-session-789"),
-            data={"warehouse_id": "abc123"},
+            data={"session_id": "test-session-789", "warehouse_id": "abc123"},
         )
 
     def test_close_session_invalid_id_type(self, sea_client):
@@ -451,8 +453,6 @@ def test_cancel_command(self, sea_client, mock_http_client, sea_command_id):
         assert kwargs["path"] == sea_client.CANCEL_STATEMENT_PATH_WITH_ID.format(
             "test-statement-123"
         )
-        assert "warehouse_id" in kwargs["data"]
-        assert kwargs["data"]["warehouse_id"] == "abc123"
 
     def test_close_command(self, sea_client, mock_http_client, sea_command_id):
         """Test closing a command."""
@@ -469,8 +469,6 @@ def test_close_command(self, sea_client, mock_http_client, sea_command_id):
         assert kwargs["path"] == sea_client.STATEMENT_PATH_WITH_ID.format(
             "test-statement-123"
         )
-        assert "warehouse_id" in kwargs["data"]
-        assert kwargs["data"]["warehouse_id"] == "abc123"
 
     def test_get_query_state(self, sea_client, mock_http_client, sea_command_id):
         """Test getting the state of a query."""
@@ -493,40 +491,50 @@ def test_get_query_state(self, sea_client, mock_http_client, sea_command_id):
         assert kwargs["path"] == sea_client.STATEMENT_PATH_WITH_ID.format(
             "test-statement-123"
         )
-        assert "warehouse_id" in kwargs["data"]
-        assert kwargs["data"]["warehouse_id"] == "abc123"
 
     def test_get_execution_result(
         self, sea_client, mock_http_client, mock_cursor, sea_command_id
     ):
         """Test getting the result of a command execution."""
         # Set up mock response
-        mock_http_client._make_request.return_value = {
+        sea_response = {
             "statement_id": "test-statement-123",
             "status": {"state": "SUCCEEDED"},
             "manifest": {
-                "schema": [
-                    {
-                        "name": "col1",
-                        "type_name": "STRING",
-                        "type_text": "string",
-                        "nullable": True,
-                    }
-                ],
+                "format": "JSON_ARRAY",
+                "schema": {
+                    "column_count": 1,
+                    "columns": [
+                        {
+                            "name": "test_value",
+                            "type_text": "INT",
+                            "type_name": "INT",
+                            "position": 0,
+                        }
+                    ],
+                },
+                "total_chunk_count": 1,
+                "chunks": [{"chunk_index": 0, "row_offset": 0, "row_count": 1}],
                 "total_row_count": 1,
-                "total_byte_count": 100,
+                "truncated": False,
+            },
+            "result": {
+                "chunk_index": 0,
+                "row_offset": 0,
+                "row_count": 1,
+                "data_array": [["1"]],
             },
-            "result": {"data": [["value1"]]},
         }
+        mock_http_client._make_request.return_value = sea_response
 
         # Create a real result set to verify the implementation
         result = sea_client.get_execution_result(sea_command_id, mock_cursor)
+        print(result)
 
         # Verify basic properties of the result
         assert result.statement_id == "test-statement-123"
         assert result.status.state == CommandState.SUCCEEDED
-        assert len(result.description) == 1
-        assert result.description[0][0] == "col1"  # column name
+        assert result.manifest.schema[0].name == "test_value"
 
         # Verify the HTTP request
         mock_http_client._make_request.assert_called_once()
@@ -535,8 +543,6 @@ def test_get_execution_result(
         assert kwargs["path"] == sea_client.STATEMENT_PATH_WITH_ID.format(
             "test-statement-123"
         )
-        assert "warehouse_id" in kwargs["data"]
-        assert kwargs["data"]["warehouse_id"] == "abc123"
 
     # Tests for metadata operations