diff --git a/.bot/prompts/engineer/system.md b/.bot/prompts/engineer/system.md index 2c10da0e7..27c37ccff 100644 --- a/.bot/prompts/engineer/system.md +++ b/.bot/prompts/engineer/system.md @@ -106,13 +106,11 @@ up whatever it needs. == HOW TO WORK (bug-fix flow) == 0. **Pick the BACKEND the bug is on — reproduce on that one.** The connector has - three backends and a bug on one won't reproduce on another. Choose from the + two backends and a bug on one won't reproduce on another. Choose from the issue: - **kernel** (`use_kernel=True`) — only if the issue is specifically about the Rust kernel / `use_kernel`. Repro in `tests/e2e/test_kernel_backend.py` and run it alone (it needs the real wheel; see the `-m "not realkernel"` note above). - - **SEA** (`use_sea=True`) — if the issue implicates SEA, or the area is - backend-parametrized (add the `{"use_sea": True}` param case). - **Thrift** (the default, no kwarg) — everything else; this is the common case. See CONTRIBUTING.md → "Backends and test tiers" for the full matrix (selection kwarg + where each backend's tests live). diff --git a/CHANGELOG.md b/CHANGELOG.md index b4ae65801..2e9719271 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,9 @@ # Release History # Unreleased +- **Breaking:** Remove the deprecated pure-Python SEA backend and its options; use Thrift (default) or `use_kernel=True`. +- **Breaking:** Kernel IDs now use `BackendType.KERNEL`, `SessionId.from_kernel_session_id`, and `CommandId.from_kernel_statement_id`. Remove `BackendType.SEA`, the `from_sea_*` / `to_sea_*` helpers, `CommandState.from_sea_state`, and `DatabricksClientType.SEA`. +- Warn when removed SEA options are passed. On kernel ≥ 1.1.0, use `_pool_maxsize` for the kernel HTTP pool; it does not configure the Thrift query transport. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 37736ed1a..4bdaea008 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -19,11 +19,6 @@ but ignored (or rejected) on the other is called out in the **Note** column. > does **not** splat `**kwargs`. Options implemented in the shared Python driver > layer can still apply to both backends; unforwarded backend-specific options > are ignored or rejected on the kernel path. -> -> A separate pure-Python **SEA** backend (`use_sea=True`) also exists but is -> being deprecated. Its selector and SEA-only hybrid-disposition option are -> listed for completeness, but this reference does not provide a third SEA -> support column. ## Legend @@ -60,10 +55,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | ------------------------ | ------ | :----: | :----: | ------------- | -------------------------------------------------------------------------------------------------- | -| `use_kernel` | `bool` | — | — | `False` | Routing selector: `False` selects Thrift; `True` selects Kernel. Mutually exclusive with `use_sea`. | -| `use_sea` | `bool` | — | — | `False` | Select the deprecated pure-Python SEA backend. Mutually exclusive with `use_kernel`. | -| `use_hybrid_disposition` | `bool` | ❌ | ❌ | `False` | SEA-only result-disposition option; ignored by Thrift and Kernel. | -| `max_connections` | `int` | ❌ | ❌ | `10` | Internal SEA-only HTTP-pool limit; ignored by Thrift and Kernel. | +| `use_kernel` | `bool` | — | — | `False` | Routing selector: `False` selects Thrift; `True` selects Kernel. | --- diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f744689b0..ca18f8d4c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -146,14 +146,13 @@ The suites marked `[not documented]` require additional configuration which will #### Backends and test tiers -The connector has **three execution backends**, selected per connection. When you +The connector has **two execution backends**, selected per connection. When you reproduce or fix a bug, use the backend the bug is actually on — a Thrift bug won't -reproduce on a SEA or kernel connection, and vice versa: +reproduce on a kernel connection, and vice versa: | Backend | Select via (connect kwarg / `extra_params`) | Where its tests live | | --- | --- | --- | -| **Thrift** (default) | *(nothing — the default path)* | the general `tests/e2e` suite (the `{}` parametrize case) and mocked `tests/unit` | -| **SEA** (Statement Execution API) *(deprecated — use Kernel for SEA-native connections)* | `use_sea=True` | the general `tests/e2e` suite (the `{"use_sea": True}` parametrize case, e.g. `tests/e2e/test_driver.py`) and mocked `tests/unit` | +| **Thrift** (default) | *(nothing — the default path)* | the general `tests/e2e` suite (the default connection) and mocked `tests/unit` | | **Kernel** (Rust, optional) | `use_kernel=True` | the dedicated `tests/e2e/test_kernel_backend.py` / `test_kernel_tls.py`, plus the offline routing test `tests/unit/test_session.py -m realkernel` | Notes that matter when running the suite: diff --git a/examples/experimental/sea_connector_test.py b/examples/experimental/sea_connector_test.py deleted file mode 100644 index d4fda6e13..000000000 --- a/examples/experimental/sea_connector_test.py +++ /dev/null @@ -1,126 +0,0 @@ -""" -Main script to run all SEA connector tests. - -DEPRECATED: the pure-Python SEA backend (``use_sea=True``) exercised by -these examples is incomplete (e.g. no positional ``?`` parameter binding) -and slated for removal. For a SEA-native connection use ``use_kernel=True`` -instead — install it with ``pip install 'databricks-sql-connector[kernel]'``. - -This script runs all the individual test modules and displays -a summary of test results with visual indicators. - -In order to run the script, the following environment variables need to be set: -- DATABRICKS_SERVER_HOSTNAME: The hostname of the Databricks server -- DATABRICKS_HTTP_PATH: The HTTP path of the Databricks server -- DATABRICKS_TOKEN: The token to use for authentication -""" - -import os -import sys -import logging -import subprocess -from typing import List, Tuple - -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)}" - ) - logger.error("Please set these variables before running the tests.") - sys.exit(1) - - # 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) diff --git a/examples/experimental/tests/__init__.py b/examples/experimental/tests/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/examples/experimental/tests/test_sea_async_query.py b/examples/experimental/tests/test_sea_async_query.py deleted file mode 100644 index 5bc6c6793..000000000 --- a/examples/experimental/tests/test_sea_async_query.py +++ /dev/null @@ -1,241 +0,0 @@ -""" -Test for SEA asynchronous query execution functionality. -""" -import os -import sys -import logging -import time -from databricks.sql.client import Connection -from databricks.sql.backend.types import CommandState - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - - -def test_sea_async_query_with_cloud_fetch(): - """ - Test executing a query asynchronously using the SEA backend with cloud fetch enabled. - - This function connects to a Databricks SQL endpoint using the SEA backend, - executes a simple query asynchronously with cloud fetch enabled, 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." - ) - return False - - try: - # Create connection with cloud fetch enabled - logger.info( - "Creating connection for asynchronous query execution with cloud fetch 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_query_result_lz4_compression=False, - ) - - logger.info( - f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}" - ) - - # Execute a query that generates large rows to force multiple chunks - requested_row_count = 5000 - cursor = connection.cursor() - query = f""" - SELECT - id, - concat('value_', repeat('a', 10000)) as test_value - FROM range(1, {requested_row_count} + 1) AS t(id) - """ - - logger.info( - f"Executing asynchronous query with cloud fetch to generate {requested_row_count} rows" - ) - cursor.execute_async(query) - logger.info( - "Asynchronous query submitted successfully with cloud fetch enabled" - ) - - # Check query state - logger.info("Checking query state...") - while cursor.is_query_pending(): - logger.info("Query is still pending, waiting...") - time.sleep(1) - - logger.info("Query is no longer pending, getting results...") - cursor.get_async_execution_result() - - results = [cursor.fetchone()] - results.extend(cursor.fetchmany(10)) - results.extend(cursor.fetchall()) - actual_row_count = len(results) - - logger.info( - f"Requested {requested_row_count} rows, received {actual_row_count} rows" - ) - - # Verify total row count - if actual_row_count != requested_row_count: - logger.error( - f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}" - ) - return False - - logger.info( - "PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly" - ) - - # Close resources - cursor.close() - connection.close() - logger.info("Successfully closed SEA session") - - return True - - except Exception as e: - logger.error( - f"Error during SEA asynchronous query execution test with cloud fetch: {str(e)}" - ) - import traceback - - logger.error(traceback.format_exc()) - return False - - -def test_sea_async_query_without_cloud_fetch(): - """ - Test executing a query asynchronously using the SEA backend with cloud fetch disabled. - - This function connects to a Databricks SQL endpoint using the SEA backend, - executes a simple query asynchronously with cloud fetch 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." - ) - return False - - try: - # Create connection with cloud fetch disabled - logger.info( - "Creating connection for asynchronous query execution with cloud fetch 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_query_result_lz4_compression=False, - ) - - logger.info( - f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}" - ) - - # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits - requested_row_count = 100 - cursor = connection.cursor() - query = f""" - SELECT - id, - concat('value_', repeat('a', 100)) as test_value - FROM range(1, {requested_row_count} + 1) AS t(id) - """ - - logger.info( - f"Executing asynchronous query without cloud fetch to generate {requested_row_count} rows" - ) - cursor.execute_async(query) - logger.info( - "Asynchronous query submitted successfully with cloud fetch disabled" - ) - - # Check query state - logger.info("Checking query state...") - while cursor.is_query_pending(): - logger.info("Query is still pending, waiting...") - time.sleep(1) - - logger.info("Query is no longer pending, getting results...") - cursor.get_async_execution_result() - results = [cursor.fetchone()] - results.extend(cursor.fetchmany(10)) - results.extend(cursor.fetchall()) - actual_row_count = len(results) - - logger.info( - f"Requested {requested_row_count} rows, received {actual_row_count} rows" - ) - - # Verify total row count - if actual_row_count != requested_row_count: - logger.error( - f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}" - ) - return False - - logger.info( - "PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly" - ) - - # Close resources - cursor.close() - connection.close() - logger.info("Successfully closed SEA session") - - return True - - except Exception as e: - logger.error( - f"Error during SEA asynchronous query execution test without cloud fetch: {str(e)}" - ) - import traceback - - logger.error(traceback.format_exc()) - return False - - -def test_sea_async_query_exec(): - """ - Run both asynchronous query tests and return overall success. - """ - with_cloud_fetch_success = test_sea_async_query_with_cloud_fetch() - logger.info( - f"Asynchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}" - ) - - without_cloud_fetch_success = test_sea_async_query_without_cloud_fetch() - logger.info( - f"Asynchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}" - ) - - return with_cloud_fetch_success and without_cloud_fetch_success - - -if __name__ == "__main__": - success = test_sea_async_query_exec() - sys.exit(0 if success else 1) diff --git a/examples/experimental/tests/test_sea_metadata.py b/examples/experimental/tests/test_sea_metadata.py deleted file mode 100644 index a200d97d3..000000000 --- a/examples/experimental/tests/test_sea_metadata.py +++ /dev/null @@ -1,98 +0,0 @@ -""" -Test for SEA metadata functionality. -""" -import os -import sys -import logging -from databricks.sql.client import Connection - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - - -def test_sea_metadata(): - """ - Test metadata operations using the SEA backend. - - This function connects to a Databricks SQL endpoint using the SEA backend, - and executes metadata operations like catalogs(), schemas(), tables(), and columns(). - """ - 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." - ) - return False - - if not catalog: - logger.error( - "DATABRICKS_CATALOG environment variable is required for metadata tests." - ) - return False - - try: - # Create connection - logger.info("Creating connection for metadata operations") - 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", - ) - - logger.info( - f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}" - ) - - # Test catalogs - cursor = connection.cursor() - logger.info("Fetching catalogs...") - cursor.catalogs() - logger.info("Successfully fetched catalogs") - - # Test schemas - logger.info(f"Fetching schemas for catalog '{catalog}'...") - cursor.schemas(catalog_name=catalog) - logger.info("Successfully fetched schemas") - - # Test tables - logger.info(f"Fetching tables for catalog '{catalog}', schema 'default'...") - cursor.tables(catalog_name=catalog, schema_name="default") - logger.info("Successfully fetched tables") - - # Test columns for a specific table - # Using a common table that should exist in most environments - logger.info( - f"Fetching columns for catalog '{catalog}', schema 'default', table 'customer'..." - ) - cursor.columns( - catalog_name=catalog, schema_name="default", table_name="customer" - ) - logger.info("Successfully fetched columns") - - # Close resources - cursor.close() - connection.close() - logger.info("Successfully closed SEA session") - - return True - - except Exception as e: - logger.error(f"Error during SEA metadata test: {str(e)}") - import traceback - - logger.error(traceback.format_exc()) - return False - - -if __name__ == "__main__": - success = test_sea_metadata() - sys.exit(0 if success else 1) diff --git a/examples/experimental/tests/test_sea_session.py b/examples/experimental/tests/test_sea_session.py deleted file mode 100644 index 516c1bbb8..000000000 --- a/examples/experimental/tests/test_sea_session.py +++ /dev/null @@ -1,71 +0,0 @@ -""" -Test for SEA session management functionality. -""" -import os -import sys -import logging -from databricks.sql.client import Connection - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - - -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." - ) - return False - - 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", - ) - - 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") - - return True - - except Exception as e: - logger.error(f"Error testing SEA session: {str(e)}") - import traceback - - logger.error(traceback.format_exc()) - return False - - -if __name__ == "__main__": - success = test_sea_session() - sys.exit(0 if success else 1) diff --git a/examples/experimental/tests/test_sea_sync_query.py b/examples/experimental/tests/test_sea_sync_query.py deleted file mode 100644 index 4e12d5aa4..000000000 --- a/examples/experimental/tests/test_sea_sync_query.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Test for SEA synchronous query execution functionality. -""" -import os -import sys -import logging -from databricks.sql.client import Connection - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - - -def test_sea_sync_query_with_cloud_fetch(): - """ - Test executing a query synchronously using the SEA backend with cloud fetch enabled. - - This function connects to a Databricks SQL endpoint using the SEA backend, - executes a simple query with cloud fetch enabled, 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." - ) - return False - - try: - # Create connection with cloud fetch enabled - logger.info( - "Creating connection for synchronous query execution with cloud fetch 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_query_result_lz4_compression=False, - ) - - logger.info( - f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}" - ) - - # Execute a query that generates large rows to force multiple chunks - requested_row_count = 10000 - cursor = connection.cursor() - query = f""" - SELECT - id, - concat('value_', repeat('a', 10000)) as test_value - FROM range(1, {requested_row_count} + 1) AS t(id) - """ - - logger.info( - f"Executing synchronous query with cloud fetch to generate {requested_row_count} rows" - ) - cursor.execute(query) - results = [cursor.fetchone()] - results.extend(cursor.fetchmany(10)) - results.extend(cursor.fetchall()) - actual_row_count = len(results) - logger.info( - f"{actual_row_count} rows retrieved against {requested_row_count} requested" - ) - - # Verify total row count - if actual_row_count != requested_row_count: - logger.error( - f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}" - ) - return False - - # Close resources - cursor.close() - connection.close() - logger.info("Successfully closed SEA session") - - return True - - except Exception as e: - logger.error( - f"Error during SEA synchronous query execution test with cloud fetch: {str(e)}" - ) - import traceback - - logger.error(traceback.format_exc()) - return False - - -def test_sea_sync_query_without_cloud_fetch(): - """ - Test executing a query synchronously using the SEA backend with cloud fetch disabled. - - This function connects to a Databricks SQL endpoint using the SEA backend, - executes a simple query with cloud fetch 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." - ) - return False - - try: - # Create connection with cloud fetch disabled - logger.info( - "Creating connection for synchronous query execution with cloud fetch 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_query_result_lz4_compression=False, - ) - - logger.info( - f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}" - ) - - # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits - requested_row_count = 100 - cursor = connection.cursor() - logger.info( - f"Executing synchronous query without cloud fetch: SELECT {requested_row_count} rows" - ) - cursor.execute( - "SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)" - ) - - results = [cursor.fetchone()] - results.extend(cursor.fetchmany(10)) - results.extend(cursor.fetchall()) - actual_row_count = len(results) - logger.info( - f"{actual_row_count} rows retrieved against {requested_row_count} requested" - ) - - # Verify total row count - if actual_row_count != requested_row_count: - logger.error( - f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}" - ) - return False - - # Close resources - cursor.close() - connection.close() - logger.info("Successfully closed SEA session") - - return True - - except Exception as e: - logger.error( - f"Error during SEA synchronous query execution test without cloud fetch: {str(e)}" - ) - import traceback - - logger.error(traceback.format_exc()) - return False - - -def test_sea_sync_query_exec(): - """ - Run both synchronous query tests and return overall success. - """ - with_cloud_fetch_success = test_sea_sync_query_with_cloud_fetch() - logger.info( - f"Synchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}" - ) - - without_cloud_fetch_success = test_sea_sync_query_without_cloud_fetch() - logger.info( - f"Synchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}" - ) - - return with_cloud_fetch_success and without_cloud_fetch_success - - -if __name__ == "__main__": - success = test_sea_sync_query_exec() - sys.exit(0 if success else 1) diff --git a/src/databricks/sql/backend/databricks_client.py b/src/databricks/sql/backend/databricks_client.py index 0c22b9e07..abd6c8244 100644 --- a/src/databricks/sql/backend/databricks_client.py +++ b/src/databricks/sql/backend/databricks_client.py @@ -9,8 +9,8 @@ # Type-annotation-only import (deferred by ``from __future__ import # annotations``). ``execute_command`` is typed with ``TSparkParameter`` for - # backwards compatibility, but this abstract base -- and the SEA/kernel - # implementations of it -- never import the Apache Thrift ``thrift`` package + # backwards compatibility, but this abstract base and the kernel + # implementation never import the Apache Thrift ``thrift`` package # at load time. See ``test_lazy_thrift_import``. from databricks.sql.thrift_api.TCLIService import ttypes diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 3477caf6b..26f4da8bd 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -474,7 +474,7 @@ def kernel_auth_kwargs( return kwargs # 5. Custom credentials_provider — the connector's primary M2M path - # on Thrift/SEA, but unusable on the kernel: it's an opaque token + # on Thrift, but unusable on the kernel: it's an opaque token # source with no extractable client_id/secret, so the kernel # can't own the token lifecycle. Point the caller at the raw # M2M kwargs instead. diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..4720b2a54 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -38,7 +38,6 @@ from databricks.sql.backend.kernel.result_set import KernelResultSet from databricks.sql.backend.kernel.type_mapping import bind_tspark_params from databricks.sql.backend.types import ( - BackendType, CommandId, CommandState, SessionId, @@ -491,13 +490,14 @@ def open_session( # the secret, so drop ours. self._auth_options.pop("oauth_client_secret", None) - # Use the kernel's real server-issued session id, not a - # synthetic UUID. Matches what the native SEA backend does. + # Use the kernel's real server-issued session id. # ``session_id`` is a PyO3 attribute access — also wrapped so # any conversion error surfaces as a mapped PEP 249 exception # instead of bubbling raw from the boundary. try: - session_id = SessionId.from_sea_session_id(self._kernel_session.session_id) + session_id = SessionId.from_kernel_session_id( + self._kernel_session.session_id + ) except Exception as exc: raise _wrap_kernel_exception("open_session", exc) from exc self._session_id = session_id @@ -609,7 +609,7 @@ def execute_command( bind_tspark_params(stmt, parameters) if async_op: async_exec = stmt.submit() - command_id = CommandId.from_sea_statement_id( + command_id = CommandId.from_kernel_statement_id( async_exec.statement_id ) cursor.active_command_id = command_id @@ -664,7 +664,7 @@ def execute_command( canceller.statement_id() if canceller is not None else None ) if stmt_id: - cursor.active_command_id = CommandId.from_sea_statement_id( + cursor.active_command_id = CommandId.from_kernel_statement_id( stmt_id ) except Exception: @@ -683,7 +683,7 @@ def execute_command( except Exception: pass - command_id = CommandId.from_sea_statement_id(executed.statement_id) + command_id = CommandId.from_kernel_statement_id(executed.statement_id) # Surface the affected-row count for DML (INSERT/UPDATE/DELETE/ # MERGE) as ``cursor.rowcount`` instead of the hardcoded ``-1``. # ``num_modified_rows`` is ``None`` for SELECT (and warehouses @@ -1009,7 +1009,7 @@ def _synthetic_command_id(self) -> CommandId: consumes ``cursor.query_id`` downstream (telemetry, log ingestion) sees a UUID-shaped string rather than a connector-internal magic prefix it cannot parse.""" - return CommandId.from_sea_statement_id(uuid.uuid4().hex) + return CommandId.from_kernel_statement_id(uuid.uuid4().hex) def get_catalogs( self, diff --git a/src/databricks/sql/backend/kernel/type_mapping.py b/src/databricks/sql/backend/kernel/type_mapping.py index 962e3d1a8..55d9d90f6 100644 --- a/src/databricks/sql/backend/kernel/type_mapping.py +++ b/src/databricks/sql/backend/kernel/type_mapping.py @@ -6,13 +6,6 @@ flattens the conversion so ``KernelResultSet`` and any future kernel-result wrapper share the same mapping. -The string constants come from ``SqlType`` in the SEA backend's -``conversion`` module — same single source of truth both backends -already use. The Arrow → ``SqlType`` lookup itself is kernel- -specific (SEA receives type-text from the server and normalises it; -the kernel receives Arrow schemas directly), so the mapping -function stays local but the names are shared. - Parameter binding (``TSparkParameter`` → kernel ``Statement.bind_param``) is handled by ``bind_tspark_params`` — forwards the connector's already-string-encoded form to the kernel @@ -25,7 +18,6 @@ import pyarrow -from databricks.sql.backend.sea.utils.conversion import SqlType from databricks.sql.exc import NotSupportedError if TYPE_CHECKING: @@ -44,41 +36,37 @@ def _arrow_type_to_dbapi_string(arrow_type: pyarrow.DataType) -> str: - """Map a pyarrow type to the Databricks SQL type name used in - PEP 249 ``description``. Names come from ``SqlType`` so the - kernel and SEA backends emit identical type-code strings; - consumers can branch on them identically. - """ + """Map an Arrow type to the lowercase type name used by Thrift descriptions.""" if pyarrow.types.is_boolean(arrow_type): - return SqlType.BOOLEAN + return "boolean" if pyarrow.types.is_int8(arrow_type): - return SqlType.TINYINT + return "tinyint" if pyarrow.types.is_int16(arrow_type): - return SqlType.SMALLINT + return "smallint" if pyarrow.types.is_int32(arrow_type): - return SqlType.INT + return "int" if pyarrow.types.is_int64(arrow_type): - return SqlType.BIGINT + return "bigint" if pyarrow.types.is_float32(arrow_type): - return SqlType.FLOAT + return "float" if pyarrow.types.is_float64(arrow_type): - return SqlType.DOUBLE + return "double" if pyarrow.types.is_decimal(arrow_type): - return SqlType.DECIMAL + return "decimal" if pyarrow.types.is_string(arrow_type) or pyarrow.types.is_large_string(arrow_type): - return SqlType.STRING + return "string" if pyarrow.types.is_binary(arrow_type) or pyarrow.types.is_large_binary(arrow_type): - return SqlType.BINARY + return "binary" if pyarrow.types.is_date(arrow_type): - return SqlType.DATE + return "date" if pyarrow.types.is_timestamp(arrow_type): - return SqlType.TIMESTAMP + return "timestamp" if pyarrow.types.is_list(arrow_type) or pyarrow.types.is_large_list(arrow_type): - return SqlType.ARRAY + return "array" if pyarrow.types.is_struct(arrow_type): - return SqlType.STRUCT + return "struct" if pyarrow.types.is_map(arrow_type): - return SqlType.MAP + return "map" # Fallback for types the kernel hasn't been observed to emit yet # (time32/time64, unsigned ints, dictionary, string_view, # binary_view, fixed_size_*). ``str(arrow_type)`` produces shapes @@ -175,7 +163,7 @@ def _databricks_type_for_field(field: pyarrow.Field) -> str: # bytes); compare against bytes to avoid one encode per field. type_name = md.get(b"databricks.type_name") if type_name is not None: - # Lowercase to match the canonical SqlType slugs the Thrift + # Lowercase to match the type names the Thrift # backend produces (``"array"`` / ``"map"`` / ``"struct"`` / # ``"variant"``). Other server-reported names (``"INT"`` etc.) # would also pass through this branch but we deliberately diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py deleted file mode 100644 index fff7daab4..000000000 --- a/src/databricks/sql/backend/sea/backend.py +++ /dev/null @@ -1,845 +0,0 @@ -from __future__ import annotations - -import logging -import time -import re -from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set - -from databricks.sql.backend.sea.models.base import ( - ExternalLink, - ResultManifest, - StatementStatus, -) -from databricks.sql.backend.sea.models.responses import GetChunksResponse -from databricks.sql.backend.sea.utils.constants import ( - ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP, - ResultFormat, - ResultDisposition, - ResultCompression, - WaitTimeout, - MetadataCommands, -) -from databricks.sql.backend.sea.utils.normalize import normalize_sea_type_to_thrift - -if TYPE_CHECKING: - from databricks.sql.client import Cursor - - # Type-annotation-only import (deferred by ``from __future__ import - # annotations``). ``execute_command`` accepts the Thrift-shaped - # ``TSparkParameter`` for interface compatibility, but only reads its - # attributes (duck-typed) at runtime, so the SEA backend never imports the - # Apache Thrift ``thrift`` package. See ``test_lazy_thrift_import``. - from databricks.sql.thrift_api.TCLIService import ttypes - -from databricks.sql.backend.sea.result_set import SeaResultSet - -from databricks.sql.backend.databricks_client import DatabricksClient -from databricks.sql.backend.types import ( - SessionId, - CommandId, - CommandState, - BackendType, - ExecuteResponse, -) -from databricks.sql.exc import DatabaseError, ServerOperationError -from databricks.sql.backend.sea.utils.http_client import SeaHttpClient -from databricks.sql.types import SSLOptions - -from databricks.sql.backend.sea.models import ( - ExecuteStatementRequest, - GetStatementRequest, - CancelStatementRequest, - CloseStatementRequest, - CreateSessionRequest, - DeleteSessionRequest, - StatementParameter, - ExecuteStatementResponse, - GetStatementResponse, - CreateSessionResponse, -) - -logger = logging.getLogger(__name__) - - -def _filter_session_configuration( - session_configuration: Optional[Dict[str, Any]], -) -> Dict[str, str]: - """ - Filter and normalise the provided session configuration parameters. - - The Statement Execution API supports only a subset of SQL session - configuration options. This helper validates the supplied - ``session_configuration`` dictionary against the allow-list defined in - ``ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP`` and returns a new - dictionary that contains **only** the supported parameters. - - Args: - session_configuration: Optional mapping of session configuration - names to their desired values. Key comparison is - case-insensitive. - - Returns: - Dict[str, str]: A dictionary containing only the supported - configuration parameters with lower-case keys and string values. If - *session_configuration* is ``None`` or empty, an empty dictionary is - returned. - """ - - if not session_configuration: - return {} - - filtered_session_configuration = {} - ignored_configs: Set[str] = set() - - for key, value in session_configuration.items(): - if key.upper() in ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP: - filtered_session_configuration[key.lower()] = str(value) - else: - ignored_configs.add(key) - - if ignored_configs: - logger.warning( - "Some session configurations were ignored because they are not supported: %s", - ignored_configs, - ) - logger.warning( - "Supported session configurations are: %s", - list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys()), - ) - - return filtered_session_configuration - - -class SeaDatabricksClient(DatabricksClient): - """ - Statement Execution API (SEA) implementation of the DatabricksClient interface. - """ - - # SEA API paths - BASE_PATH = "/api/2.0/sql/" - SESSION_PATH = BASE_PATH + "sessions" - SESSION_PATH_WITH_ID = SESSION_PATH + "/{}" - STATEMENT_PATH = BASE_PATH + "statements" - STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}" - CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel" - CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}" - - # SEA constants - POLL_INTERVAL_SECONDS = 0.2 - - def __init__( - self, - server_hostname: str, - port: int, - http_path: str, - http_headers: List[Tuple[str, str]], - auth_provider, - ssl_options: SSLOptions, - **kwargs, - ): - """ - Initialize the SEA backend client. - - Args: - server_hostname: Hostname of the Databricks server - port: Port number for the connection - http_path: HTTP path for the connection - http_headers: List of HTTP headers to include in requests - auth_provider: Authentication provider - ssl_options: SSL configuration options - **kwargs: Additional keyword arguments - """ - - logger.debug( - "SeaDatabricksClient.__init__(server_hostname=%s, port=%s, http_path=%s)", - server_hostname, - port, - http_path, - ) - - # The SEA backend is deprecated and incomplete (e.g. it does not - # support positional parameter binding) and is slated for removal. - # Steer users to the Rust kernel backend, which is the supported path. - logger.warning( - "The SEA backend (use_sea=True) is deprecated and incomplete and " - "should not be used in production; it is slated for removal. Use " - "the kernel backend instead by passing use_kernel=True and " - "installing the kernel extra: " - "pip install 'databricks-sql-connector[kernel]'." - ) - - self._max_download_threads = kwargs.get("max_download_threads", 10) - self._ssl_options = ssl_options - self._use_arrow_native_complex_types = kwargs.get( - "_use_arrow_native_complex_types", True - ) - - self.use_hybrid_disposition = kwargs.get("use_hybrid_disposition", False) - self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True) - - # Extract warehouse ID from http_path - self.warehouse_id = self._extract_warehouse_id(http_path) - - # Initialize HTTP client - self._http_client = SeaHttpClient( - server_hostname=server_hostname, - port=port, - http_path=http_path, - http_headers=http_headers, - auth_provider=auth_provider, - ssl_options=ssl_options, - **kwargs, - ) - - def _extract_warehouse_id(self, http_path: str) -> str: - """ - Extract the warehouse ID from the HTTP path. - - Args: - http_path: The HTTP path from which to extract the warehouse ID - - Returns: - The extracted warehouse ID - - Raises: - ValueError: If the warehouse ID cannot be extracted from the path - """ - - # [^?&]+ stops at query params (e.g. ?o= for SPOG routing) - warehouse_pattern = re.compile(r".*/warehouses/([^?&]+)") - endpoint_pattern = re.compile(r".*/endpoints/([^?&]+)") - - for pattern in [warehouse_pattern, endpoint_pattern]: - match = pattern.match(http_path) - if not match: - continue - warehouse_id = match.group(1) - logger.debug( - f"Extracted warehouse ID: {warehouse_id} from path: {http_path}" - ) - return warehouse_id - - # If no match found, raise error - error_message = ( - f"Could not extract warehouse ID from http_path: {http_path}. " - f"Expected format: /path/to/warehouses/{{warehouse_id}} or " - f"/path/to/endpoints/{{warehouse_id}}." - f"Note: SEA only works for warehouses." - ) - logger.error(error_message) - raise ValueError(error_message) - - @property - def max_download_threads(self) -> int: - """Get the maximum number of download threads for cloud fetch operations.""" - return self._max_download_threads - - def open_session( - self, - session_configuration: Optional[Dict[str, Any]], - catalog: Optional[str], - schema: Optional[str], - ) -> SessionId: - """ - Opens a new session with the Databricks SQL service using SEA. - - Args: - session_configuration: Optional dictionary of configuration parameters for the session. - Only specific parameters are supported as documented at: - https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-parameters - catalog: Optional catalog name to use as the initial catalog for the session - schema: Optional schema name to use as the initial schema for the session - - Returns: - SessionId: A session identifier object that can be used for subsequent operations - - Raises: - Error: If the session configuration is invalid - OperationalError: If there's an error establishing the session - """ - - logger.debug( - "SeaDatabricksClient.open_session(session_configuration=%s, catalog=%s, schema=%s)", - session_configuration, - catalog, - schema, - ) - - session_configuration = _filter_session_configuration(session_configuration) - - request_data = CreateSessionRequest( - warehouse_id=self.warehouse_id, - session_confs=session_configuration, - catalog=catalog, - schema=schema, - ) - - response = self._http_client._make_request( - method="POST", path=self.SESSION_PATH, data=request_data.to_dict() - ) - - session_response = CreateSessionResponse.from_dict(response) - session_id = session_response.session_id - if not session_id: - raise ServerOperationError( - "Failed to create session: No session ID returned", - { - "operation-id": None, - "diagnostic-info": None, - }, - ) - - return SessionId.from_sea_session_id(session_id) - - def close_session(self, session_id: SessionId) -> None: - """ - Closes an existing session with the Databricks SQL service. - - Args: - session_id: The session identifier returned by open_session() - - Raises: - ValueError: If the session ID is invalid - OperationalError: If there's an error closing the session - """ - - logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id) - - if session_id.backend_type != BackendType.SEA: - raise ValueError("Not a valid SEA session ID") - sea_session_id = session_id.to_sea_session_id() - - request_data = DeleteSessionRequest( - warehouse_id=self.warehouse_id, - session_id=sea_session_id, - ) - - self._http_client._make_request( - method="DELETE", - path=self.SESSION_PATH_WITH_ID.format(sea_session_id), - data=request_data.to_dict(), - ) - - def _extract_description_from_manifest( - self, manifest: ResultManifest - ) -> List[Tuple]: - """ - Extract column description from a manifest object, in the format defined by - the spec: https://peps.python.org/pep-0249/#description - - Args: - manifest: The ResultManifest object containing schema information - - Returns: - Optional[List]: A list of column tuples or None if no columns are found - """ - - schema_data = manifest.schema - columns_data = schema_data.get("columns", []) - - columns = [] - for col_data in columns_data: - # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok) - name = col_data.get("name", "") - type_name = col_data.get("type_name", "") - - # Normalize SEA type to Thrift conventions before any processing - type_name = normalize_sea_type_to_thrift(type_name, col_data) - - # Now strip _TYPE suffix and convert to lowercase - type_name = ( - type_name[:-5] if type_name.endswith("_TYPE") else type_name - ).lower() - precision = col_data.get("type_precision") - scale = col_data.get("type_scale") - - columns.append( - ( - name, # name - type_name, # type_code - None, # display_size (not provided by SEA) - None, # internal_size (not provided by SEA) - precision, # precision - scale, # scale - None, # null_ok - ) - ) - - return columns - - def _results_message_to_execute_response( - self, response: Union[ExecuteStatementResponse, GetStatementResponse] - ) -> ExecuteResponse: - """ - Convert a SEA response to an ExecuteResponse and extract result data. - - Args: - sea_response: The response from the SEA API - command_id: The command ID - - Returns: - ExecuteResponse: The normalized execute response - """ - - # Extract description from manifest schema - description = self._extract_description_from_manifest(response.manifest) - - # Check for compression - lz4_compressed = ( - response.manifest.result_compression == ResultCompression.LZ4_FRAME.value - ) - - execute_response = ExecuteResponse( - command_id=CommandId.from_sea_statement_id(response.statement_id), - status=response.status.state, - description=description, - has_been_closed_server_side=False, - lz4_compressed=lz4_compressed, - is_staging_operation=response.manifest.is_volume_operation, - arrow_schema_bytes=None, - result_format=response.manifest.format, - ) - - return execute_response - - def _response_to_result_set( - self, - response: Union[ExecuteStatementResponse, GetStatementResponse], - cursor: Cursor, - ) -> SeaResultSet: - """ - Convert a SEA response to a SeaResultSet. - """ - - execute_response = self._results_message_to_execute_response(response) - - return SeaResultSet( - connection=cursor.connection, - execute_response=execute_response, - sea_client=self, - result_data=response.result, - manifest=response.manifest, - buffer_size_bytes=cursor.buffer_size_bytes, - arraysize=cursor.arraysize, - ) - - def _check_command_not_in_failed_or_closed_state( - self, status: StatementStatus, command_id: CommandId - ) -> None: - state = status.state - if state == CommandState.CLOSED: - raise DatabaseError( - "Command {} unexpectedly closed server side".format(command_id), - { - "operation-id": command_id, - }, - ) - if state == CommandState.FAILED: - error = status.error - error_code = error.error_code if error else "UNKNOWN_ERROR_CODE" - error_message = error.message if error else "UNKNOWN_ERROR_MESSAGE" - raise ServerOperationError( - "Command failed: {} - {}".format(error_code, error_message), - { - "operation-id": command_id, - }, - ) - - def _wait_until_command_done( - self, response: ExecuteStatementResponse - ) -> Union[ExecuteStatementResponse, GetStatementResponse]: - """ - Wait until a command is done. - """ - - final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response - command_id = CommandId.from_sea_statement_id(final_response.statement_id) - - while final_response.status.state in [ - CommandState.PENDING, - CommandState.RUNNING, - ]: - time.sleep(self.POLL_INTERVAL_SECONDS) - final_response = self._poll_query(command_id) - - self._check_command_not_in_failed_or_closed_state( - final_response.status, command_id - ) - - return final_response - - def execute_command( - self, - operation: str, - session_id: SessionId, - max_rows: int, - max_bytes: int, - lz4_compression: bool, - cursor: Cursor, - use_cloud_fetch: bool, - parameters: List[ttypes.TSparkParameter], - async_op: bool, - enforce_embedded_schema_correctness: bool, - row_limit: Optional[int] = None, - query_tags: Optional[Dict[str, Optional[str]]] = None, - ) -> Union[SeaResultSet, None]: - """ - Execute a SQL command using the SEA backend. - - Args: - operation: SQL command to execute - session_id: Session identifier - max_rows: Maximum number of rows to fetch - max_bytes: Maximum number of bytes to fetch - lz4_compression: Whether to use LZ4 compression - cursor: Cursor executing the command - use_cloud_fetch: Whether to use cloud fetch - parameters: SQL parameters - async_op: Whether to execute asynchronously - enforce_embedded_schema_correctness: Whether to enforce schema correctness - - Returns: - ResultSet: A SeaResultSet instance for the executed command - """ - - if session_id.backend_type != BackendType.SEA: - raise ValueError("Not a valid SEA session ID") - - sea_session_id = session_id.to_sea_session_id() - - # Convert parameters to StatementParameter objects - sea_parameters = [] - if parameters: - for param in parameters: - sea_parameters.append( - StatementParameter( - name=param.name, - value=( - param.value.stringValue if param.value is not None else None - ), - type=param.type, - ) - ) - - format = ( - ResultFormat.ARROW_STREAM if use_cloud_fetch else ResultFormat.JSON_ARRAY - ).value - disposition = ( - ( - ResultDisposition.HYBRID - if self.use_hybrid_disposition - else ResultDisposition.EXTERNAL_LINKS - ) - if use_cloud_fetch - else ResultDisposition.INLINE - ).value - result_compression = ( - ResultCompression.LZ4_FRAME if lz4_compression else ResultCompression.NONE - ).value - - request = ExecuteStatementRequest( - warehouse_id=self.warehouse_id, - session_id=sea_session_id, - statement=operation, - disposition=disposition, - format=format, - wait_timeout=(WaitTimeout.ASYNC if async_op else WaitTimeout.SYNC).value, - on_wait_timeout="CONTINUE", - row_limit=row_limit, - parameters=sea_parameters if sea_parameters else None, - result_compression=result_compression, - query_tags=query_tags, - ) - - response_data = self._http_client._make_request( - method="POST", path=self.STATEMENT_PATH, data=request.to_dict() - ) - response = ExecuteStatementResponse.from_dict(response_data) - statement_id = response.statement_id - - command_id = CommandId.from_sea_statement_id(statement_id) - - # Store the command ID in the cursor - cursor.active_command_id = command_id - - # If async operation, return and let the client poll for results - if async_op: - return None - - final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response - if response.status.state != CommandState.SUCCEEDED: - final_response = self._wait_until_command_done(response) - - return self._response_to_result_set(final_response, cursor) - - def cancel_command(self, command_id: CommandId) -> None: - """ - Cancel a running command. - - Args: - command_id: Command identifier to cancel - - Raises: - ValueError: If the command ID is invalid - """ - - if command_id.backend_type != BackendType.SEA: - raise ValueError("Not a valid SEA command ID") - - sea_statement_id = command_id.to_sea_statement_id() - - request = CancelStatementRequest(statement_id=sea_statement_id) - self._http_client._make_request( - method="POST", - path=self.CANCEL_STATEMENT_PATH_WITH_ID.format(sea_statement_id), - data=request.to_dict(), - ) - - def close_command(self, command_id: CommandId) -> None: - """ - Close a command and release resources. - - Args: - command_id: Command identifier to close - - Raises: - ValueError: If the command ID is invalid - """ - - if command_id.backend_type != BackendType.SEA: - raise ValueError("Not a valid SEA command ID") - - sea_statement_id = command_id.to_sea_statement_id() - - request = CloseStatementRequest(statement_id=sea_statement_id) - self._http_client._make_request( - method="DELETE", - path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id), - data=request.to_dict(), - ) - - def _poll_query(self, command_id: CommandId) -> GetStatementResponse: - """ - Poll for the current command info. - """ - - if command_id.backend_type != BackendType.SEA: - raise ValueError("Not a valid SEA command ID") - - sea_statement_id = command_id.to_sea_statement_id() - - request = GetStatementRequest(statement_id=sea_statement_id) - response_data = self._http_client._make_request( - method="GET", - path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id), - data=request.to_dict(), - ) - response = GetStatementResponse.from_dict(response_data) - - return response - - def get_query_state(self, command_id: CommandId) -> CommandState: - """ - Get the state of a running query. - - Args: - command_id: Command identifier - - Returns: - CommandState: The current state of the command - - Raises: - ValueError: If the command ID is invalid - """ - - response = self._poll_query(command_id) - return response.status.state - - def get_execution_result( - self, - command_id: CommandId, - cursor: Cursor, - ) -> SeaResultSet: - """ - Get the result of a command execution. - - Args: - command_id: Command identifier - cursor: Cursor executing the command - - Returns: - SeaResultSet: A SeaResultSet instance with the execution results - - Raises: - ValueError: If the command ID is invalid - """ - - response = self._poll_query(command_id) - return self._response_to_result_set(response, cursor) - - def get_chunk_links( - self, statement_id: str, chunk_index: int - ) -> List[ExternalLink]: - """ - Get links for chunks starting from the specified index. - Args: - statement_id: The statement ID - chunk_index: The starting chunk index - Returns: - ExternalLink: External link for the chunk - """ - - response_data = self._http_client._make_request( - method="GET", - path=self.CHUNK_PATH_WITH_ID_AND_INDEX.format(statement_id, chunk_index), - ) - response = GetChunksResponse.from_dict(response_data) - - links = response.external_links or [] - return links - - # == Metadata Operations == - - def get_catalogs( - self, - session_id: SessionId, - max_rows: int, - max_bytes: int, - cursor: Cursor, - ) -> SeaResultSet: - """Get available catalogs by executing 'SHOW CATALOGS'.""" - result = self.execute_command( - operation=MetadataCommands.SHOW_CATALOGS.value, - session_id=session_id, - max_rows=max_rows, - max_bytes=max_bytes, - lz4_compression=False, - cursor=cursor, - use_cloud_fetch=self.use_cloud_fetch, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result is not None, "execute_command returned None in synchronous mode" - return result - - def get_schemas( - self, - session_id: SessionId, - max_rows: int, - max_bytes: int, - cursor: Cursor, - catalog_name: Optional[str] = None, - schema_name: Optional[str] = None, - ) -> SeaResultSet: - """Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'.""" - if not catalog_name: - raise DatabaseError("Catalog name is required for get_schemas") - - operation = MetadataCommands.SHOW_SCHEMAS.value.format(catalog_name) - - if schema_name: - operation += MetadataCommands.LIKE_PATTERN.value.format(schema_name) - - result = self.execute_command( - operation=operation, - session_id=session_id, - max_rows=max_rows, - max_bytes=max_bytes, - lz4_compression=False, - cursor=cursor, - use_cloud_fetch=self.use_cloud_fetch, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result is not None, "execute_command returned None in synchronous mode" - return result - - def get_tables( - self, - session_id: SessionId, - max_rows: int, - max_bytes: int, - cursor: Cursor, - catalog_name: Optional[str] = None, - schema_name: Optional[str] = None, - table_name: Optional[str] = None, - table_types: Optional[List[str]] = None, - ) -> SeaResultSet: - """Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'.""" - operation = ( - MetadataCommands.SHOW_TABLES_ALL_CATALOGS.value - if catalog_name in [None, "*", "%"] - else MetadataCommands.SHOW_TABLES.value.format( - MetadataCommands.CATALOG_SPECIFIC.value.format(catalog_name) - ) - ) - - if schema_name: - operation += MetadataCommands.SCHEMA_LIKE_PATTERN.value.format(schema_name) - - if table_name: - operation += MetadataCommands.LIKE_PATTERN.value.format(table_name) - - result = self.execute_command( - operation=operation, - session_id=session_id, - max_rows=max_rows, - max_bytes=max_bytes, - lz4_compression=False, - cursor=cursor, - use_cloud_fetch=self.use_cloud_fetch, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result is not None, "execute_command returned None in synchronous mode" - - # Apply client-side filtering by table_types - from databricks.sql.backend.sea.utils.filters import ResultSetFilter - - result = ResultSetFilter.filter_tables_by_type(result, table_types) - - return result - - def get_columns( - self, - session_id: SessionId, - max_rows: int, - max_bytes: int, - cursor: Cursor, - catalog_name: Optional[str] = None, - schema_name: Optional[str] = None, - table_name: Optional[str] = None, - column_name: Optional[str] = None, - ) -> SeaResultSet: - """Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'.""" - if not catalog_name: - raise DatabaseError("Catalog name is required for get_columns") - - operation = MetadataCommands.SHOW_COLUMNS.value.format(catalog_name) - - if schema_name: - operation += MetadataCommands.SCHEMA_LIKE_PATTERN.value.format(schema_name) - - if table_name: - operation += MetadataCommands.TABLE_LIKE_PATTERN.value.format(table_name) - - if column_name: - operation += MetadataCommands.LIKE_PATTERN.value.format(column_name) - - result = self.execute_command( - operation=operation, - session_id=session_id, - max_rows=max_rows, - max_bytes=max_bytes, - lz4_compression=False, - cursor=cursor, - use_cloud_fetch=self.use_cloud_fetch, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result is not None, "execute_command returned None in synchronous mode" - return result diff --git a/src/databricks/sql/backend/sea/models/__init__.py b/src/databricks/sql/backend/sea/models/__init__.py deleted file mode 100644 index 8450ec85d..000000000 --- a/src/databricks/sql/backend/sea/models/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -""" -Models for the SEA (Statement Execution API) backend. - -This package contains data models for SEA API requests and responses. -""" - -from databricks.sql.backend.sea.models.base import ( - ServiceError, - StatementStatus, - ExternalLink, - ResultData, - ResultManifest, -) - -from databricks.sql.backend.sea.models.requests import ( - StatementParameter, - ExecuteStatementRequest, - GetStatementRequest, - CancelStatementRequest, - CloseStatementRequest, - CreateSessionRequest, - DeleteSessionRequest, -) - -from databricks.sql.backend.sea.models.responses import ( - ExecuteStatementResponse, - GetStatementResponse, - CreateSessionResponse, - GetChunksResponse, -) - -__all__ = [ - # Base models - "ServiceError", - "StatementStatus", - "ExternalLink", - "ResultData", - "ResultManifest", - # Request models - "StatementParameter", - "ExecuteStatementRequest", - "GetStatementRequest", - "CancelStatementRequest", - "CloseStatementRequest", - "CreateSessionRequest", - "DeleteSessionRequest", - # Response models - "ExecuteStatementResponse", - "GetStatementResponse", - "CreateSessionResponse", - "GetChunksResponse", -] diff --git a/src/databricks/sql/backend/sea/models/base.py b/src/databricks/sql/backend/sea/models/base.py deleted file mode 100644 index 3eacc8887..000000000 --- a/src/databricks/sql/backend/sea/models/base.py +++ /dev/null @@ -1,82 +0,0 @@ -""" -Base models for the SEA (Statement Execution API) backend. - -These models define the common structures used in SEA API requests and responses. -""" - -from typing import Dict, List, Any, Optional, Union -from dataclasses import dataclass, field - -from databricks.sql.backend.types import CommandState - - -@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: CommandState - error: Optional[ServiceError] = None - sql_state: Optional[str] = None - - -@dataclass -class ExternalLink: - """External link information for result data.""" - - external_link: str - expiration: str - chunk_index: int - byte_count: int = 0 - row_count: int = 0 - row_offset: int = 0 - next_chunk_index: Optional[int] = None - next_chunk_internal_link: Optional[str] = None - 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 -class ResultManifest: - """Manifest information for a result set.""" - - format: str - schema: Dict[str, Any] - total_row_count: int - total_byte_count: int - total_chunk_count: int - truncated: bool = False - chunks: Optional[List[ChunkInfo]] = None - result_compression: Optional[str] = None - is_volume_operation: bool = False diff --git a/src/databricks/sql/backend/sea/models/requests.py b/src/databricks/sql/backend/sea/models/requests.py deleted file mode 100644 index eb156fb1a..000000000 --- a/src/databricks/sql/backend/sea/models/requests.py +++ /dev/null @@ -1,141 +0,0 @@ -""" -Request models for the SEA (Statement Execution API) backend. - -These models define the structures used in SEA API requests. -""" - -from typing import Dict, List, Any, Optional, Union -from dataclasses import dataclass, field - - -@dataclass -class StatementParameter: - """Representation of a parameter for a SQL statement.""" - - name: str - value: Optional[str] = None - type: Optional[str] = None - - -@dataclass -class ExecuteStatementRequest: - """Representation of a request to execute a SQL statement.""" - - session_id: str - statement: str - warehouse_id: str - disposition: str = "EXTERNAL_LINKS" - format: str = "JSON_ARRAY" - result_compression: Optional[str] = None - parameters: Optional[List[StatementParameter]] = None - wait_timeout: str = "10s" - on_wait_timeout: str = "CONTINUE" - row_limit: Optional[int] = None - query_tags: Optional[Dict[str, Optional[str]]] = None - - def to_dict(self) -> Dict[str, Any]: - """Convert the request to a dictionary for JSON serialization.""" - result: Dict[str, Any] = { - "warehouse_id": self.warehouse_id, - "session_id": self.session_id, - "statement": self.statement, - "disposition": self.disposition, - "format": self.format, - "wait_timeout": self.wait_timeout, - "on_wait_timeout": self.on_wait_timeout, - } - - if self.row_limit is not None and self.row_limit > 0: - result["row_limit"] = self.row_limit - - if self.result_compression: - result["result_compression"] = self.result_compression - - if self.parameters: - result["parameters"] = [ - { - "name": param.name, - "value": param.value, - "type": param.type, - } - for param in self.parameters - ] - - # SEA API expects query_tags as an array of {key, value} objects. - # None/empty values are left to the server to handle as key-only tags. - if self.query_tags: - result["query_tags"] = [ - {"key": k, "value": v} for k, v in self.query_tags.items() - ] - - return result - - -@dataclass -class GetStatementRequest: - """Representation of a request to get information about a statement.""" - - statement_id: str - - def to_dict(self) -> Dict[str, Any]: - """Convert the request to a dictionary for JSON serialization.""" - return {"statement_id": self.statement_id} - - -@dataclass -class CancelStatementRequest: - """Representation of a request to cancel a statement.""" - - statement_id: str - - def to_dict(self) -> Dict[str, Any]: - """Convert the request to a dictionary for JSON serialization.""" - return {"statement_id": self.statement_id} - - -@dataclass -class CloseStatementRequest: - """Representation of a request to close a statement.""" - - statement_id: str - - def to_dict(self) -> Dict[str, Any]: - """Convert the request to a dictionary for JSON serialization.""" - return {"statement_id": self.statement_id} - - -@dataclass -class CreateSessionRequest: - """Representation of a request to create a new session.""" - - warehouse_id: str - session_confs: Optional[Dict[str, str]] = None - catalog: Optional[str] = None - schema: Optional[str] = None - - def to_dict(self) -> Dict[str, Any]: - """Convert the request to a dictionary for JSON serialization.""" - result: Dict[str, Any] = {"warehouse_id": self.warehouse_id} - - if self.session_confs: - result["session_confs"] = self.session_confs - - if self.catalog: - result["catalog"] = self.catalog - - if self.schema: - result["schema"] = self.schema - - return result - - -@dataclass -class DeleteSessionRequest: - """Representation of a request to delete a session.""" - - warehouse_id: str - session_id: str - - def to_dict(self) -> Dict[str, str]: - """Convert the request to a dictionary for JSON serialization.""" - return {"warehouse_id": self.warehouse_id, "session_id": self.session_id} diff --git a/src/databricks/sql/backend/sea/models/responses.py b/src/databricks/sql/backend/sea/models/responses.py deleted file mode 100644 index 5a5580481..000000000 --- a/src/databricks/sql/backend/sea/models/responses.py +++ /dev/null @@ -1,196 +0,0 @@ -""" -Response models for the SEA (Statement Execution API) backend. - -These models define the structures used in SEA API responses. -""" - -import base64 -from typing import Dict, Any, List, Optional -from dataclasses import dataclass - -from databricks.sql.backend.types import CommandState -from databricks.sql.backend.sea.models.base import ( - StatementStatus, - ResultManifest, - ResultData, - ServiceError, - ExternalLink, - ChunkInfo, -) - - -def _parse_status(data: Dict[str, Any]) -> StatementStatus: - """Parse status from response data.""" - status_data = data.get("status", {}) - error = None - if "error" in status_data: - error_data = status_data["error"] - error = ServiceError( - message=error_data.get("message", ""), - error_code=error_data.get("error_code"), - ) - - state = CommandState.from_sea_state(status_data.get("state", "")) - if state is None: - raise ValueError(f"Invalid state: {status_data.get('state', '')}") - - return StatementStatus( - state=state, - error=error, - sql_state=status_data.get("sql_state"), - ) - - -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=chunks, - result_compression=manifest_data.get("result_compression"), - is_volume_operation=manifest_data.get("is_volume_operation", False), - ) - - -def _parse_result(data: Dict[str, Any]) -> ResultData: - """Parse result data from response data.""" - result_data = data.get("result", {}) - external_links = None - - if "external_links" in result_data: - external_links = [] - for link_data in result_data["external_links"]: - external_links.append( - ExternalLink( - external_link=link_data.get("external_link", ""), - expiration=link_data.get("expiration", ""), - chunk_index=link_data.get("chunk_index", 0), - byte_count=link_data.get("byte_count", 0), - row_count=link_data.get("row_count", 0), - row_offset=link_data.get("row_offset", 0), - next_chunk_index=link_data.get("next_chunk_index"), - next_chunk_internal_link=link_data.get("next_chunk_internal_link"), - http_headers=link_data.get("http_headers"), - ) - ) - - # Handle attachment field - decode from base64 if present - attachment = result_data.get("attachment") - if attachment is not None: - attachment = base64.b64decode(attachment) - - 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=attachment, - ) - - -@dataclass -class ExecuteStatementResponse: - """Representation of the response from executing a SQL statement.""" - - statement_id: str - status: StatementStatus - manifest: ResultManifest - result: ResultData - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse": - """Create an ExecuteStatementResponse from a dictionary.""" - return cls( - statement_id=data.get("statement_id", ""), - status=_parse_status(data), - manifest=_parse_manifest(data), - result=_parse_result(data), - ) - - -@dataclass -class GetStatementResponse: - """Representation of the response from getting information about a statement.""" - - statement_id: str - status: StatementStatus - manifest: ResultManifest - result: ResultData - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse": - """Create a GetStatementResponse from a dictionary.""" - return cls( - statement_id=data.get("statement_id", ""), - status=_parse_status(data), - manifest=_parse_manifest(data), - result=_parse_result(data), - ) - - -@dataclass -class CreateSessionResponse: - """Representation of the response from creating a new session.""" - - session_id: str - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> "CreateSessionResponse": - """Create a CreateSessionResponse from a dictionary.""" - return cls(session_id=data.get("session_id", "")) - - -@dataclass -class GetChunksResponse: - """ - Response from getting chunks for a statement. - - The response model can be found in the docs, here: - https://docs.databricks.com/api/workspace/statementexecution/getstatementresultchunkn - """ - - 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 - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> "GetChunksResponse": - """Create a GetChunksResponse from a dictionary.""" - result = _parse_result({"result": data}) - return cls( - data=result.data, - external_links=result.external_links, - byte_count=result.byte_count, - chunk_index=result.chunk_index, - next_chunk_index=result.next_chunk_index, - next_chunk_internal_link=result.next_chunk_internal_link, - row_count=result.row_count, - row_offset=result.row_offset, - ) diff --git a/src/databricks/sql/backend/sea/queue.py b/src/databricks/sql/backend/sea/queue.py deleted file mode 100644 index 5ca070bc3..000000000 --- a/src/databricks/sql/backend/sea/queue.py +++ /dev/null @@ -1,402 +0,0 @@ -from __future__ import annotations - -from abc import ABC -import threading -from typing import Dict, List, Optional, Tuple, Union, TYPE_CHECKING - -from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager -from databricks.sql.telemetry.models.enums import StatementType - -from databricks.sql.cloudfetch.downloader import ResultSetDownloadHandler - -try: - import pyarrow -except ImportError: - pyarrow = None - -import dateutil - -if TYPE_CHECKING: - from databricks.sql.backend.sea.backend import SeaDatabricksClient - from databricks.sql.backend.sea.models.base import ( - ExternalLink, - ResultData, - ResultManifest, - ) - - # Type-annotation-only import (deferred by ``from __future__ import - # annotations``). The SEA backend reuses the Thrift ``TSparkArrowResultLink`` - # only as the payload the shared cloud-fetch download manager expects; it is - # constructed via a function-local import in ``_convert_to_thrift_link`` so - # importing the SEA backend never imports the Apache Thrift ``thrift`` - # package. See ``test_lazy_thrift_import``. - from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink -from databricks.sql.backend.sea.utils.constants import ResultFormat -from databricks.sql.exc import ProgrammingError, ServerOperationError -from databricks.sql.types import SSLOptions -from databricks.sql.utils import ( - ArrowQueue, - CloudFetchQueue, - ResultSetQueue, - create_arrow_table_from_arrow_file, -) - -import logging - -logger = logging.getLogger(__name__) - - -class SeaResultSetQueueFactory(ABC): - @staticmethod - def build_queue( - result_data: ResultData, - manifest: ResultManifest, - statement_id: str, - ssl_options: SSLOptions, - description: List[Tuple], - max_download_threads: int, - sea_client: SeaDatabricksClient, - lz4_compressed: bool, - http_client, - ) -> ResultSetQueue: - """ - Factory method to build a result set queue for SEA backend. - - Args: - result_data (ResultData): Result data from SEA response - manifest (ResultManifest): Manifest from SEA response - statement_id (str): Statement ID for the query - description (List[List[Any]]): Column descriptions - max_download_threads (int): Maximum number of download threads - sea_client (SeaDatabricksClient): SEA client for fetching additional links - lz4_compressed (bool): Whether the data is LZ4 compressed - - Returns: - ResultSetQueue: The appropriate queue for the result data - """ - - if manifest.format == ResultFormat.JSON_ARRAY.value: - # INLINE disposition with JSON_ARRAY format - return JsonQueue(result_data.data) - elif manifest.format == ResultFormat.ARROW_STREAM.value: - if result_data.attachment is not None: - # direct results from Hybrid disposition - arrow_file = ( - ResultSetDownloadHandler._decompress_data(result_data.attachment) - if lz4_compressed - else result_data.attachment - ) - arrow_table = create_arrow_table_from_arrow_file( - arrow_file, description - ) - logger.debug(f"Created arrow table with {arrow_table.num_rows} rows") - return ArrowQueue(arrow_table, manifest.total_row_count) - - # EXTERNAL_LINKS disposition - return SeaCloudFetchQueue( - result_data=result_data, - max_download_threads=max_download_threads, - ssl_options=ssl_options, - sea_client=sea_client, - statement_id=statement_id, - total_chunk_count=manifest.total_chunk_count, - lz4_compressed=lz4_compressed, - description=description, - http_client=http_client, - ) - raise ProgrammingError("Invalid result format") - - -class JsonQueue(ResultSetQueue): - """Queue implementation for JSON_ARRAY format data.""" - - def __init__(self, data_array: Optional[List[List[str]]]): - """Initialize with JSON array data.""" - self.data_array = data_array or [] - self.cur_row_index = 0 - self.num_rows = len(self.data_array) - - def next_n_rows(self, num_rows: int) -> List[List[str]]: - """Get the next n rows from the data array.""" - length = min(num_rows, self.num_rows - self.cur_row_index) - slice = self.data_array[self.cur_row_index : self.cur_row_index + length] - self.cur_row_index += length - return slice - - def remaining_rows(self) -> List[List[str]]: - """Get all remaining rows from the data array.""" - slice = self.data_array[self.cur_row_index :] - self.cur_row_index += len(slice) - return slice - - def close(self): - return - - -class LinkFetcher: - """ - Background helper that incrementally retrieves *external links* for a - result set produced by the SEA backend and feeds them to a - :class:`databricks.sql.cloudfetch.download_manager.ResultFileDownloadManager`. - - The SEA backend splits large result sets into *chunks*. Each chunk is - stored remotely (e.g., in object storage) and exposed via a signed URL - encapsulated by an :class:`ExternalLink`. Only the first batch of links is - returned with the initial query response. The remaining links must be - pulled on demand using the *next-chunk* token embedded in each - :pyattr:`ExternalLink.next_chunk_index`. - - LinkFetcher takes care of this choreography so callers (primarily - ``SeaCloudFetchQueue``) can simply ask for the link of a specific - ``chunk_index`` and block until it becomes available. - - Key responsibilities: - - • Maintain an in-memory mapping from ``chunk_index`` → ``ExternalLink``. - • Launch a background worker thread that continuously requests the next - batch of links from the backend until all chunks have been discovered or - an unrecoverable error occurs. - • Bridge SEA link objects to the Thrift representation expected by the - existing download manager. - • Provide a synchronous API (`get_chunk_link`) that blocks until the desired - link is present in the cache. - """ - - def __init__( - self, - download_manager: ResultFileDownloadManager, - backend: SeaDatabricksClient, - statement_id: str, - initial_links: List[ExternalLink], - total_chunk_count: int, - ): - self.download_manager = download_manager - self.backend = backend - self._statement_id = statement_id - - self._shutdown_event = threading.Event() - - self._link_data_update = threading.Condition() - self._error: Optional[Exception] = None - self.chunk_index_to_link: Dict[int, ExternalLink] = {} - - self._add_links(initial_links) - self.total_chunk_count = total_chunk_count - - # DEBUG: capture initial state for observability - logger.debug( - "LinkFetcher[%s]: initialized with %d initial link(s); expecting %d total chunk(s)", - statement_id, - len(initial_links), - total_chunk_count, - ) - - def _add_links(self, links: List[ExternalLink]): - """Cache *links* locally and enqueue them with the download manager.""" - logger.debug( - "LinkFetcher[%s]: caching %d link(s) – chunks %s", - self._statement_id, - len(links), - ", ".join(str(l.chunk_index) for l in links) if links else "", - ) - for link in links: - self.chunk_index_to_link[link.chunk_index] = link - self.download_manager.add_link(LinkFetcher._convert_to_thrift_link(link)) - - def _get_next_chunk_index(self) -> Optional[int]: - """Return the next *chunk_index* that should be requested from the backend, or ``None`` if we have them all.""" - with self._link_data_update: - max_chunk_index = max(self.chunk_index_to_link.keys(), default=None) - if max_chunk_index is None: - return 0 - max_link = self.chunk_index_to_link[max_chunk_index] - return max_link.next_chunk_index - - def _trigger_next_batch_download(self) -> bool: - """Fetch the next batch of links from the backend and return *True* on success.""" - logger.debug( - "LinkFetcher[%s]: requesting next batch of links", self._statement_id - ) - next_chunk_index = self._get_next_chunk_index() - if next_chunk_index is None: - return False - - try: - links = self.backend.get_chunk_links(self._statement_id, next_chunk_index) - with self._link_data_update: - self._add_links(links) - self._link_data_update.notify_all() - except Exception as e: - logger.error( - f"LinkFetcher: Error fetching links for chunk {next_chunk_index}: {e}" - ) - with self._link_data_update: - self._error = e - self._link_data_update.notify_all() - return False - - logger.debug( - "LinkFetcher[%s]: received %d new link(s)", - self._statement_id, - len(links), - ) - return True - - def get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]: - """Return (blocking) the :class:`ExternalLink` associated with *chunk_index*.""" - logger.debug( - "LinkFetcher[%s]: waiting for link of chunk %d", - self._statement_id, - chunk_index, - ) - if chunk_index >= self.total_chunk_count: - return None - - with self._link_data_update: - while chunk_index not in self.chunk_index_to_link: - if self._error: - raise self._error - if self._shutdown_event.is_set(): - raise ProgrammingError( - "LinkFetcher is shutting down without providing link for chunk index {}".format( - chunk_index - ) - ) - self._link_data_update.wait() - - return self.chunk_index_to_link[chunk_index] - - @staticmethod - def _convert_to_thrift_link(link: ExternalLink) -> TSparkArrowResultLink: - """Convert SEA external links to Thrift format for compatibility with existing download manager.""" - from databricks.sql.thrift_api.TCLIService.ttypes import ( - TSparkArrowResultLink, - ) - - # Parse the ISO format expiration time - expiry_time = int(dateutil.parser.parse(link.expiration).timestamp()) - return TSparkArrowResultLink( - fileLink=link.external_link, - expiryTime=expiry_time, - rowCount=link.row_count, - bytesNum=link.byte_count, - startRowOffset=link.row_offset, - httpHeaders=link.http_headers or {}, - ) - - def _worker_loop(self): - """Entry point for the background thread.""" - logger.debug("LinkFetcher[%s]: worker thread started", self._statement_id) - while not self._shutdown_event.is_set(): - links_downloaded = self._trigger_next_batch_download() - if not links_downloaded: - self._shutdown_event.set() - logger.debug("LinkFetcher[%s]: worker thread exiting", self._statement_id) - with self._link_data_update: - self._link_data_update.notify_all() - - def start(self): - """Spawn the worker thread.""" - logger.debug("LinkFetcher[%s]: starting worker thread", self._statement_id) - self._worker_thread = threading.Thread( - target=self._worker_loop, name=f"LinkFetcher-{self._statement_id}" - ) - self._worker_thread.start() - - def stop(self): - """Signal the worker thread to stop and wait for its termination.""" - logger.debug("LinkFetcher[%s]: stopping worker thread", self._statement_id) - self._shutdown_event.set() - self._worker_thread.join() - logger.debug("LinkFetcher[%s]: worker thread stopped", self._statement_id) - - -class SeaCloudFetchQueue(CloudFetchQueue): - """Queue implementation for EXTERNAL_LINKS disposition with ARROW format for SEA backend.""" - - def __init__( - self, - result_data: ResultData, - max_download_threads: int, - ssl_options: SSLOptions, - sea_client: SeaDatabricksClient, - statement_id: str, - total_chunk_count: int, - http_client, - lz4_compressed: bool = False, - description: List[Tuple] = [], - ): - """ - Initialize the SEA CloudFetchQueue. - - Args: - initial_links: Initial list of external links to download - schema_bytes: Arrow schema bytes - max_download_threads: Maximum number of download threads - ssl_options: SSL options for downloads - sea_client: SEA client for fetching additional links - statement_id: Statement ID for the query - total_chunk_count: Total number of chunks in the result set - lz4_compressed: Whether the data is LZ4 compressed - description: Column descriptions - """ - - super().__init__( - max_download_threads=max_download_threads, - ssl_options=ssl_options, - statement_id=statement_id, - schema_bytes=None, - lz4_compressed=lz4_compressed, - description=description, - # TODO: fix these arguments when telemetry is implemented in SEA - session_id_hex=None, - chunk_id=0, - http_client=http_client, - ) - - logger.debug( - "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format( - statement_id, total_chunk_count - ) - ) - - initial_links = result_data.external_links or [] - - # Track the current chunk we're processing - self._current_chunk_index = 0 - - self.link_fetcher = None # for empty responses, we do not need a link fetcher - if total_chunk_count > 0: - self.link_fetcher = LinkFetcher( - download_manager=self.download_manager, - backend=sea_client, - statement_id=statement_id, - initial_links=initial_links, - total_chunk_count=total_chunk_count, - ) - self.link_fetcher.start() - - # Initialize table and position - self.table = self._create_next_table() - - def _create_next_table(self) -> Union["pyarrow.Table", None]: - """Create next table by retrieving the logical next downloaded file.""" - if self.link_fetcher is None: - return None - - chunk_link = self.link_fetcher.get_chunk_link(self._current_chunk_index) - if chunk_link is None: - return None - - row_offset = chunk_link.row_offset - # NOTE: link has already been submitted to download manager at this point - arrow_table = self._create_table_at_offset(row_offset) - - self._current_chunk_index += 1 - - return arrow_table - - def close(self): - super().close() - if self.link_fetcher: - self.link_fetcher.stop() diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py deleted file mode 100644 index 17838ed81..000000000 --- a/src/databricks/sql/backend/sea/result_set.py +++ /dev/null @@ -1,266 +0,0 @@ -from __future__ import annotations - -from typing import Any, List, Optional, TYPE_CHECKING - -import logging - -from databricks.sql.backend.sea.models.base import ResultData, ResultManifest -from databricks.sql.backend.sea.utils.conversion import SqlTypeConverter - -try: - import pyarrow -except ImportError: - pyarrow = None - -if TYPE_CHECKING: - from databricks.sql.client import Connection - from databricks.sql.backend.sea.backend import SeaDatabricksClient -from databricks.sql.types import Row -from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory -from databricks.sql.backend.types import ExecuteResponse -from databricks.sql.result_set import ResultSet - -logger = logging.getLogger(__name__) - - -class SeaResultSet(ResultSet): - """ResultSet implementation for SEA backend.""" - - def __init__( - self, - connection: Connection, - execute_response: ExecuteResponse, - sea_client: SeaDatabricksClient, - result_data: ResultData, - manifest: ResultManifest, - buffer_size_bytes: int = 104857600, - arraysize: int = 10000, - ): - """ - Initialize a SeaResultSet with the response from a SEA query execution. - - Args: - connection: The parent connection - execute_response: Response from the execute command - sea_client: The SeaDatabricksClient instance for direct access - buffer_size_bytes: Buffer size for fetching results - arraysize: Default number of rows to fetch - result_data: Result data from SEA response - manifest: Manifest from SEA response - """ - - self.manifest = manifest - - statement_id = execute_response.command_id.to_sea_statement_id() - if statement_id is None: - raise ValueError("Command ID is not a SEA statement ID") - - results_queue = SeaResultSetQueueFactory.build_queue( - result_data, - self.manifest, - statement_id, - ssl_options=connection.session.ssl_options, - description=execute_response.description, - max_download_threads=sea_client.max_download_threads, - sea_client=sea_client, - lz4_compressed=execute_response.lz4_compressed, - http_client=connection.session.http_client, - ) - - # Call parent constructor with common attributes - super().__init__( - connection=connection, - backend=sea_client, - arraysize=arraysize, - buffer_size_bytes=buffer_size_bytes, - command_id=execute_response.command_id, - status=execute_response.status, - has_been_closed_server_side=execute_response.has_been_closed_server_side, - results_queue=results_queue, - description=execute_response.description, - is_staging_operation=execute_response.is_staging_operation, - lz4_compressed=execute_response.lz4_compressed, - arrow_schema_bytes=execute_response.arrow_schema_bytes, - ) - - def _convert_json_types(self, row: List[str]) -> List[Any]: - """ - Convert string values in the row to appropriate Python types based on column metadata. - """ - - # JSON + INLINE gives us string values, so we convert them to appropriate - # types based on column metadata - converted_row = [] - - for i, value in enumerate(row): - column_name = self.description[i][0] - column_type = self.description[i][1] - precision = self.description[i][4] - scale = self.description[i][5] - - converted_value = SqlTypeConverter.convert_value( - value, - column_type, - column_name=column_name, - precision=precision, - scale=scale, - ) - converted_row.append(converted_value) - - return converted_row - - def _convert_json_to_arrow_table(self, rows: List[List[str]]) -> "pyarrow.Table": - """ - Convert raw data rows to Arrow table. - - Args: - rows: List of raw data rows - - Returns: - PyArrow Table containing the converted values - """ - - if not rows: - return pyarrow.Table.from_pydict({}) - - # create a generator for row conversion - converted_rows_iter = (self._convert_json_types(row) for row in rows) - cols = list(map(list, zip(*converted_rows_iter))) - - names = [col[0] for col in self.description] - return pyarrow.Table.from_arrays(cols, names=names) - - def _create_json_table(self, rows: List[List[str]]) -> List[Row]: - """ - Convert raw data rows to Row objects with named columns based on description. - - Args: - rows: List of raw data rows - Returns: - List of Row objects with named columns and converted values - """ - - ResultRow = Row(*[col[0] for col in self.description]) - return [ResultRow(*self._convert_json_types(row)) for row in rows] - - def fetchmany_json(self, size: int) -> List[List[str]]: - """ - Fetch the next set of rows as a columnar table. - - Args: - size: Number of rows to fetch - - Returns: - Columnar table containing the fetched rows - - Raises: - ValueError: If size is negative - """ - - if size < 0: - raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") - - results = self.results.next_n_rows(size) - self._next_row_index += len(results) - - return results - - def fetchall_json(self) -> List[List[str]]: - """ - Fetch all remaining rows as a columnar table. - - Returns: - Columnar table containing all remaining rows - """ - - results = self.results.remaining_rows() - self._next_row_index += len(results) - - return results - - def fetchmany_arrow(self, size: int) -> "pyarrow.Table": - """ - Fetch the next set of rows as an Arrow table. - - Args: - size: Number of rows to fetch - - Returns: - PyArrow Table containing the fetched rows - - Raises: - ImportError: If PyArrow is not installed - ValueError: If size is negative - """ - - if size < 0: - raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") - - results = self.results.next_n_rows(size) - if isinstance(self.results, JsonQueue): - results = self._convert_json_to_arrow_table(results) - - self._next_row_index += results.num_rows - - return results - - def fetchall_arrow(self) -> "pyarrow.Table": - """ - Fetch all remaining rows as an Arrow table. - """ - - results = self.results.remaining_rows() - if isinstance(self.results, JsonQueue): - results = self._convert_json_to_arrow_table(results) - - self._next_row_index += results.num_rows - - return results - - def fetchone(self) -> Optional[Row]: - """ - Fetch the next row of a query result set, returning a single sequence, - or None when no more data is available. - - Returns: - A single Row object or None if no more rows are available - """ - - if isinstance(self.results, JsonQueue): - res = self._create_json_table(self.fetchmany_json(1)) - else: - res = self._convert_arrow_table(self.fetchmany_arrow(1)) - - return res[0] if res else None - - def fetchmany(self, size: int) -> List[Row]: - """ - Fetch the next set of rows of a query result, returning a list of rows. - - Args: - size: Number of rows to fetch (defaults to arraysize if None) - - Returns: - List of Row objects - - Raises: - ValueError: If size is negative - """ - - if isinstance(self.results, JsonQueue): - return self._create_json_table(self.fetchmany_json(size)) - else: - return self._convert_arrow_table(self.fetchmany_arrow(size)) - - def fetchall(self) -> List[Row]: - """ - Fetch all remaining rows of a query result, returning them as a list of rows. - - Returns: - List of Row objects containing all remaining rows - """ - - if isinstance(self.results, JsonQueue): - return self._create_json_table(self.fetchall_json()) - else: - return self._convert_arrow_table(self.fetchall_arrow()) diff --git a/src/databricks/sql/backend/sea/utils/constants.py b/src/databricks/sql/backend/sea/utils/constants.py deleted file mode 100644 index 61ecf969e..000000000 --- a/src/databricks/sql/backend/sea/utils/constants.py +++ /dev/null @@ -1,68 +0,0 @@ -""" -Constants for the Statement Execution API (SEA) backend. -""" - -from typing import Dict -from enum import Enum - -# from https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-parameters -ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP: Dict[str, str] = { - "ANSI_MODE": "true", - "ENABLE_PHOTON": "true", - "LEGACY_TIME_PARSER_POLICY": "Exception", - "MAX_FILE_PARTITION_BYTES": "128m", - "READ_ONLY_EXTERNAL_METASTORE": "false", - "STATEMENT_TIMEOUT": "0", - "TIMEZONE": "UTC", - "USE_CACHED_RESULT": "true", - "QUERY_TAGS": "", -} - - -class ResultFormat(Enum): - """Enum for result format values.""" - - ARROW_STREAM = "ARROW_STREAM" - JSON_ARRAY = "JSON_ARRAY" - - -class ResultDisposition(Enum): - """Enum for result disposition values.""" - - HYBRID = "INLINE_OR_EXTERNAL_LINKS" - EXTERNAL_LINKS = "EXTERNAL_LINKS" - INLINE = "INLINE" - - -class ResultCompression(Enum): - """Enum for result compression values.""" - - LZ4_FRAME = "LZ4_FRAME" - NONE = None - - -class WaitTimeout(Enum): - """Enum for wait timeout values.""" - - ASYNC = "0s" - SYNC = "10s" - - -class MetadataCommands(Enum): - """SQL commands used in the SEA backend. - - These constants are used for metadata operations and other SQL queries - to ensure consistency and avoid string literal duplication. - """ - - SHOW_CATALOGS = "SHOW CATALOGS" - SHOW_SCHEMAS = "SHOW SCHEMAS IN {}" - SHOW_TABLES = "SHOW TABLES IN {}" - SHOW_TABLES_ALL_CATALOGS = "SHOW TABLES IN ALL CATALOGS" - SHOW_COLUMNS = "SHOW COLUMNS IN CATALOG {}" - - LIKE_PATTERN = " LIKE '{}'" - SCHEMA_LIKE_PATTERN = " SCHEMA" + LIKE_PATTERN - TABLE_LIKE_PATTERN = " TABLE" + LIKE_PATTERN - - CATALOG_SPECIFIC = "CATALOG {}" diff --git a/src/databricks/sql/backend/sea/utils/conversion.py b/src/databricks/sql/backend/sea/utils/conversion.py deleted file mode 100644 index 69c6dfbe2..000000000 --- a/src/databricks/sql/backend/sea/utils/conversion.py +++ /dev/null @@ -1,173 +0,0 @@ -""" -Type conversion utilities for the Databricks SQL Connector. - -This module provides functionality to convert string values from SEA Inline results -to appropriate Python types based on column metadata. -""" - -import datetime -import decimal -import logging -from dateutil import parser -from typing import Callable, Dict, Optional - -logger = logging.getLogger(__name__) - - -def _convert_decimal( - value: str, precision: Optional[int] = None, scale: Optional[int] = None -) -> decimal.Decimal: - """ - Convert a string value to a decimal with optional precision and scale. - - Args: - value: The string value to convert - precision: Optional precision (total number of significant digits) for the decimal - scale: Optional scale (number of decimal places) for the decimal - - Returns: - A decimal.Decimal object with appropriate precision and scale - """ - - # First create the decimal from the string value - result = decimal.Decimal(value) - - # Apply scale (quantize to specific number of decimal places) if specified - quantizer = None - if scale is not None: - quantizer = decimal.Decimal(f'0.{"0" * scale}') - - # Apply precision (total number of significant digits) if specified - context = None - if precision is not None: - context = decimal.Context(prec=precision) - - if quantizer is not None: - result = result.quantize(quantizer, context=context) - - return result - - -class SqlType: - """ - SQL type constants based on Thrift TTypeId values. - - These correspond to the normalized type names that come from the SEA backend - after normalize_sea_type_to_thrift processing (lowercase, without _TYPE suffix). - """ - - # Numeric types - TINYINT = "tinyint" # Maps to TTypeId.TINYINT_TYPE - SMALLINT = "smallint" # Maps to TTypeId.SMALLINT_TYPE - INT = "int" # Maps to TTypeId.INT_TYPE - BIGINT = "bigint" # Maps to TTypeId.BIGINT_TYPE - FLOAT = "float" # Maps to TTypeId.FLOAT_TYPE - DOUBLE = "double" # Maps to TTypeId.DOUBLE_TYPE - DECIMAL = "decimal" # Maps to TTypeId.DECIMAL_TYPE - - # Boolean type - BOOLEAN = "boolean" # Maps to TTypeId.BOOLEAN_TYPE - - # Date/Time types - DATE = "date" # Maps to TTypeId.DATE_TYPE - TIMESTAMP = "timestamp" # Maps to TTypeId.TIMESTAMP_TYPE - INTERVAL_YEAR_MONTH = ( - "interval_year_month" # Maps to TTypeId.INTERVAL_YEAR_MONTH_TYPE - ) - INTERVAL_DAY_TIME = "interval_day_time" # Maps to TTypeId.INTERVAL_DAY_TIME_TYPE - - # String types - CHAR = "char" # Maps to TTypeId.CHAR_TYPE - VARCHAR = "varchar" # Maps to TTypeId.VARCHAR_TYPE - STRING = "string" # Maps to TTypeId.STRING_TYPE - - # Binary type - BINARY = "binary" # Maps to TTypeId.BINARY_TYPE - - # Complex types - ARRAY = "array" # Maps to TTypeId.ARRAY_TYPE - MAP = "map" # Maps to TTypeId.MAP_TYPE - STRUCT = "struct" # Maps to TTypeId.STRUCT_TYPE - - # Other types - NULL = "null" # Maps to TTypeId.NULL_TYPE - UNION = "union" # Maps to TTypeId.UNION_TYPE - USER_DEFINED = "user_defined" # Maps to TTypeId.USER_DEFINED_TYPE - - -class SqlTypeConverter: - """ - Utility class for converting SQL types to Python types. - Based on the Thrift TTypeId types after normalization. - """ - - # SQL type to conversion function mapping - # TODO: complex types - TYPE_MAPPING: Dict[str, Callable] = { - # Numeric types - SqlType.TINYINT: lambda v: int(v), - SqlType.SMALLINT: lambda v: int(v), - SqlType.INT: lambda v: int(v), - SqlType.BIGINT: lambda v: int(v), - SqlType.FLOAT: lambda v: float(v), - SqlType.DOUBLE: lambda v: float(v), - SqlType.DECIMAL: _convert_decimal, - # Boolean type - SqlType.BOOLEAN: lambda v: v.lower() in ("true", "t", "1", "yes", "y"), - # Date/Time types - SqlType.DATE: lambda v: datetime.date.fromisoformat(v), - SqlType.TIMESTAMP: lambda v: parser.parse(v), - SqlType.INTERVAL_YEAR_MONTH: lambda v: v, # Keep as string for now - SqlType.INTERVAL_DAY_TIME: lambda v: v, # Keep as string for now - # String types - no conversion needed - SqlType.CHAR: lambda v: v, - SqlType.VARCHAR: lambda v: v, - SqlType.STRING: lambda v: v, - # Binary type - SqlType.BINARY: lambda v: bytes.fromhex(v), - # Other types - SqlType.NULL: lambda v: None, - # Complex types and user-defined types return as-is - SqlType.USER_DEFINED: lambda v: v, - } - - @staticmethod - def convert_value( - value: str, - sql_type: str, - column_name: Optional[str], - **kwargs, - ) -> object: - """ - Convert a string value to the appropriate Python type based on SQL type. - - Args: - value: The string value to convert - sql_type: The SQL type (e.g., 'tinyint', 'decimal') - column_name: The name of the column being converted - **kwargs: Additional keyword arguments for the conversion function - - Returns: - The converted value in the appropriate Python type - """ - - sql_type = sql_type.lower().strip() - - if sql_type not in SqlTypeConverter.TYPE_MAPPING: - return value - - converter_func = SqlTypeConverter.TYPE_MAPPING[sql_type] - try: - if sql_type == SqlType.DECIMAL: - precision = kwargs.get("precision", None) - scale = kwargs.get("scale", None) - return converter_func(value, precision, scale) - else: - return converter_func(value) - except Exception as e: - warning_message = f"Error converting value '{value}' to {sql_type}" - if column_name: - warning_message += f" in column {column_name}" - warning_message += f": {e}" - logger.warning(warning_message) - return value diff --git a/src/databricks/sql/backend/sea/utils/filters.py b/src/databricks/sql/backend/sea/utils/filters.py deleted file mode 100644 index f125c409a..000000000 --- a/src/databricks/sql/backend/sea/utils/filters.py +++ /dev/null @@ -1,287 +0,0 @@ -""" -Client-side filtering utilities for Databricks SQL connector. - -This module provides filtering capabilities for result sets returned by different backends. -""" - -from __future__ import annotations - -import io -import logging -from typing import ( - List, - Optional, - Any, - cast, - TYPE_CHECKING, -) - -if TYPE_CHECKING: - from databricks.sql.backend.sea.result_set import SeaResultSet - -from databricks.sql.backend.types import ExecuteResponse -from databricks.sql.backend.sea.models.base import ResultData -from databricks.sql.backend.sea.backend import SeaDatabricksClient -from databricks.sql.utils import CloudFetchQueue, ArrowQueue - -try: - import pyarrow - import pyarrow.compute as pc -except ImportError: - pyarrow = None - pc = None - -logger = logging.getLogger(__name__) - - -class ResultSetFilter: - """ - A general-purpose filter for result sets. - """ - - @staticmethod - def _create_execute_response(result_set: SeaResultSet) -> ExecuteResponse: - """ - Create an ExecuteResponse with parameters from the original result set. - - Args: - result_set: Original result set to copy parameters from - - Returns: - ExecuteResponse: New execute response object - """ - return ExecuteResponse( - command_id=result_set.command_id, - status=result_set.status, - description=result_set.description, - has_been_closed_server_side=result_set.has_been_closed_server_side, - lz4_compressed=result_set.lz4_compressed, - arrow_schema_bytes=result_set._arrow_schema_bytes, - is_staging_operation=False, - ) - - @staticmethod - def _update_manifest(result_set: SeaResultSet, new_row_count: int): - """ - Create a copy of the manifest with updated row count. - - Args: - result_set: Original result set to copy manifest from - new_row_count: New total row count for filtered data - - Returns: - Updated manifest copy - """ - filtered_manifest = result_set.manifest - filtered_manifest.total_row_count = new_row_count - return filtered_manifest - - @staticmethod - def _create_filtered_result_set( - result_set: SeaResultSet, - result_data: ResultData, - row_count: int, - ) -> "SeaResultSet": - """ - Create a new filtered SeaResultSet with the provided data. - - Args: - result_set: Original result set to copy parameters from - result_data: New result data for the filtered set - row_count: Number of rows in the filtered data - - Returns: - New filtered SeaResultSet - """ - from databricks.sql.backend.sea.result_set import SeaResultSet - - execute_response = ResultSetFilter._create_execute_response(result_set) - filtered_manifest = ResultSetFilter._update_manifest(result_set, row_count) - - return SeaResultSet( - connection=result_set.connection, - execute_response=execute_response, - sea_client=cast(SeaDatabricksClient, result_set.backend), - result_data=result_data, - manifest=filtered_manifest, - buffer_size_bytes=result_set.buffer_size_bytes, - arraysize=result_set.arraysize, - ) - - @staticmethod - def _filter_arrow_table( - table: Any, # pyarrow.Table - column_name: str, - allowed_values: List[str], - case_sensitive: bool = True, - ) -> Any: # returns pyarrow.Table - """ - Filter a PyArrow table by column values. - - Args: - table: The PyArrow table to filter - column_name: The name of the column to filter on - allowed_values: List of allowed values for the column - case_sensitive: Whether to perform case-sensitive comparison - - Returns: - A filtered PyArrow table - """ - if not pyarrow: - raise ImportError("PyArrow is required for Arrow table filtering") - - if table.num_rows == 0: - return table - - # Handle case-insensitive filtering by normalizing both column and allowed values - if not case_sensitive: - # Convert allowed values to uppercase - allowed_values = [v.upper() for v in allowed_values] - # Get column values as uppercase - column = pc.utf8_upper(table[column_name]) - else: - # Use column as-is - column = table[column_name] - - # Convert allowed_values to PyArrow Array - allowed_array = pyarrow.array(allowed_values) - - # Construct a boolean mask: True where column is in allowed_list - mask = pc.is_in(column, value_set=allowed_array) - return table.filter(mask) - - @staticmethod - def _filter_arrow_result_set( - result_set: SeaResultSet, - column_index: int, - allowed_values: List[str], - case_sensitive: bool = True, - ) -> SeaResultSet: - """ - Filter a SEA result set that contains Arrow tables. - - Args: - result_set: The SEA result set to filter (containing Arrow data) - column_index: The index of the column to filter on - allowed_values: List of allowed values for the column - case_sensitive: Whether to perform case-sensitive comparison - - Returns: - A filtered SEA result set - """ - # Validate column index and get column name - if column_index >= len(result_set.description): - raise ValueError(f"Column index {column_index} is out of bounds") - column_name = result_set.description[column_index][0] - - # Get all remaining rows as Arrow table and filter it - arrow_table = result_set.results.remaining_rows() - filtered_table = ResultSetFilter._filter_arrow_table( - arrow_table, column_name, allowed_values, case_sensitive - ) - - # Convert the filtered table to Arrow stream format for ResultData - sink = io.BytesIO() - with pyarrow.ipc.new_stream(sink, filtered_table.schema) as writer: - writer.write_table(filtered_table) - arrow_stream_bytes = sink.getvalue() - - # Create ResultData with attachment containing the filtered data - result_data = ResultData( - data=None, # No JSON data - external_links=None, # No external links - attachment=arrow_stream_bytes, # Arrow data as attachment - ) - - return ResultSetFilter._create_filtered_result_set( - result_set, result_data, filtered_table.num_rows - ) - - @staticmethod - def _filter_json_result_set( - result_set: SeaResultSet, - column_index: int, - allowed_values: List[str], - case_sensitive: bool = False, - ) -> SeaResultSet: - """ - Filter a result set by values in a specific column. - - Args: - result_set: The result set to filter - column_index: The index of the column to filter on - allowed_values: List of allowed values for the column - case_sensitive: Whether to perform case-sensitive comparison - - Returns: - A filtered result set - """ - # Validate column index (optional - not in arrow version but good practice) - if column_index >= len(result_set.description): - raise ValueError(f"Column index {column_index} is out of bounds") - - # Extract rows - all_rows = result_set.results.remaining_rows() - - # Convert allowed values if case-insensitive - if not case_sensitive: - allowed_values = [v.upper() for v in allowed_values] - # Helper lambda to get column value based on case sensitivity - get_column_value = lambda row: ( - row[column_index].upper() if not case_sensitive else row[column_index] - ) - - # Filter rows based on allowed values - filtered_rows = [ - row - for row in all_rows - if len(row) > column_index and get_column_value(row) in allowed_values - ] - - # Create filtered result set - result_data = ResultData(data=filtered_rows, external_links=None) - - return ResultSetFilter._create_filtered_result_set( - result_set, result_data, len(filtered_rows) - ) - - @staticmethod - def filter_tables_by_type( - result_set: SeaResultSet, table_types: Optional[List[str]] = None - ) -> SeaResultSet: - """ - Filter a result set of tables by the specified table types. - - This is a client-side filter that processes the result set after it has been - retrieved from the server. It filters out tables whose type does not match - any of the types in the table_types list. - - Args: - result_set: The original result set containing tables - table_types: List of table types to include (e.g., ["TABLE", "VIEW"]) - - Returns: - A filtered result set containing only tables of the specified types - """ - # Default table types if none specified - DEFAULT_TABLE_TYPES = ["TABLE", "VIEW", "SYSTEM TABLE"] - valid_types = table_types if table_types else DEFAULT_TABLE_TYPES - - # Check if we have an Arrow table (cloud fetch) or JSON data - # Table type is the 6th column (index 5) - if isinstance(result_set.results, (CloudFetchQueue, ArrowQueue)): - # For Arrow tables, we need to handle filtering differently - return ResultSetFilter._filter_arrow_result_set( - result_set, - column_index=5, - allowed_values=valid_types, - case_sensitive=True, - ) - else: - # For JSON data, use the existing filter method - return ResultSetFilter._filter_json_result_set( - result_set, - column_index=5, - allowed_values=valid_types, - case_sensitive=True, - ) diff --git a/src/databricks/sql/backend/sea/utils/http_client.py b/src/databricks/sql/backend/sea/utils/http_client.py deleted file mode 100644 index 8ef6cec96..000000000 --- a/src/databricks/sql/backend/sea/utils/http_client.py +++ /dev/null @@ -1,305 +0,0 @@ -import json -import logging -import ssl -import urllib.parse -import urllib.request -from typing import Dict, Any, Optional, List, Tuple, Union - -from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager -from urllib3.util import make_headers -from urllib3.exceptions import MaxRetryError - -from databricks.sql.auth.authenticators import AuthProvider -from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy -from databricks.sql.types import SSLOptions -from databricks.sql.exc import ( - RequestError, -) -from databricks.sql.common.http_utils import ( - detect_and_parse_proxy, -) -from databricks.sql.common.url_utils import normalize_host_with_protocol - -logger = logging.getLogger(__name__) - - -class SeaHttpClient: - """ - HTTP client for Statement Execution API (SEA). - - This client uses urllib3 for robust HTTP communication with retry policies - and connection pooling. - """ - - retry_policy: Union[DatabricksRetryPolicy, int] - _pool: Optional[Union[HTTPConnectionPool, HTTPSConnectionPool]] - proxy_uri: Optional[str] - proxy_auth: Optional[Dict[str, str]] - realhost: Optional[str] - realport: Optional[int] - - def __init__( - self, - server_hostname: str, - port: int, - http_path: str, - http_headers: List[Tuple[str, str]], - auth_provider: AuthProvider, - ssl_options: SSLOptions, - **kwargs, - ): - """ - Initialize the SEA HTTP client. - - Args: - server_hostname: Hostname of the Databricks server - port: Port number for the connection - http_path: HTTP path for the connection - http_headers: List of HTTP headers to include in requests - auth_provider: Authentication provider - ssl_options: SSL configuration options - **kwargs: Additional keyword arguments including retry policy settings - """ - - self.server_hostname = server_hostname - self.port = port or 443 - self.http_path = http_path - self.auth_provider = auth_provider - self.ssl_options = ssl_options - - # Build base URL using url_utils for consistent normalization - normalized_host = normalize_host_with_protocol(server_hostname) - self.base_url = f"{normalized_host}:{self.port}" - - # Parse URL for proxy handling - parsed_url = urllib.parse.urlparse(self.base_url) - self.scheme = parsed_url.scheme - self.host = parsed_url.hostname - self.port = parsed_url.port or (443 if self.scheme == "https" else 80) - - # Setup headers - self.headers: Dict[str, str] = dict(http_headers) - self.headers.update({"Content-Type": "application/json"}) - - # Extract retry policy settings - self._retry_delay_min = kwargs.get("_retry_delay_min", 1.0) - self._retry_delay_max = kwargs.get("_retry_delay_max", 60.0) - self._retry_stop_after_attempts_count = kwargs.get( - "_retry_stop_after_attempts_count", 30 - ) - self._retry_stop_after_attempts_duration = kwargs.get( - "_retry_stop_after_attempts_duration", 900.0 - ) - self._retry_delay_default = kwargs.get("_retry_delay_default", 5.0) - self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) - self._respect_server_retry_after_header = kwargs.get( - "_respect_server_retry_after_header", False - ) - - # Connection pooling settings - self.max_connections = kwargs.get("max_connections", 10) - - # Setup retry policy - self.enable_v3_retries = kwargs.get("_enable_v3_retries", True) - - if self.enable_v3_retries: - urllib3_kwargs = {"allowed_methods": ["GET", "POST", "DELETE"]} - _max_redirects = kwargs.get("_retry_max_redirects") - if _max_redirects: - if _max_redirects > self._retry_stop_after_attempts_count: - logger.warning( - "_retry_max_redirects > _retry_stop_after_attempts_count so it will have no affect!" - ) - urllib3_kwargs["redirect"] = _max_redirects - - self.retry_policy = DatabricksRetryPolicy( - delay_min=self._retry_delay_min, - delay_max=self._retry_delay_max, - stop_after_attempts_count=self._retry_stop_after_attempts_count, - stop_after_attempts_duration=self._retry_stop_after_attempts_duration, - delay_default=self._retry_delay_default, - force_dangerous_codes=self.force_dangerous_codes, - respect_server_retry_after_header=self._respect_server_retry_after_header, - urllib3_kwargs=urllib3_kwargs, - ) - else: - # Legacy behavior - no automatic retries - logger.warning( - "Legacy retry behavior is enabled for this connection." - " This behaviour is not supported for the SEA backend." - ) - self.retry_policy = 0 - - # Handle proxy settings using shared utility - proxy_auth_method = kwargs.get("_proxy_auth_method") - proxy_uri, proxy_auth = detect_and_parse_proxy( - self.scheme, self.host, proxy_auth_method=proxy_auth_method - ) - - if proxy_uri: - parsed_proxy = urllib.parse.urlparse(proxy_uri) - self.realhost = self.host - self.realport = self.port - self.proxy_uri = proxy_uri - self.host = parsed_proxy.hostname - self.port = parsed_proxy.port or (443 if self.scheme == "https" else 80) - self.proxy_auth = proxy_auth - else: - self.realhost = self.realport = self.proxy_auth = self.proxy_uri = None - - # Initialize connection pool - self._pool = None - self._open() - - def _open(self): - """Initialize the connection pool.""" - pool_kwargs = {"maxsize": self.max_connections} - - if self.scheme == "http": - pool_class = HTTPConnectionPool - else: # https - pool_class = HTTPSConnectionPool - pool_kwargs.update( - { - "cert_reqs": ( - ssl.CERT_REQUIRED - if self.ssl_options.tls_verify - else ssl.CERT_NONE - ), - "ca_certs": self.ssl_options.tls_trusted_ca_file, - "cert_file": self.ssl_options.tls_client_cert_file, - "key_file": self.ssl_options.tls_client_cert_key_file, - "key_password": self.ssl_options.tls_client_cert_key_password, - } - ) - - if self.using_proxy(): - proxy_manager = ProxyManager( - self.proxy_uri, - num_pools=1, - proxy_headers=self.proxy_auth, - ) - self._pool = proxy_manager.connection_from_host( - host=self.realhost, - port=self.realport, - scheme=self.scheme, - pool_kwargs=pool_kwargs, - ) - else: - self._pool = pool_class(self.host, self.port, **pool_kwargs) - - def close(self): - """Close the connection pool.""" - if self._pool: - self._pool.clear() - - def using_proxy(self) -> bool: - """Check if proxy is being used.""" - return self.realhost is not None - - def set_retry_command_type(self, command_type: CommandType): - """Set the command type for retry policy decision making.""" - if isinstance(self.retry_policy, DatabricksRetryPolicy): - self.retry_policy.command_type = command_type - - def start_retry_timer(self): - """Start the retry timer for duration-based retry limits.""" - if isinstance(self.retry_policy, DatabricksRetryPolicy): - self.retry_policy.start_retry_timer() - - def _get_auth_headers(self) -> Dict[str, str]: - """Get authentication headers from the auth provider.""" - headers: Dict[str, str] = {} - self.auth_provider.add_headers(headers) - return headers - - def _make_request( - self, - method: str, - path: str, - data: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """ - Make an HTTP request to the SEA endpoint. - - Args: - method: HTTP method (GET, POST, DELETE) - path: API endpoint path - data: Request payload data - - Returns: - Dict[str, Any]: Response data parsed from JSON - - Raises: - RequestError: If the request fails after retries - """ - - # Prepare headers - headers = {**self.headers, **self._get_auth_headers()} - - # Prepare request body - body = json.dumps(data).encode("utf-8") if data else b"" - if body: - headers["Content-Length"] = str(len(body)) - - # Set command type for retry policy - command_type = self._get_command_type_from_path(path, method) - self.set_retry_command_type(command_type) - self.start_retry_timer() - - logger.debug(f"Making {method} request to {path}") - - if self._pool is None: - raise RequestError("Connection pool not initialized", None) - - try: - with self._pool.request( - method=method.upper(), - url=path, - body=body, - headers=headers, - preload_content=False, - retries=self.retry_policy, - ) as response: - # Handle successful responses - if 200 <= response.status < 300: - if response.data: - return json.loads(response.data.decode()) - else: - return {} - - error_message = f"SEA HTTP request failed with status {response.status}" - raise Exception(error_message) - except MaxRetryError as e: - logger.error(f"SEA HTTP request failed with MaxRetryError: {e}") - raise - except Exception as e: - logger.error(f"SEA HTTP request failed with exception: {e}") - error_message = f"Error during request to server. {e}" - raise RequestError(error_message, None, None, e) - - def _get_command_type_from_path(self, path: str, method: str) -> CommandType: - """ - Determine the command type based on the API path and method. - - This helps the retry policy make appropriate decisions for different - types of SEA operations. - """ - - path = path.lower() - method = method.upper() - - if "/statements" in path: - if method == "POST" and path.endswith("/statements"): - return CommandType.EXECUTE_STATEMENT - elif "/cancel" in path: - return CommandType.OTHER # Cancel operation - elif method == "DELETE": - return CommandType.CLOSE_OPERATION - elif method == "GET": - return CommandType.GET_OPERATION_STATUS - elif "/sessions" in path: - if method == "DELETE": - return CommandType.CLOSE_SESSION - - return CommandType.OTHER diff --git a/src/databricks/sql/backend/sea/utils/normalize.py b/src/databricks/sql/backend/sea/utils/normalize.py deleted file mode 100644 index d725d294b..000000000 --- a/src/databricks/sql/backend/sea/utils/normalize.py +++ /dev/null @@ -1,50 +0,0 @@ -""" -Type normalization utilities for SEA backend. - -This module provides functionality to normalize SEA type names to match -Thrift type naming conventions. -""" - -from typing import Dict, Any - -# SEA types that need to be translated to Thrift types -# The list of all SEA types is available in the REST reference at: -# https://docs.databricks.com/api/workspace/statementexecution/executestatement -# The list of all Thrift types can be found in the ttypes.TTypeId definition -# The SEA types that do not align with Thrift are explicitly mapped below -SEA_TO_THRIFT_TYPE_MAP = { - "BYTE": "TINYINT", - "SHORT": "SMALLINT", - "LONG": "BIGINT", - "INTERVAL": "INTERVAL", # Default mapping, will be overridden if type_interval_type is present -} - - -def normalize_sea_type_to_thrift(type_name: str, col_data: Dict[str, Any]) -> str: - """ - Normalize SEA type names to match Thrift type naming conventions. - - Args: - type_name: The type name from SEA (e.g., "BYTE", "LONG", "INTERVAL") - col_data: The full column data dictionary from manifest (for accessing type_interval_type) - - Returns: - Normalized type name matching Thrift conventions - """ - # Early return if type doesn't need mapping - if type_name not in SEA_TO_THRIFT_TYPE_MAP: - return type_name - - normalized_type = SEA_TO_THRIFT_TYPE_MAP[type_name] - - # Special handling for interval types - if type_name == "INTERVAL": - type_interval_type = col_data.get("type_interval_type") - if type_interval_type: - return ( - "INTERVAL_YEAR_MONTH" - if any(t in type_interval_type.upper() for t in ["YEAR", "MONTH"]) - else "INTERVAL_DAY_TIME" - ) - - return normalized_type diff --git a/src/databricks/sql/backend/types.py b/src/databricks/sql/backend/types.py index 402cf6beb..d727343f5 100644 --- a/src/databricks/sql/backend/types.py +++ b/src/databricks/sql/backend/types.py @@ -72,7 +72,7 @@ def from_thrift_state( # Function-local import: this classmethod is only ever called from the # Thrift backend, so deferring the import keeps ``thrift`` out of the - # SEA/kernel load path. + # kernel load path. from databricks.sql.thrift_api.TCLIService import ttypes if state in ( @@ -97,26 +97,6 @@ def from_thrift_state( else: return None - @classmethod - def from_sea_state(cls, state: str) -> Optional["CommandState"]: - """ - Map SEA state string to CommandState enum. - Args: - state: SEA state string - Returns: - CommandState: The corresponding CommandState enum value - """ - state_mapping = { - "PENDING": cls.PENDING, - "RUNNING": cls.RUNNING, - "SUCCEEDED": cls.SUCCEEDED, - "FAILED": cls.FAILED, - "CLOSED": cls.CLOSED, - "CANCELED": cls.CANCELLED, - } - - return state_mapping.get(state, None) - class BackendType(Enum): """ @@ -124,15 +104,15 @@ class BackendType(Enum): """ THRIFT = "thrift" - SEA = "sea" + KERNEL = "kernel" class SessionId: """ - A normalized session identifier that works with both Thrift and SEA backends. + A normalized session identifier that works with both Thrift and kernel backends. This class abstracts away the differences between Thrift's TSessionHandle and - SEA's session ID string, providing a consistent interface for the connector. + the kernel's session ID string, providing a consistent interface for the connector. """ def __init__( @@ -146,7 +126,7 @@ def __init__( Initialize a SessionId. Args: - backend_type: The type of backend (THRIFT or SEA) + backend_type: The type of backend (THRIFT or KERNEL) guid: The primary identifier for the session secret: The secret part of the identifier (only used for Thrift) properties: Additional information about the session @@ -161,14 +141,14 @@ def __str__(self) -> str: """ Return a string representation of the SessionId. - For SEA backend, returns the guid. + For kernel backend, returns the guid. For Thrift backend, returns a format like "guid|secret". Returns: A string representation of the session ID """ - if self.backend_type == BackendType.SEA: + if self.backend_type == BackendType.KERNEL: return str(self.guid) elif self.backend_type == BackendType.THRIFT: secret_hex = ( @@ -207,20 +187,20 @@ def from_thrift_handle( return cls(BackendType.THRIFT, guid_bytes, secret_bytes, properties) @classmethod - def from_sea_session_id( + def from_kernel_session_id( cls, session_id: str, properties: Optional[Dict[str, Any]] = None ): """ - Create a SessionId from a SEA session ID. + Create a SessionId from a kernel session ID. Args: - session_id: The SEA session ID string + session_id: The kernel session ID string Returns: A SessionId instance """ - return cls(BackendType.SEA, session_id, properties=properties) + return cls(BackendType.KERNEL, session_id, properties=properties) def to_thrift_handle(self): """ @@ -241,19 +221,6 @@ def to_thrift_handle(self): sessionId=handle_identifier, serverProtocolVersion=server_protocol_version ) - def to_sea_session_id(self): - """ - Get the SEA session ID string. - - Returns: - The session ID string or None if this is not a SEA session ID - """ - - if self.backend_type != BackendType.SEA: - return None - - return self.guid - @property def hex_guid(self) -> str: """ @@ -275,7 +242,7 @@ def protocol_version(self): Returns: The server protocol version or None if it does not exist - It is not expected to exist for SEA sessions. + It is not expected to exist for kernel sessions. """ return self.properties.get("serverProtocolVersion") @@ -283,10 +250,10 @@ def protocol_version(self): class CommandId: """ - A normalized command identifier that works with both Thrift and SEA backends. + A normalized command identifier that works with both Thrift and kernel backends. This class abstracts away the differences between Thrift's TOperationHandle and - SEA's statement ID string, providing a consistent interface for the connector. + the kernel's statement ID string, providing a consistent interface for the connector. """ def __init__( @@ -302,7 +269,7 @@ def __init__( Initialize a CommandId. Args: - backend_type: The type of backend (THRIFT or SEA) + backend_type: The type of backend (THRIFT or KERNEL) guid: The primary identifier for the command secret: The secret part of the identifier (only used for Thrift) operation_type: The operation type (only used for Thrift) @@ -321,14 +288,14 @@ def __str__(self) -> str: """ Return a string representation of the CommandId. - For SEA backend, returns the guid. + For kernel backend, returns the guid. For Thrift backend, returns a format like "guid|secret". Returns: A string representation of the command ID """ - if self.backend_type == BackendType.SEA: + if self.backend_type == BackendType.KERNEL: return str(self.guid) elif self.backend_type == BackendType.THRIFT: secret_hex = ( @@ -367,18 +334,18 @@ def from_thrift_handle(cls, operation_handle): ) @classmethod - def from_sea_statement_id(cls, statement_id: str): + def from_kernel_statement_id(cls, statement_id: str): """ - Create a CommandId from a SEA statement ID. + Create a CommandId from a kernel statement ID. Args: - statement_id: The SEA statement ID string + statement_id: The kernel statement ID string Returns: A CommandId instance """ - return cls(BackendType.SEA, statement_id) + return cls(BackendType.KERNEL, statement_id) def to_thrift_handle(self): """ @@ -401,19 +368,6 @@ def to_thrift_handle(self): modifiedRowCount=self.modified_row_count, ) - def to_sea_statement_id(self): - """ - Get the SEA statement ID string. - - Returns: - The statement ID string or None if this is not a SEA statement ID - """ - - if self.backend_type != BackendType.SEA: - return None - - return self.guid - def to_hex_guid(self) -> str: """ Get a hexadecimal string representation of the command ID. diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 940b12775..b9989a022 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -81,7 +81,7 @@ # annotations``). ``get_protocol_version`` and ``_prepare_native_parameters`` # are typed with these Thrift-generated types, but the Thrift backend and # its result set are imported lazily (only on the Thrift connect path), so - # importing this module -- and connecting with the SEA or kernel backend -- + # importing this module -- and connecting with the kernel backend -- # never imports the Apache Thrift ``thrift`` package. See # ``test_lazy_thrift_import``. from databricks.sql.thrift_api.TCLIService.ttypes import ( @@ -127,7 +127,7 @@ def __getattr__(name: str) -> Any: ``from databricks.sql.client import `` keeps working for any existing caller -- and the ``patch("databricks.sql.client.ThriftDatabricksClient")`` test seam is preserved -- without importing the Apache Thrift ``thrift`` - package at module load, which is what keeps the SEA/kernel connect path + package at module load, which is what keeps the kernel connect path Thrift-free (see ``test_lazy_thrift_import``). """ if name == "ThriftDatabricksClient": @@ -164,15 +164,6 @@ def __init__( Connect to a Databricks SQL endpoint or a Databricks cluster. Parameters: - :param use_sea: `bool`, optional (default is False) - Use the native pure-Python SEA backend instead of - the Thrift backend. - - Deprecated and incomplete — this backend has feature - gaps (e.g. it does not support positional ``?`` - parameter binding) and is slated for removal. For a - SEA-native connection use ``use_kernel=True`` instead, - which is the supported path. :param use_kernel: `bool`, optional (default is False) Route the connection through the Rust kernel (``databricks-sql-kernel`` via PyO3), a SEA-native @@ -183,16 +174,16 @@ def __init__( is a no-op and ``use_kernel=True`` raises a clear ``ImportError``. Supports PAT, OAuth M2M, and OAuth U2M auth, and native (positional and named) parameter - binding. Mutually exclusive with ``use_sea``. + binding. Telemetry note: on the kernel path telemetry is strictly opt-in. An explicit ``enable_telemetry=True`` turns telemetry on unconditionally; it is NOT gated by the server-side ``enableTelemetryForPythonDriver`` - feature flag that applies on the Thrift/SEA paths. When + feature flag that applies on the Thrift path. When ``enable_telemetry`` is unset the kernel owns the enable decision. This is an intentional divergence from the - Thrift/SEA paths, where an explicit ``True`` can still + Thrift path, where an explicit ``True`` can still be suppressed by the feature flag. :param enable_geospatial_support: `bool`, optional (default is False) Kernel backend only. Controls the public representation of @@ -201,8 +192,6 @@ def __init__( strings (for example ``"SRID=4326;POINT(1 2)"``). The conversion is local to the kernel/driver and this option is never sent to the SQL Execution API. - :param use_hybrid_disposition: `bool`, optional (default is False) - Use the hybrid disposition instead of the inline disposition. :param server_hostname: Databricks instance host name. :param http_path: Http path either to a DBSQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef) or to a DBR interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123) @@ -378,6 +367,26 @@ def read(self) -> Optional[OAuthToken]: # use_cloud_fetch # Enable use of cloud fetch to extract large query results in parallel via cloud storage + if kwargs.pop("use_sea", False): + raise ValueError( + "The SEA backend has been removed. Use use_kernel=True and " + "install 'databricks-sql-connector[kernel]'." + ) + + if kwargs.pop("use_hybrid_disposition", False): + logger.warning( + "Parameter 'use_hybrid_disposition' was removed with the SEA " + "backend and is ignored." + ) + + if "max_connections" in kwargs: + kwargs.pop("max_connections") + logger.warning( + "Parameter 'max_connections' was removed with the SEA backend " + "and is ignored. On kernel >= 1.1.0, use '_pool_maxsize' for the " + "kernel HTTP pool; it does not configure the Thrift query transport." + ) + logger.debug( "Connection.__init__(server_hostname=%s, http_path=%s)", server_hostname, @@ -499,11 +508,7 @@ def read(self) -> Optional[OAuthToken]: driver_connection_params = DriverConnectionParameters( http_path=http_path, - mode=( - DatabricksClientType.SEA - if self.session.use_sea - else DatabricksClientType.THRIFT - ), + mode=DatabricksClientType.THRIFT, host_info=HostDetails(host_url=server_hostname, port=self.session.port), auth_mech=TelemetryHelper.get_auth_mechanism(self.session.auth_provider), auth_flow=TelemetryHelper.get_auth_flow(self.session.auth_provider), @@ -520,7 +525,6 @@ def read(self) -> Optional[OAuthToken]: use_system_trust_store=True, # Python uses system SSL by default enable_arrow=pyarrow is not None, enable_direct_results=True, # Always enabled in Python - enable_sea_hybrid_results=kwargs.get("use_hybrid_disposition", False), http_connection_pool_size=kwargs.get("pool_maxsize", None), rows_fetched_per_block=DEFAULT_ARRAY_SIZE, async_poll_interval_millis=2000, # Default polling interval @@ -551,10 +555,9 @@ def _open_session_with_reyden_fallback( when a Reyden / Real-Time warehouse rejects the default Thrift protocol. Reyden warehouses reject a Thrift ``OpenSession`` (SQLSTATE ``KP001``); - the kernel (SEA) backend is the supported path. Auto-recovery applies - only when the caller did not pick a backend explicitly (neither - ``use_kernel`` nor ``use_sea``). On a rejection the warehouse is - remembered so later connections skip the doomed Thrift attempt. + the kernel backend is the supported path. Auto-recovery applies only + when the caller did not set ``use_kernel=True``. On a rejection the + warehouse is remembered so later connections skip the Thrift attempt. """ def build_session(session_kwargs: dict) -> Session: @@ -598,10 +601,7 @@ def kernel_recovery_kwargs() -> dict: # An explicit backend choice is always honored — auto-recovery engages # only on the default (Thrift) path. - explicit_backend = kwargs.get("use_kernel", False) or kwargs.get( - "use_sea", False - ) - if explicit_backend: + if kwargs.get("use_kernel", False): return build_session(kwargs) warehouse_id = extract_warehouse_id(http_path) @@ -671,7 +671,7 @@ def __exit__(self, exc_type, exc_value, traceback): self.close() def __del__(self): - if self.open: + if hasattr(self, "session") and self.open: logger.debug( "Closing unclosed connection for session " "{}".format(self.get_session_id_hex()) @@ -1954,9 +1954,7 @@ def cancel(self) -> None: # backend) may still have a server statement in flight. Such a # backend exposes ``cancel_running_cursor(cursor)`` -> bool to # cancel it; it returns True if something was actually - # cancelled. Opt-in via getattr so the Thrift / SEA backends - # (which set ``active_command_id`` before blocking) are - # unaffected. + # cancelled. Thrift sets ``active_command_id`` before blocking. cancel_running_cursor = getattr(self.backend, "cancel_running_cursor", None) if cancel_running_cursor is not None and cancel_running_cursor(self): return diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index ece7406b7..b02cec567 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: # Type-annotation-only import; see the note in downloader.py. Keeping the - # ``thrift`` package out of this module lets the SEA/kernel backends use the + # ``thrift`` package out of this module lets the kernel backend use the # cloud-fetch download manager without importing Apache Thrift. from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 295f147cd..926ce48d0 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -17,7 +17,7 @@ # Imported for type annotations only. ``from __future__ import annotations`` # makes every annotation a string, so this import is never evaluated at # runtime -- which keeps the (Apache Thrift) ``thrift`` package out of the - # cloud-fetch code path used by the SEA and kernel backends. See the + # cloud-fetch code path used by the kernel backend. See the # ``test_lazy_thrift_import`` regression test. from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink diff --git a/src/databricks/sql/parameters/native.py b/src/databricks/sql/parameters/native.py index d0d811940..27c92f730 100644 --- a/src/databricks/sql/parameters/native.py +++ b/src/databricks/sql/parameters/native.py @@ -12,7 +12,7 @@ # these out of the runtime import graph; the ``TSparkParameter*`` objects # are constructed via function-local imports inside the ``as_tspark_param`` # / ``_tspark_*`` helpers below, which only run on the Thrift execute path. - # This keeps the Apache Thrift ``thrift`` package out of the SEA/kernel + # This keeps the Apache Thrift ``thrift`` package out of the kernel # load path (see ``test_lazy_thrift_import``). from databricks.sql.thrift_api.TCLIService.ttypes import ( TSparkParameter, @@ -26,7 +26,7 @@ # Thrift-generated ``ttypes``. They are resolved lazily via ``__getattr__`` # (PEP 562) so ``from databricks.sql.parameters.native import TSparkParameter`` # keeps working for existing callers without importing the Apache Thrift -# ``thrift`` package at module load -- which is what keeps the SEA/kernel path +# ``thrift`` package at module load -- which is what keeps the kernel path # Thrift-free. See ``test_lazy_thrift_import``. _LAZY_THRIFT_REEXPORTS = frozenset( {"TSparkParameter", "TSparkParameterValue", "TSparkParameterValueArg"} diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 365b06f8f..78a560b2f 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -1,7 +1,7 @@ import logging import re from functools import cached_property -from typing import Dict, Tuple, List, Optional, Any, Type, TYPE_CHECKING +from typing import Dict, Tuple, List, Optional, Any, TYPE_CHECKING from databricks.sql.types import SSLOptions from databricks.sql.auth.auth import get_python_sql_connector_auth_provider @@ -11,7 +11,7 @@ from databricks.sql import __version__ from databricks.sql import USER_AGENT_NAME from databricks.sql.backend.databricks_client import DatabricksClient -from databricks.sql.backend.types import SessionId, BackendType +from databricks.sql.backend.types import SessionId from databricks.sql.common.unified_http_client import UnifiedHttpClient from databricks.sql.common.agent import detect as detect_agent from databricks.sql.common.feature_flag import ( @@ -22,7 +22,6 @@ if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient - from databricks.sql.backend.sea.backend import SeaDatabricksClient logger = logging.getLogger(__name__) @@ -30,7 +29,7 @@ # The backend client classes are resolved lazily (PEP 562) rather than imported # at module load. ``ThriftDatabricksClient`` pulls in the Apache Thrift # ``thrift`` package, so importing it eagerly would drag ``thrift`` into the -# SEA and kernel connect paths (breaking build systems that ship their own +# kernel connect path (breaking build systems that ship their own # ``thrift``, e.g. Buck). Exposing them as module attributes -- rather than as # function-local imports inside ``open`` -- also keeps the long-standing test # seam ``patch("databricks.sql.session.ThriftDatabricksClient")`` working. @@ -40,10 +39,6 @@ def __getattr__(name: str) -> Any: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient return ThriftDatabricksClient - if name == "SeaDatabricksClient": - from databricks.sql.backend.sea.backend import SeaDatabricksClient - - return SeaDatabricksClient raise AttributeError(f"module {__name__!r} has no attribute {name!r}") @@ -103,8 +98,8 @@ def __init__( # Extract workspace context from http_path for SPOG routing. # On SPOG hosts, the http_path can contain either ?o= or an - # all-purpose-compute /o// path segment. For SEA, telemetry, - # and feature flags, we inject x-databricks-org-id as an HTTP header. + # all-purpose-compute /o// path segment. For non-Thrift + # requests, we inject x-databricks-org-id as an HTTP header. self._spog_headers = self._extract_spog_headers(http_path, all_headers) if self._spog_headers: all_headers = all_headers + list(self._spog_headers.items()) @@ -138,7 +133,7 @@ def __init__( # So for ``use_kernel`` we hand the bridge only a minimal PAT # provider when an ``access_token`` is present, and ``None`` # otherwise (OAuth M2M/U2M resolve purely from the raw kwargs - # the bridge reads). The Thrift / SEA backends are unchanged. + # the bridge reads). The Thrift backend is unchanged. self.use_kernel = kwargs.get("use_kernel", False) if self.use_kernel: access_token = kwargs.get("access_token") @@ -184,13 +179,6 @@ def _create_backend( kwargs: dict, ) -> DatabricksClient: """Create and return the appropriate backend client.""" - self.use_sea = kwargs.get("use_sea", False) - - if self.use_kernel and self.use_sea: - raise ValueError( - "use_kernel and use_sea are mutually exclusive — pick one." - ) - if self.use_kernel: # Lazy import so the connector doesn't ImportError at # startup when the kernel wheel isn't installed — the @@ -205,8 +193,7 @@ def _create_backend( # NOT build the connector's own OAuth provider (see __init__ # above), so these raw kwargs are the only source of the # OAuth client id/secret and optional federation client id. - # These are kernel-only; the Thrift / SEA backends are - # unaffected. + # These are kernel-only; the Thrift backend is unaffected. kernel_auth_options = { "auth_type": kwargs.get("auth_type"), "oauth_client_id": kwargs.get("oauth_client_id"), @@ -244,7 +231,7 @@ def _create_backend( # tenant and workspace-resource-id are optional (the kernel # auto-discovers the tenant; the resource id adds the # management-token resource-id header for an RBAC-only SP). - # Kernel-only; Thrift / SEA are unaffected. + # Kernel-only; Thrift is unaffected. "azure_client_id": kwargs.get("azure_client_id"), "azure_client_secret": kwargs.get("azure_client_secret"), "azure_tenant_id": kwargs.get("azure_tenant_id"), @@ -258,7 +245,7 @@ def _create_backend( # are passed; `_retry_delay_default` is intentionally omitted # (the kernel's no-Retry-After backoff is exponential from # its min-wait, so a flat default delay has no equivalent). - # Kernel-only; Thrift / SEA are unaffected. + # Kernel-only; Thrift is unaffected. kernel_retry_options = { "retry_delay_min": kwargs.get("_retry_delay_min"), "retry_delay_max": kwargs.get("_retry_delay_max"), @@ -272,13 +259,13 @@ def _create_backend( # Forward the binding/runtime identity and telemetry knobs # added by kernel telemetry phase 7. Python-side telemetry # still owns feature-flag evaluation and event export for the - # Thrift/SEA paths; the kernel path needs the same driver + # Thrift path; the kernel path needs the same driver # identity at Session construction time so kernel-owned # telemetry can populate its system configuration. kernel_telemetry_options = { # Preserve the caller's explicit telemetry choice. When unset, # leave it as None so the kernel owns the enable decision -- - # this is an INTENTIONAL divergence from the Thrift/SEA path, + # this is an INTENTIONAL divergence from the Thrift path, # not an oversight. There, an unset enable_telemetry defaults to # True (client.py) but only actually emits when the # `enableTelemetryForPythonDriver` server feature flag is on @@ -294,7 +281,7 @@ def _create_backend( # Match the connector's default batch size (client.py forwards # the same TelemetryClientFactory.DEFAULT_BATCH_SIZE fallback) # so an unset telemetry_batch_size resolves to the same value - # on the kernel path as on the Thrift/SEA path, rather than + # on the kernel path as on the Thrift path, rather than # letting the kernel silently pick its own internal default. "telemetry_batch_size": kwargs.get( "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE @@ -304,14 +291,14 @@ def _create_backend( # mirroring the enable_telemetry handling above rather than # forcing a default. This is deliberately NOT defaulted to True: # although ClientContext's signature default is True - # (auth/common.py:55), the Thrift/SEA path never reaches it -- + # (auth/common.py:55), the Thrift path never reaches it -- # build_client_context (utils.py:1018) always passes # _telemetry_circuit_breaker_enabled explicitly (None when the # caller leaves it unset), and ClientContext coerces it with # bool(None) -> False (auth/common.py:89). So the *effective* - # Thrift/SEA default when unset is False, not True; forwarding + # Thrift default when unset is False, not True; forwarding # True here would turn the circuit breaker on for an - # unconfigured connection while Thrift/SEA leaves it off. + # unconfigured connection while Thrift leaves it off. # Passing None instead defers to the kernel's own default; # only an explicit caller value overrides it. "telemetry_circuit_breaker_enabled": kwargs.get( @@ -339,20 +326,10 @@ def _create_backend( escape_catalog_name=kwargs.get("escape_catalog_name", False), ) - # These reference the lazily-resolved module attributes defined via - # ``__getattr__`` above (or a test's ``patch(...)`` of them), so neither - # backend class -- and in particular ``thrift`` -- is imported until the - # branch that actually needs it runs. + # Resolve the Thrift client lazily to keep it off the kernel import path. import databricks.sql.session as _session_module - databricks_client_class: Type[DatabricksClient] - if self.use_sea: - logger.debug("Creating SEA backend client") - databricks_client_class = _session_module.SeaDatabricksClient - else: - logger.debug("Creating Thrift backend client") - databricks_client_class = _session_module.ThriftDatabricksClient - + logger.debug("Creating Thrift backend client") common_args = { "server_hostname": server_hostname, "port": self.port, @@ -364,7 +341,7 @@ def _create_backend( "_use_arrow_native_complex_types": _use_arrow_native_complex_types, **kwargs, } - return databricks_client_class(**common_args) + return _session_module.ThriftDatabricksClient(**common_args) # All-purpose-compute Thrift http_path: # [/]sql/protocolv1/o//[/...][?...] @@ -386,7 +363,7 @@ def _extract_spog_headers(http_path, existing_headers): both. Returns an empty dict when no workspace ID can be determined. On SPOG (Custom URL) hosts this header is required for non-Thrift - endpoints — telemetry, feature flags, SEA — to be routed to the right + endpoints — telemetry and feature flags — to be routed to the right workspace. Without it, PoPP falls back to default routing and workspace-scoped requests are redirected to ``/login``. """ @@ -458,7 +435,7 @@ def server_parameterized_queries_enabled(protocolVersion): # Function-local import: the protocol-version constant lives in the # Thrift-generated ttypes, but this check only ever runs with a # Thrift-negotiated protocol version, so deferring keeps ``thrift`` out - # of the SEA/kernel load path. + # of the kernel load path. from databricks.sql.thrift_api.TCLIService import ttypes if ( diff --git a/src/databricks/sql/telemetry/models/enums.py b/src/databricks/sql/telemetry/models/enums.py index dd8f26eb0..d183f2e3d 100644 --- a/src/databricks/sql/telemetry/models/enums.py +++ b/src/databricks/sql/telemetry/models/enums.py @@ -16,7 +16,6 @@ class AuthMech(Enum): class DatabricksClientType(Enum): - SEA = "SEA" THRIFT = "THRIFT" diff --git a/src/databricks/sql/telemetry/models/event.py b/src/databricks/sql/telemetry/models/event.py index 4d5a45038..ff02268b0 100644 --- a/src/databricks/sql/telemetry/models/event.py +++ b/src/databricks/sql/telemetry/models/event.py @@ -50,7 +50,6 @@ class DriverConnectionParameters(JsonSerializableMixin): use_system_trust_store (bool): Whether system trust store is used enable_arrow (bool): Whether Arrow format is enabled enable_direct_results (bool): Whether direct results are enabled - enable_sea_hybrid_results (bool): Whether SEA hybrid results are enabled http_connection_pool_size (int): HTTP connection pool size rows_fetched_per_block (int): Number of rows fetched per block async_poll_interval_millis (int): Async polling interval in milliseconds @@ -78,7 +77,6 @@ class DriverConnectionParameters(JsonSerializableMixin): use_system_trust_store: Optional[bool] = None enable_arrow: Optional[bool] = None enable_direct_results: Optional[bool] = None - enable_sea_hybrid_results: Optional[bool] = None http_connection_pool_size: Optional[int] = None rows_fetched_per_block: Optional[int] = None async_poll_interval_millis: Optional[int] = None diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index fed2001cb..201e38cb4 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -757,7 +757,7 @@ def connection_failure_log( ) telemetry_client._driver_connection_params = DriverConnectionParameters( http_path=http_path, - mode=DatabricksClientType.THRIFT, # TODO: Add SEA mode + mode=DatabricksClientType.THRIFT, host_info=HostDetails(host_url=host_url, port=port), ) telemetry_client._user_agent = user_agent diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 4e82370b9..554e9aad6 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -27,7 +27,7 @@ # annotations``). ``TSparkRowSetType`` is used at runtime only inside # ``ThriftResultSetQueueFactory.build_queue``, which is a Thrift-only code # path and imports it locally. This keeps the Apache Thrift ``thrift`` - # package out of the SEA/kernel load path (see ``test_lazy_thrift_import``). + # package out of the kernel load path (see ``test_lazy_thrift_import``). from databricks.sql.thrift_api.TCLIService.ttypes import ( TRowSet, TSparkArrowResultLink, @@ -116,7 +116,7 @@ def build_queue( """ # Function-local import: this factory is only used by the Thrift - # backend, so deferring keeps ``thrift`` out of the SEA/kernel path. + # backend, so deferring keeps ``thrift`` out of the kernel path. from databricks.sql.thrift_api.TCLIService.ttypes import TSparkRowSetType if row_set_type == TSparkRowSetType.ARROW_BASED_SET: diff --git a/tests/e2e/common/retry_test_mixins.py b/tests/e2e/common/retry_test_mixins.py index bd4d669b4..195bafa53 100755 --- a/tests/e2e/common/retry_test_mixins.py +++ b/tests/e2e/common/retry_test_mixins.py @@ -77,6 +77,7 @@ def _noop_export(self, event): return None try: + def _noop_send_with_unified_client(self, *args, **kwargs): return None @@ -84,9 +85,7 @@ def _noop_send_with_unified_client(self, *args, **kwargs): TelemetryClientFactory, "initialize_telemetry_client", staticmethod(_install_noop), - ), patch.object( - TelemetryClient, "_send_telemetry", _noop_send - ), patch.object( + ), patch.object(TelemetryClient, "_send_telemetry", _noop_send), patch.object( TelemetryClient, "_export_event", _noop_export ), patch.object( # Backstop layer 4: no-op the actual HTTP boundary. _send_telemetry @@ -109,31 +108,17 @@ def _noop_send_with_unified_client(self, *args, **kwargs): class Client429ResponseMixin: - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_client_should_retry_automatically_when_getting_429(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_client_should_retry_automatically_when_getting_429(self): + with self.cursor({}) as cursor: for _ in range(10): cursor.execute("SELECT 1") rows = cursor.fetchall() self.assertEqual(len(rows), 1) self.assertEqual(rows[0][0], 1) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_client_should_not_retry_429_if_RateLimitRetry_is_0(self, extra_params): + def test_client_should_not_retry_429_if_RateLimitRetry_is_0(self): with pytest.raises(self.error_type) as cm: - extra_params = {**extra_params, **self.conf_to_disable_rate_limit_retries} + extra_params = {**self.conf_to_disable_rate_limit_retries} with self.cursor(extra_params) as cursor: for _ in range(10): cursor.execute("SELECT 1") @@ -153,48 +138,32 @@ def test_client_should_not_retry_429_if_RateLimitRetry_is_0(self, extra_params): class Client503ResponseMixin: - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_wait_cluster_startup(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_wait_cluster_startup(self): + with self.cursor({}) as cursor: cursor.execute("SELECT 1") cursor.fetchall() - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def _test_retry_disabled_with_message( - self, error_msg_substring, exception_type, extra_params - ): + def _test_retry_disabled_with_message(self, error_msg_substring, exception_type): with pytest.raises(exception_type) as cm: - with self.connection( - self.conf_to_disable_temporarily_unavailable_retries, extra_params - ): + with self.connection(self.conf_to_disable_temporarily_unavailable_retries): pass assert error_msg_substring in str(cm.exception) class SimpleHttpResponse: """A simple HTTP response mock that works with both urllib3 v1.x and v2.x""" - - def __init__(self, status: int, headers: dict, redirect_location: Optional[str] = None): + + def __init__( + self, status: int, headers: dict, redirect_location: Optional[str] = None + ): # Import the correct HTTP message type that urllib3 v1.x expects try: from http.client import HTTPMessage except ImportError: from httplib import HTTPMessage - + self.status = status - # Create proper HTTPMessage for urllib3 v1.x compatibility + # Create proper HTTPMessage for urllib3 v1.x compatibility self.headers = HTTPMessage() for key, value in headers.items(): self.headers[key] = str(value) @@ -211,41 +180,41 @@ def __init__(self, status: int, headers: dict, redirect_location: Optional[str] self._body = b"" self._fp = io.BytesIO(self._body) self._url = "https://example.com" - + def get_redirect_location(self, *args, **kwargs): """Return the redirect location or False""" return False if self._redirect_location is None else self._redirect_location - + def read(self, amt=None): """Mock read method for file-like behavior""" return self._body - + def close(self): """Mock close method""" pass - + def drain_conn(self): """Mock drain_conn method for urllib3 v2.x""" pass - + def isclosed(self): """Mock isclosed method for urllib3 v1.x""" return False - + def release_conn(self): """Mock release_conn method for thrift HTTP client""" pass - + @property def data(self): """Mock data property for urllib3 v2.x""" return self._body - + @property def url(self): """Mock url property""" return self._url - + @url.setter def url(self, value): """Mock url setter""" @@ -257,7 +226,7 @@ def mocked_server_response( status: int = 200, headers: dict = {}, redirect_location: Optional[str] = None ): """Context manager for patching urllib3 responses with version compatibility""" - + mock_response = SimpleHttpResponse(status, headers, redirect_location) with patch("urllib3.connectionpool.HTTPSConnectionPool._get_conn") as getconn_mock: @@ -283,7 +252,7 @@ def mock_sequential_server_responses(responses: List[dict]): SimpleHttpResponse( status=resp["status"], headers=resp["headers"], - redirect_location=resp["redirect_location"] + redirect_location=resp["redirect_location"], ) for resp in responses ] @@ -324,17 +293,8 @@ class PySQLRetryTestsMixin: "enable_telemetry": False, } - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry") - def test_retry_urllib3_settings_are_honored( - self, mock_send_telemetry, extra_params - ): + def test_retry_urllib3_settings_are_honored(self, mock_send_telemetry): """Databricks overrides some of urllib3's configuration. This tests confirms that what configuration we DON'T override is preserved in urllib3's internals """ @@ -354,36 +314,22 @@ def test_retry_urllib3_settings_are_honored( assert rp.read == 11 assert rp.redirect == 12 - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry") - def test_oserror_retries(self, mock_send_telemetry, extra_params): + def test_oserror_retries(self, mock_send_telemetry): """If a network error occurs during make_request, the request is retried according to policy""" with _isolated_from_telemetry(), patch( "urllib3.connectionpool.HTTPSConnectionPool._validate_conn", ) as mock_validate_conn: mock_validate_conn.side_effect = OSError("Some arbitrary network error") with pytest.raises(MaxRetryError) as cm: - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: pass assert mock_validate_conn.call_count == 6 - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry") - def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params): + def test_retry_max_count_not_exceeded(self, mock_send_telemetry): """GIVEN the max_attempts_count is 5 WHEN the server sends nothing but 429 responses THEN the connector issues six request (original plus five retries) @@ -393,20 +339,13 @@ def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params): status=429, headers={"Retry-After": "0"} ) as mock_obj: with pytest.raises(MaxRetryError) as cm: - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: pass assert mock_obj.return_value.getresponse.call_count == 6 - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry") - def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params): + def test_retry_exponential_backoff(self, mock_send_telemetry): """GIVEN the retry policy is configured for reasonable exponential backoff WHEN the server sends nothing but 429 responses with retry-afters THEN the connector will use those retry-afters values as floor @@ -419,7 +358,7 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params): status=429, headers={"Retry-After": "8"} ) as mock_obj: with pytest.raises(RequestError) as cm: - extra_params = {**extra_params, **retry_policy} + extra_params = {**retry_policy} with self.connection(extra_params=extra_params) as conn: pass @@ -436,33 +375,19 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params): # Should be less than 26, but this is a safe margin for CI/CD slowness assert duration < 30 - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_max_duration_not_exceeded(self, extra_params): + def test_retry_max_duration_not_exceeded(self): """GIVEN the max attempt duration of 10 seconds WHEN the server sends a Retry-After header of 60 seconds THEN the connector raises a MaxRetryDurationError """ with mocked_server_response(status=429, headers={"Retry-After": "60"}): with pytest.raises(RequestError) as cm: - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: pass assert isinstance(cm.value.args[1], MaxRetryDurationError) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_abort_non_recoverable_error(self, extra_params): + def test_retry_abort_non_recoverable_error(self): """GIVEN the server returns a code 501 WHEN the connector receives this response THEN nothing is retried and an exception is raised @@ -471,24 +396,17 @@ def test_retry_abort_non_recoverable_error(self, extra_params): # Code 501 is a Not Implemented error with mocked_server_response(status=501): with pytest.raises(RequestError) as cm: - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: pass assert isinstance(cm.value.args[1], NonRecoverableNetworkError) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_abort_unsafe_execute_statement_retry_condition(self, extra_params): + def test_retry_abort_unsafe_execute_statement_retry_condition(self): """GIVEN the server sends a code other than 429 or 503 WHEN the connector sent an ExecuteStatement command THEN nothing is retried because it's idempotent """ - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: with conn.cursor() as cursor: # Code 502 is a Bad Gateway, which we commonly see in production under heavy load @@ -497,14 +415,7 @@ def test_retry_abort_unsafe_execute_statement_retry_condition(self, extra_params cursor.execute("Not a real query") assert isinstance(cm.value.args[1], UnsafeToRetryError) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_dangerous_codes(self, extra_params): + def test_retry_dangerous_codes(self): """GIVEN the server sends a dangerous code and the user forced this to be retryable WHEN the connector sent an ExecuteStatement command THEN the command is retried @@ -520,7 +431,7 @@ def test_retry_dangerous_codes(self, extra_params): } # Prove that these codes are not retried by default - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: with conn.cursor() as cursor: for dangerous_code in DANGEROUS_CODES: @@ -539,14 +450,7 @@ def test_retry_dangerous_codes(self, extra_params): with pytest.raises(MaxRetryError) as cm: cursor.execute("Not a real query") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_safe_execute_statement_retry_condition(self, extra_params): + def test_retry_safe_execute_statement_retry_condition(self): """GIVEN the server sends either code 429 or 503 WHEN the connector sent an ExecuteStatement command THEN the request is retried because these are idempotent @@ -559,7 +463,6 @@ def test_retry_safe_execute_statement_retry_condition(self, extra_params): with self.connection( extra_params={ - **extra_params, **self._retry_policy, "_retry_stop_after_attempts_count": 1, } @@ -571,14 +474,7 @@ def test_retry_safe_execute_statement_retry_condition(self, extra_params): cursor.execute("This query never reaches the server") assert mock_obj.return_value.getresponse.call_count == 2 - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_abort_close_session_on_404(self, extra_params, caplog): + def test_retry_abort_close_session_on_404(self, caplog): """GIVEN the connector sends a CloseSession command WHEN server sends a 404 (which is not retried since commit 41b28159) THEN nothing is retried because 404 is globally non-retryable @@ -590,20 +486,13 @@ def test_retry_abort_close_session_on_404(self, extra_params, caplog): {"status": 404, "headers": {}, "redirect_location": None}, ] - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: with mock_sequential_server_responses(responses): # Should not raise an exception, the error is caught internally conn.close() - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_abort_close_operation_on_404(self, extra_params, caplog): + def test_retry_abort_close_operation_on_404(self, caplog): """GIVEN the connector sends a CancelOperation command WHEN server sends a 404 (which is not retried since commit 41b28159) THEN nothing is retried because 404 is globally non-retryable @@ -615,7 +504,7 @@ def test_retry_abort_close_operation_on_404(self, extra_params, caplog): {"status": 404, "headers": {}, "redirect_location": None}, ] - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: with conn.cursor() as curs: with patch( @@ -629,17 +518,8 @@ def test_retry_abort_close_operation_on_404(self, extra_params, caplog): # Should not raise an exception, the error is caught internally curs.close() - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry") - def test_3xx_redirect_codes_are_not_retried( - self, mock_send_telemetry, extra_params - ): + def test_3xx_redirect_codes_are_not_retried(self, mock_send_telemetry): """GIVEN the connector is configured with a custom max_redirects WHEN the DatabricksRetryPolicy receives a 302 redirect THEN the connector does not retry since 3xx codes are not retried per policy @@ -648,15 +528,14 @@ def test_3xx_redirect_codes_are_not_retried( max_redirects, expected_call_count = 1, 1 # Code 302 is a redirect, but 3xx codes are not retried per policy - # Note: We don't set redirect_location because that would cause urllib3 v2.x + # Note: We don't set redirect_location because that would cause urllib3 v2.x # to follow redirects internally, bypassing our retry policy test - with mocked_server_response( - status=302, redirect_location=None - ) as mock_obj: - with pytest.raises(RequestError): # Should get RequestError, not MaxRetryError + with mocked_server_response(status=302, redirect_location=None) as mock_obj: + with pytest.raises( + RequestError + ): # Should get RequestError, not MaxRetryError with self.connection( extra_params={ - **extra_params, **self._retry_policy, "_retry_max_redirects": max_redirects, } @@ -665,16 +544,9 @@ def test_3xx_redirect_codes_are_not_retried( # Total call count should be 1 (original only, no retries for 3xx codes) assert mock_obj.return_value.getresponse.call_count == expected_call_count - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry") def test_3xx_codes_not_retried_regardless_of_max_redirects_setting( - self, mock_send_telemetry, extra_params + self, mock_send_telemetry ): """GIVEN the connector is configured without a custom max_redirects WHEN the DatabricksRetryPolicy receives a 302 redirect @@ -684,15 +556,14 @@ def test_3xx_codes_not_retried_regardless_of_max_redirects_setting( according to the DatabricksRetryPolicy regardless of redirect settings. """ # Code 302 is a redirect, but 3xx codes are not retried per policy - # Note: We don't set redirect_location because that would cause urllib3 v2.x + # Note: We don't set redirect_location because that would cause urllib3 v2.x # to follow redirects internally, bypassing our retry policy test - with mocked_server_response( - status=302, redirect_location=None - ) as mock_obj: - with pytest.raises(RequestError): # Should get RequestError, not MaxRetryError + with mocked_server_response(status=302, redirect_location=None) as mock_obj: + with pytest.raises( + RequestError + ): # Should get RequestError, not MaxRetryError with self.connection( extra_params={ - **extra_params, **self._retry_policy, } ): @@ -701,16 +572,7 @@ def test_3xx_codes_not_retried_regardless_of_max_redirects_setting( # Total call count should be 1 (original only, no retries for 3xx codes) assert mock_obj.return_value.getresponse.call_count == 1 - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_3xx_codes_stop_request_immediately_no_retry_attempts( - self, extra_params - ): + def test_3xx_codes_stop_request_immediately_no_retry_attempts(self): # Since 3xx codes are not retried per policy, we only ever see the first 302 response responses = [ {"status": 302, "headers": {}, "redirect_location": "/foo.bar"}, @@ -726,7 +588,6 @@ def test_3xx_codes_stop_request_immediately_no_retry_attempts( with mock_sequential_server_responses(responses): with self.connection( extra_params={ - **extra_params, **self._retry_policy, **additional_settings, } @@ -737,19 +598,9 @@ def test_3xx_codes_stop_request_immediately_no_retry_attempts( assert "too many redirects" not in str(cm.value.message) assert "Error during request to server" in str(cm.value.message) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_max_redirects_exceeds_max_attempts_count_warns_user( - self, extra_params, caplog - ): + def test_retry_max_redirects_exceeds_max_attempts_count_warns_user(self, caplog): with self.connection( extra_params={ - **extra_params, **self._retry_policy, **{ "_retry_max_redirects": 100, @@ -759,17 +610,9 @@ def test_retry_max_redirects_exceeds_max_attempts_count_warns_user( ): assert "it will have no affect!" in caplog.text - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_retry_legacy_behavior_warns_user(self, extra_params, caplog): + def test_retry_legacy_behavior_warns_user(self, caplog): with self.connection( extra_params={ - **extra_params, **self._retry_policy, "_enable_v3_retries": False, } @@ -778,14 +621,7 @@ def test_retry_legacy_behavior_warns_user(self, extra_params, caplog): "Legacy retry behavior is enabled for this connection." in caplog.text ) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_403_not_retried(self, extra_params): + def test_403_not_retried(self): """GIVEN the server returns a code 403 WHEN the connector receives this response THEN nothing is retried and an exception is raised @@ -794,19 +630,12 @@ def test_403_not_retried(self, extra_params): # Code 403 is a Forbidden error with mocked_server_response(status=403): with pytest.raises(RequestError) as cm: - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params) as conn: pass assert isinstance(cm.value.args[1], NonRecoverableNetworkError) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - {"use_sea": True}, - ], - ) - def test_401_not_retried(self, extra_params): + def test_401_not_retried(self): """GIVEN the server returns a code 401 WHEN the connector receives this response THEN nothing is retried and an exception is raised @@ -815,7 +644,7 @@ def test_401_not_retried(self, extra_params): # Code 401 is an Unauthorized error with mocked_server_response(status=401): with pytest.raises(RequestError) as cm: - extra_params = {**extra_params, **self._retry_policy} + extra_params = {**self._retry_policy} with self.connection(extra_params=extra_params): pass assert isinstance(cm.value.args[1], NonRecoverableNetworkError) diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 319b7de5e..bd3355014 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -139,16 +139,15 @@ def assertEqualRowValues(self, actual, expected): class TestPySQLLargeWideResultSet(PySQLPytestTestCase): - @pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}]) @pytest.mark.parametrize("lz4_compression", [False, True]) - def test_query_with_large_wide_result_set(self, extra_params, lz4_compression): + def test_query_with_large_wide_result_set(self, lz4_compression): resultSize = 100 * 1000 * 1000 # 100 MB width = 8192 # B rows = resultSize // width cols = width // 36 fetchmany_size = 10 * 1024 * 1024 // width self.arraysize = 1000 - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.connection.lz4_compression = lz4_compression uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)]) cursor.execute( @@ -163,14 +162,13 @@ def test_query_with_large_wide_result_set(self, extra_params, lz4_compression): class TestPySQLLargeNarrowResultSet(PySQLPytestTestCase): - @pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}]) - def test_query_with_large_narrow_result_set(self, extra_params): + def test_query_with_large_narrow_result_set(self): resultSize = 100 * 1000 * 1000 # 100 MB width = 8 # sizeof(long) rows = resultSize / width fetchmany_size = 10 * 1024 * 1024 // width self.arraysize = 10000000 - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.execute("SELECT * FROM RANGE({rows})".format(rows=rows)) for row_id, row in enumerate( fetch_rows(self, cursor, rows, fetchmany_size) @@ -179,8 +177,7 @@ def test_query_with_large_narrow_result_set(self, extra_params): class TestPySQLLongRunningQuery(PySQLPytestTestCase): - @pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}]) - def test_long_running_query(self, extra_params): + def test_long_running_query(self): """Incrementally increase query size until it takes at least 1 minute, and asserts that the query completes successfully. """ @@ -191,19 +188,15 @@ def test_long_running_query(self, extra_params): duration = -1 scale0 = 10000 scale_factor = 50 - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: while duration < min_duration: assert scale_factor < 4096, "Detected infinite loop" start = time.time() - cursor.execute( - """SELECT count(*) + cursor.execute("""SELECT count(*) FROM RANGE({scale}) x JOIN RANGE({scale0}) y ON from_unixtime(x.id * y.id, "yyyy-MM-dd") LIKE "%not%a%date%" - """.format( - scale=scale_factor * scale0, scale0=scale0 - ) - ) + """.format(scale=scale_factor * scale0, scale0=scale0)) (n,) = cursor.fetchone() assert n == 0 duration = time.time() - start @@ -248,19 +241,10 @@ def test_cloud_fetch(self): class TestPySQLAsyncQueriesSuite(PySQLPytestTestCase): - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - }, - ], - ) - def test_execute_async__long_running(self, extra_params): + def test_execute_async__long_running(self): long_running_query = "SELECT COUNT(*) FROM RANGE(10000 * 16) x JOIN RANGE(10000) y ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'" - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.execute_async(long_running_query) ## Polling after every POLLING_INTERVAL seconds @@ -273,21 +257,10 @@ def test_execute_async__long_running(self, extra_params): assert result[0].asDict() == {"count(1)": 0} - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - ], - ) - def test_execute_async__small_result(self, extra_params): + def test_execute_async__small_result(self): small_result_query = "SELECT 1" - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.execute_async(small_result_query) ## Fake sleep for 5 secs @@ -303,16 +276,7 @@ def test_execute_async__small_result(self, extra_params): assert result[0].asDict() == {"1": 1} - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - }, - ], - ) - def test_execute_async__large_result(self, extra_params): + def test_execute_async__large_result(self): x_dimension = 1000 y_dimension = 1000 large_result_query = f""" @@ -326,7 +290,7 @@ def test_execute_async__large_result(self, extra_params): RANGE({y_dimension}) y """ - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.execute_async(large_result_query) ## Fake sleep for 5 secs @@ -426,22 +390,8 @@ def test_incorrect_query_throws_exception(self): cursor.execute("CREATE TABLE IF NOT EXISTS TABLE table_234234234") assert "table_234234234" in str(cm.value) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_create_table_will_return_empty_result_set(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_create_table_will_return_empty_result_set(self): + with self.cursor({}) as cursor: table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_")) try: cursor.execute( @@ -697,24 +647,10 @@ def test_get_catalogs(self): ] @skipUnless(pysql_supports_arrow(), "arrow test need arrow support") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_get_arrow(self, extra_params): + def test_get_arrow(self): # These tests are quite light weight as the arrow fetch methods are used internally # by everything else - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.execute("SELECT * FROM range(10)") table_1 = cursor.fetchmany_arrow(1).to_pydict() assert table_1 == OrderedDict([("id", [0])]) @@ -722,20 +658,9 @@ def test_get_arrow(self, extra_params): table_2 = cursor.fetchall_arrow().to_pydict() assert table_2 == OrderedDict([("id", [1, 2, 3, 4, 5, 6, 7, 8, 9])]) - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - ], - ) - def test_unicode(self, extra_params): + def test_unicode(self): unicode_str = "数据砖" - with self.cursor(extra_params) as cursor: + with self.cursor({}) as cursor: cursor.execute("SELECT '{}'".format(unicode_str)) results = cursor.fetchall() assert len(results) == 1 and len(results[0]) == 1 @@ -773,22 +698,8 @@ def execute_really_long_query(): assert len(cursor.fetchall()) == 3 @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_can_execute_command_after_failure(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_can_execute_command_after_failure(self): + with self.cursor({}) as cursor: with pytest.raises(DatabaseError): cursor.execute("this is a sytnax error") @@ -798,22 +709,8 @@ def test_can_execute_command_after_failure(self, extra_params): self.assertEqualRowValues(res, [[1]]) @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_can_execute_command_after_success(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_can_execute_command_after_success(self): + with self.cursor({}) as cursor: cursor.execute("SELECT 1;") cursor.execute("SELECT 2;") @@ -825,22 +722,8 @@ def generate_multi_row_query(self): return query @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_fetchone(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_fetchone(self): + with self.cursor({}) as cursor: query = self.generate_multi_row_query() cursor.execute(query) @@ -851,19 +734,8 @@ def test_fetchone(self, extra_params): assert cursor.fetchone() == None @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - ], - ) - def test_fetchall(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_fetchall(self): + with self.cursor({}) as cursor: query = self.generate_multi_row_query() cursor.execute(query) @@ -872,22 +744,8 @@ def test_fetchall(self, extra_params): assert cursor.fetchone() == None @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_fetchmany_when_stride_fits(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_fetchmany_when_stride_fits(self): + with self.cursor({}) as cursor: query = "SELECT * FROM range(4)" cursor.execute(query) @@ -895,22 +753,8 @@ def test_fetchmany_when_stride_fits(self, extra_params): self.assertEqualRowValues(cursor.fetchmany(2), [[2], [3]]) @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_fetchmany_in_excess(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_fetchmany_in_excess(self): + with self.cursor({}) as cursor: query = "SELECT * FROM range(4)" cursor.execute(query) @@ -918,22 +762,8 @@ def test_fetchmany_in_excess(self, extra_params): self.assertEqualRowValues(cursor.fetchmany(3), [[3]]) @skipIf(pysql_has_version("<", "2"), "requires pysql v2") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_iterator_api(self, extra_params): - with self.cursor(extra_params) as cursor: + def test_iterator_api(self): + with self.cursor({}) as cursor: query = "SELECT * FROM range(4)" cursor.execute(query) @@ -1006,28 +836,13 @@ def test_timestamps_arrow(self): ), "timestamp {} did not match {}".format(timestamp, expected) @skipUnless(pysql_supports_arrow(), "arrow test needs arrow support") - @pytest.mark.parametrize( - "extra_params", - [ - {}, - { - "use_sea": True, - "use_cloud_fetch": False, - "enable_query_result_lz4_compression": False, - }, - { - "use_sea": True, - }, - ], - ) - def test_multi_timestamps_arrow(self, extra_params): + def test_multi_timestamps_arrow(self): with self.cursor( { "session_configuration": { "ansi_mode": False, "query_tags": "test:multi-timestamps,driver:python", }, - **extra_params, } ) as cursor: query, expected = self.multi_query() diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 66a3722ec..1b1e4ab79 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -421,8 +421,8 @@ def test_cancel_command_calls_the_backend(self): def test_cancel_command_will_issue_warning_for_cancel_with_no_executing_command( self, logger_instance ): - # Backends like Thrift/SEA set active_command_id before blocking - # and do NOT define ``cancel_running_cursor``. Use ``spec`` so the + # Thrift sets active_command_id before blocking and does not + # define ``cancel_running_cursor``. Use ``spec`` so the # mock doesn't auto-advertise that opt-in hook (a bare ``Mock()`` # returns a truthy Mock for any attribute). mock_thrift_backend = Mock(spec=["cancel_command"]) diff --git a/tests/unit/test_feature_flag_lifecycle.py b/tests/unit/test_feature_flag_lifecycle.py index 98dff5851..6ed732dcd 100644 --- a/tests/unit/test_feature_flag_lifecycle.py +++ b/tests/unit/test_feature_flag_lifecycle.py @@ -8,9 +8,8 @@ from databricks.sql.telemetry.telemetry_client import TelemetryHelper -@pytest.mark.parametrize("use_sea", [False, True]) @pytest.mark.parametrize("oauth", [False, True]) -def test_flags_available_with_telemetry_disabled(session_feature_flags, use_sea, oauth): +def test_flags_available_with_telemetry_disabled(session_feature_flags, oauth): session_feature_flags.get_instance.side_effect = ( FeatureFlagsContextFactory.get_instance ) @@ -31,19 +30,17 @@ def create_backend(**kwargs): assert kwargs["auth_provider"] is auth return backend - backend_name = "SeaDatabricksClient" if use_sea else "ThriftDatabricksClient" try: with patch("databricks.sql.client.UnifiedHttpClient", return_value=http), patch( "databricks.sql.session.get_python_sql_connector_auth_provider", return_value=auth, ) as auth_factory, patch( - f"databricks.sql.session.{backend_name}", side_effect=create_backend + "databricks.sql.session.ThriftDatabricksClient", side_effect=create_backend ): conn = sql.connect( "flags.example", "/sql/1.0/warehouses/test?o=123", enable_telemetry=False, - use_sea=use_sea, **( {"auth_type": "databricks-oauth"} if oauth diff --git a/tests/unit/test_filters.py b/tests/unit/test_filters.py deleted file mode 100644 index 4efe51f3e..000000000 --- a/tests/unit/test_filters.py +++ /dev/null @@ -1,157 +0,0 @@ -""" -Tests for the ResultSetFilter class. -""" - -import unittest -from unittest.mock import MagicMock, patch - -from databricks.sql.backend.sea.utils.filters import ResultSetFilter - - -class TestResultSetFilter(unittest.TestCase): - """Tests for the ResultSetFilter class.""" - - def setUp(self): - """Set up test fixtures.""" - # Create a mock SeaResultSet - self.mock_sea_result_set = MagicMock() - - # Set up the remaining_rows method on the results attribute - self.mock_sea_result_set.results = MagicMock() - self.mock_sea_result_set.results.remaining_rows.return_value = [ - ["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""], - ["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""], - [ - "catalog1", - "schema1", - "table3", - "owner1", - "2023-01-01", - "SYSTEM TABLE", - "", - ], - [ - "catalog1", - "schema1", - "table4", - "owner1", - "2023-01-01", - "EXTERNAL TABLE", - "", - ], - ] - - # Set up the connection and other required attributes - self.mock_sea_result_set.connection = MagicMock() - self.mock_sea_result_set.backend = MagicMock() - self.mock_sea_result_set.buffer_size_bytes = 1000 - self.mock_sea_result_set.arraysize = 100 - self.mock_sea_result_set.statement_id = "test-statement-id" - self.mock_sea_result_set.lz4_compressed = False - - # Create a mock CommandId - from databricks.sql.backend.types import CommandId, BackendType - - mock_command_id = CommandId(BackendType.SEA, "test-statement-id") - self.mock_sea_result_set.command_id = mock_command_id - - self.mock_sea_result_set.status = MagicMock() - self.mock_sea_result_set.description = [ - ("catalog_name", "string", None, None, None, None, True), - ("schema_name", "string", None, None, None, None, True), - ("table_name", "string", None, None, None, None, True), - ("owner", "string", None, None, None, None, True), - ("creation_time", "string", None, None, None, None, True), - ("table_type", "string", None, None, None, None, True), - ("remarks", "string", None, None, None, None, True), - ] - self.mock_sea_result_set.has_been_closed_server_side = False - self.mock_sea_result_set._arrow_schema_bytes = None - - def test__filter_json_result_set(self): - """Test filtering by column values with various options.""" - # Case 1: Case-sensitive filtering - allowed_values = ["table1", "table3"] - - with patch( - "databricks.sql.backend.sea.utils.filters.isinstance", return_value=True - ): - with patch( - "databricks.sql.backend.sea.result_set.SeaResultSet" - ) as mock_sea_result_set_class: - mock_instance = MagicMock() - mock_sea_result_set_class.return_value = mock_instance - - # Call _filter_json_result_set on the table_name column (index 2) - result = ResultSetFilter._filter_json_result_set( - self.mock_sea_result_set, 2, allowed_values, case_sensitive=True - ) - - # Verify the filter was applied correctly - mock_sea_result_set_class.assert_called_once() - - # Check the filtered data passed to the constructor - args, kwargs = mock_sea_result_set_class.call_args - result_data = kwargs.get("result_data") - self.assertIsNotNone(result_data) - self.assertEqual(len(result_data.data), 2) - self.assertIn(result_data.data[0][2], allowed_values) - self.assertIn(result_data.data[1][2], allowed_values) - - # Case 2: Case-insensitive filtering - mock_sea_result_set_class.reset_mock() - with patch( - "databricks.sql.backend.sea.utils.filters.isinstance", return_value=True - ): - with patch( - "databricks.sql.backend.sea.result_set.SeaResultSet" - ) as mock_sea_result_set_class: - mock_instance = MagicMock() - mock_sea_result_set_class.return_value = mock_instance - - # Call _filter_json_result_set with case-insensitive matching - result = ResultSetFilter._filter_json_result_set( - self.mock_sea_result_set, - 2, - ["TABLE1", "TABLE3"], - case_sensitive=False, - ) - mock_sea_result_set_class.assert_called_once() - - def test_filter_tables_by_type(self): - """Test filtering tables by type with various options.""" - # Case 1: Specific table types - table_types = ["TABLE", "VIEW"] - - # Mock results as JsonQueue (not CloudFetchQueue or ArrowQueue) - from databricks.sql.backend.sea.queue import JsonQueue - - self.mock_sea_result_set.results = JsonQueue([]) - - with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter: - ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, table_types) - args, kwargs = mock_filter.call_args - self.assertEqual(args[0], self.mock_sea_result_set) - self.assertEqual(kwargs.get("column_index"), 5) # Table type column index - self.assertEqual(kwargs.get("allowed_values"), table_types) - self.assertEqual(kwargs.get("case_sensitive"), True) - - # Case 2: Default table types (None or empty list) - with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter: - # Test with None - ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None) - args, kwargs = mock_filter.call_args - self.assertEqual( - kwargs.get("allowed_values"), ["TABLE", "VIEW", "SYSTEM TABLE"] - ) - - # Test with empty list - ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, []) - args, kwargs = mock_filter.call_args - self.assertEqual( - kwargs.get("allowed_values"), ["TABLE", "VIEW", "SYSTEM TABLE"] - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index 000515915..be73da4b9 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -169,7 +169,7 @@ def test_generic_oauth_provider_raises_not_supported(self): def test_external_credentials_provider_raises_not_supported(self): """A user-supplied ``credentials_provider`` is the connector's - primary M2M path on Thrift/SEA, but it's an opaque token source + primary M2M path on Thrift, but it's an opaque token source with no extractable raw creds — the kernel can't own the token lifecycle, so the bridge rejects it and points at the ``oauth_client_id`` / ``oauth_client_secret`` M2M kwargs.""" diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..0f1538fee 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -290,7 +290,7 @@ def test_no_open_session_guards_raise_interface_error(): # The async-state methods attach to the kernel session by id; they # must also guard a closed/absent session before the kernel call. - cid = CommandId.from_sea_statement_id("any-id") + cid = CommandId.from_kernel_statement_id("any-id") with pytest.raises(InterfaceError): c.get_query_state(cid) with pytest.raises(InterfaceError): @@ -1076,13 +1076,13 @@ def test_cancel_command_tolerant_when_handle_missing(): """``cancel_command`` is documented to be a no-op when there's no tracked async handle (matches Thrift's tolerance).""" c = _make_client() - fake_command_id = CommandId.from_sea_statement_id("not-tracked") + fake_command_id = CommandId.from_kernel_statement_id("not-tracked") c.cancel_command(fake_command_id) # must not raise def test_close_command_tolerant_when_handle_missing(): c = _make_client() - fake_command_id = CommandId.from_sea_statement_id("not-tracked") + fake_command_id = CommandId.from_kernel_statement_id("not-tracked") c.close_command(fake_command_id) # must not raise @@ -1112,7 +1112,7 @@ def test_get_query_state_returns_succeeded_when_server_404s(): polling loop terminates cleanly.""" c = _make_client() _attach_returns(c, status_error=_FakeKernelError(code="NotFound")) - cid = CommandId.from_sea_statement_id("sync-only") + cid = CommandId.from_kernel_statement_id("sync-only") assert c.get_query_state(cid) == CommandState.SUCCEEDED @@ -1121,7 +1121,7 @@ def test_independent_async_read_requires_reader_api(read_result): c = _make_client() c._kernel_session = MagicMock() owner = MagicMock(spec=["status", "await_result"]) - cid = CommandId.from_sea_statement_id("tracked-async") + cid = CommandId.from_kernel_statement_id("tracked-async") c._async_handles[cid.guid] = owner c._async_result_stream_started.add(cid.guid) @@ -1142,7 +1142,7 @@ def test_get_query_state_returns_closed_from_server(): closed-state bookkeeping.""" c = _make_client() _attach_returns(c, status=("Closed", None)) - cid = CommandId.from_sea_statement_id("closed-async") + cid = CommandId.from_kernel_statement_id("closed-async") assert c.get_query_state(cid) == CommandState.CLOSED @@ -1152,7 +1152,7 @@ def test_get_query_state_propagates_non_not_found_error(): PEP 249 exception so the caller can retry / surface it.""" c = _make_client() _attach_returns(c, status_error=_FakeKernelError(code="Unavailable")) - cid = CommandId.from_sea_statement_id("flaky") + cid = CommandId.from_kernel_statement_id("flaky") with pytest.raises(DatabaseError): c.get_query_state(cid) @@ -1165,7 +1165,7 @@ def test_get_query_state_uses_retained_owning_handle_before_result_stream(): c._kernel_session = MagicMock() handle = MagicMock() handle.status.return_value = ("Running", None) - cid = CommandId.from_sea_statement_id("async-status-owning") + cid = CommandId.from_kernel_statement_id("async-status-owning") c._async_handles[cid.guid] = handle assert c.get_query_state(cid) == CommandState.RUNNING @@ -1187,7 +1187,7 @@ def test_get_query_state_concurrent_poll_uses_reader(): reader = MagicMock() reader.status.return_value = ("Running", None) owning_handle.reader.return_value = reader - cid = CommandId.from_sea_statement_id("async-status-concurrent") + cid = CommandId.from_kernel_statement_id("async-status-concurrent") c._async_handles[cid.guid] = owning_handle # Simulate the first poll being mid-flight: its reservation is set. @@ -1221,7 +1221,7 @@ def test_get_query_state_uses_reader_after_result_stream_started(): reader = MagicMock() reader.status.return_value = ("Succeeded", None) owning_handle.reader.return_value = reader - cid = CommandId.from_sea_statement_id("async-status-attached") + cid = CommandId.from_kernel_statement_id("async-status-attached") c._async_handles[cid.guid] = owning_handle c._async_result_stream_started.add(cid.guid) @@ -1246,7 +1246,7 @@ def test_get_execution_result_uses_retained_owning_handle_first(): cursor.arraysize = 100 cursor.buffer_size_bytes = 1024 cursor.row_limit = 5 - cid = CommandId.from_sea_statement_id("async-1") + cid = CommandId.from_kernel_statement_id("async-1") c._async_handles[cid.guid] = handle rs = c.get_execution_result(cid, cursor=cursor) @@ -1267,7 +1267,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle(): cursor = MagicMock() cursor.arraysize = 100 cursor.buffer_size_bytes = 1024 - cid = CommandId.from_sea_statement_id("async-1") + cid = CommandId.from_kernel_statement_id("async-1") rs = c.get_execution_result(cid, cursor=cursor) @@ -1284,7 +1284,7 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle(): c._kernel_session = MagicMock() handle = MagicMock() handle.await_result.side_effect = _FakeKernelError(code="Unavailable") - cid = CommandId.from_sea_statement_id("async-retry-owning") + cid = CommandId.from_kernel_statement_id("async-retry-owning") c._async_handles[cid.guid] = handle with pytest.raises(OperationalError): @@ -1309,7 +1309,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade owning_stream.arrow_schema.side_effect = _FakeKernelError(code="Internal") owning_handle = MagicMock() owning_handle.await_result.return_value = owning_stream - cid = CommandId.from_sea_statement_id("async-construct-fail") + cid = CommandId.from_kernel_statement_id("async-construct-fail") c._async_handles[cid.guid] = owning_handle with pytest.raises(DatabaseError): @@ -1346,7 +1346,7 @@ def test_get_execution_result_maps_not_found_to_programming_error(): mapped PEP 249 exception rather than a raw error.""" c = _make_client() _attach_returns(c, status_error=_FakeKernelError(code="NotFound")) - cid = CommandId.from_sea_statement_id("gone") + cid = CommandId.from_kernel_statement_id("gone") with pytest.raises(DatabaseError): c.get_execution_result(cid, cursor=MagicMock()) @@ -1355,7 +1355,7 @@ def test_cancel_command_reraises_kernel_error(): c = _make_client() fake_handle = MagicMock() fake_handle.cancel.side_effect = _FakeKernelError(code="Unavailable") - cid = CommandId.from_sea_statement_id("abc") + cid = CommandId.from_kernel_statement_id("abc") c._async_handles[cid.guid] = fake_handle with pytest.raises(OperationalError): c.cancel_command(cid) @@ -1365,7 +1365,7 @@ def test_close_command_reraises_kernel_error(): c = _make_client() fake_handle = MagicMock() fake_handle.close.side_effect = _FakeKernelError(code="Internal") - cid = CommandId.from_sea_statement_id("abc") + cid = CommandId.from_kernel_statement_id("abc") c._async_handles[cid.guid] = fake_handle with pytest.raises(DatabaseError): c.close_command(cid) @@ -1379,7 +1379,7 @@ def test_get_query_state_raises_on_failed_state_with_failure(): _attach_returns( c, status=("Failed", _FakeKernelError(code="SqlError", message="bad")) ) - cid = CommandId.from_sea_statement_id("abc") + cid = CommandId.from_kernel_statement_id("abc") with pytest.raises(DatabaseError, match="bad"): c.get_query_state(cid) @@ -1395,7 +1395,7 @@ def test_get_query_state_handles_non_baseexception_failure(): # ``failure`` is a plain dict (not BaseException) — simulates a # kernel binding that exposes the failure as a structured value. _attach_returns(c, status=("Failed", {"code": "Internal", "msg": "weird"})) - cid = CommandId.from_sea_statement_id("xyz") + cid = CommandId.from_kernel_statement_id("xyz") # Must surface as a PEP 249 exception (OperationalError via the # wrap helper's fallback path), not TypeError. with pytest.raises(OperationalError): @@ -1405,7 +1405,7 @@ def test_get_query_state_handles_non_baseexception_failure(): def test_get_query_state_returns_state_when_no_failure(): c = _make_client() _attach_returns(c, status=("Running", None)) - cid = CommandId.from_sea_statement_id("abc") + cid = CommandId.from_kernel_statement_id("abc") assert c.get_query_state(cid) == CommandState.RUNNING @@ -1457,7 +1457,7 @@ def test_close_session_closes_and_drops_swept_handles(): server.""" c = _make_client() handle = MagicMock() - cid = CommandId.from_sea_statement_id("xyz") + cid = CommandId.from_kernel_statement_id("xyz") c._async_handles[cid.guid] = handle c._kernel_session = MagicMock() c.close_session(MagicMock()) @@ -1479,7 +1479,7 @@ def test_get_query_state_returns_closed_after_close_command(): c = _make_client() c._kernel_session = MagicMock() handle = MagicMock() - cid = CommandId.from_sea_statement_id("async-1") + cid = CommandId.from_kernel_statement_id("async-1") c._async_handles[cid.guid] = handle c.close_command(cid) assert handle.close.called @@ -1584,7 +1584,7 @@ def test_get_execution_result_is_re_callable(): reader = MagicMock() reader.await_result.return_value = fake_stream owning_handle.reader.return_value = reader - cid = CommandId.from_sea_statement_id("async-recall-twice") + cid = CommandId.from_kernel_statement_id("async-recall-twice") c._async_handles[cid.guid] = owning_handle cursor = MagicMock() cursor.arraysize = 100 diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index cfa2d3f2e..2baaccc78 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -54,7 +54,7 @@ def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet: connection=connection, backend=backend, kernel_handle=handle, - command_id=CommandId.from_sea_statement_id("smoke-test"), + command_id=CommandId.from_kernel_statement_id("smoke-test"), arraysize=100, buffer_size_bytes=1024, ) @@ -232,7 +232,7 @@ def test_close_skips_kernel_call_when_connection_already_closed(int_schema): connection=connection, backend=backend, kernel_handle=handle, - command_id=CommandId.from_sea_statement_id("conn-closed-test"), + command_id=CommandId.from_kernel_statement_id("conn-closed-test"), arraysize=100, buffer_size_bytes=1024, ) diff --git a/tests/unit/test_lazy_thrift_import.py b/tests/unit/test_lazy_thrift_import.py index 5763a91c6..1b039314f 100644 --- a/tests/unit/test_lazy_thrift_import.py +++ b/tests/unit/test_lazy_thrift_import.py @@ -4,7 +4,7 @@ eagerly whenever ``connect()`` (i.e. ``databricks.sql.client``) was imported, regardless of which backend the caller selected. Build systems that vendor their own ``thrift`` (e.g. Meta's Buck) then hit a namespace collision even on -the SEA / kernel code paths, which never speak Thrift on the wire. +the kernel code path, which never speaks Thrift on the wire. These tests assert that importing the connector and the non-Thrift backends does NOT import ``thrift``, while the Thrift backend still does. The invariant @@ -94,7 +94,7 @@ def _probe_import(module_name: str) -> _ImportProbeResult: # Modules on the connect()/execute() path that must stay Thrift-free so the -# SEA and kernel backends can be used without the ``thrift`` package present. +# kernel backend can be used without the ``thrift`` package present. THRIFT_FREE_MODULES = [ "databricks.sql", "databricks.sql.client", @@ -103,8 +103,6 @@ def _probe_import(module_name: str) -> _ImportProbeResult: "databricks.sql.parameters.native", "databricks.sql.backend.types", "databricks.sql.backend.databricks_client", - "databricks.sql.backend.sea.backend", - "databricks.sql.backend.sea.queue", "databricks.sql.backend.kernel.type_mapping", "databricks.sql.backend.kernel.client", "databricks.sql.cloudfetch.downloader", @@ -118,7 +116,7 @@ def test_module_does_not_import_thrift(module_name): Apache Thrift runtime. This is what unblocks callers (e.g. Buck-based builds) that ship their own - ``thrift`` package and use only the SEA or kernel backend. + ``thrift`` package and use only the kernel backend. """ result = _probe_import(module_name) diff --git a/tests/unit/test_sea_backend.py b/tests/unit/test_sea_backend.py deleted file mode 100644 index 513c83c8a..000000000 --- a/tests/unit/test_sea_backend.py +++ /dev/null @@ -1,1209 +0,0 @@ -""" -Tests for the SEA (Statement Execution API) backend implementation. - -This module contains tests for the SeaDatabricksClient class, which implements -the Databricks SQL connector's SEA backend functionality. -""" - -import pytest -from unittest.mock import patch, MagicMock, Mock - -from databricks.sql.backend.sea.backend import ( - SeaDatabricksClient, - _filter_session_configuration, -) -from databricks.sql.backend.sea.models.base import ServiceError, StatementStatus -from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType -from databricks.sql.parameters.native import IntegerParameter, TDbsqlParameter -from databricks.sql.thrift_api.TCLIService import ttypes -from databricks.sql.types import SSLOptions -from databricks.sql.auth.authenticators import AuthProvider -from databricks.sql.exc import ( - Error, - NotSupportedError, - ProgrammingError, - ServerOperationError, - DatabaseError, -) - - -class TestSeaBackend: - """Test suite for the SeaDatabricksClient class.""" - - @pytest.fixture - def mock_http_client(self): - """Create a mock HTTP client.""" - with patch( - "databricks.sql.backend.sea.backend.SeaHttpClient" - ) as mock_client_class: - mock_client = mock_client_class.return_value - yield mock_client - - @pytest.fixture - def sea_client(self, mock_http_client): - """Create a SeaDatabricksClient instance with mocked dependencies.""" - server_hostname = "test-server.databricks.com" - port = 443 - http_path = "/sql/warehouses/abc123" - http_headers = [("header1", "value1"), ("header2", "value2")] - auth_provider = AuthProvider() - ssl_options = SSLOptions() - - client = SeaDatabricksClient( - server_hostname=server_hostname, - port=port, - http_path=http_path, - http_headers=http_headers, - auth_provider=auth_provider, - ssl_options=ssl_options, - use_cloud_fetch=False, - ) - - return client - - @pytest.fixture - def sea_client_cloud_fetch(self, mock_http_client): - """Create a SeaDatabricksClient instance with cloud fetch enabled.""" - server_hostname = "test-server.databricks.com" - port = 443 - http_path = "/sql/warehouses/abc123" - http_headers = [("header1", "value1"), ("header2", "value2")] - auth_provider = AuthProvider() - ssl_options = SSLOptions() - - client = SeaDatabricksClient( - server_hostname=server_hostname, - port=port, - http_path=http_path, - http_headers=http_headers, - auth_provider=auth_provider, - ssl_options=ssl_options, - use_cloud_fetch=True, - ) - - return client - - @pytest.fixture - def sea_session_id(self): - """Create a SEA session ID.""" - return SessionId.from_sea_session_id("test-session-123") - - @pytest.fixture - def sea_command_id(self): - """Create a SEA command ID.""" - return CommandId.from_sea_statement_id("test-statement-123") - - @pytest.fixture - def mock_cursor(self): - """Create a mock cursor.""" - cursor = Mock() - cursor.active_command_id = None - cursor.buffer_size_bytes = 1000 - cursor.arraysize = 100 - return cursor - - @pytest.fixture - def thrift_session_id(self): - """Create a Thrift session ID (not SEA).""" - mock_thrift_handle = MagicMock() - mock_thrift_handle.sessionId.guid = b"guid" - mock_thrift_handle.sessionId.secret = b"secret" - return SessionId.from_thrift_handle(mock_thrift_handle) - - @pytest.fixture - def thrift_command_id(self): - """Create a Thrift command ID (not SEA).""" - mock_thrift_operation_handle = MagicMock() - mock_thrift_operation_handle.operationId.guid = b"guid" - mock_thrift_operation_handle.operationId.secret = b"secret" - return CommandId.from_thrift_handle(mock_thrift_operation_handle) - - def test_initialization(self, mock_http_client): - """Test client initialization and warehouse ID extraction.""" - # Test with warehouses format - client1 = SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/warehouses/abc123", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - assert client1.warehouse_id == "abc123" - assert client1.max_download_threads == 10 # Default value - - # Test with endpoints format - client2 = SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/endpoints/def456", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - assert client2.warehouse_id == "def456" - - # Test with SPOG query param ?o= in http_path - client_spog = SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/1.0/warehouses/abc123?o=6051921418418893", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - assert client_spog.warehouse_id == "abc123" - - # Test with SPOG query param on endpoints path - client_spog_ep = SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/1.0/endpoints/def456?o=6051921418418893", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - assert client_spog_ep.warehouse_id == "def456" - - # Test with multiple query params - client_spog_multi = SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/1.0/warehouses/abc123?o=123&extra=val", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - assert client_spog_multi.warehouse_id == "abc123" - - # Test with custom max_download_threads - client3 = SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/warehouses/abc123", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - max_download_threads=5, - ) - assert client3.max_download_threads == 5 - - # Test with invalid HTTP path - with pytest.raises(ValueError) as excinfo: - SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/invalid/path", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - assert "Could not extract warehouse ID" in str(excinfo.value) - - def test_initialization_warns_backend_incomplete(self, mock_http_client, caplog): - """Constructing a SEA client emits a warning steering users to the - kernel backend, since the SEA path is incomplete and slated for - deprecation.""" - import logging - - with caplog.at_level( - logging.WARNING, logger="databricks.sql.backend.sea.backend" - ): - SeaDatabricksClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/warehouses/abc123", - http_headers=[], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - ) - - warnings = [r.message for r in caplog.records if r.levelno == logging.WARNING] - assert any( - "incomplete" in m and "use_kernel=True" in m and "[kernel]" in m - for m in warnings - ), warnings - - def test_session_management(self, sea_client, mock_http_client, thrift_session_id): - """Test session management methods.""" - # Test open_session with minimal parameters - mock_http_client._make_request.return_value = {"session_id": "test-session-123"} - session_id = sea_client.open_session(None, None, None) - assert isinstance(session_id, SessionId) - assert session_id.backend_type == BackendType.SEA - assert session_id.guid == "test-session-123" - mock_http_client._make_request.assert_called_with( - method="POST", path=sea_client.SESSION_PATH, data={"warehouse_id": "abc123"} - ) - - # Test open_session with all parameters - mock_http_client.reset_mock() - mock_http_client._make_request.return_value = {"session_id": "test-session-456"} - session_config = { - "ANSI_MODE": "FALSE", # Supported parameter - "STATEMENT_TIMEOUT": "3600", # Supported parameter - "QUERY_TAGS": "team:marketing,dashboard:abc123", # Supported parameter - "unsupported_param": "value", # Unsupported parameter - } - catalog = "test_catalog" - schema = "test_schema" - session_id = sea_client.open_session(session_config, catalog, schema) - assert session_id.guid == "test-session-456" - expected_data = { - "warehouse_id": "abc123", - "session_confs": { - "ansi_mode": "FALSE", - "statement_timeout": "3600", - "query_tags": "team:marketing,dashboard:abc123", - }, - "catalog": catalog, - "schema": schema, - } - mock_http_client._make_request.assert_called_with( - method="POST", path=sea_client.SESSION_PATH, data=expected_data - ) - - # Test open_session error handling - mock_http_client.reset_mock() - mock_http_client._make_request.return_value = {} - with pytest.raises(Error) as excinfo: - sea_client.open_session(None, None, None) - assert "Failed to create session" in str(excinfo.value) - - # Test close_session with valid ID - mock_http_client.reset_mock() - session_id = SessionId.from_sea_session_id("test-session-789") - sea_client.close_session(session_id) - mock_http_client._make_request.assert_called_with( - method="DELETE", - path=sea_client.SESSION_PATH_WITH_ID.format("test-session-789"), - data={"session_id": "test-session-789", "warehouse_id": "abc123"}, - ) - - # Test close_session with invalid ID type - with pytest.raises(ValueError) as excinfo: - sea_client.close_session(thrift_session_id) - assert "Not a valid SEA session ID" in str(excinfo.value) - - def test_command_execution_sync( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """Test synchronous command execution.""" - # Test synchronous execution - execute_response = { - "statement_id": "test-statement-123", - "status": {"state": "SUCCEEDED"}, - "manifest": { - "schema": [ - { - "name": "col1", - "type_name": "STRING", - "type_text": "string", - "nullable": True, - } - ], - "total_row_count": 1, - "total_byte_count": 100, - }, - "result": {"data": [["value1"]]}, - } - mock_http_client._make_request.return_value = execute_response - - with patch.object( - sea_client, "_response_to_result_set", return_value="mock_result_set" - ) as mock_get_result: - result = sea_client.execute_command( - operation="SELECT 1", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result == "mock_result_set" - - # Test with invalid session ID - with pytest.raises(ValueError) as excinfo: - mock_thrift_handle = MagicMock() - mock_thrift_handle.sessionId.guid = b"guid" - mock_thrift_handle.sessionId.secret = b"secret" - thrift_session_id = SessionId.from_thrift_handle(mock_thrift_handle) - - sea_client.execute_command( - operation="SELECT 1", - session_id=thrift_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert "Not a valid SEA session ID" in str(excinfo.value) - - def test_command_execution_async( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """Test asynchronous command execution.""" - # Test asynchronous execution - execute_response = { - "statement_id": "test-statement-456", - "status": {"state": "PENDING"}, - } - mock_http_client._make_request.return_value = execute_response - - result = sea_client.execute_command( - operation="SELECT 1", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=True, - enforce_embedded_schema_correctness=False, - ) - assert result is None - assert isinstance(mock_cursor.active_command_id, CommandId) - assert mock_cursor.active_command_id.guid == "test-statement-456" - - def test_command_execution_advanced( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """Test advanced command execution scenarios.""" - # Test with polling - initial_response = { - "statement_id": "test-statement-789", - "status": {"state": "RUNNING"}, - } - poll_response = { - "statement_id": "test-statement-789", - "status": {"state": "SUCCEEDED"}, - "manifest": {"schema": [], "total_row_count": 0, "total_byte_count": 0}, - "result": {"data": []}, - } - mock_http_client._make_request.side_effect = [initial_response, poll_response] - - with patch.object( - sea_client, "_response_to_result_set", return_value="mock_result_set" - ) as mock_get_result: - with patch("time.sleep"): - result = sea_client.execute_command( - operation="SELECT * FROM large_table", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result == "mock_result_set" - - # Test with parameters - mock_http_client.reset_mock() - mock_http_client._make_request.side_effect = None # Reset side_effect - execute_response = { - "statement_id": "test-statement-123", - "status": {"state": "SUCCEEDED"}, - } - mock_http_client._make_request.return_value = execute_response - dbsql_param = IntegerParameter(name="param1", value=1) - param = dbsql_param.as_tspark_param(named=True) - - with patch.object(sea_client, "_response_to_result_set"): - sea_client.execute_command( - operation="SELECT * FROM table WHERE col = :param1", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[param], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - args, kwargs = mock_http_client._make_request.call_args - assert "parameters" in kwargs["data"] - assert len(kwargs["data"]["parameters"]) == 1 - assert kwargs["data"]["parameters"][0]["name"] == "param1" - assert kwargs["data"]["parameters"][0]["value"] == "1" - assert kwargs["data"]["parameters"][0]["type"] == "INT" - - # Test execution failure - mock_http_client.reset_mock() - error_response = { - "statement_id": "test-statement-123", - "status": { - "state": "FAILED", - "error": { - "message": "Syntax error in SQL", - "error_code": "SYNTAX_ERROR", - }, - }, - } - mock_http_client._make_request.return_value = error_response - - with patch("time.sleep"): - with patch.object( - sea_client, "get_query_state", return_value=CommandState.FAILED - ): - with pytest.raises(Error) as excinfo: - sea_client.execute_command( - operation="SELECT * FROM nonexistent_table", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert "Command failed" in str(excinfo.value) - - def _execute_response(self): - return { - "statement_id": "test-statement-123", - "status": {"state": "SUCCEEDED"}, - "manifest": {"schema": [], "total_row_count": 0, "total_byte_count": 0}, - "result": {"data": []}, - } - - def _run_execute_command(self, sea_client, sea_session_id, mock_cursor, **kwargs): - """Helper to invoke execute_command with default args.""" - return sea_client.execute_command( - operation="SELECT 1", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - **kwargs, - ) - - def test_execute_command_query_tags_string_values( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """query_tags with string values are included in the request payload.""" - mock_http_client._make_request.return_value = self._execute_response() - with patch.object(sea_client, "_response_to_result_set"): - self._run_execute_command( - sea_client, - sea_session_id, - mock_cursor, - query_tags={"env": "prod", "team": "data"}, - ) - _, kwargs = mock_http_client._make_request.call_args - assert kwargs["data"]["query_tags"] == [ - {"key": "env", "value": "prod"}, - {"key": "team", "value": "data"}, - ] - - def test_execute_command_query_tags_none_value( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """query_tags with a None value omit the value field (key-only tag).""" - mock_http_client._make_request.return_value = self._execute_response() - with patch.object(sea_client, "_response_to_result_set"): - self._run_execute_command( - sea_client, - sea_session_id, - mock_cursor, - query_tags={"env": "prod", "team": None}, - ) - _, kwargs = mock_http_client._make_request.call_args - assert kwargs["data"]["query_tags"] == [ - {"key": "env", "value": "prod"}, - {"key": "team", "value": None}, - ] - - def test_execute_command_no_query_tags_omitted( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """query_tags field is absent from the request when not provided.""" - mock_http_client._make_request.return_value = self._execute_response() - with patch.object(sea_client, "_response_to_result_set"): - self._run_execute_command(sea_client, sea_session_id, mock_cursor) - _, kwargs = mock_http_client._make_request.call_args - assert "query_tags" not in kwargs["data"] - - def test_execute_command_empty_query_tags_omitted( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """Empty query_tags dict is treated as absent — field omitted from request.""" - mock_http_client._make_request.return_value = self._execute_response() - with patch.object(sea_client, "_response_to_result_set"): - self._run_execute_command( - sea_client, sea_session_id, mock_cursor, query_tags={} - ) - _, kwargs = mock_http_client._make_request.call_args - assert "query_tags" not in kwargs["data"] - - def test_execute_command_async_query_tags( - self, sea_client, mock_http_client, mock_cursor, sea_session_id - ): - """query_tags are included in async execute requests (execute_async path).""" - mock_http_client._make_request.return_value = { - "statement_id": "test-statement-async", - "status": {"state": "PENDING"}, - } - sea_client.execute_command( - operation="SELECT 1", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=True, - enforce_embedded_schema_correctness=False, - query_tags={"job": "nightly-etl"}, - ) - _, kwargs = mock_http_client._make_request.call_args - assert kwargs["data"]["query_tags"] == [{"key": "job", "value": "nightly-etl"}] - - def test_command_management( - self, - sea_client, - mock_http_client, - sea_command_id, - thrift_command_id, - mock_cursor, - ): - """Test command management methods.""" - # Test cancel_command - mock_http_client._make_request.return_value = {} - sea_client.cancel_command(sea_command_id) - mock_http_client._make_request.assert_called_with( - method="POST", - path=sea_client.CANCEL_STATEMENT_PATH_WITH_ID.format("test-statement-123"), - data={"statement_id": "test-statement-123"}, - ) - - # Test cancel_command with invalid ID - with pytest.raises(ValueError) as excinfo: - sea_client.cancel_command(thrift_command_id) - assert "Not a valid SEA command ID" in str(excinfo.value) - - # Test close_command - mock_http_client.reset_mock() - sea_client.close_command(sea_command_id) - mock_http_client._make_request.assert_called_with( - method="DELETE", - path=sea_client.STATEMENT_PATH_WITH_ID.format("test-statement-123"), - data={"statement_id": "test-statement-123"}, - ) - - # Test close_command with invalid ID - with pytest.raises(ValueError) as excinfo: - sea_client.close_command(thrift_command_id) - assert "Not a valid SEA command ID" in str(excinfo.value) - - # Test get_query_state - mock_http_client.reset_mock() - mock_http_client._make_request.return_value = { - "statement_id": "test-statement-123", - "status": {"state": "RUNNING"}, - } - state = sea_client.get_query_state(sea_command_id) - assert state == CommandState.RUNNING - mock_http_client._make_request.assert_called_with( - method="GET", - path=sea_client.STATEMENT_PATH_WITH_ID.format("test-statement-123"), - data={"statement_id": "test-statement-123"}, - ) - - # Test get_query_state with invalid ID - with pytest.raises(ValueError) as excinfo: - sea_client.get_query_state(thrift_command_id) - assert "Not a valid SEA command ID" in str(excinfo.value) - - # Test get_execution_result - mock_http_client.reset_mock() - sea_response = { - "statement_id": "test-statement-123", - "status": {"state": "SUCCEEDED"}, - "manifest": { - "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, - "truncated": False, - }, - "result": { - "chunk_index": 0, - "row_offset": 0, - "row_count": 1, - "data_array": [["1"]], - }, - } - mock_http_client._make_request.return_value = sea_response - result = sea_client.get_execution_result(sea_command_id, mock_cursor) - assert result.command_id.to_sea_statement_id() == "test-statement-123" - assert result.status == CommandState.SUCCEEDED - - # Test get_execution_result with invalid ID - with pytest.raises(ValueError) as excinfo: - sea_client.get_execution_result(thrift_command_id, mock_cursor) - assert "Not a valid SEA command ID" in str(excinfo.value) - - def test_check_command_state(self, sea_client, sea_command_id): - """Test _check_command_not_in_failed_or_closed_state method.""" - # Test with RUNNING state (should not raise) - sea_client._check_command_not_in_failed_or_closed_state( - StatementStatus(state=CommandState.RUNNING), sea_command_id - ) - - # Test with SUCCEEDED state (should not raise) - sea_client._check_command_not_in_failed_or_closed_state( - StatementStatus(state=CommandState.SUCCEEDED), sea_command_id - ) - - # Test with CLOSED state (should raise DatabaseError) - with pytest.raises(DatabaseError) as excinfo: - sea_client._check_command_not_in_failed_or_closed_state( - StatementStatus(state=CommandState.CLOSED), sea_command_id - ) - assert "Command test-statement-123 unexpectedly closed server side" in str( - excinfo.value - ) - - # Test with FAILED state (should raise ServerOperationError) - with pytest.raises(ServerOperationError) as excinfo: - sea_client._check_command_not_in_failed_or_closed_state( - StatementStatus( - state=CommandState.FAILED, - error=ServiceError(message="Test error", error_code="TEST_ERROR"), - ), - sea_command_id, - ) - assert "Command failed" in str(excinfo.value) - - def test_extract_description_from_manifest(self, sea_client): - """Test _extract_description_from_manifest.""" - manifest_obj = MagicMock() - manifest_obj.schema = { - "columns": [ - { - "name": "col1", - "type_name": "STRING", - "type_precision": 10, - "type_scale": 2, - }, - { - "name": "col2", - "type_name": "INT", - "nullable": False, - }, - ] - } - - description = sea_client._extract_description_from_manifest(manifest_obj) - assert description is not None - assert len(description) == 2 - assert description[0][0] == "col1" # name - assert description[0][1] == "string" # type_code - assert description[0][4] == 10 # precision - assert description[0][5] == 2 # scale - assert description[0][6] is None # null_ok - assert description[1][0] == "col2" # name - assert description[1][1] == "int" # type_code - assert description[1][6] is None # null_ok - - def test_extract_description_from_manifest_with_type_normalization( - self, sea_client - ): - """Test _extract_description_from_manifest with SEA to Thrift type normalization.""" - manifest_obj = MagicMock() - manifest_obj.schema = { - "columns": [ - { - "name": "byte_col", - "type_name": "BYTE", - }, - { - "name": "short_col", - "type_name": "SHORT", - }, - { - "name": "long_col", - "type_name": "LONG", - }, - { - "name": "interval_ym_col", - "type_name": "INTERVAL", - "type_interval_type": "YEAR TO MONTH", - }, - { - "name": "interval_dt_col", - "type_name": "INTERVAL", - "type_interval_type": "DAY TO SECOND", - }, - { - "name": "interval_default_col", - "type_name": "INTERVAL", - # No type_interval_type field - }, - ] - } - - description = sea_client._extract_description_from_manifest(manifest_obj) - assert description is not None - assert len(description) == 6 - - # Check normalized types - assert description[0][0] == "byte_col" - assert description[0][1] == "tinyint" # BYTE -> tinyint - - assert description[1][0] == "short_col" - assert description[1][1] == "smallint" # SHORT -> smallint - - assert description[2][0] == "long_col" - assert description[2][1] == "bigint" # LONG -> bigint - - assert description[3][0] == "interval_ym_col" - assert description[3][1] == "interval_year_month" # INTERVAL with YEAR/MONTH - - assert description[4][0] == "interval_dt_col" - assert description[4][1] == "interval_day_time" # INTERVAL with DAY/TIME - - assert description[5][0] == "interval_default_col" - assert description[5][1] == "interval" # INTERVAL without subtype - - def test_filter_session_configuration(self): - """Test that _filter_session_configuration converts all values to strings.""" - session_config = { - "ANSI_MODE": True, - "statement_timeout": 3600, - "TIMEZONE": "UTC", - "enable_photon": False, - "MAX_FILE_PARTITION_BYTES": 128.5, - "QUERY_TAGS": "team:engineering,project:data-pipeline", - "unsupported_param": "value", - "ANOTHER_UNSUPPORTED": 42, - } - - result = _filter_session_configuration(session_config) - - # Verify result is not None - assert result is not None - - # Verify all returned values are strings - for key, value in result.items(): - assert isinstance( - value, str - ), f"Value for key '{key}' is not a string: {type(value)}" - - # Verify specific conversions - expected_result = { - "ansi_mode": "True", # boolean True -> "True", key lowercased - "statement_timeout": "3600", # int -> "3600", key lowercased - "timezone": "UTC", # string -> "UTC", key lowercased - "enable_photon": "False", # boolean False -> "False", key lowercased - "max_file_partition_bytes": "128.5", # float -> "128.5", key lowercased - "query_tags": "team:engineering,project:data-pipeline", - } - - assert result == expected_result - - # Test with None input - assert _filter_session_configuration(None) == {} - - # Test with only unsupported parameters - unsupported_config = { - "unsupported_param1": "value1", - "unsupported_param2": 123, - } - result = _filter_session_configuration(unsupported_config) - assert result == {} - - # Test case insensitivity for keys - case_insensitive_config = { - "ansi_mode": "false", # lowercase key - "STATEMENT_TIMEOUT": 7200, # uppercase key - "TiMeZoNe": "America/New_York", # mixed case key - "QueRy_TaGs": "team:marketing,test:case-insensitive", - } - result = _filter_session_configuration(case_insensitive_config) - expected_case_result = { - "ansi_mode": "false", - "statement_timeout": "7200", - "timezone": "America/New_York", - "query_tags": "team:marketing,test:case-insensitive", - } - assert result == expected_case_result - - # Verify all values are strings in case insensitive test - for key, value in result.items(): - assert isinstance( - value, str - ), f"Value for key '{key}' is not a string: {type(value)}" - - def test_results_message_to_execute_response_is_staging_operation(self, sea_client): - """Test that is_staging_operation is correctly set from manifest.is_volume_operation.""" - # Test when is_volume_operation is True - response = MagicMock() - response.statement_id = "test-statement-123" - response.status.state = CommandState.SUCCEEDED - response.manifest.is_volume_operation = True - response.manifest.result_compression = "NONE" - response.manifest.format = "JSON_ARRAY" - - # Mock the _extract_description_from_manifest method to return None - with patch.object( - sea_client, "_extract_description_from_manifest", return_value=None - ): - result = sea_client._results_message_to_execute_response(response) - assert result.is_staging_operation is True - - # Test when is_volume_operation is False - response.manifest.is_volume_operation = False - with patch.object( - sea_client, "_extract_description_from_manifest", return_value=None - ): - result = sea_client._results_message_to_execute_response(response) - assert result.is_staging_operation is False - - def test_get_catalogs(self, sea_client, sea_session_id, mock_cursor): - """Test the get_catalogs method.""" - # Mock the execute_command method - mock_result_set = Mock() - with patch.object( - sea_client, "execute_command", return_value=mock_result_set - ) as mock_execute: - # Call get_catalogs - result = sea_client.get_catalogs( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - ) - - # Verify execute_command was called with the correct parameters - mock_execute.assert_called_once_with( - operation="SHOW CATALOGS", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - - # Verify the result is correct - assert result == mock_result_set - - def test_get_schemas(self, sea_client, sea_session_id, mock_cursor): - """Test the get_schemas method with various parameter combinations.""" - # Mock the execute_command method - mock_result_set = Mock() - with patch.object( - sea_client, "execute_command", return_value=mock_result_set - ) as mock_execute: - # Case 1: With catalog name only - result = sea_client.get_schemas( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - ) - - mock_execute.assert_called_with( - operation="SHOW SCHEMAS IN test_catalog", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - - # Case 2: With catalog and schema names - result = sea_client.get_schemas( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - schema_name="test_schema", - ) - - mock_execute.assert_called_with( - operation="SHOW SCHEMAS IN test_catalog LIKE 'test_schema'", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - - # Case 3: Without catalog name (should raise ValueError) - with pytest.raises(DatabaseError) as excinfo: - sea_client.get_schemas( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - ) - assert "Catalog name is required for get_schemas" in str(excinfo.value) - - def test_get_tables(self, sea_client, sea_session_id, mock_cursor): - """Test the get_tables method with various parameter combinations.""" - # Mock the execute_command method - from databricks.sql.backend.sea.result_set import SeaResultSet - - mock_result_set = Mock(spec=SeaResultSet) - - with patch.object( - sea_client, "execute_command", return_value=mock_result_set - ) as mock_execute: - # Mock the filter_tables_by_type method - with patch( - "databricks.sql.backend.sea.utils.filters.ResultSetFilter.filter_tables_by_type", - return_value=mock_result_set, - ) as mock_filter: - # Case 1: With catalog name only - result = sea_client.get_tables( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - ) - - mock_execute.assert_called_with( - operation="SHOW TABLES IN CATALOG test_catalog", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - mock_filter.assert_called_with(mock_result_set, None) - - # Case 2: With all parameters - table_types = ["TABLE", "VIEW"] - result = sea_client.get_tables( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - schema_name="test_schema", - table_name="test_table", - table_types=table_types, - ) - - mock_execute.assert_called_with( - operation="SHOW TABLES IN CATALOG test_catalog SCHEMA LIKE 'test_schema' LIKE 'test_table'", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - mock_filter.assert_called_with(mock_result_set, table_types) - - # Case 3: With wildcard catalog - result = sea_client.get_tables( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="*", - ) - - mock_execute.assert_called_with( - operation="SHOW TABLES IN ALL CATALOGS", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - - def test_get_columns(self, sea_client, sea_session_id, mock_cursor): - """Test the get_columns method with various parameter combinations.""" - # Mock the execute_command method - mock_result_set = Mock() - with patch.object( - sea_client, "execute_command", return_value=mock_result_set - ) as mock_execute: - # Case 1: With catalog name only - result = sea_client.get_columns( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - ) - - mock_execute.assert_called_with( - operation="SHOW COLUMNS IN CATALOG test_catalog", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - - # Case 2: With all parameters - result = sea_client.get_columns( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - schema_name="test_schema", - table_name="test_table", - column_name="test_column", - ) - - mock_execute.assert_called_with( - operation="SHOW COLUMNS IN CATALOG test_catalog SCHEMA LIKE 'test_schema' TABLE LIKE 'test_table' LIKE 'test_column'", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=False, - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - - # Case 3: Without catalog name (should raise ValueError) - with pytest.raises(DatabaseError) as excinfo: - sea_client.get_columns( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - ) - assert "Catalog name is required for get_columns" in str(excinfo.value) - - def test_get_tables_with_cloud_fetch( - self, sea_client_cloud_fetch, sea_session_id, mock_cursor - ): - """Test the get_tables method with cloud fetch enabled.""" - # Mock the execute_command method and ResultSetFilter - mock_result_set = Mock() - - with patch.object( - sea_client_cloud_fetch, "execute_command", return_value=mock_result_set - ) as mock_execute: - with patch( - "databricks.sql.backend.sea.utils.filters.ResultSetFilter" - ) as mock_filter: - mock_filter.filter_tables_by_type.return_value = mock_result_set - - # Call get_tables - result = sea_client_cloud_fetch.get_tables( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - ) - - # Verify execute_command was called with use_cloud_fetch=True - mock_execute.assert_called_with( - operation="SHOW TABLES IN CATALOG test_catalog", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=True, # Should use True since client was created with use_cloud_fetch=True - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result == mock_result_set - - def test_get_schemas_with_cloud_fetch( - self, sea_client_cloud_fetch, sea_session_id, mock_cursor - ): - """Test the get_schemas method with cloud fetch enabled.""" - # Mock the execute_command method - mock_result_set = Mock() - with patch.object( - sea_client_cloud_fetch, "execute_command", return_value=mock_result_set - ) as mock_execute: - # Test with catalog name - result = sea_client_cloud_fetch.get_schemas( - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - cursor=mock_cursor, - catalog_name="test_catalog", - ) - - mock_execute.assert_called_with( - operation="SHOW SCHEMAS IN test_catalog", - session_id=sea_session_id, - max_rows=100, - max_bytes=1000, - lz4_compression=False, - cursor=mock_cursor, - use_cloud_fetch=True, # Should use True since client was created with use_cloud_fetch=True - parameters=[], - async_op=False, - enforce_embedded_schema_correctness=False, - ) - assert result == mock_result_set diff --git a/tests/unit/test_sea_conversion.py b/tests/unit/test_sea_conversion.py deleted file mode 100644 index 234cca868..000000000 --- a/tests/unit/test_sea_conversion.py +++ /dev/null @@ -1,149 +0,0 @@ -""" -Tests for the conversion module in the SEA backend. - -This module contains tests for the SqlType and SqlTypeConverter classes. -""" - -import pytest -import datetime -import decimal -from unittest.mock import Mock, patch - -from databricks.sql.backend.sea.utils.conversion import SqlType, SqlTypeConverter - - -class TestSqlTypeConverter: - """Test suite for the SqlTypeConverter class.""" - - def test_convert_numeric_types(self): - """Test converting numeric types.""" - # Test integer types - assert SqlTypeConverter.convert_value("123", SqlType.TINYINT, None) == 123 - assert SqlTypeConverter.convert_value("456", SqlType.SMALLINT, None) == 456 - assert SqlTypeConverter.convert_value("789", SqlType.INT, None) == 789 - assert ( - SqlTypeConverter.convert_value("1234567890", SqlType.BIGINT, None) - == 1234567890 - ) - - # Test floating point types - assert SqlTypeConverter.convert_value("123.45", SqlType.FLOAT, None) == 123.45 - assert SqlTypeConverter.convert_value("678.90", SqlType.DOUBLE, None) == 678.90 - - # Test decimal type - decimal_value = SqlTypeConverter.convert_value("123.45", SqlType.DECIMAL, None) - assert isinstance(decimal_value, decimal.Decimal) - assert decimal_value == decimal.Decimal("123.45") - - # Test decimal with precision and scale - decimal_value = SqlTypeConverter.convert_value( - "123.45", SqlType.DECIMAL, None, precision=5, scale=2 - ) - assert isinstance(decimal_value, decimal.Decimal) - assert decimal_value == decimal.Decimal("123.45") - - # Test invalid numeric input - result = SqlTypeConverter.convert_value("not_a_number", SqlType.INT, None) - assert result == "not_a_number" # Returns original value on error - - def test_convert_boolean_type(self): - """Test converting boolean types.""" - # True values - assert SqlTypeConverter.convert_value("true", SqlType.BOOLEAN, None) is True - assert SqlTypeConverter.convert_value("True", SqlType.BOOLEAN, None) is True - assert SqlTypeConverter.convert_value("t", SqlType.BOOLEAN, None) is True - assert SqlTypeConverter.convert_value("1", SqlType.BOOLEAN, None) is True - assert SqlTypeConverter.convert_value("yes", SqlType.BOOLEAN, None) is True - assert SqlTypeConverter.convert_value("y", SqlType.BOOLEAN, None) is True - - # False values - assert SqlTypeConverter.convert_value("false", SqlType.BOOLEAN, None) is False - assert SqlTypeConverter.convert_value("False", SqlType.BOOLEAN, None) is False - assert SqlTypeConverter.convert_value("f", SqlType.BOOLEAN, None) is False - assert SqlTypeConverter.convert_value("0", SqlType.BOOLEAN, None) is False - assert SqlTypeConverter.convert_value("no", SqlType.BOOLEAN, None) is False - assert SqlTypeConverter.convert_value("n", SqlType.BOOLEAN, None) is False - - def test_convert_datetime_types(self): - """Test converting datetime types.""" - # Test date type - date_value = SqlTypeConverter.convert_value("2023-01-15", SqlType.DATE, None) - assert isinstance(date_value, datetime.date) - assert date_value == datetime.date(2023, 1, 15) - - # Test timestamp type - timestamp_value = SqlTypeConverter.convert_value( - "2023-01-15T12:30:45", SqlType.TIMESTAMP, None - ) - assert isinstance(timestamp_value, datetime.datetime) - assert timestamp_value.year == 2023 - assert timestamp_value.month == 1 - assert timestamp_value.day == 15 - assert timestamp_value.hour == 12 - assert timestamp_value.minute == 30 - assert timestamp_value.second == 45 - - # Test interval types (currently return as string) - interval_ym_value = SqlTypeConverter.convert_value( - "1-6", SqlType.INTERVAL_YEAR_MONTH, None - ) - assert interval_ym_value == "1-6" - - interval_dt_value = SqlTypeConverter.convert_value( - "1 day 2 hours", SqlType.INTERVAL_DAY_TIME, None - ) - assert interval_dt_value == "1 day 2 hours" - - # Test invalid date input - result = SqlTypeConverter.convert_value("not_a_date", SqlType.DATE, None) - assert result == "not_a_date" # Returns original value on error - - def test_convert_string_types(self): - """Test converting string types.""" - # String types don't need conversion, they should be returned as-is - assert ( - SqlTypeConverter.convert_value("test string", SqlType.STRING, None) - == "test string" - ) - assert ( - SqlTypeConverter.convert_value("test char", SqlType.CHAR, None) - == "test char" - ) - assert ( - SqlTypeConverter.convert_value("test varchar", SqlType.VARCHAR, None) - == "test varchar" - ) - - def test_convert_binary_type(self): - """Test converting binary type.""" - # Test valid hex string - binary_value = SqlTypeConverter.convert_value( - "48656C6C6F", SqlType.BINARY, None - ) - assert isinstance(binary_value, bytes) - assert binary_value == b"Hello" - - # Test invalid binary input - result = SqlTypeConverter.convert_value("not_hex", SqlType.BINARY, None) - assert result == "not_hex" # Returns original value on error - - def test_convert_unsupported_type(self): - """Test converting an unsupported type.""" - # Should return the original value - assert ( - SqlTypeConverter.convert_value("test", "unsupported_type", None) == "test" - ) - - # Complex types should return as-is (not yet implemented in TYPE_MAPPING) - assert ( - SqlTypeConverter.convert_value("complex_value", SqlType.ARRAY, None) - == "complex_value" - ) - assert ( - SqlTypeConverter.convert_value("complex_value", SqlType.MAP, None) - == "complex_value" - ) - assert ( - SqlTypeConverter.convert_value("complex_value", SqlType.STRUCT, None) - == "complex_value" - ) diff --git a/tests/unit/test_sea_http_client.py b/tests/unit/test_sea_http_client.py deleted file mode 100644 index 5100a5cb0..000000000 --- a/tests/unit/test_sea_http_client.py +++ /dev/null @@ -1,235 +0,0 @@ -import json -import unittest -from unittest.mock import patch, Mock, MagicMock -import pytest - -from databricks.sql.backend.sea.utils.http_client import SeaHttpClient -from databricks.sql.auth.retry import CommandType -from databricks.sql.auth.authenticators import AuthProvider -from databricks.sql.types import SSLOptions -from databricks.sql.exc import RequestError - - -class TestSeaHttpClient: - @pytest.fixture - def mock_auth_provider(self): - auth_provider = Mock(spec=AuthProvider) - auth_provider.add_headers = Mock(return_value=None) - return auth_provider - - @pytest.fixture - def ssl_options(self): - return SSLOptions( - tls_verify=True, - tls_trusted_ca_file=None, - tls_client_cert_file=None, - tls_client_cert_key_file=None, - tls_client_cert_key_password=None, - ) - - @pytest.fixture - def sea_http_client(self, mock_auth_provider, ssl_options): - with patch( - "databricks.sql.backend.sea.utils.http_client.HTTPSConnectionPool" - ) as mock_pool: - client = SeaHttpClient( - server_hostname="test-server.databricks.com", - port=443, - http_path="/sql/1.0/warehouses/abc123", - http_headers=[("User-Agent", "test-agent")], - auth_provider=mock_auth_provider, - ssl_options=ssl_options, - ) - # Replace the real pool with a mock - client._pool = Mock() - return client - - @pytest.mark.parametrize( - "server_hostname,port,expected_base_url", - [ - # Basic hostname without protocol - ("myserver.com", 443, "https://myserver.com:443"), - # Hostname with trailing slash - ("myserver.com/", 443, "https://myserver.com:443"), - # Hostname with https:// protocol - ("https://myserver.com", 443, "https://myserver.com:443"), - # Hostname with http:// protocol (preserved as-is) - ("http://myserver.com", 443, "http://myserver.com:443"), - # Hostname with protocol and trailing slash - ("https://myserver.com/", 443, "https://myserver.com:443"), - # Custom port - ("myserver.com", 8080, "https://myserver.com:8080"), - # Protocol with custom port - ("https://myserver.com", 8080, "https://myserver.com:8080"), - ], - ) - def test_base_url_construction( - self, server_hostname, port, expected_base_url, mock_auth_provider, ssl_options - ): - """Test that base_url is constructed correctly from various hostname inputs.""" - with patch("databricks.sql.backend.sea.utils.http_client.HTTPSConnectionPool"): - client = SeaHttpClient( - server_hostname=server_hostname, - port=port, - http_path="/sql/1.0/warehouses/test", - http_headers=[], - auth_provider=mock_auth_provider, - ssl_options=ssl_options, - ) - assert client.base_url == expected_base_url - - def test_get_command_type_from_path(self, sea_http_client): - """Test the _get_command_type_from_path method with various paths and methods.""" - # Test statement execution - assert ( - sea_http_client._get_command_type_from_path("/statements", "POST") - == CommandType.EXECUTE_STATEMENT - ) - - # Test statement cancellation - assert ( - sea_http_client._get_command_type_from_path( - "/statements/123/cancel", "POST" - ) - == CommandType.OTHER - ) - - # Test statement deletion (close operation) - assert ( - sea_http_client._get_command_type_from_path("/statements/123", "DELETE") - == CommandType.CLOSE_OPERATION - ) - - # Test get statement status - assert ( - sea_http_client._get_command_type_from_path("/statements/123", "GET") - == CommandType.GET_OPERATION_STATUS - ) - - # Test session close - assert ( - sea_http_client._get_command_type_from_path("/sessions/456", "DELETE") - == CommandType.CLOSE_SESSION - ) - - # Test other paths - assert ( - sea_http_client._get_command_type_from_path("/other/endpoint", "GET") - == CommandType.OTHER - ) - assert ( - sea_http_client._get_command_type_from_path("/other/endpoint", "POST") - == CommandType.OTHER - ) - - @patch( - "databricks.sql.backend.sea.utils.http_client.SeaHttpClient._get_auth_headers" - ) - def test_make_request_success(self, mock_get_auth_headers, sea_http_client): - """Test successful _make_request calls.""" - # Setup mock response - mock_response = Mock() - mock_response.status = 200 - # Mock response.data.decode() to return a valid JSON string - mock_response.data.decode.return_value = '{"result": "success"}' - mock_response.__enter__ = Mock(return_value=mock_response) - mock_response.__exit__ = Mock(return_value=None) - - # Setup mock auth headers - mock_get_auth_headers.return_value = {"Authorization": "Bearer test-token"} - - # Configure the pool's request method to return our mock response - sea_http_client._pool.request.return_value = mock_response - - # Test GET request without data - result = sea_http_client._make_request("GET", "/test/path") - - # Verify the request was made correctly - sea_http_client._pool.request.assert_called_with( - method="GET", - url="/test/path", - body=b"", - headers={ - "Content-Type": "application/json", - "User-Agent": "test-agent", - "Authorization": "Bearer test-token", - }, - preload_content=False, - retries=sea_http_client.retry_policy, - ) - - # Check the result - assert result == {"result": "success"} - - # Test POST request with data - test_data = {"query": "SELECT * FROM test"} - result = sea_http_client._make_request("POST", "/statements", test_data) - - # Verify the request was made with the correct body - expected_body = json.dumps(test_data).encode("utf-8") - sea_http_client._pool.request.assert_called_with( - method="POST", - url="/statements", - body=expected_body, - headers={ - "Content-Type": "application/json", - "User-Agent": "test-agent", - "Authorization": "Bearer test-token", - "Content-Length": str(len(expected_body)), - }, - preload_content=False, - retries=sea_http_client.retry_policy, - ) - - @patch( - "databricks.sql.backend.sea.utils.http_client.SeaHttpClient._get_auth_headers" - ) - def test_make_request_error_response(self, mock_get_auth_headers, sea_http_client): - """Test _make_request with error HTTP status.""" - # Setup mock response with error status - mock_response = Mock() - mock_response.status = 400 - mock_response.__enter__ = Mock(return_value=mock_response) - mock_response.__exit__ = Mock(return_value=None) - - # Setup mock auth headers - mock_get_auth_headers.return_value = {"Authorization": "Bearer test-token"} - - # Configure the pool's request method to return our mock response - sea_http_client._pool.request.return_value = mock_response - - # Test request with error response - with pytest.raises(Exception) as excinfo: - sea_http_client._make_request("GET", "/test/path") - - assert "SEA HTTP request failed with status 400" in str(excinfo.value) - - @patch( - "databricks.sql.backend.sea.utils.http_client.SeaHttpClient._get_auth_headers" - ) - def test_make_request_connection_error( - self, mock_get_auth_headers, sea_http_client - ): - """Test _make_request with connection error.""" - # Setup mock auth headers - mock_get_auth_headers.return_value = {"Authorization": "Bearer test-token"} - - # Configure the pool's request to raise an exception - sea_http_client._pool.request.side_effect = Exception("Connection error") - - # Test request with connection error - with pytest.raises(RequestError) as excinfo: - sea_http_client._make_request("GET", "/test/path") - - assert "Error during request to server" in str(excinfo.value) - - def test_make_request_no_pool(self, sea_http_client): - """Test _make_request when pool is not initialized.""" - # Set pool to None to simulate uninitialized pool - sea_http_client._pool = None - - # Test request with no pool - with pytest.raises(RequestError) as excinfo: - sea_http_client._make_request("GET", "/test/path") - - assert "Connection pool not initialized" in str(excinfo.value) diff --git a/tests/unit/test_sea_queue.py b/tests/unit/test_sea_queue.py deleted file mode 100644 index 6471cb4fd..000000000 --- a/tests/unit/test_sea_queue.py +++ /dev/null @@ -1,735 +0,0 @@ -""" -Tests for SEA-related queue classes. - -This module contains tests for the JsonQueue, SeaResultSetQueueFactory, and SeaCloudFetchQueue classes. -It also tests the Hybrid disposition which can create either ArrowQueue or SeaCloudFetchQueue based on -whether attachment is set. -""" - -import pytest -from unittest.mock import Mock, patch, MagicMock - -from databricks.sql.backend.sea.queue import ( - JsonQueue, - LinkFetcher, - SeaResultSetQueueFactory, - SeaCloudFetchQueue, -) -from databricks.sql.backend.sea.models.base import ( - ResultData, - ResultManifest, - ExternalLink, -) -from databricks.sql.backend.sea.utils.constants import ResultFormat -from databricks.sql.exc import ProgrammingError, ServerOperationError -from databricks.sql.types import SSLOptions -from databricks.sql.utils import ArrowQueue -import threading -import time - - -class TestJsonQueue: - """Test suite for the JsonQueue class.""" - - @pytest.fixture - def sample_data(self): - """Create sample data for testing.""" - return [ - ["value1", 1, True], - ["value2", 2, False], - ["value3", 3, True], - ["value4", 4, False], - ["value5", 5, True], - ] - - def test_init(self, sample_data): - """Test initialization of JsonQueue.""" - queue = JsonQueue(sample_data) - assert queue.data_array == sample_data - assert queue.cur_row_index == 0 - assert queue.num_rows == len(sample_data) - - def test_init_with_none(self): - """Test initialization with None data.""" - queue = JsonQueue(None) - assert queue.data_array == [] - assert queue.cur_row_index == 0 - assert queue.num_rows == 0 - - def test_next_n_rows_partial(self, sample_data): - """Test fetching a subset of rows.""" - queue = JsonQueue(sample_data) - result = queue.next_n_rows(2) - assert result == sample_data[:2] - assert queue.cur_row_index == 2 - - def test_next_n_rows_all(self, sample_data): - """Test fetching all rows.""" - queue = JsonQueue(sample_data) - result = queue.next_n_rows(len(sample_data)) - assert result == sample_data - assert queue.cur_row_index == len(sample_data) - - def test_next_n_rows_more_than_available(self, sample_data): - """Test fetching more rows than available.""" - queue = JsonQueue(sample_data) - result = queue.next_n_rows(len(sample_data) + 10) - assert result == sample_data - assert queue.cur_row_index == len(sample_data) - - def test_next_n_rows_zero(self, sample_data): - """Test fetching zero rows.""" - queue = JsonQueue(sample_data) - result = queue.next_n_rows(0) - assert result == [] - assert queue.cur_row_index == 0 - - def test_remaining_rows(self, sample_data): - """Test fetching all remaining rows.""" - queue = JsonQueue(sample_data) - - # Fetch some rows first - queue.next_n_rows(2) - - # Now fetch remaining - result = queue.remaining_rows() - assert result == sample_data[2:] - assert queue.cur_row_index == len(sample_data) - - def test_remaining_rows_all(self, sample_data): - """Test fetching all remaining rows from the start.""" - queue = JsonQueue(sample_data) - result = queue.remaining_rows() - assert result == sample_data - assert queue.cur_row_index == len(sample_data) - - def test_remaining_rows_empty(self, sample_data): - """Test fetching remaining rows when none are left.""" - queue = JsonQueue(sample_data) - - # Fetch all rows first - queue.next_n_rows(len(sample_data)) - - # Now fetch remaining (should be empty) - result = queue.remaining_rows() - assert result == [] - assert queue.cur_row_index == len(sample_data) - - -class TestSeaResultSetQueueFactory: - """Test suite for the SeaResultSetQueueFactory class.""" - - @pytest.fixture - def json_manifest(self): - """Create a JSON manifest for testing.""" - return ResultManifest( - format=ResultFormat.JSON_ARRAY.value, - schema={}, - total_row_count=5, - total_byte_count=1000, - total_chunk_count=1, - ) - - @pytest.fixture - def arrow_manifest(self): - """Create an Arrow manifest for testing.""" - return ResultManifest( - format=ResultFormat.ARROW_STREAM.value, - schema={}, - total_row_count=5, - total_byte_count=1000, - total_chunk_count=1, - ) - - @pytest.fixture - def invalid_manifest(self): - """Create an invalid manifest for testing.""" - return ResultManifest( - format="INVALID_FORMAT", - schema={}, - total_row_count=5, - total_byte_count=1000, - total_chunk_count=1, - ) - - @pytest.fixture - def sample_data(self): - """Create sample result data.""" - return [ - ["value1", "1", "true"], - ["value2", "2", "false"], - ] - - @pytest.fixture - def ssl_options(self): - """Create SSL options for testing.""" - return SSLOptions(tls_verify=True) - - @pytest.fixture - def mock_sea_client(self): - """Create a mock SEA client.""" - client = Mock() - client.max_download_threads = 10 - return client - - @pytest.fixture - def description(self): - """Create column descriptions.""" - return [ - ("col1", "string", None, None, None, None, None), - ("col2", "int", None, None, None, None, None), - ("col3", "boolean", None, None, None, None, None), - ] - - def test_build_queue_json_array(self, json_manifest, sample_data): - """Test building a JSON array queue.""" - result_data = ResultData(data=sample_data) - mock_http_client = MagicMock() - - queue = SeaResultSetQueueFactory.build_queue( - result_data=result_data, - manifest=json_manifest, - statement_id="test-statement", - ssl_options=SSLOptions(), - description=[], - max_download_threads=10, - sea_client=Mock(), - lz4_compressed=False, - http_client=mock_http_client, - ) - - assert isinstance(queue, JsonQueue) - assert queue.data_array == sample_data - - def test_build_queue_arrow_stream( - self, arrow_manifest, ssl_options, mock_sea_client, description - ): - """Test building an Arrow stream queue.""" - external_links = [ - ExternalLink( - external_link="https://example.com/data/chunk0", - expiration="2025-07-03T05:51:18.118009", - row_count=100, - byte_count=1024, - row_offset=0, - chunk_index=0, - next_chunk_index=1, - http_headers={"Authorization": "Bearer token123"}, - ) - ] - result_data = ResultData(data=None, external_links=external_links) - - mock_http_client = MagicMock() - - with patch( - "databricks.sql.backend.sea.queue.ResultFileDownloadManager" - ), patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None): - queue = SeaResultSetQueueFactory.build_queue( - result_data=result_data, - manifest=arrow_manifest, - statement_id="test-statement", - ssl_options=ssl_options, - description=description, - max_download_threads=10, - sea_client=mock_sea_client, - lz4_compressed=False, - http_client=mock_http_client, - ) - - assert isinstance(queue, SeaCloudFetchQueue) - - def test_build_queue_invalid_format(self, invalid_manifest): - """Test building a queue with invalid format.""" - result_data = ResultData(data=[]) - mock_http_client = MagicMock() - - with pytest.raises(ProgrammingError, match="Invalid result format"): - SeaResultSetQueueFactory.build_queue( - result_data=result_data, - manifest=invalid_manifest, - statement_id="test-statement", - ssl_options=SSLOptions(), - description=[], - max_download_threads=10, - sea_client=Mock(), - lz4_compressed=False, - http_client=mock_http_client, - ) - - -class TestSeaCloudFetchQueue: - """Test suite for the SeaCloudFetchQueue class.""" - - @pytest.fixture - def ssl_options(self): - """Create SSL options for testing.""" - return SSLOptions(tls_verify=True) - - @pytest.fixture - def mock_sea_client(self): - """Create a mock SEA client.""" - client = Mock() - client.max_download_threads = 10 - return client - - @pytest.fixture - def description(self): - """Create column descriptions.""" - return [ - ("col1", "string", None, None, None, None, None), - ("col2", "int", None, None, None, None, None), - ("col3", "boolean", None, None, None, None, None), - ] - - @pytest.fixture - def sample_external_link(self): - """Create a sample external link.""" - return ExternalLink( - external_link="https://example.com/data/chunk0", - expiration="2025-07-03T05:51:18.118009", - row_count=100, - byte_count=1024, - row_offset=0, - chunk_index=0, - next_chunk_index=1, - http_headers={"Authorization": "Bearer token123"}, - ) - - @pytest.fixture - def sample_external_link_no_headers(self): - """Create a sample external link without headers.""" - return ExternalLink( - external_link="https://example.com/data/chunk0", - expiration="2025-07-03T05:51:18.118009", - row_count=100, - byte_count=1024, - row_offset=0, - chunk_index=0, - next_chunk_index=1, - http_headers=None, - ) - - def test_convert_to_thrift_link(self, sample_external_link): - """Test conversion of ExternalLink to TSparkArrowResultLink.""" - # Call the method directly - result = LinkFetcher._convert_to_thrift_link(sample_external_link) - - # Verify the conversion - assert result.fileLink == sample_external_link.external_link - assert result.rowCount == sample_external_link.row_count - assert result.bytesNum == sample_external_link.byte_count - assert result.startRowOffset == sample_external_link.row_offset - assert result.httpHeaders == sample_external_link.http_headers - - def test_convert_to_thrift_link_no_headers(self, sample_external_link_no_headers): - """Test conversion of ExternalLink with no headers to TSparkArrowResultLink.""" - # Call the method directly - result = LinkFetcher._convert_to_thrift_link(sample_external_link_no_headers) - - # Verify the conversion - assert result.fileLink == sample_external_link_no_headers.external_link - assert result.rowCount == sample_external_link_no_headers.row_count - assert result.bytesNum == sample_external_link_no_headers.byte_count - assert result.startRowOffset == sample_external_link_no_headers.row_offset - assert result.httpHeaders == {} - - @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager") - @patch("databricks.sql.backend.sea.queue.logger") - def test_init_with_valid_initial_link( - self, - mock_logger, - mock_download_manager_class, - mock_sea_client, - ssl_options, - description, - sample_external_link, - ): - """Test initialization with valid initial link.""" - # Create a queue with valid initial link - mock_http_client = MagicMock() - with patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None): - queue = SeaCloudFetchQueue( - result_data=ResultData(external_links=[sample_external_link]), - max_download_threads=5, - ssl_options=ssl_options, - sea_client=mock_sea_client, - statement_id="test-statement-123", - total_chunk_count=1, - lz4_compressed=False, - description=description, - http_client=mock_http_client, - ) - - # Verify attributes - assert queue._current_chunk_index == 0 - assert queue.link_fetcher is not None - - @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager") - @patch("databricks.sql.backend.sea.queue.logger") - def test_init_no_initial_links( - self, - mock_logger, - mock_download_manager_class, - mock_sea_client, - ssl_options, - description, - ): - """Test initialization with no initial links.""" - # Create a queue with empty initial links - mock_http_client = MagicMock() - queue = SeaCloudFetchQueue( - result_data=ResultData(external_links=[]), - max_download_threads=5, - ssl_options=ssl_options, - sea_client=mock_sea_client, - statement_id="test-statement-123", - total_chunk_count=0, - lz4_compressed=False, - description=description, - http_client=mock_http_client, - ) - assert queue.table is None - - @patch("databricks.sql.backend.sea.queue.logger") - def test_create_next_table_success(self, mock_logger): - """Test _create_next_table with successful table creation.""" - # Create a queue instance without initializing - queue = Mock(spec=SeaCloudFetchQueue) - queue._current_chunk_index = 0 - queue.download_manager = Mock() - queue.link_fetcher = Mock() - - # Mock the dependencies - mock_table = Mock() - mock_chunk_link = Mock() - queue.link_fetcher.get_chunk_link = Mock(return_value=mock_chunk_link) - queue._create_table_at_offset = Mock(return_value=mock_table) - - # Call the method directly - SeaCloudFetchQueue._create_next_table(queue) - - # Verify the chunk index was incremented - assert queue._current_chunk_index == 1 - - # Verify the chunk link was retrieved - queue.link_fetcher.get_chunk_link.assert_called_once_with(0) - - # Verify the table was created from the link - queue._create_table_at_offset.assert_called_once_with( - mock_chunk_link.row_offset - ) - - -class TestHybridDisposition: - """Test suite for the Hybrid disposition handling in SeaResultSetQueueFactory.""" - - @pytest.fixture - def arrow_manifest(self): - """Create an Arrow manifest for testing.""" - return ResultManifest( - format=ResultFormat.ARROW_STREAM.value, - schema={}, - total_row_count=5, - total_byte_count=1000, - total_chunk_count=1, - ) - - @pytest.fixture - def description(self): - """Create column descriptions.""" - return [ - ("col1", "string", None, None, None, None, None), - ("col2", "int", None, None, None, None, None), - ("col3", "boolean", None, None, None, None, None), - ] - - @pytest.fixture - def ssl_options(self): - """Create SSL options for testing.""" - return SSLOptions(tls_verify=True) - - @pytest.fixture - def mock_sea_client(self): - """Create a mock SEA client.""" - client = Mock() - client.max_download_threads = 10 - return client - - @patch("databricks.sql.backend.sea.queue.create_arrow_table_from_arrow_file") - def test_hybrid_disposition_with_attachment( - self, - mock_create_table, - arrow_manifest, - description, - ssl_options, - mock_sea_client, - ): - """Test that ArrowQueue is created when attachment is present.""" - # Create mock arrow table - mock_arrow_table = Mock() - mock_arrow_table.num_rows = 5 - mock_create_table.return_value = mock_arrow_table - - # Create result data with attachment - attachment_data = b"mock_arrow_data" - result_data = ResultData(attachment=attachment_data) - mock_http_client = MagicMock() - # Build queue - queue = SeaResultSetQueueFactory.build_queue( - result_data=result_data, - manifest=arrow_manifest, - statement_id="test-statement", - ssl_options=ssl_options, - description=description, - max_download_threads=10, - sea_client=mock_sea_client, - lz4_compressed=False, - http_client=mock_http_client, - ) - - # Verify ArrowQueue was created - assert isinstance(queue, ArrowQueue) - mock_create_table.assert_called_once_with(attachment_data, description) - - @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager") - @patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None) - def test_hybrid_disposition_with_external_links( - self, - mock_create_table, - mock_download_manager, - arrow_manifest, - description, - ssl_options, - mock_sea_client, - ): - """Test that SeaCloudFetchQueue is created when attachment is None but external links are present.""" - # Create external links - external_links = [ - ExternalLink( - external_link="https://example.com/data/chunk0", - expiration="2025-07-03T05:51:18.118009", - row_count=100, - byte_count=1024, - row_offset=0, - chunk_index=0, - next_chunk_index=1, - http_headers={"Authorization": "Bearer token123"}, - ) - ] - - # Create result data with external links but no attachment - result_data = ResultData(external_links=external_links, attachment=None) - - # Build queue - mock_http_client = MagicMock() - queue = SeaResultSetQueueFactory.build_queue( - result_data=result_data, - manifest=arrow_manifest, - statement_id="test-statement", - ssl_options=ssl_options, - description=description, - max_download_threads=10, - sea_client=mock_sea_client, - lz4_compressed=False, - http_client=mock_http_client, - ) - - # Verify SeaCloudFetchQueue was created - assert isinstance(queue, SeaCloudFetchQueue) - mock_create_table.assert_called_once() - - @patch("databricks.sql.backend.sea.queue.ResultSetDownloadHandler._decompress_data") - @patch("databricks.sql.backend.sea.queue.create_arrow_table_from_arrow_file") - def test_hybrid_disposition_with_compressed_attachment( - self, - mock_create_table, - mock_decompress, - arrow_manifest, - description, - ssl_options, - mock_sea_client, - ): - """Test that ArrowQueue is created with decompressed data when attachment is present and lz4_compressed is True.""" - # Create mock arrow table - mock_arrow_table = Mock() - mock_arrow_table.num_rows = 5 - mock_create_table.return_value = mock_arrow_table - - # Setup decompression mock - compressed_data = b"compressed_data" - decompressed_data = b"decompressed_data" - mock_decompress.return_value = decompressed_data - - # Create result data with attachment - result_data = ResultData(attachment=compressed_data) - mock_http_client = MagicMock() - # Build queue with lz4_compressed=True - queue = SeaResultSetQueueFactory.build_queue( - result_data=result_data, - manifest=arrow_manifest, - statement_id="test-statement", - ssl_options=ssl_options, - description=description, - max_download_threads=10, - sea_client=mock_sea_client, - lz4_compressed=True, - http_client=mock_http_client, - ) - - # Verify ArrowQueue was created with decompressed data - assert isinstance(queue, ArrowQueue) - mock_decompress.assert_called_once_with(compressed_data) - mock_create_table.assert_called_once_with(decompressed_data, description) - - -class TestLinkFetcher: - """Unit tests for the LinkFetcher helper class.""" - - @pytest.fixture - def sample_links(self): - """Provide a pair of ExternalLink objects forming two sequential chunks.""" - link0 = ExternalLink( - external_link="https://example.com/data/chunk0", - expiration="2030-01-01T00:00:00.000000", - row_count=100, - byte_count=1024, - row_offset=0, - chunk_index=0, - next_chunk_index=1, - http_headers={"Authorization": "Bearer token0"}, - ) - - link1 = ExternalLink( - external_link="https://example.com/data/chunk1", - expiration="2030-01-01T00:00:00.000000", - row_count=100, - byte_count=1024, - row_offset=100, - chunk_index=1, - next_chunk_index=None, - http_headers={"Authorization": "Bearer token1"}, - ) - - return link0, link1 - - def _create_fetcher( - self, - initial_links, - backend_mock=None, - download_manager_mock=None, - total_chunk_count=10, - ): - """Helper to create a LinkFetcher instance with supplied mocks.""" - if backend_mock is None: - backend_mock = Mock() - if download_manager_mock is None: - download_manager_mock = Mock() - - return ( - LinkFetcher( - download_manager=download_manager_mock, - backend=backend_mock, - statement_id="statement-123", - initial_links=list(initial_links), - total_chunk_count=total_chunk_count, - ), - backend_mock, - download_manager_mock, - ) - - def test_add_links_and_get_next_chunk_index(self, sample_links): - """Verify that initial links are stored and next chunk index is computed correctly.""" - link0, link1 = sample_links - - fetcher, _backend, download_manager = self._create_fetcher([link0]) - - # add_link should have been called for the initial link - download_manager.add_link.assert_called_once() - - # Internal mapping should contain the link - assert fetcher.chunk_index_to_link[0] == link0 - - # The next chunk index should be 1 (from link0.next_chunk_index) - assert fetcher._get_next_chunk_index() == 1 - - # Add second link and validate it is present - fetcher._add_links([link1]) - assert fetcher.chunk_index_to_link[1] == link1 - - def test_trigger_next_batch_download_success(self, sample_links): - """Check that _trigger_next_batch_download fetches and stores new links.""" - link0, link1 = sample_links - - backend_mock = Mock() - backend_mock.get_chunk_links = Mock(return_value=[link1]) - - fetcher, backend, download_manager = self._create_fetcher( - [link0], backend_mock=backend_mock - ) - - # Trigger download of the next chunk (index 1) - success = fetcher._trigger_next_batch_download() - - assert success is True - backend.get_chunk_links.assert_called_once_with("statement-123", 1) - assert fetcher.chunk_index_to_link[1] == link1 - # Two calls to add_link: one for initial link, one for new link - assert download_manager.add_link.call_count == 2 - - def test_trigger_next_batch_download_error(self, sample_links): - """Ensure that errors from backend are captured and surfaced.""" - link0, _link1 = sample_links - - backend_mock = Mock() - backend_mock.get_chunk_links.side_effect = ServerOperationError( - "Backend failure" - ) - - fetcher, backend, download_manager = self._create_fetcher( - [link0], backend_mock=backend_mock - ) - - success = fetcher._trigger_next_batch_download() - - assert success is False - assert fetcher._error is not None - - def test_get_chunk_link_waits_until_available(self, sample_links): - """Validate that get_chunk_link blocks until the requested link is available and then returns it.""" - link0, link1 = sample_links - - backend_mock = Mock() - # Configure backend to return link1 when requested for chunk index 1 - backend_mock.get_chunk_links = Mock(return_value=[link1]) - - fetcher, backend, download_manager = self._create_fetcher( - [link0], backend_mock=backend_mock, total_chunk_count=2 - ) - - # Holder to capture the link returned from the background thread - result_container = {} - - def _worker(): - result_container["link"] = fetcher.get_chunk_link(1) - - thread = threading.Thread(target=_worker) - thread.start() - - # Give the thread a brief moment to start and attempt to fetch (and therefore block) - time.sleep(0.1) - - # Trigger the backend fetch which will add link1 and notify waiting threads - fetcher._trigger_next_batch_download() - - thread.join(timeout=2) - - # The thread should have finished and captured link1 - assert result_container.get("link") == link1 - - def test_get_chunk_link_out_of_range_returns_none(self, sample_links): - """Requesting a chunk index >= total_chunk_count should immediately return None.""" - link0, _ = sample_links - - fetcher, _backend, _dm = self._create_fetcher([link0], total_chunk_count=1) - - assert fetcher.get_chunk_link(10) is None diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py deleted file mode 100644 index 1c3e3b5b4..000000000 --- a/tests/unit/test_sea_result_set.py +++ /dev/null @@ -1,567 +0,0 @@ -""" -Tests for the SeaResultSet class. - -This module contains tests for the SeaResultSet class, which implements -the result set functionality for the SEA (Statement Execution API) backend. -""" - -import pytest -from unittest.mock import Mock, patch - -try: - import pyarrow -except ImportError: - pyarrow = None - -from databricks.sql.backend.sea.result_set import SeaResultSet, Row -from databricks.sql.backend.sea.queue import JsonQueue -from databricks.sql.backend.sea.utils.constants import ResultFormat -from databricks.sql.backend.types import CommandId, CommandState -from databricks.sql.backend.sea.models.base import ResultData, ResultManifest - - -class TestSeaResultSet: - """Test suite for the SeaResultSet class.""" - - @pytest.fixture - def mock_connection(self): - """Create a mock connection.""" - connection = Mock() - connection.open = True - connection.session = Mock() - connection.session.ssl_options = Mock() - return connection - - @pytest.fixture - def mock_sea_client(self): - """Create a mock SEA client.""" - client = Mock() - client.max_download_threads = 10 - return client - - @pytest.fixture - def execute_response(self): - """Create a sample execute response.""" - mock_response = Mock() - mock_response.command_id = CommandId.from_sea_statement_id("test-statement-123") - mock_response.status = CommandState.SUCCEEDED - mock_response.has_been_closed_server_side = False - mock_response.has_more_rows = False - mock_response.results_queue = None - mock_response.description = [ - ("col1", "string", None, None, None, None, None), - ("col2", "int", None, None, None, None, None), - ("col3", "boolean", None, None, None, None, None), - ] - mock_response.is_staging_operation = False - mock_response.lz4_compressed = False - mock_response.arrow_schema_bytes = None - return mock_response - - @pytest.fixture - def sample_data(self): - """Create sample data for testing.""" - return [ - ["value1", "1", "true"], - ["value2", "2", "false"], - ["value3", "3", "true"], - ["value4", "4", "false"], - ["value5", "5", "true"], - ] - - def _create_empty_manifest(self, format: ResultFormat): - """Create an empty manifest.""" - return ResultManifest( - format=format.value, - schema={}, - total_row_count=-1, - total_byte_count=-1, - total_chunk_count=-1, - ) - - @pytest.fixture - def result_set_with_data( - self, mock_connection, mock_sea_client, execute_response, sample_data - ): - """Create a SeaResultSet with sample data.""" - # Create ResultData with inline data - result_data = ResultData( - data=sample_data, external_links=None, row_count=len(sample_data) - ) - - # Initialize SeaResultSet with result data - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", - return_value=JsonQueue(sample_data), - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=result_data, - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - - return result_set - - @pytest.fixture - def mock_arrow_queue(self): - """Create a mock Arrow queue.""" - queue = Mock() - if pyarrow is not None: - queue.next_n_rows.return_value = Mock(spec=pyarrow.Table) - queue.next_n_rows.return_value.num_rows = 0 - queue.remaining_rows.return_value = Mock(spec=pyarrow.Table) - queue.remaining_rows.return_value.num_rows = 0 - return queue - - @pytest.fixture - def mock_json_queue(self): - """Create a mock JSON queue.""" - queue = Mock(spec=JsonQueue) - queue.next_n_rows.return_value = [] - queue.remaining_rows.return_value = [] - return queue - - @pytest.fixture - def result_set_with_arrow_queue( - self, mock_connection, mock_sea_client, execute_response, mock_arrow_queue - ): - """Create a SeaResultSet with an Arrow queue.""" - # Create ResultData with external links - result_data = ResultData(data=None, external_links=[], row_count=0) - - # Initialize SeaResultSet with result data - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", - return_value=mock_arrow_queue, - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=result_data, - manifest=ResultManifest( - format=ResultFormat.ARROW_STREAM.value, - schema={}, - total_row_count=0, - total_byte_count=0, - total_chunk_count=0, - ), - buffer_size_bytes=1000, - arraysize=100, - ) - - return result_set - - @pytest.fixture - def result_set_with_json_queue( - self, mock_connection, mock_sea_client, execute_response, mock_json_queue - ): - """Create a SeaResultSet with a JSON queue.""" - # Create ResultData with inline data - result_data = ResultData(data=[], external_links=None, row_count=0) - - # Initialize SeaResultSet with result data - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", - return_value=mock_json_queue, - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=result_data, - manifest=ResultManifest( - format=ResultFormat.JSON_ARRAY.value, - schema={}, - total_row_count=0, - total_byte_count=0, - total_chunk_count=0, - ), - buffer_size_bytes=1000, - arraysize=100, - ) - - return result_set - - def test_init_with_execute_response( - self, mock_connection, mock_sea_client, execute_response - ): - """Test initializing SeaResultSet with an execute response.""" - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue" - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=ResultData(data=[]), - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - - # Verify basic properties - assert result_set.command_id == execute_response.command_id - assert result_set.status == CommandState.SUCCEEDED - assert result_set.connection == mock_connection - assert result_set.backend == mock_sea_client - assert result_set.buffer_size_bytes == 1000 - assert result_set.arraysize == 100 - assert result_set.description == execute_response.description - - def test_init_with_invalid_command_id( - self, mock_connection, mock_sea_client, execute_response - ): - """Test initializing SeaResultSet with invalid command ID.""" - # Mock the command ID to return None - mock_command_id = Mock() - mock_command_id.to_sea_statement_id.return_value = None - execute_response.command_id = mock_command_id - - with pytest.raises(ValueError, match="Command ID is not a SEA statement ID"): - SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=ResultData(data=[]), - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - - def test_close(self, mock_connection, mock_sea_client, execute_response): - """Test closing a result set.""" - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue" - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=ResultData(data=[]), - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - - # Close the result set - result_set.close() - - # Verify the backend's close_command was called - mock_sea_client.close_command.assert_called_once_with(result_set.command_id) - assert result_set.has_been_closed_server_side is True - assert result_set.status == CommandState.CLOSED - - def test_close_when_already_closed_server_side( - self, mock_connection, mock_sea_client, execute_response - ): - """Test closing a result set that has already been closed server-side.""" - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue" - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=ResultData(data=[]), - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - result_set.has_been_closed_server_side = True - - # Close the result set - result_set.close() - - # Verify the backend's close_command was NOT called - mock_sea_client.close_command.assert_not_called() - assert result_set.has_been_closed_server_side is True - assert result_set.status == CommandState.CLOSED - - def test_close_when_connection_closed( - self, mock_connection, mock_sea_client, execute_response - ): - """Test closing a result set when the connection is closed.""" - mock_connection.open = False - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue" - ): - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=ResultData(data=[]), - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - - # Close the result set - result_set.close() - - # Verify the backend's close_command was NOT called - mock_sea_client.close_command.assert_not_called() - assert result_set.has_been_closed_server_side is True - assert result_set.status == CommandState.CLOSED - - def test_convert_json_types(self, result_set_with_data, sample_data): - """Test the _convert_json_types method.""" - # Call _convert_json_types - converted_row = result_set_with_data._convert_json_types(sample_data[0]) - - # Verify the conversion - assert converted_row[0] == "value1" # string stays as string - assert converted_row[1] == 1 # "1" converted to int - assert converted_row[2] is True # "true" converted to boolean - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_convert_json_to_arrow_table(self, result_set_with_data, sample_data): - """Test the _convert_json_to_arrow_table method.""" - # Call _convert_json_to_arrow_table - result_table = result_set_with_data._convert_json_to_arrow_table(sample_data) - - # Verify the result - assert isinstance(result_table, pyarrow.Table) - assert result_table.num_rows == len(sample_data) - assert result_table.num_columns == 3 - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_convert_json_to_arrow_table_empty(self, result_set_with_data): - """Test the _convert_json_to_arrow_table method with empty data.""" - # Call _convert_json_to_arrow_table with empty data - result_table = result_set_with_data._convert_json_to_arrow_table([]) - - # Verify the result - assert isinstance(result_table, pyarrow.Table) - assert result_table.num_rows == 0 - - def test_create_json_table(self, result_set_with_data, sample_data): - """Test the _create_json_table method.""" - # Call _create_json_table - result_rows = result_set_with_data._create_json_table(sample_data) - - # Verify the result - assert len(result_rows) == len(sample_data) - assert isinstance(result_rows[0], Row) - assert result_rows[0].col1 == "value1" - assert result_rows[0].col2 == 1 - assert result_rows[0].col3 is True - - def test_fetchmany_json(self, result_set_with_data): - """Test the fetchmany_json method.""" - # Test fetching a subset of rows - result = result_set_with_data.fetchmany_json(2) - assert len(result) == 2 - assert result_set_with_data._next_row_index == 2 - - # Test fetching the next subset - result = result_set_with_data.fetchmany_json(2) - assert len(result) == 2 - assert result_set_with_data._next_row_index == 4 - - # Test fetching more than available - result = result_set_with_data.fetchmany_json(10) - assert len(result) == 1 # Only one row left - assert result_set_with_data._next_row_index == 5 - - def test_fetchmany_json_negative_size(self, result_set_with_data): - """Test the fetchmany_json method with negative size.""" - with pytest.raises( - ValueError, match="size argument for fetchmany is -1 but must be >= 0" - ): - result_set_with_data.fetchmany_json(-1) - - def test_fetchall_json(self, result_set_with_data, sample_data): - """Test the fetchall_json method.""" - # Test fetching all rows - result = result_set_with_data.fetchall_json() - assert result == sample_data - assert result_set_with_data._next_row_index == len(sample_data) - - # Test fetching again (should return empty) - result = result_set_with_data.fetchall_json() - assert result == [] - assert result_set_with_data._next_row_index == len(sample_data) - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_fetchmany_arrow(self, result_set_with_data, sample_data): - """Test the fetchmany_arrow method.""" - # Test with JSON queue (should convert to Arrow) - result = result_set_with_data.fetchmany_arrow(2) - assert isinstance(result, pyarrow.Table) - assert result.num_rows == 2 - assert result_set_with_data._next_row_index == 2 - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_fetchmany_arrow_negative_size(self, result_set_with_data): - """Test the fetchmany_arrow method with negative size.""" - with pytest.raises( - ValueError, match="size argument for fetchmany is -1 but must be >= 0" - ): - result_set_with_data.fetchmany_arrow(-1) - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_fetchall_arrow(self, result_set_with_data, sample_data): - """Test the fetchall_arrow method.""" - # Test with JSON queue (should convert to Arrow) - result = result_set_with_data.fetchall_arrow() - assert isinstance(result, pyarrow.Table) - assert result.num_rows == len(sample_data) - assert result_set_with_data._next_row_index == len(sample_data) - - def test_fetchone(self, result_set_with_data): - """Test the fetchone method.""" - # Test fetching one row at a time - row1 = result_set_with_data.fetchone() - assert isinstance(row1, Row) - assert row1.col1 == "value1" - assert row1.col2 == 1 - assert row1.col3 is True - assert result_set_with_data._next_row_index == 1 - - row2 = result_set_with_data.fetchone() - assert isinstance(row2, Row) - assert row2.col1 == "value2" - assert row2.col2 == 2 - assert row2.col3 is False - assert result_set_with_data._next_row_index == 2 - - # Fetch the rest - result_set_with_data.fetchall() - - # Test fetching when no more rows - row_none = result_set_with_data.fetchone() - assert row_none is None - - def test_fetchmany(self, result_set_with_data): - """Test the fetchmany method.""" - # Test fetching multiple rows - rows = result_set_with_data.fetchmany(2) - assert len(rows) == 2 - assert isinstance(rows[0], Row) - assert rows[0].col1 == "value1" - assert rows[0].col2 == 1 - assert rows[0].col3 is True - assert rows[1].col1 == "value2" - assert rows[1].col2 == 2 - assert rows[1].col3 is False - assert result_set_with_data._next_row_index == 2 - - # Test with invalid size - with pytest.raises( - ValueError, match="size argument for fetchmany is -1 but must be >= 0" - ): - result_set_with_data.fetchmany(-1) - - def test_fetchall(self, result_set_with_data, sample_data): - """Test the fetchall method.""" - # Test fetching all rows - rows = result_set_with_data.fetchall() - assert len(rows) == len(sample_data) - assert isinstance(rows[0], Row) - assert rows[0].col1 == "value1" - assert rows[0].col2 == 1 - assert rows[0].col3 is True - assert result_set_with_data._next_row_index == len(sample_data) - - # Test fetching again (should return empty) - rows = result_set_with_data.fetchall() - assert len(rows) == 0 - - def test_iteration(self, result_set_with_data, sample_data): - """Test iterating over the result set.""" - # Test iteration - rows = list(result_set_with_data) - assert len(rows) == len(sample_data) - assert isinstance(rows[0], Row) - assert rows[0].col1 == "value1" - assert rows[0].col2 == 1 - assert rows[0].col3 is True - - def test_is_staging_operation( - self, mock_connection, mock_sea_client, execute_response - ): - """Test the is_staging_operation property.""" - # Set is_staging_operation to True - execute_response.is_staging_operation = True - - with patch( - "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue" - ): - # Create a result set - result_set = SeaResultSet( - connection=mock_connection, - execute_response=execute_response, - sea_client=mock_sea_client, - result_data=ResultData(data=[]), - manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), - buffer_size_bytes=1000, - arraysize=100, - ) - - # Test the property - assert result_set.is_staging_operation is True - - # Edge case tests - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_fetchone_empty_arrow_queue(self, result_set_with_arrow_queue): - """Test fetchone with an empty Arrow queue.""" - # Setup _convert_arrow_table to return empty list - result_set_with_arrow_queue._convert_arrow_table = Mock(return_value=[]) - - # Call fetchone - result = result_set_with_arrow_queue.fetchone() - - # Verify result is None - assert result is None - - # Verify _convert_arrow_table was called - result_set_with_arrow_queue._convert_arrow_table.assert_called_once() - - def test_fetchone_empty_json_queue(self, result_set_with_json_queue): - """Test fetchone with an empty JSON queue.""" - # Setup _create_json_table to return empty list - result_set_with_json_queue._create_json_table = Mock(return_value=[]) - - # Call fetchone - result = result_set_with_json_queue.fetchone() - - # Verify result is None - assert result is None - - # Verify _create_json_table was called - result_set_with_json_queue._create_json_table.assert_called_once() - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_fetchmany_empty_arrow_queue(self, result_set_with_arrow_queue): - """Test fetchmany with an empty Arrow queue.""" - # Setup _convert_arrow_table to return empty list - result_set_with_arrow_queue._convert_arrow_table = Mock(return_value=[]) - - # Call fetchmany - result = result_set_with_arrow_queue.fetchmany(10) - - # Verify result is an empty list - assert result == [] - - # Verify _convert_arrow_table was called - result_set_with_arrow_queue._convert_arrow_table.assert_called_once() - - @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") - def test_fetchall_empty_arrow_queue(self, result_set_with_arrow_queue): - """Test fetchall with an empty Arrow queue.""" - # Setup _convert_arrow_table to return empty list - result_set_with_arrow_queue._convert_arrow_table = Mock(return_value=[]) - - # Call fetchall - result = result_set_with_arrow_queue.fetchall() - - # Verify result is an empty list - assert result == [] - - # Verify _convert_arrow_table was called - result_set_with_arrow_queue._convert_arrow_table.assert_called_once() diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index a9f774208..9725b1f10 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -33,6 +33,76 @@ class TestSession: "enable_telemetry": False, } + @pytest.mark.parametrize("use_kernel", [False, True]) + @patch("databricks.sql.client.UnifiedHttpClient") + @patch("databricks.sql.client.Session") + def test_removed_sea_backend_rejected_before_connecting( + self, mock_session, mock_http_client, use_kernel + ): + gc.collect() + with patch("sys.unraisablehook") as unraisable_hook: + with pytest.raises( + ValueError, match="SEA backend has been removed.*use_kernel=True" + ): + databricks.sql.connect( + **self.DUMMY_CONNECTION_ARGS, use_sea=True, use_kernel=use_kernel + ) + gc.collect() + unraisable_hook.assert_not_called() + mock_session.assert_not_called() + mock_http_client.assert_not_called() + + @pytest.mark.parametrize("use_kernel", [False, True]) + @pytest.mark.parametrize( + "options, expected_warning", + [ + ({}, None), + ({"use_sea": False}, None), + ({"use_hybrid_disposition": False}, None), + ({"use_hybrid_disposition": True}, "use_hybrid_disposition"), + ({"max_connections": 10}, "max_connections"), + ({"max_connections": 0}, "max_connections"), + ({"max_connections": None}, "max_connections"), + ( + {"max_connections": 10, "_pool_maxsize": 42}, + "max_connections", + ), + ], + ) + @patch("databricks.sql.session.Session._create_backend") + def test_removed_backend_options_warn_and_are_not_forwarded( + self, mock_create_backend, options, expected_warning, use_kernel, caplog + ): + backend_type = BackendType.KERNEL if use_kernel else BackendType.THRIFT + mock_create_backend.return_value.open_session.return_value = SessionId( + backend_type, "session-id" + ) + with caplog.at_level("WARNING", logger="databricks.sql.client"): + with databricks.sql.connect( + **self.DUMMY_CONNECTION_ARGS, use_kernel=use_kernel, **options + ): + pass + forwarded = mock_create_backend.call_args.args[-1] + assert ( + not {"use_sea", "use_hybrid_disposition", "max_connections"} + & forwarded.keys() + ) + if "_pool_maxsize" in options: + assert forwarded["_pool_maxsize"] == options["_pool_maxsize"] + messages = [ + record.getMessage() + for record in caplog.records + if record.name == "databricks.sql.client" + ] + if expected_warning: + assert len(messages) == 1 + assert expected_warning in messages[0] + if expected_warning == "max_connections": + assert "On kernel >= 1.1.0, use '_pool_maxsize'" in messages[0] + assert "does not configure the Thrift query transport" in messages[0] + else: + assert messages == [] + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) def test_close_uses_the_correct_session_id(self, mock_client_class): instance = mock_client_class.return_value @@ -469,7 +539,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self): ): instance = mock_kernel_client.return_value instance.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) conn = databricks.sql.connect( @@ -537,7 +607,7 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self): ): instance = mock_kernel_client.return_value instance.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) conn = databricks.sql.connect( @@ -586,7 +656,7 @@ def test_oauth_token_cache_enabled_threaded_into_kernel_auth_options(self): ): instance = mock_kernel_client.return_value instance.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) conn = databricks.sql.connect( @@ -632,7 +702,7 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self): ): instance = mock_kernel_client.return_value instance.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) conn = databricks.sql.connect( @@ -674,7 +744,7 @@ def test_telemetry_enabled_defaults_none_for_kernel_client(self): ): instance = mock_kernel_client.return_value instance.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) conn = databricks.sql.connect( @@ -719,7 +789,7 @@ def test_user_agent_entry_reaches_kernel_client_http_headers(self): ): instance = mock_kernel_client.return_value instance.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) conn = databricks.sql.connect( @@ -782,7 +852,7 @@ def test_connect_use_kernel_instantiates_real_kernel_backend(self): with patch.object( KernelDatabricksClient, "open_session", - return_value=SessionId(BackendType.SEA, "sess-id", None), + return_value=SessionId(BackendType.KERNEL, "sess-id", None), ): conn = databricks.sql.connect( server_hostname="foo.cloud.databricks.com", @@ -857,7 +927,7 @@ def _fake_kernel(self): "databricks.sql.backend.kernel.client.KernelDatabricksClient" ) as mock_kernel: mock_kernel.return_value.open_session.return_value = SessionId( - BackendType.SEA, "sess-id", None + BackendType.KERNEL, "sess-id", None ) yield mock_kernel @@ -960,15 +1030,6 @@ def test_non_reyden_thrift_error_not_recovered(self, mock_thrift): self._connect() assert not isinstance(excinfo.value, ReydenThriftUnsupportedError) - @patch("%s.session.SeaDatabricksClient" % PACKAGE) - def test_explicit_use_sea_not_recovered(self, mock_sea): - # An explicit backend choice is always honored — even the (contrived) - # case of SEA surfacing the marker must not trigger Thrift→kernel - # recovery. - mock_sea.return_value.open_session.side_effect = self._reject() - with pytest.raises(ReydenThriftUnsupportedError): - self._connect(use_sea=True) - @patch("%s.session.ThriftDatabricksClient" % PACKAGE) def test_kernel_retry_failure_chains_both_errors(self, mock_thrift): mock_thrift.return_value.open_session.side_effect = self._reject() diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index f000eaa43..ab7c56d65 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -685,7 +685,7 @@ def test_driver_connection_parameters_all_fields(self): params = DriverConnectionParameters( http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.SEA, + mode=DatabricksClientType.THRIFT, host_info=host_info, auth_mech=AuthMech.OAUTH, auth_flow=AuthFlow.BROWSER_BASED_AUTHENTICATION, @@ -702,7 +702,6 @@ def test_driver_connection_parameters_all_fields(self): use_system_trust_store=True, enable_arrow=True, enable_direct_results=True, - enable_sea_hybrid_results=True, http_connection_pool_size=100, rows_fetched_per_block=100000, async_poll_interval_millis=2000, @@ -718,7 +717,7 @@ def test_driver_connection_parameters_all_fields(self): # Verify all new fields are in JSON assert json_dict["http_path"] == "/sql/1.0/warehouses/abc123" - assert json_dict["mode"] == "SEA" + assert json_dict["mode"] == "THRIFT" assert json_dict["host_info"]["host_url"] == "workspace.databricks.com" assert json_dict["auth_mech"] == "OAUTH" assert json_dict["auth_flow"] == "BROWSER_BASED_AUTHENTICATION" @@ -735,7 +734,6 @@ def test_driver_connection_parameters_all_fields(self): assert json_dict["use_system_trust_store"] is True assert json_dict["enable_arrow"] is True assert json_dict["enable_direct_results"] is True - assert json_dict["enable_sea_hybrid_results"] is True assert json_dict["http_connection_pool_size"] == 100 assert json_dict["rows_fetched_per_block"] == 100000 assert json_dict["async_poll_interval_millis"] == 2000 @@ -804,7 +802,7 @@ def test_telemetry_event_complete_serialization(self): connection_params = DriverConnectionParameters( http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.SEA, + mode=DatabricksClientType.THRIFT, host_info=host_info, auth_mech=AuthMech.OAUTH, use_proxy=True, @@ -864,7 +862,7 @@ def test_json_serialization_excludes_none_values(self): params = DriverConnectionParameters( http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.SEA, + mode=DatabricksClientType.THRIFT, host_info=host_info, # All optional fields left as None ) @@ -894,7 +892,6 @@ def test_connection_with_proxy_populates_telemetry(self, mock_setup_pools, mock_ mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False - mock_session_instance.use_sea = True mock_session_instance.port = 443 mock_session_instance.host = "workspace.databricks.com" mock_session.return_value = mock_session_instance @@ -919,7 +916,7 @@ def test_connection_with_proxy_populates_telemetry(self, mock_setup_pools, mock_ # Verify fields are populated assert driver_params.http_path == "/sql/1.0/warehouses/test" - assert driver_params.mode == DatabricksClientType.SEA + assert driver_params.mode == DatabricksClientType.THRIFT assert driver_params.host_info.host_url == "workspace.databricks.com" assert driver_params.host_info.port == 443 @@ -930,7 +927,6 @@ def test_connection_with_azure_params_populates_telemetry(self, mock_setup_pools mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False - mock_session_instance.use_sea = False mock_session_instance.port = 443 mock_session_instance.host = "workspace.azuredatabricks.net" mock_session.return_value = mock_session_instance @@ -960,7 +956,6 @@ def test_connection_populates_arrow_and_performance_params(self, mock_setup_pool mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False - mock_session_instance.use_sea = True mock_session_instance.port = 443 mock_session_instance.host = "workspace.databricks.com" mock_session.return_value = mock_session_instance @@ -1005,7 +1000,6 @@ def test_federated_pat_populates_telemetry_as_pat(self, mock_setup_pools, mock_s mock_session_instance.use_kernel = False mock_session_instance.auth_provider = federated_pat mock_session_instance.is_open = False - mock_session_instance.use_sea = False mock_session_instance.port = 443 mock_session_instance.host = "workspace.databricks.com" mock_session.return_value = mock_session_instance @@ -1033,7 +1027,6 @@ def test_cf_proxy_fields_default_to_false_none(self, mock_setup_pools, mock_sess mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False - mock_session_instance.use_sea = True mock_session_instance.port = 443 mock_session_instance.host = "workspace.databricks.com" mock_session.return_value = mock_session_instance