Skip to content

Commit 5d993da

Browse files
committed
Added ability to specify custom SQL in make_query
1 parent 53034fd commit 5d993da

File tree

2 files changed

+17
-9
lines changed

2 files changed

+17
-9
lines changed

README.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -913,6 +913,9 @@ See also [`ws`](#ws-fixture).
913913
### `make_query` fixture
914914
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.
915915

916+
Keyword Arguments:
917+
- `query`: The query to be stored. Default is `SELECT * FROM <newly created random table>`.
918+
916919
Usage:
917920
```python
918921
from databricks.sdk.service.sql import PermissionLevel

src/databricks/labs/pytester/fixtures/redash.py

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,25 @@
11
from collections.abc import Generator
22

33
from pytest import fixture
4-
from databricks.sdk.errors import DatabricksError
54
from databricks.sdk.service.sql import LegacyQuery
65

76
from databricks.labs.pytester.fixtures.baseline import factory
87

98

109
@fixture
1110
def make_query(
12-
ws, make_table, make_random, log_workspace_link, watchdog_purge_suffix
11+
ws,
12+
make_table,
13+
make_random,
14+
log_workspace_link,
15+
watchdog_purge_suffix,
1316
) -> Generator[LegacyQuery, None, None]:
1417
"""
1518
Create a query and remove it after the test is done. Returns the `databricks.sdk.service.sql.LegacyQuery` object.
1619
20+
Keyword Arguments:
21+
- `query`: The query to be stored. Default is `SELECT * FROM <newly created random table>`.
22+
1723
Usage:
1824
```python
1925
from databricks.sdk.service.sql import PermissionLevel
@@ -33,22 +39,21 @@ def test_permissions_for_redash(
3339
```
3440
"""
3541

36-
def create() -> LegacyQuery:
37-
table = make_table()
42+
def create(sql_query: str | None = None) -> LegacyQuery:
43+
if sql_query is None:
44+
table = make_table()
45+
sql_query = f"SELECT * FROM {table.catalog_name}.{table.schema_name}.{table.name}"
3846
query_name = f"dummy_query_Q{make_random(4)}_{watchdog_purge_suffix}"
3947
query = ws.queries_legacy.create(
4048
name=query_name,
4149
description="TEST QUERY FOR UCX",
42-
query=f"SELECT * FROM {table.schema_name}.{table.name}",
50+
query=sql_query,
4351
tags=["original_query_tag"],
4452
)
4553
log_workspace_link(f"{query_name} query", f'sql/editor/{query.id}')
4654
return query
4755

4856
def remove(query: LegacyQuery):
49-
try:
50-
ws.queries_legacy.delete(query_id=query.id)
51-
except DatabricksError:
52-
pass
57+
ws.queries_legacy.delete(query_id=query.id)
5358

5459
yield from factory("query", create, remove)

0 commit comments

Comments
 (0)