From 5d993dac32a93efd647f9c4bea47b4efdebda02b Mon Sep 17 00:00:00 2001
From: Serge Smertin <serge.smertin@databricks.com>
Date: Mon, 16 Sep 2024 13:47:36 +0200
Subject: [PATCH] Added ability to specify custom SQL in `make_query`

---
 README.md                                     |  3 +++
 .../labs/pytester/fixtures/redash.py          | 23 +++++++++++--------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 9f8cd25..f31fd16 100644
--- a/README.md
+++ b/README.md
@@ -913,6 +913,9 @@ See also [`ws`](#ws-fixture).
 ### `make_query` fixture
 Create a query and remove it after the test is done. Returns the [`LegacyQuery`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/sql.html#databricks.sdk.service.sql.LegacyQuery) object.
 
+Keyword Arguments:
+- `query`: The query to be stored. Default is `SELECT * FROM <newly created random table>`.
+
 Usage:
 ```python
 from databricks.sdk.service.sql import PermissionLevel
diff --git a/src/databricks/labs/pytester/fixtures/redash.py b/src/databricks/labs/pytester/fixtures/redash.py
index 9f55819..b041671 100644
--- a/src/databricks/labs/pytester/fixtures/redash.py
+++ b/src/databricks/labs/pytester/fixtures/redash.py
@@ -1,7 +1,6 @@
 from collections.abc import Generator
 
 from pytest import fixture
-from databricks.sdk.errors import DatabricksError
 from databricks.sdk.service.sql import LegacyQuery
 
 from databricks.labs.pytester.fixtures.baseline import factory
@@ -9,11 +8,18 @@
 
 @fixture
 def make_query(
-    ws, make_table, make_random, log_workspace_link, watchdog_purge_suffix
+    ws,
+    make_table,
+    make_random,
+    log_workspace_link,
+    watchdog_purge_suffix,
 ) -> Generator[LegacyQuery, None, None]:
     """
     Create a query and remove it after the test is done. Returns the `databricks.sdk.service.sql.LegacyQuery` object.
 
+    Keyword Arguments:
+    - `query`: The query to be stored. Default is `SELECT * FROM <newly created random table>`.
+
     Usage:
     ```python
     from databricks.sdk.service.sql import PermissionLevel
@@ -33,22 +39,21 @@ def test_permissions_for_redash(
     ```
     """
 
-    def create() -> LegacyQuery:
-        table = make_table()
+    def create(sql_query: str | None = None) -> LegacyQuery:
+        if sql_query is None:
+            table = make_table()
+            sql_query = f"SELECT * FROM {table.catalog_name}.{table.schema_name}.{table.name}"
         query_name = f"dummy_query_Q{make_random(4)}_{watchdog_purge_suffix}"
         query = ws.queries_legacy.create(
             name=query_name,
             description="TEST QUERY FOR UCX",
-            query=f"SELECT * FROM {table.schema_name}.{table.name}",
+            query=sql_query,
             tags=["original_query_tag"],
         )
         log_workspace_link(f"{query_name} query", f'sql/editor/{query.id}')
         return query
 
     def remove(query: LegacyQuery):
-        try:
-            ws.queries_legacy.delete(query_id=query.id)
-        except DatabricksError:
-            pass
+        ws.queries_legacy.delete(query_id=query.id)
 
     yield from factory("query", create, remove)