Updates to use the Databricks Python sdk 0.30.0 (#247)

These changes contains the update to use the Databricks Python sdk
0.30.0

Author: JCZuurmond

src/databricks/labs/lsql/core.py

@@ -8,7 +8,7 @@
 import types
 from collections.abc import Callable, Iterator
 from datetime import timedelta
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import requests
 import sqlglot
@@ -17,7 +17,6 @@
 from databricks.sdk.service.sql import (
     ColumnInfoTypeName,
     Disposition,
-    ExecuteStatementResponse,
     Format,
     ServiceError,
     ServiceErrorCode,
@@ -26,6 +25,19 @@
     StatementStatus,
 )
 
+if TYPE_CHECKING:  # MyPy complains about the double import
+    from databricks.sdk.service.sql import StatementResponse
+else:
+    try:
+        # databricks-sdk>=0.30.0
+        from databricks.sdk.service.sql import StatementResponse
+    except ImportError:
+        # databricks-sdk<0.30.0
+        from databricks.sdk.service.sql import (
+            ExecuteStatementResponse as StatementResponse,
+        )
+
+
 MAX_SLEEP_PER_ATTEMPT = 10
 
 MAX_PLATFORM_TIMEOUT = 50
@@ -179,7 +191,7 @@ def execute(
         catalog: str | None = None,
         schema: str | None = None,
         timeout: timedelta | None = None,
-    ) -> ExecuteStatementResponse:
+    ) -> StatementResponse:
         """Execute a SQL statement and block until results are ready, including starting
         the warehouse if needed.
 
@@ -214,7 +226,7 @@ def execute(
             it is handled on the server side. If the timeout is greater than 50 seconds,
             Databricks SDK for Python cancels the statement execution and throws `TimeoutError`.
             If not given, it will use the timeout specified in the constructor.
-        :return: ExecuteStatementResponse
+        :return: StatementResponse
         """
         # The wait_timeout field must be 0 seconds (disables wait),
         # or between 5 seconds and 50 seconds.
@@ -259,7 +271,7 @@ def execute(
             if not state:
                 state = StatementState.FAILED
             if state == StatementState.SUCCEEDED:
-                return ExecuteStatementResponse(
+                return StatementResponse(
                     manifest=res.manifest, result=res.result, statement_id=statement_id, status=result_status
                 )
             status_message = f"current status: {state.value}"
@@ -378,7 +390,7 @@ def fetch_value(self, statement: str, **kwargs) -> Any | None:
             return v
         return None
 
-    def _result_converter(self, execute_response: ExecuteStatementResponse):
+    def _result_converter(self, execute_response: StatementResponse):
         """Get the result schema from the execute response."""
         manifest = execute_response.manifest
         if not manifest:

tests/unit/test_backends.py

@@ -16,11 +16,11 @@
 from databricks.sdk.service.sql import (
     ColumnInfo,
     ColumnInfoTypeName,
-    ExecuteStatementResponse,
     Format,
     ResultData,
     ResultManifest,
     ResultSchema,
+    StatementResponse,
     StatementState,
     StatementStatus,
 )
@@ -56,7 +56,7 @@ class Bar:
 
 def test_statement_execution_backend_execute_happy():
     ws = create_autospec(WorkspaceClient)
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED)
     )
 
@@ -78,7 +78,7 @@ def test_statement_execution_backend_execute_happy():
 
 def test_statement_execution_backend_with_overrides():
     ws = create_autospec(WorkspaceClient)
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED)
     )
 
@@ -101,7 +101,7 @@ def test_statement_execution_backend_with_overrides():
 def test_statement_execution_backend_fetch_happy():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         result=ResultData(data_array=[["1"], ["2"], ["3"]]),
@@ -117,7 +117,7 @@ def test_statement_execution_backend_fetch_happy():
 
 def test_statement_execution_backend_save_table_overwrite_empty_table():
     ws = create_autospec(WorkspaceClient)
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED)
     )
     seb = StatementExecutionBackend(ws, "abc")
@@ -161,7 +161,7 @@ def test_statement_execution_backend_save_table_overwrite_empty_table():
 def test_statement_execution_backend_save_table_empty_records():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED)
     )
 
@@ -185,7 +185,7 @@ def test_statement_execution_backend_save_table_empty_records():
 def test_statement_execution_backend_save_table_two_records():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED)
     )
 
@@ -222,7 +222,7 @@ def test_statement_execution_backend_save_table_two_records():
 def test_statement_execution_backend_save_table_in_batches_of_two(mocker):
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED)
     )
 

tests/unit/test_core.py

