Changes to work with Databricks SDK v0.38.0 (#350)

Changes to work with Databricks SDK [`v0.38.0`](https://github.com/databricks/databricks-sdk-py/releases/tag/v0.38.0) Resolves #349 Resolves #348 Resolves #347 Resolves #346 Resolves #345 Resolves #344 Resolves #343 Resolves #342 Resolves #341 Resolves #340 Resolves #339 Resolves #338 Resolves #337 Resolves #336 Resolves #335 Resolves #334 Resolves #333 Resolves #332
databrickslabs · Nov 19, 2024 · 348cb43 · 348cb43
1 parent 55699cb
commit 348cb43
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 48 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -21,7 +21,7 @@ classifiers = [
 ]
 dependencies = [
   "databricks-labs-blueprint[yaml]>=0.4.2",
-  "databricks-sdk~=0.37",
+  "databricks-sdk~=0.38",
   "sqlglot>=22.3.1"
 ]
 

diff --git a/src/databricks/labs/lsql/dashboards.py b/src/databricks/labs/lsql/dashboards.py
@@ -1133,9 +1133,9 @@ def create_dashboard(
             warehouse_id=warehouse_id,
         )
         if dashboard_id is not None:
-            sdk_dashboard = self._ws.lakeview.update(dashboard_id, dashboard=dashboard_to_create.as_dict())  # type: ignore
+            sdk_dashboard = self._ws.lakeview.update(dashboard_id, dashboard=dashboard_to_create)
         else:
-            sdk_dashboard = self._ws.lakeview.create(dashboard=dashboard_to_create.as_dict())  # type: ignore
+            sdk_dashboard = self._ws.lakeview.create(dashboard=dashboard_to_create)
         if publish:
             assert sdk_dashboard.dashboard_id is not None
             self._ws.lakeview.publish(sdk_dashboard.dashboard_id, warehouse_id=warehouse_id)

diff --git a/tests/integration/test_core.py b/tests/integration/test_core.py
@@ -17,49 +17,44 @@ def test_sql_execution_chunked(ws, disposition):
     assert total == 1999999000000
 
 
-def test_sql_execution(ws, env_or_skip):
-    results = []
+NYC_TAXI_TRIPS_LIMITED = """
+WITH zipcodes AS (
+  SELECT DISTINCT pickup_zip, dropoff_zip 
+  FROM samples.nyctaxi.trips 
+  WHERE pickup_zip = 10282 AND dropoff_zip <= 10005
+)
+
+SELECT 
+  trips.pickup_zip, 
+  trips.dropoff_zip, 
+  trips.tpep_pickup_datetime, 
+  trips.tpep_dropoff_datetime 
+FROM 
+    zipcodes 
+  JOIN 
+    samples.nyctaxi.trips AS trips 
+    ON zipcodes.pickup_zip = trips.pickup_zip AND zipcodes.dropoff_zip = trips.dropoff_zip
+ORDER BY trips.dropoff_zip, trips.tpep_pickup_datetime, trips.tpep_dropoff_datetime 
+"""
+
+
+def test_sql_execution(ws, env_or_skip) -> None:
     see = StatementExecutionExt(ws, warehouse_id=env_or_skip("TEST_DEFAULT_WAREHOUSE_ID"))
-    for pickup_zip, dropoff_zip in see.fetch_all(
-        "SELECT pickup_zip, dropoff_zip FROM nyctaxi.trips LIMIT 10", catalog="samples"
-    ):
-        results.append((pickup_zip, dropoff_zip))
-    assert results == [
-        (10282, 10171),
-        (10110, 10110),
-        (10103, 10023),
-        (10022, 10017),
-        (10110, 10282),
-        (10009, 10065),
-        (10153, 10199),
-        (10112, 10069),
-        (10023, 10153),
-        (10012, 10003),
-    ]
-
-
-def test_sql_execution_partial(ws, env_or_skip):
-    results = []
+
+    records = see.fetch_all(NYC_TAXI_TRIPS_LIMITED, catalog="samples")
+
+    assert len([True for _ in records]) > 1
+
+
+def test_sql_execution_as_iterator(ws, env_or_skip) -> None:
+    number_of_records = 0
     see = StatementExecutionExt(ws, warehouse_id=env_or_skip("TEST_DEFAULT_WAREHOUSE_ID"), catalog="samples")
-    for row in see("SELECT * FROM nyctaxi.trips LIMIT 10"):
-        pickup_time, dropoff_time = row[0], row[1]
-        pickup_zip = row.pickup_zip
-        dropoff_zip = row["dropoff_zip"]
+    for row in see(NYC_TAXI_TRIPS_LIMITED):
+        pickup_zip, dropoff_zip, pickup_time, dropoff_time = row[0], row[1], row[2], row[3]
         all_fields = row.asDict()
         logger.info(f"{pickup_zip}@{pickup_time} -> {dropoff_zip}@{dropoff_time}: {all_fields}")
-        results.append((pickup_zip, dropoff_zip))
-    assert results == [
-        (10282, 10171),
-        (10110, 10110),
-        (10103, 10023),
-        (10022, 10017),
-        (10110, 10282),
-        (10009, 10065),
-        (10153, 10199),
-        (10112, 10069),
-        (10023, 10153),
-        (10012, 10003),
-    ]
+        number_of_records += 1
+    assert number_of_records > 1
 
 
 def test_fetch_one(ws):
@@ -73,9 +68,10 @@ def test_fetch_one_fails_if_limit_is_bigger(ws):
         see.fetch_one("SELECT * FROM samples.nyctaxi.trips LIMIT 100")
 
 
-def test_fetch_one_works(ws):
+def test_fetch_one_works(ws) -> None:
     see = StatementExecutionExt(ws)
-    row = see.fetch_one("SELECT * FROM samples.nyctaxi.trips LIMIT 1")
+    row = see.fetch_one("SELECT pickup_zip FROM samples.nyctaxi.trips WHERE pickup_zip == 10282 LIMIT 1")
+    assert row is not None
     assert row.pickup_zip == 10282
 
 

diff --git a/tests/integration/test_dashboards.py b/tests/integration/test_dashboards.py
@@ -58,7 +58,7 @@ def create(*, display_name: str = "") -> SDKDashboard:
             display_name = f"created_by_lsql_{make_random()}"
         else:
             display_name = f"{display_name} ({make_random()})"
-        dashboard = ws.lakeview.create(dashboard=SDKDashboard(display_name=display_name).as_dict())
+        dashboard = ws.lakeview.create(dashboard=SDKDashboard(display_name=display_name))
         if is_in_debug():
             dashboard_url = f"{ws.config.host}/sql/dashboardsv3/{dashboard.dashboard_id}"
             webbrowser.open(dashboard_url)
@@ -117,7 +117,7 @@ def test_dashboards_creates_exported_dashboard_definition(ws, make_dashboard) ->
     dashboard_content = (Path(__file__).parent / "dashboards" / "dashboard.lvdash.json").read_text()
 
     dashboard_to_create = dataclasses.replace(sdk_dashboard, serialized_dashboard=dashboard_content)
-    ws.lakeview.update(sdk_dashboard.dashboard_id, dashboard=dashboard_to_create.as_dict())
+    ws.lakeview.update(sdk_dashboard.dashboard_id, dashboard=dashboard_to_create)
     lakeview_dashboard = Dashboard.from_dict(json.loads(dashboard_content))
     new_dashboard = dashboards.get_dashboard(sdk_dashboard.path)
 

diff --git a/tests/unit/test_dashboards.py b/tests/unit/test_dashboards.py
@@ -1478,7 +1478,7 @@ def test_dashboards_calls_create_without_dashboard_id() -> None:
 
     dashboards.create_dashboard(dashboard_metadata, parent_path="/non/existing/path", warehouse_id="warehouse")
 
-    ws.lakeview.create.assert_called_with(dashboard=sdk_dashboard.as_dict())
+    ws.lakeview.create.assert_called_with(dashboard=sdk_dashboard)
     ws.lakeview.update.assert_not_called()
     ws.lakeview.publish.assert_not_called()
 
@@ -1498,7 +1498,7 @@ def test_dashboards_calls_update_with_dashboard_id() -> None:
     dashboards.create_dashboard(dashboard_metadata, dashboard_id="id", warehouse_id="warehouse")
 
     ws.lakeview.create.assert_not_called()
-    ws.lakeview.update.assert_called_with("id", dashboard=sdk_dashboard.as_dict())
+    ws.lakeview.update.assert_called_with("id", dashboard=sdk_dashboard)
     ws.lakeview.publish.assert_not_called()