move sea stuff into sea/ dir

varun-edachali-dbx · varun-edachali-dbx · commit cdc7f4247b8d · 2025-06-09T09:30:18.000Z
Signed-off-by: varun-edachali-dbx &lt;varun.edachali@databricks.com&gt;
diff --git a/examples/experimental/sea_connector_test.py b/examples/experimental/sea_connector_test.py
@@ -450,7 +450,7 @@ def test_sea_result_set_with_multiple_chunks():
         cursor.execute(query)
         
         # Attempt to access the manifest to check for multiple chunks
-        from databricks.sql.backend.sea_backend import SeaDatabricksClient
+        from databricks.sql.backend.sea.backend import SeaDatabricksClient
         if isinstance(connection.session.backend, SeaDatabricksClient):
             # Get the statement ID from the cursor's active result set
             statement_id = cursor.active_result_set.statement_id
diff --git a/src/databricks/sql/backend/filters.py b/src/databricks/sql/backend/filters.py
@@ -19,7 +19,7 @@
 
 from databricks.sql.utils import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.types import ExecuteResponse, CommandId
-from databricks.sql.backend.models.base import ResultData
+from databricks.sql.backend.sea.models.base import ResultData
 
 if TYPE_CHECKING:
     from databricks.sql.result_set import ResultSet, SeaResultSet
diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py
@@ -6,7 +6,7 @@
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
     from databricks.sql.result_set import ResultSet
-    from databricks.sql.backend.models.responses import GetChunksResponse
+    from databricks.sql.backend.sea.models.responses import GetChunksResponse
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (
@@ -17,13 +17,17 @@
     ExecuteResponse,
 )
 from databricks.sql.exc import Error, NotSupportedError, ServerOperationError
-from databricks.sql.backend.utils.http_client import CustomHttpClient
+from databricks.sql.backend.sea.utils.http_client import CustomHttpClient
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 from databricks.sql.utils import SeaResultSetQueueFactory
-from databricks.sql.backend.models.base import ResultData, ExternalLink, ResultManifest
+from databricks.sql.backend.sea.models.base import (
+    ResultData,
+    ExternalLink,
+    ResultManifest,
+)
 
-from databricks.sql.backend.models import (
+from databricks.sql.backend.sea.models import (
     ExecuteStatementRequest,
     GetStatementRequest,
     CancelStatementRequest,
@@ -254,8 +258,8 @@ def fetch_chunk_links(
         Raises:
             Error: If there's an error fetching the chunk links
         """
-        from databricks.sql.backend.models.responses import GetChunksResponse
-        from databricks.sql.backend.models.base import ExternalLink
+        from databricks.sql.backend.sea.models.responses import GetChunksResponse
+        from databricks.sql.backend.sea.models.base import ExternalLink
 
         logger.info(f"Fetching chunk {chunk_index} links for statement {statement_id}")
 
@@ -304,7 +308,7 @@ def get_chunk_links(
         Returns:
             GetChunksResponse: Response containing external links
         """
-        from databricks.sql.backend.models.responses import GetChunksResponse
+        from databricks.sql.backend.sea.models.responses import GetChunksResponse
 
         params = {"chunk_index": chunk_index}
 
diff --git a/src/databricks/sql/backend/sea/models/__init__.py b/src/databricks/sql/backend/sea/models/__init__.py
@@ -4,7 +4,7 @@
 This package contains data models for SEA API requests and responses.
 """
 
-from databricks.sql.backend.models.base import (
+from databricks.sql.backend.sea.models.base import (
     ServiceError,
     StatementStatus,
     ExternalLink,
@@ -13,7 +13,7 @@
     ResultManifest,
 )
 
-from databricks.sql.backend.models.requests import (
+from databricks.sql.backend.sea.models.requests import (
     StatementParameter,
     ExecuteStatementRequest,
     GetStatementRequest,
@@ -23,7 +23,7 @@
     DeleteSessionRequest,
 )
 
-from databricks.sql.backend.models.responses import (
+from databricks.sql.backend.sea.models.responses import (
     ExecuteStatementResponse,
     GetStatementResponse,
     CreateSessionResponse,
diff --git a/src/databricks/sql/backend/sea/models/base.py b/src/databricks/sql/backend/sea/models/base.py
diff --git a/src/databricks/sql/backend/sea/models/requests.py b/src/databricks/sql/backend/sea/models/requests.py
diff --git a/src/databricks/sql/backend/sea/models/responses.py b/src/databricks/sql/backend/sea/models/responses.py
@@ -8,7 +8,7 @@
 from dataclasses import dataclass, field
 
 from databricks.sql.backend.types import CommandState
-from databricks.sql.backend.models.base import (
+from databricks.sql.backend.sea.models.base import (
     StatementStatus,
     ResultManifest,
     ResultData,
diff --git a/src/databricks/sql/backend/sea/utils/http_client.py b/src/databricks/sql/backend/sea/utils/http_client.py
diff --git a/src/databricks/sql/cloud_fetch_queue.py b/src/databricks/sql/cloud_fetch_queue.py
@@ -9,7 +9,7 @@
 from typing import Any, List, Optional, Union, TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from databricks.sql.backend.sea_backend import SeaDatabricksClient
+    from databricks.sql.backend.sea.backend import SeaDatabricksClient
     from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 
 from abc import ABC, abstractmethod
@@ -25,7 +25,7 @@
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.types import SSLOptions
-from databricks.sql.backend.models.base import ExternalLink
+from databricks.sql.backend.sea.models.base import ExternalLink
 from databricks.sql.utils import ResultSetQueue
 
 logger = logging.getLogger(__name__)
diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
@@ -8,7 +8,7 @@
 from databricks.sql import __version__
 from databricks.sql import USER_AGENT_NAME
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
-from databricks.sql.backend.sea_backend import SeaDatabricksClient
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import SessionId, BackendType
 
diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from databricks.sql.backend.sea_backend import SeaDatabricksClient
+    from databricks.sql.backend.sea.backend import SeaDatabricksClient
 
 from abc import ABC, abstractmethod
 from collections import OrderedDict, namedtuple
@@ -28,7 +28,11 @@
     TSparkRowSetType,
 )
 from databricks.sql.types import SSLOptions
-from databricks.sql.backend.models.base import ResultData, ExternalLink, ResultManifest
+from databricks.sql.backend.sea.models.base import (
+    ResultData,
+    ExternalLink,
+    ResultManifest,
+)
 from databricks.sql.parameters.native import ParameterStructure, TDbsqlParameter
 
 import logging
diff --git a/tests/unit/test_result_set_queue_factories.py b/tests/unit/test_result_set_queue_factories.py
@@ -9,7 +9,7 @@
     SeaResultSetQueueFactory,
     JsonQueue,
 )
-from databricks.sql.backend.models.base import ResultData, ResultManifest
+from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
 
 
 class TestResultSetQueueFactories(unittest.TestCase):
diff --git a/tests/unit/test_sea_backend.py b/tests/unit/test_sea_backend.py
@@ -9,7 +9,7 @@
 import pytest
 from unittest.mock import patch, MagicMock, Mock
 
-from databricks.sql.backend.sea_backend import SeaDatabricksClient
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.result_set import SeaResultSet
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
 from databricks.sql.types import SSLOptions
@@ -24,7 +24,7 @@ class TestSeaBackend:
     def mock_http_client(self):
         """Create a mock HTTP client."""
         with patch(
-            "databricks.sql.backend.sea_backend.CustomHttpClient"
+            "databricks.sql.backend.sea.backend.CustomHttpClient"
         ) as mock_client_class:
             mock_client = mock_client_class.return_value
             yield mock_client