@@ -8,16 +8,15 @@
     ColumnInfo,
     ColumnInfoTypeName,
     EndpointInfo,
-    ExecuteStatementResponse,
     ExternalLink,
     Format,
-    GetStatementResponse,
     ResultData,
     ResultManifest,
     ResultSchema,
     ServiceError,
     ServiceErrorCode,
     State,
+    StatementResponse,
     StatementState,
     StatementStatus,
     timedelta,
@@ -67,7 +66,7 @@ def test_row_factory_with_generator():
 def test_selects_warehouse_from_config():
     ws = create_autospec(WorkspaceClient)
     ws.config.warehouse_id = "abc"
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         statement_id="bcd",
     )
@@ -96,7 +95,7 @@ def test_selects_warehouse_from_existing_first_running():
         EndpointInfo(id="cde", state=State.RUNNING),
     ]
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
     )
 
@@ -125,7 +124,7 @@ def test_selects_warehouse_from_existing_not_running():
         EndpointInfo(id="cde", state=State.STARTING),
     ]
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
     )
 
@@ -155,11 +154,11 @@ def test_no_warehouse_given():
 
 def test_execute_poll_succeeds():
     ws = create_autospec(WorkspaceClient)
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.PENDING),
         statement_id="bcd",
     )
-    ws.statement_execution.get_statement.return_value = GetStatementResponse(
+    ws.statement_execution.get_statement.return_value = StatementResponse(
         manifest=ResultManifest(),
         result=ResultData(byte_count=100500),
         statement_id="bcd",
@@ -202,7 +201,7 @@ def test_execute_poll_succeeds():
 def test_execute_fails(status_error, platform_error_type):
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.FAILED, error=status_error),
         statement_id="bcd",
     )
@@ -216,14 +215,14 @@ def test_execute_fails(status_error, platform_error_type):
 def test_execute_poll_waits():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.PENDING),
         statement_id="bcd",
     )
 
     ws.statement_execution.get_statement.side_effect = [
-        GetStatementResponse(status=StatementStatus(state=StatementState.RUNNING), statement_id="bcd"),
-        GetStatementResponse(
+        StatementResponse(status=StatementStatus(state=StatementState.RUNNING), statement_id="bcd"),
+        StatementResponse(
             manifest=ResultManifest(),
             result=ResultData(byte_count=100500),
             statement_id="bcd",
@@ -242,12 +241,12 @@ def test_execute_poll_waits():
 def test_execute_poll_timeouts_on_client():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.PENDING),
         statement_id="bcd",
     )
 
-    ws.statement_execution.get_statement.return_value = GetStatementResponse(
+    ws.statement_execution.get_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.RUNNING),
         statement_id="bcd",
     )
@@ -262,7 +261,7 @@ def test_execute_poll_timeouts_on_client():
 def test_fetch_all_no_chunks():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(
             schema=ResultSchema(
@@ -298,7 +297,7 @@ def test_fetch_all_no_chunks():
 def test_fetch_all_no_chunks_no_converter():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(
             schema=ResultSchema(
@@ -323,7 +322,7 @@ def test_fetch_all_no_chunks_no_converter():
 def test_fetch_all_two_chunks():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(
             schema=ResultSchema(
@@ -361,7 +360,7 @@ def test_fetch_all_two_chunks():
 def test_fetch_one():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         result=ResultData(data_array=[["4"]]),
@@ -389,7 +388,7 @@ def test_fetch_one():
 def test_fetch_one_none():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         statement_id="bcd",
@@ -405,7 +404,7 @@ def test_fetch_one_none():
 def test_fetch_one_disable_magic():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         result=ResultData(data_array=[["4"], ["5"], ["6"]]),
@@ -433,7 +432,7 @@ def test_fetch_one_disable_magic():
 def test_fetch_value():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         result=ResultData(data_array=[["4"]]),
@@ -450,7 +449,7 @@ def test_fetch_value():
 def test_fetch_value_none():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         statement_id="bcd",
@@ -466,7 +465,7 @@ def test_fetch_value_none():
 def test_callable_returns_iterator():
     ws = create_autospec(WorkspaceClient)
 
-    ws.statement_execution.execute_statement.return_value = ExecuteStatementResponse(
+    ws.statement_execution.execute_statement.return_value = StatementResponse(
         status=StatementStatus(state=StatementState.SUCCEEDED),
         manifest=ResultManifest(schema=ResultSchema(columns=[ColumnInfo(name="id", type_name=ColumnInfoTypeName.INT)])),
         result=ResultData(data_array=[["4"], ["5"], ["6"]]),