Skip to content

Commit

Permalink
Added make_query fixture
Browse files Browse the repository at this point in the history
  • Loading branch information
nfx committed Sep 12, 2024
1 parent e4a1643 commit 5e0aeb7
Show file tree
Hide file tree
Showing 4 changed files with 99 additions and 3 deletions.
33 changes: 30 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
* [`log_workspace_link` fixture](#log_workspace_link-fixture)
* [`make_dashboard_permissions` fixture](#make_dashboard_permissions-fixture)
* [`make_alert_permissions` fixture](#make_alert_permissions-fixture)
* [`make_query` fixture](#make_query-fixture)
* [`make_query_permissions` fixture](#make_query_permissions-fixture)
* [`make_registered_model_permissions` fixture](#make_registered_model_permissions-fixture)
* [`make_serving_endpoint_permissions` fixture](#make_serving_endpoint_permissions-fixture)
Expand Down Expand Up @@ -204,7 +205,7 @@ def test_workspace_operations(ws):
assert len(clusters) >= 0
```

See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`sql_backend`](#sql_backend-fixture), [`workspace_library`](#workspace_library-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture).
See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`sql_backend`](#sql_backend-fixture), [`workspace_library`](#workspace_library-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture).


[[back to top](#python-testing-for-databricks)]
Expand All @@ -230,7 +231,7 @@ random_string = make_random(k=8)
assert len(random_string) == 8
```

See also [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`workspace_library`](#workspace_library-fixture).
See also [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`workspace_library`](#workspace_library-fixture).


[[back to top](#python-testing-for-databricks)]
Expand Down Expand Up @@ -705,7 +706,7 @@ def test_catalog_fixture(make_catalog, make_schema, make_table):
logger.info(f"Created new schema: {from_table_1}")
```

See also [`ws`](#ws-fixture), [`sql_backend`](#sql_backend-fixture), [`make_schema`](#make_schema-fixture), [`make_random`](#make_random-fixture).
See also [`make_query`](#make_query-fixture), [`ws`](#ws-fixture), [`sql_backend`](#sql_backend-fixture), [`make_schema`](#make_schema-fixture), [`make_random`](#make_random-fixture).


[[back to top](#python-testing-for-databricks)]
Expand Down Expand Up @@ -851,6 +852,32 @@ _No description yet._
See also [`ws`](#ws-fixture).


[[back to top](#python-testing-for-databricks)]

### `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.

Usage:
```python
from databricks.sdk.service.sql import PermissionLevel

def test_permissions_for_redash(
make_user,
make_query,
make_query_permissions,
):
user = make_user()
query = make_query()
make_query_permissions(
object_id=query.id,
permission_level=PermissionLevel.CAN_EDIT,
user_name=user.display_name,
)
```

See also [`ws`](#ws-fixture), [`make_table`](#make_table-fixture), [`make_random`](#make_random-fixture).


[[back to top](#python-testing-for-databricks)]

### `make_query_permissions` fixture
Expand Down
2 changes: 2 additions & 0 deletions src/databricks/labs/pytester/fixtures/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from databricks.labs.pytester.fixtures.wheel import workspace_library
from databricks.labs.pytester.fixtures.environment import debug_env, debug_env_name, env_or_skip
from databricks.labs.pytester.fixtures.ml import make_experiment, make_model
from databricks.labs.pytester.fixtures.redash import make_query

__all__ = [
'debug_env_name',
Expand Down Expand Up @@ -94,6 +95,7 @@
'log_workspace_link',
'make_dashboard_permissions',
'make_alert_permissions',
'make_query',
'make_query_permissions',
'make_registered_model_permissions',
'make_serving_endpoint_permissions',
Expand Down
52 changes: 52 additions & 0 deletions src/databricks/labs/pytester/fixtures/redash.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
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 get_purge_suffix, factory


@fixture
def make_query(ws, make_table, make_random, log_workspace_link) -> Generator[LegacyQuery, None, None]:
"""
Create a query and remove it after the test is done. Returns the `databricks.sdk.service.sql.LegacyQuery` object.
Usage:
```python
from databricks.sdk.service.sql import PermissionLevel
def test_permissions_for_redash(
make_user,
make_query,
make_query_permissions,
):
user = make_user()
query = make_query()
make_query_permissions(
object_id=query.id,
permission_level=PermissionLevel.CAN_EDIT,
user_name=user.display_name,
)
```
"""

def create() -> LegacyQuery:
table = make_table()
query_name = f"dummy_query_Q{make_random(4)}_{get_purge_suffix()}"
query = ws.queries_legacy.create(
name=query_name,
description="TEST QUERY FOR UCX",
query=f"SELECT * FROM {table.schema_name}.{table.name}",
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

yield from factory("query", create, remove)
15 changes: 15 additions & 0 deletions tests/integration/fixtures/test_redash.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from databricks.sdk.service.sql import PermissionLevel


def test_permissions_for_redash(
make_user,
make_query,
make_query_permissions,
):
user = make_user()
query = make_query()
make_query_permissions(
object_id=query.id,
permission_level=PermissionLevel.CAN_EDIT,
user_name=user.display_name,
)

0 comments on commit 5e0aeb7

Please sign in to comment.