use tags instead of name suffix for queries (#47)

## Changes Current implementation uses name suffix, which makes query names less readable This PR implements watchdog query tags instead ### Linked issues None ### Tests - [x] ran unit tests --------- Co-authored-by: Eric Vergnaud <[email protected]>
databrickslabs · Sep 19, 2024 · 53f69c0 · 53f69c0
1 parent 0fa2127
commit 53f69c0
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -1033,7 +1033,7 @@ def test_permissions_for_redash(
     )
 ```
 
-See also [`ws`](#ws-fixture), [`make_table`](#make_table-fixture), [`make_random`](#make_random-fixture), [`watchdog_purge_suffix`](#watchdog_purge_suffix-fixture).
+See also [`ws`](#ws-fixture), [`make_table`](#make_table-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]
@@ -1104,15 +1104,15 @@ See also [`ws`](#ws-fixture).
 ### `watchdog_remove_after` fixture
 Purge time for test objects, representing the (UTC-based) hour from which objects may be purged.
 
-See also [`make_cluster`](#make_cluster-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_schema`](#make_schema-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_warehouse`](#make_warehouse-fixture), [`watchdog_purge_suffix`](#watchdog_purge_suffix-fixture).
+See also [`make_cluster`](#make_cluster-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_schema`](#make_schema-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_warehouse`](#make_warehouse-fixture), [`watchdog_purge_suffix`](#watchdog_purge_suffix-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]
 
 ### `watchdog_purge_suffix` fixture
 HEX-encoded purge time suffix for test objects.
 
-See also [`make_acc_group`](#make_acc_group-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_user`](#make_user-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
+See also [`make_acc_group`](#make_acc_group-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_repo`](#make_repo-fixture), [`make_user`](#make_user-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]

diff --git a/src/databricks/labs/pytester/fixtures/redash.py b/src/databricks/labs/pytester/fixtures/redash.py
@@ -1,3 +1,4 @@
+import json
 from collections.abc import Generator
 
 from pytest import fixture
@@ -12,7 +13,7 @@ def make_query(
     make_table,
     make_random,
     log_workspace_link,
-    watchdog_purge_suffix,
+    watchdog_remove_after,
 ) -> Generator[LegacyQuery, None, None]:
     """
     Create a query and remove it after the test is done. Returns the `databricks.sdk.service.sql.LegacyQuery` object.
@@ -39,16 +40,23 @@ def test_permissions_for_redash(
     ```
     """
 
-    def create(sql_query: str | None = None) -> LegacyQuery:
+    def create(sql_query: str | None = None, **kwargs) -> 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}"
+        # add RemoveAfter tag for watchdog
+        remove_after_tag = json.dumps({"key": "RemoveAfter", "value": watchdog_remove_after})
+        tags: list[str] | None = kwargs.get('tags', None)
+        if tags:
+            tags.append(remove_after_tag)
+        else:
+            tags = [remove_after_tag]
+        query_name = f"dummy_query_Q{make_random(4)}"
         query = ws.queries_legacy.create(
             name=query_name,
             description="TEST QUERY FOR UCX",
             query=sql_query,
-            tags=["original_query_tag"],
+            tags=tags,
         )
         log_workspace_link(f"{query_name} query", f'sql/editor/{query.id}')
         return query