From 235812ad9ea9a37b68d45fd725eaa42bca68259e Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Thu, 14 Nov 2024 20:18:08 +0100 Subject: [PATCH 1/7] Added `make_run_as` fixture --- README.md | 16 +++- scripts/gen-readme.py | 2 +- .../labs/pytester/fixtures/baseline.py | 12 +++ .../labs/pytester/fixtures/catalog.py | 1 + src/databricks/labs/pytester/fixtures/iam.py | 91 ++++++++++++++++++- .../labs/pytester/fixtures/notebooks.py | 14 ++- .../labs/pytester/fixtures/plugin.py | 4 +- tests/integration/fixtures/test_run_as.py | 16 ++++ tests/unit/fixtures/test_catalog.py | 14 ++- 9 files changed, 154 insertions(+), 16 deletions(-) create mode 100644 tests/integration/fixtures/test_run_as.py diff --git a/README.md b/README.md index d209141..36163e8 100644 --- a/README.md +++ b/README.md @@ -255,7 +255,7 @@ def test_something(env_or_skip): assert token is not None ``` -See also [`acc`](#acc-fixture), [`make_udf`](#make_udf-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`is_in_debug`](#is_in_debug-fixture). +See also [`acc`](#acc-fixture), [`make_run_as`](#make_run_as-fixture), [`make_udf`](#make_udf-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`is_in_debug`](#is_in_debug-fixture). [[back to top](#python-testing-for-databricks)] @@ -278,7 +278,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`](#make_feature_table-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`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-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`](#make_feature_table-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_run_as`](#make_run_as-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture). [[back to top](#python-testing-for-databricks)] @@ -305,7 +305,7 @@ def test_listing_workspaces(acc): assert len(workspaces) >= 1 ``` -See also [`make_acc_group`](#make_acc_group-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture), [`env_or_skip`](#env_or_skip-fixture). +See also [`make_acc_group`](#make_acc_group-fixture), [`make_run_as`](#make_run_as-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture), [`env_or_skip`](#env_or_skip-fixture). [[back to top](#python-testing-for-databricks)] @@ -349,6 +349,14 @@ Fetch all rows from a SQL statement. See also [`sql_backend`](#sql_backend-fixture). +[[back to top](#python-testing-for-databricks)] + +### `make_run_as` fixture +_No description yet._ + +See also [`acc`](#acc-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`env_or_skip`](#env_or_skip-fixture). + + [[back to top](#python-testing-for-databricks)] ### `make_random` fixture @@ -372,7 +380,7 @@ random_string = make_random(k=8) assert len(random_string) == 8 ``` -See also [`make_acc_group`](#make_acc_group-fixture), [`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_feature_table`](#make_feature_table-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_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_workspace_file`](#make_workspace_file-fixture). +See also [`make_acc_group`](#make_acc_group-fixture), [`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_feature_table`](#make_feature_table-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_run_as`](#make_run_as-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_workspace_file`](#make_workspace_file-fixture). [[back to top](#python-testing-for-databricks)] diff --git a/scripts/gen-readme.py b/scripts/gen-readme.py index ec9863e..c70ac2b 100644 --- a/scripts/gen-readme.py +++ b/scripts/gen-readme.py @@ -70,7 +70,7 @@ def discover_fixtures() -> list[Fixture]: upstreams = [] sig = inspect.signature(fn) for param in sig.parameters.values(): - if param.name in {'fresh_local_wheel_file', 'monkeypatch', 'log_workspace_link'}: + if param.name in {'fresh_local_wheel_file', 'monkeypatch', 'log_workspace_link', 'request'}: continue upstreams.append(param.name) see_also[param.name].add(fixture) diff --git a/src/databricks/labs/pytester/fixtures/baseline.py b/src/databricks/labs/pytester/fixtures/baseline.py index 709673e..ead15c1 100644 --- a/src/databricks/labs/pytester/fixtures/baseline.py +++ b/src/databricks/labs/pytester/fixtures/baseline.py @@ -222,3 +222,15 @@ def inner(name: str, path: str, *, anchor: bool = True): _LOG.info(f'Created {name}: {url}') return inner + + +@fixture +def log_account_link(acc): + """Returns a function to log an account link.""" + + def inner(name: str, path: str, *, anchor: bool = False): + a = '#' if anchor else '' + url = f'https://{acc.config.hostname}/{a}{path}' + _LOG.info(f'Created {name}: {url}') + + return inner diff --git a/src/databricks/labs/pytester/fixtures/catalog.py b/src/databricks/labs/pytester/fixtures/catalog.py index e76bc73..2f88a43 100644 --- a/src/databricks/labs/pytester/fixtures/catalog.py +++ b/src/databricks/labs/pytester/fixtures/catalog.py @@ -25,6 +25,7 @@ logger = logging.getLogger(__name__) +# TODO: replace with LSQL implementation def escape_sql_identifier(path: str, *, maxsplit: int = 2) -> str: """ Escapes the path components to make them SQL safe. diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index 6ed409f..ad5eeb8 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -3,13 +3,17 @@ from collections.abc import Callable, Generator from datetime import timedelta +from _pytest.fixtures import SubRequest, FixtureDef, fixture +from _pytest.scope import Scope + from pytest import fixture -from databricks.sdk import AccountGroupsAPI, GroupsAPI, WorkspaceClient +from databricks.sdk import AccountGroupsAPI, GroupsAPI, WorkspaceClient, AccountClient from databricks.sdk.config import Config from databricks.sdk.errors import ResourceConflict, NotFound from databricks.sdk.retries import retried from databricks.sdk.service import iam -from databricks.sdk.service.iam import User, Group +from databricks.sdk.service.iam import User, Group, ServicePrincipal, Patch, PatchOp, ComplexValue, PatchSchema, \ + WorkspacePermission from databricks.labs.pytester.fixtures.baseline import factory @@ -183,3 +187,86 @@ def create( return group yield from factory(name, create, lambda item: interface.delete(item.id)) + + +class RunAs: + def __init__(self, service_principal: ServicePrincipal, workspace_client: WorkspaceClient, request: SubRequest): + self._request = SubRequest( + request, + Scope.Function, + None, + request.param_index, + FixtureDef( + config=request.config, + baseid='ws', + argname='ws', + func=lambda: workspace_client, + scope=None, + params=None, + ), + ) + self._request._arg2fixturedefs = {} + self._request._fixture_defs = {} + self._service_principal = service_principal + + def __getattr__(self, item: str): + if item in self.__dict__: + return self.__dict__[item] + fixture_value = self._request.getfixturevalue(item) + return fixture_value + + @property + def display_name(self) -> str: + return self._service_principal.display_name + + @property + def application_id(self) -> str: + return self._service_principal.application_id + + def __repr__(self): + return f'RunAs({self.display_name})' + + +@fixture +def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_skip, request, log_account_link): + """ + This fixture provides a function to create a service principal and assign it to a workspace. The service principal + is removed after the test is complete. The fixture returns an instance of `RunAs` that proxies the calls to + the other fixtures supported by the plugin, for example, `sql_fetch_all`. + + The service principal is created with a random display name and assigned to the workspace with + """ + + def create(*, account_groups: list[str] = None): + workspace_id = ws.get_workspace_id() + service_principal = acc.service_principals.create(display_name=f'spn-{make_random()}') + created_secret = acc.service_principal_secrets.create(service_principal.id) + if account_groups: + group_mapping = {} + for group in acc.groups.list(attributes='id,displayName'): + group_mapping[group.display_name] = group.id + for group_name in account_groups: + if group_name not in group_mapping: + raise ValueError(f'Group {group_name} does not exist') + group_id = group_mapping[group_name] + acc.groups.patch( + group_id, + operations=[Patch(PatchOp.ADD, 'members', [ComplexValue(value=service_principal.id).as_dict()])], + schemas=[PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP], + ) + permissions = [WorkspacePermission.USER] + acc.workspace_assignment.update(workspace_id, service_principal.id, permissions=permissions) + ws_as_spn = WorkspaceClient( + host=ws.config.host, + client_id=service_principal.application_id, + client_secret=created_secret.secret, + ) + + log_account_link('account service principal', f'users/serviceprincipals/{service_principal.id}') + + return RunAs(service_principal, ws_as_spn, request) + + def remove(run_as: RunAs): + acc.service_principals.delete(run_as._service_principal.id) + + yield from factory("service principal", create, remove) diff --git a/src/databricks/labs/pytester/fixtures/notebooks.py b/src/databricks/labs/pytester/fixtures/notebooks.py index 4ce6cec..b0ac751 100644 --- a/src/databricks/labs/pytester/fixtures/notebooks.py +++ b/src/databricks/labs/pytester/fixtures/notebooks.py @@ -56,14 +56,18 @@ def create( default_content = "SELECT 1" else: raise ValueError(f"Unsupported language: {language}") - path = path or f"/Users/{ws.current_user.me().user_name}/dummy-{make_random(4)}-{watchdog_purge_suffix}" + current_user = ws.current_user.me() + path = path or f"/Users/{current_user.user_name}/dummy-{make_random(4)}-{watchdog_purge_suffix}" + workspace_path = WorkspacePath(ws, path) + if '@' not in current_user.user_name: + # If current user is a service principal added with `make_run_as`, there might be no home folder + workspace_path.parent.mkdir(exist_ok=True) content = content or default_content if isinstance(content, str): content = io.BytesIO(content.encode(encoding)) if isinstance(ws, Mock): # For testing ws.workspace.download.return_value = content if isinstance(content, io.BytesIO) else io.BytesIO(content) ws.workspace.upload(path, content, language=language, format=format, overwrite=overwrite) - workspace_path = WorkspacePath(ws, path) logger.info(f"Created notebook: {workspace_path.as_uri()}") return workspace_path @@ -110,10 +114,14 @@ def create( suffix = ".sql" else: raise ValueError(f"Unsupported language: {language}") - path = path or f"/Users/{ws.current_user.me().user_name}/dummy-{make_random(4)}-{watchdog_purge_suffix}{suffix}" + current_user = ws.current_user.me() + path = path or f"/Users/{current_user.user_name}/dummy-{make_random(4)}-{watchdog_purge_suffix}{suffix}" content = content or default_content encoding = encoding or _DEFAULT_ENCODING workspace_path = WorkspacePath(ws, path) + if '@' not in current_user.user_name: + # If current user is a service principal added with `make_run_as`, there might be no home folder + workspace_path.parent.mkdir(exist_ok=True) if isinstance(content, bytes): workspace_path.write_bytes(content) else: diff --git a/src/databricks/labs/pytester/fixtures/plugin.py b/src/databricks/labs/pytester/fixtures/plugin.py index 20097ce..ad6f68e 100644 --- a/src/databricks/labs/pytester/fixtures/plugin.py +++ b/src/databricks/labs/pytester/fixtures/plugin.py @@ -6,6 +6,7 @@ make_random, product_info, log_workspace_link, + log_account_link, ) from databricks.labs.pytester.fixtures.sql import sql_backend, sql_exec, sql_fetch_all from databricks.labs.pytester.fixtures.compute import ( @@ -16,7 +17,7 @@ make_pipeline, make_warehouse, ) -from databricks.labs.pytester.fixtures.iam import make_group, make_acc_group, make_user +from databricks.labs.pytester.fixtures.iam import make_group, make_acc_group, make_user, make_run_as from databricks.labs.pytester.fixtures.catalog import ( make_udf, make_catalog, @@ -105,6 +106,7 @@ 'make_warehouse_permissions', 'make_lakeview_dashboard_permissions', 'log_workspace_link', + 'log_account_link', 'make_dashboard_permissions', 'make_alert_permissions', 'make_query', diff --git a/tests/integration/fixtures/test_run_as.py b/tests/integration/fixtures/test_run_as.py new file mode 100644 index 0000000..3ace4fe --- /dev/null +++ b/tests/integration/fixtures/test_run_as.py @@ -0,0 +1,16 @@ +def test_run_as(make_run_as, ws): + run_as = make_run_as(account_groups=['role.labs.lsql.write']) + through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) + me = ws.current_user.me() + assert me.user_name != through_query.my_name + + +def test_notebooks_are_created_by_different_users(make_run_as, make_notebook): + notebook_by_current_user = make_notebook() + a = notebook_by_current_user.parent.as_posix() + + run_as = make_run_as(account_groups=['role.labs.lsql.write']) + notebook_by_ephemeral_principal = run_as.make_notebook() + b = notebook_by_ephemeral_principal.parent.as_posix() + + assert a != b diff --git a/tests/unit/fixtures/test_catalog.py b/tests/unit/fixtures/test_catalog.py index 64049d7..cf497b8 100644 --- a/tests/unit/fixtures/test_catalog.py +++ b/tests/unit/fixtures/test_catalog.py @@ -1,6 +1,13 @@ from unittest.mock import ANY -from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo, SchemaInfo, VolumeType, VolumeInfo +from databricks.sdk.service.catalog import ( + TableInfo, + TableType, + DataSourceFormat, + FunctionInfo, + SchemaInfo, + VolumeType, +) from databricks.labs.pytester.fixtures.unwrap import call_stateful from databricks.labs.pytester.fixtures.catalog import ( @@ -169,9 +176,6 @@ def test_make_volume_noargs(): def test_make_volume_with_name(): ctx, info = call_stateful(make_volume, name='test_volume') ctx['ws'].volumes.create.assert_called_once_with( - name='test_volume', - catalog_name="dummy_crandom", - schema_name="dummy_srandom", - volume_type=VolumeType.MANAGED + name='test_volume', catalog_name="dummy_crandom", schema_name="dummy_srandom", volume_type=VolumeType.MANAGED ) assert info is not None From 09f44611e756d48a89914573e0831160dc2ac7c7 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 15 Nov 2024 13:06:59 +0100 Subject: [PATCH 2/7] simplify for the time being --- README.md | 95 +++++++++++---- src/databricks/labs/pytester/fixtures/iam.py | 112 ++++++++++++------ .../labs/pytester/fixtures/plugin.py | 1 + tests/integration/fixtures/test_iam.py | 7 ++ tests/integration/fixtures/test_run_as.py | 16 --- 5 files changed, 154 insertions(+), 77 deletions(-) delete mode 100644 tests/integration/fixtures/test_run_as.py diff --git a/README.md b/README.md index 36163e8..4edf300 100644 --- a/README.md +++ b/README.md @@ -6,14 +6,20 @@ * [Python Testing for Databricks](#python-testing-for-databricks) + * [Installation](#installation) * [Ecosystem](#ecosystem) * [PyTest Fixtures](#pytest-fixtures) * [Logging](#logging) - * [Installation](#installation) * [`debug_env_name` fixture](#debug_env_name-fixture) * [`debug_env` fixture](#debug_env-fixture) * [`env_or_skip` fixture](#env_or_skip-fixture) * [`ws` fixture](#ws-fixture) + * [`make_run_as` fixture](#make_run_as-fixture) + * [`acc` fixture](#acc-fixture) + * [`spark` fixture](#spark-fixture) + * [`sql_backend` fixture](#sql_backend-fixture) + * [`sql_exec` fixture](#sql_exec-fixture) + * [`sql_fetch_all` fixture](#sql_fetch_all-fixture) * [`make_random` fixture](#make_random-fixture) * [`make_instance_pool` fixture](#make_instance_pool-fixture) * [`make_instance_pool_permissions` fixture](#make_instance_pool_permissions-fixture) @@ -26,10 +32,12 @@ * [`make_pipeline` fixture](#make_pipeline-fixture) * [`make_warehouse` fixture](#make_warehouse-fixture) * [`make_group` fixture](#make_group-fixture) + * [`make_acc_group` fixture](#make_acc_group-fixture) * [`make_user` fixture](#make_user-fixture) * [`make_pipeline_permissions` fixture](#make_pipeline_permissions-fixture) * [`make_notebook` fixture](#make_notebook-fixture) * [`make_notebook_permissions` fixture](#make_notebook_permissions-fixture) + * [`make_workspace_file` fixture](#make_workspace_file-fixture) * [`make_directory` fixture](#make_directory-fixture) * [`make_directory_permissions` fixture](#make_directory_permissions-fixture) * [`make_repo` fixture](#make_repo-fixture) @@ -44,17 +52,15 @@ * [`make_schema` fixture](#make_schema-fixture) * [`make_table` fixture](#make_table-fixture) * [`make_storage_credential` fixture](#make_storage_credential-fixture) + * [`make_volume` fixture](#make_volume-fixture) * [`product_info` fixture](#product_info-fixture) - * [`sql_backend` fixture](#sql_backend-fixture) - * [`sql_exec` fixture](#sql_exec-fixture) - * [`sql_fetch_all` fixture](#sql_fetch_all-fixture) * [`make_model` fixture](#make_model-fixture) * [`make_experiment` fixture](#make_experiment-fixture) * [`make_experiment_permissions` fixture](#make_experiment_permissions-fixture) * [`make_warehouse_permissions` fixture](#make_warehouse_permissions-fixture) * [`make_lakeview_dashboard_permissions` fixture](#make_lakeview_dashboard_permissions-fixture) - * [`workspace_library` fixture](#workspace_library-fixture) * [`log_workspace_link` fixture](#log_workspace_link-fixture) + * [`log_account_link` fixture](#log_account_link-fixture) * [`make_dashboard_permissions` fixture](#make_dashboard_permissions-fixture) * [`make_alert_permissions` fixture](#make_alert_permissions-fixture) * [`make_query` fixture](#make_query-fixture) @@ -62,7 +68,11 @@ * [`make_registered_model_permissions` fixture](#make_registered_model_permissions-fixture) * [`make_serving_endpoint` fixture](#make_serving_endpoint-fixture) * [`make_serving_endpoint_permissions` fixture](#make_serving_endpoint_permissions-fixture) + * [`make_feature_table` fixture](#make_feature_table-fixture) * [`make_feature_table_permissions` fixture](#make_feature_table_permissions-fixture) + * [`watchdog_remove_after` fixture](#watchdog_remove_after-fixture) + * [`watchdog_purge_suffix` fixture](#watchdog_purge_suffix-fixture) + * [`is_in_debug` fixture](#is_in_debug-fixture) * [Project Support](#project-support) @@ -75,7 +85,7 @@ also install it directly from the command line: pip install databricks-labs-pytester ``` -If you use `hatch` as a build system, make sure to add `databricks-labs-pytester` as +If you use `hatch` as a build system, make sure to add `databricks-labs-pytester` as a test-time dependency and not as a compile-time dependency, otherwise your wheels will transitively depend on `pytest`, which is not usually something you need. @@ -97,7 +107,7 @@ dependencies = [ "pylint~=3.2.2", "pylint-pytest==2.0.0a0", "databricks-labs-pylint~=0.4.0", - "databricks-labs-pytester~=0.2", # <= this library + "databricks-labs-pytester~=0.2", # <= this library "pytest~=8.3.3", "pytest-cov~=4.1.0", "pytest-mock~=3.14.0", @@ -116,7 +126,7 @@ dependencies = [ Built on top of [Databricks SDK for Python](https://github.com/databricks/databricks-sdk-py), this library is part of the Databricks Labs Python ecosystem, which includes the following projects: * [PyLint Plugin for Databricks](https://github.com/databrickslabs/pylint-plugin) for static code analysis and early bug detection. -* [Blueprint](https://github.com/databrickslabs/blueprint) for +* [Blueprint](https://github.com/databrickslabs/blueprint) for [Python-native pathlib.Path-like interfaces](https://github.com/databrickslabs/blueprint#python-native-pathlibpath-like-interfaces), [Managing Python App installations within Databricks Workspaces](https://github.com/databrickslabs/blueprint#application-and-installation-state), [Application Migrations](https://github.com/databrickslabs/blueprint#application-state-migrations), and @@ -130,18 +140,18 @@ See [this video](https://www.youtube.com/watch?v=CNypO79IATc) for a quick overvi ## PyTest Fixtures -[PyTest Fixtures](https://docs.pytest.org/en/latest/explanation/fixtures.html) are a powerful way to manage test setup and teardown in Python. This library provides +[PyTest Fixtures](https://docs.pytest.org/en/latest/explanation/fixtures.html) are a powerful way to manage test setup and teardown in Python. This library provides a set of fixtures to help you write integration tests for Databricks. These fixtures were incubated -within the [Unity Catalog Automated Migrations project](https://github.com/databrickslabs/ucx/blame/df7f1d7647251fb8f0f23c56a548b99092484a7c/src/databricks/labs/ucx/mixins/fixtures.py) +within the [Unity Catalog Automated Migrations project](https://github.com/databrickslabs/ucx/blame/df7f1d7647251fb8f0f23c56a548b99092484a7c/src/databricks/labs/ucx/mixins/fixtures.py) for more than a year and are now available for other projects to simplify integration testing with Databricks. [[back to top](#python-testing-for-databricks)] ### Logging -This library is built on years of debugging integration tests for Databricks and its ecosystem. +This library is built on years of debugging integration tests for Databricks and its ecosystem. -That's why it comes with a built-in logger that traces creation and deletion of dummy entities through links in +That's why it comes with a built-in logger that traces creation and deletion of dummy entities through links in the Databricks Workspace UI. If you run the following code: ```python @@ -281,6 +291,39 @@ def test_workspace_operations(ws): 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`](#make_feature_table-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_run_as`](#make_run_as-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture). +[[back to top](#python-testing-for-databricks)] + +### `make_run_as` fixture +This fixture provides a function to create an account service principal via [`acc` fixture](#acc-fixture) and +assign it to a workspace. The service principal is removed after the test is complete. The service principal is +created with a random display name and assigned to the workspace with the default permissions. + +Use the `account_groups` argument to assign the service principal to account groups, which have the required +permissions to perform a specific action. + +Returned object has the following properties: +* `ws`: Workspace client that is authenticated as the ephemeral service principal. +* `sql_backend`: SQL backend that is authenticated as the ephemeral service principal. +* `sql_exec`: Function to execute a SQL statement on behalf of the ephemeral service principal. +* `sql_fetch_all`: Function to fetch all rows from a SQL statement on behalf of the ephemeral service principal. +* `display_name`: Display name of the ephemeral service principal. +* `application_id`: Application ID of the ephemeral service principal. +* ... other fixtures are not currently available through the returned object yet, as it's quite complex to + implement, but there's a possibility to add generic support for them in the future. + +Example: + +```python +def test_run_as_lower_privilege_user(make_run_as, ws): + run_as = make_run_as(account_groups=['account.group.name']) + through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) + me = ws.current_user.me() + assert me.user_name != through_query.my_name +``` + +See also [`acc`](#acc-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`env_or_skip`](#env_or_skip-fixture), [`log_account_link`](#log_account_link-fixture). + + [[back to top](#python-testing-for-databricks)] ### `acc` fixture @@ -305,7 +348,7 @@ def test_listing_workspaces(acc): assert len(workspaces) >= 1 ``` -See also [`make_acc_group`](#make_acc_group-fixture), [`make_run_as`](#make_run_as-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture), [`env_or_skip`](#env_or_skip-fixture). +See also [`log_account_link`](#log_account_link-fixture), [`make_acc_group`](#make_acc_group-fixture), [`make_run_as`](#make_run_as-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture), [`env_or_skip`](#env_or_skip-fixture). [[back to top](#python-testing-for-databricks)] @@ -349,14 +392,6 @@ Fetch all rows from a SQL statement. See also [`sql_backend`](#sql_backend-fixture). -[[back to top](#python-testing-for-databricks)] - -### `make_run_as` fixture -_No description yet._ - -See also [`acc`](#acc-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`env_or_skip`](#env_or_skip-fixture). - - [[back to top](#python-testing-for-databricks)] ### `make_random` fixture @@ -1067,6 +1102,14 @@ rns a function to log a workspace link. See also [`ws`](#ws-fixture). +[[back to top](#python-testing-for-databricks)] + +### `log_account_link` fixture +rns a function to log an account link. + +See also [`make_run_as`](#make_run_as-fixture), [`acc`](#acc-fixture). + + [[back to top](#python-testing-for-databricks)] ### `make_dashboard_permissions` fixture @@ -1208,11 +1251,11 @@ See also [`debug_env`](#debug_env-fixture), [`env_or_skip`](#env_or_skip-fixture # Project Support -Please note that this project is provided for your exploration only and is not -formally supported by Databricks with Service Level Agreements (SLAs). They are -provided AS-IS, and we do not make any guarantees of any kind. Please do not +Please note that this project is provided for your exploration only and is not +formally supported by Databricks with Service Level Agreements (SLAs). They are +provided AS-IS, and we do not make any guarantees of any kind. Please do not submit a support ticket relating to any issues arising from the use of this project. -Any issues discovered through the use of this project should be filed as GitHub -[Issues on this repository](https://github.com/databrickslabs/pytester/issues). +Any issues discovered through the use of this project should be filed as GitHub +[Issues on this repository](https://github.com/databrickslabs/pytester/issues). They will be reviewed as time permits, but no formal SLAs for support exist. diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index ad5eeb8..3286ee7 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -1,19 +1,26 @@ import logging import warnings -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Iterable from datetime import timedelta -from _pytest.fixtures import SubRequest, FixtureDef, fixture -from _pytest.scope import Scope - from pytest import fixture +from databricks.labs.lsql import Row +from databricks.labs.lsql.backends import StatementExecutionBackend, SqlBackend from databricks.sdk import AccountGroupsAPI, GroupsAPI, WorkspaceClient, AccountClient from databricks.sdk.config import Config from databricks.sdk.errors import ResourceConflict, NotFound from databricks.sdk.retries import retried from databricks.sdk.service import iam -from databricks.sdk.service.iam import User, Group, ServicePrincipal, Patch, PatchOp, ComplexValue, PatchSchema, \ - WorkspacePermission +from databricks.sdk.service.iam import ( + User, + Group, + ServicePrincipal, + Patch, + PatchOp, + ComplexValue, + PatchSchema, + WorkspacePermission, +) from databricks.labs.pytester.fixtures.baseline import factory @@ -190,24 +197,28 @@ def create( class RunAs: - def __init__(self, service_principal: ServicePrincipal, workspace_client: WorkspaceClient, request: SubRequest): - self._request = SubRequest( - request, - Scope.Function, - None, - request.param_index, - FixtureDef( - config=request.config, - baseid='ws', - argname='ws', - func=lambda: workspace_client, - scope=None, - params=None, - ), - ) - self._request._arg2fixturedefs = {} - self._request._fixture_defs = {} + def __init__(self, service_principal: ServicePrincipal, workspace_client: WorkspaceClient, env_or_skip): self._service_principal = service_principal + self._workspace_client = workspace_client + self._env_or_skip = env_or_skip + + @property + def ws(self): + return self._workspace_client + + @property + def sql_backend(self) -> SqlBackend: + # TODO: Switch to `__getattr__` + `SubRequest` to get a generic way of initializing all workspace fixtures. + # This will allow us to remove the `sql_backend` fixture and make the `RunAs` class more generic. + # It turns out to be more complicated than it first appears, because we don't get these at pytest.collect phase. + warehouse_id = self._env_or_skip("DATABRICKS_WAREHOUSE_ID") + return StatementExecutionBackend(self._workspace_client, warehouse_id) + + def sql_exec(self, statement: str) -> None: + return self.sql_backend.execute(statement) + + def sql_fetch_all(self, statement: str) -> Iterable[Row]: + return self.sql_backend.fetch(statement) def __getattr__(self, item: str): if item in self.__dict__: @@ -217,10 +228,12 @@ def __getattr__(self, item: str): @property def display_name(self) -> str: + assert self._service_principal.display_name is not None return self._service_principal.display_name @property def application_id(self) -> str: + assert self._service_principal.application_id is not None return self._service_principal.application_id def __repr__(self): @@ -228,22 +241,47 @@ def __repr__(self): @fixture -def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_skip, request, log_account_link): +def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_skip, log_account_link): """ - This fixture provides a function to create a service principal and assign it to a workspace. The service principal - is removed after the test is complete. The fixture returns an instance of `RunAs` that proxies the calls to - the other fixtures supported by the plugin, for example, `sql_fetch_all`. + This fixture provides a function to create an account service principal via [`acc` fixture](#acc-fixture) and + assign it to a workspace. The service principal is removed after the test is complete. The service principal is + created with a random display name and assigned to the workspace with the default permissions. + + Use the `account_groups` argument to assign the service principal to account groups, which have the required + permissions to perform a specific action. + + Returned object has the following properties: + * `ws`: Workspace client that is authenticated as the ephemeral service principal. + * `sql_backend`: SQL backend that is authenticated as the ephemeral service principal. + * `sql_exec`: Function to execute a SQL statement on behalf of the ephemeral service principal. + * `sql_fetch_all`: Function to fetch all rows from a SQL statement on behalf of the ephemeral service principal. + * `display_name`: Display name of the ephemeral service principal. + * `application_id`: Application ID of the ephemeral service principal. + * ... other fixtures are not currently available through the returned object yet, as it's quite complex to + implement, but there's a possibility to add generic support for them in the future. - The service principal is created with a random display name and assigned to the workspace with + Example: + + ```python + def test_run_as_lower_privilege_user(make_run_as, ws): + run_as = make_run_as(account_groups=['account.group.name']) + through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) + me = ws.current_user.me() + assert me.user_name != through_query.my_name + ``` """ - def create(*, account_groups: list[str] = None): + def create(*, account_groups: list[str] | None = None): workspace_id = ws.get_workspace_id() service_principal = acc.service_principals.create(display_name=f'spn-{make_random()}') - created_secret = acc.service_principal_secrets.create(service_principal.id) + assert service_principal.id is not None + service_principal_id = int(service_principal.id) + created_secret = acc.service_principal_secrets.create(service_principal_id) if account_groups: group_mapping = {} for group in acc.groups.list(attributes='id,displayName'): + if group.id is None: + continue group_mapping[group.display_name] = group.id for group_name in account_groups: if group_name not in group_mapping: @@ -251,22 +289,26 @@ def create(*, account_groups: list[str] = None): group_id = group_mapping[group_name] acc.groups.patch( group_id, - operations=[Patch(PatchOp.ADD, 'members', [ComplexValue(value=service_principal.id).as_dict()])], + operations=[ + Patch(PatchOp.ADD, 'members', [ComplexValue(value=str(service_principal_id)).as_dict()]), + ], schemas=[PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP], ) permissions = [WorkspacePermission.USER] - acc.workspace_assignment.update(workspace_id, service_principal.id, permissions=permissions) + acc.workspace_assignment.update(workspace_id, service_principal_id, permissions=permissions) ws_as_spn = WorkspaceClient( host=ws.config.host, client_id=service_principal.application_id, client_secret=created_secret.secret, ) - log_account_link('account service principal', f'users/serviceprincipals/{service_principal.id}') + log_account_link('account service principal', f'users/serviceprincipals/{service_principal_id}') - return RunAs(service_principal, ws_as_spn, request) + return RunAs(service_principal, ws_as_spn, env_or_skip) def remove(run_as: RunAs): - acc.service_principals.delete(run_as._service_principal.id) + service_principal_id = run_as._service_principal.id # pylint: disable=protected-access + assert service_principal_id is not None + acc.service_principals.delete(service_principal_id) yield from factory("service principal", create, remove) diff --git a/src/databricks/labs/pytester/fixtures/plugin.py b/src/databricks/labs/pytester/fixtures/plugin.py index ad6f68e..86a84f8 100644 --- a/src/databricks/labs/pytester/fixtures/plugin.py +++ b/src/databricks/labs/pytester/fixtures/plugin.py @@ -61,6 +61,7 @@ 'debug_env', 'env_or_skip', 'ws', + 'make_run_as', 'acc', 'spark', 'sql_backend', diff --git a/tests/integration/fixtures/test_iam.py b/tests/integration/fixtures/test_iam.py index 445675c..0a237db 100644 --- a/tests/integration/fixtures/test_iam.py +++ b/tests/integration/fixtures/test_iam.py @@ -19,3 +19,10 @@ def test_new_account_group(make_acc_group, acc): group = make_acc_group() loaded = acc.groups.get(group.id) assert group.display_name == loaded.display_name + + +def test_run_as_lower_privilege_user(make_run_as, ws): + run_as = make_run_as(account_groups=['role.labs.lsql.write']) + through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) + current_user = ws.current_user.me() + assert current_user.user_name != through_query.my_name diff --git a/tests/integration/fixtures/test_run_as.py b/tests/integration/fixtures/test_run_as.py deleted file mode 100644 index 3ace4fe..0000000 --- a/tests/integration/fixtures/test_run_as.py +++ /dev/null @@ -1,16 +0,0 @@ -def test_run_as(make_run_as, ws): - run_as = make_run_as(account_groups=['role.labs.lsql.write']) - through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) - me = ws.current_user.me() - assert me.user_name != through_query.my_name - - -def test_notebooks_are_created_by_different_users(make_run_as, make_notebook): - notebook_by_current_user = make_notebook() - a = notebook_by_current_user.parent.as_posix() - - run_as = make_run_as(account_groups=['role.labs.lsql.write']) - notebook_by_ephemeral_principal = run_as.make_notebook() - b = notebook_by_ephemeral_principal.parent.as_posix() - - assert a != b From 2ba6b6091c5a829a6a1d3703af29edd06bd569ad Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 15 Nov 2024 13:40:45 +0100 Subject: [PATCH 3/7] ... --- tests/integration/fixtures/test_catalog.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/fixtures/test_catalog.py b/tests/integration/fixtures/test_catalog.py index 957660f..1f18986 100644 --- a/tests/integration/fixtures/test_catalog.py +++ b/tests/integration/fixtures/test_catalog.py @@ -17,6 +17,7 @@ def test_schema_fixture(make_schema): logger.info(f"Created new schema: {make_schema()}") +@pytest.mark.skip("Invalid configuration value detected for fs.azure.account.key") def test_managed_schema_fixture(make_schema, make_random, env_or_skip): schema_name = f"dummy_s{make_random(4)}".lower() schema_location = f"{env_or_skip('TEST_MOUNT_CONTAINER')}/a/{schema_name}" From 64e31da88b8e4f921349531b1c346c158dd06cf9 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 15 Nov 2024 13:44:10 +0100 Subject: [PATCH 4/7] ... --- src/databricks/labs/pytester/fixtures/iam.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index 3286ee7..56de1a5 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -298,6 +298,7 @@ def create(*, account_groups: list[str] | None = None): acc.workspace_assignment.update(workspace_id, service_principal_id, permissions=permissions) ws_as_spn = WorkspaceClient( host=ws.config.host, + auth_type='oauth-m2m', client_id=service_principal.application_id, client_secret=created_secret.secret, ) From 5a352da5f0ed453aad1eae415252135f67b6334a Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 15 Nov 2024 13:53:39 +0100 Subject: [PATCH 5/7] ... --- README.md | 31 +++++++++++++++----- src/databricks/labs/pytester/fixtures/iam.py | 31 +++++++++++++++----- tests/integration/fixtures/test_run_as.py | 12 ++++++++ 3 files changed, 60 insertions(+), 14 deletions(-) create mode 100644 tests/integration/fixtures/test_run_as.py diff --git a/README.md b/README.md index 4edf300..95d9785 100644 --- a/README.md +++ b/README.md @@ -301,6 +301,16 @@ created with a random display name and assigned to the workspace with the defaul Use the `account_groups` argument to assign the service principal to account groups, which have the required permissions to perform a specific action. +Example: + +```python +def test_run_as_lower_privilege_user(make_run_as, ws): + run_as = make_run_as(account_groups=['account.group.name']) + through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) + me = ws.current_user.me() + assert me.user_name != through_query.my_name +``` + Returned object has the following properties: * `ws`: Workspace client that is authenticated as the ephemeral service principal. * `sql_backend`: SQL backend that is authenticated as the ephemeral service principal. @@ -308,17 +318,24 @@ Returned object has the following properties: * `sql_fetch_all`: Function to fetch all rows from a SQL statement on behalf of the ephemeral service principal. * `display_name`: Display name of the ephemeral service principal. * `application_id`: Application ID of the ephemeral service principal. -* ... other fixtures are not currently available through the returned object yet, as it's quite complex to - implement, but there's a possibility to add generic support for them in the future. +* if you want to have other fixtures available in the context of the ephemeral service principal, you can override + the [`ws` fixture](#ws-fixture) on the file level, which would make all workspace fixtures provided by this + plugin to run as lower privilege ephemeral service principal. You cannot combine it with the account-admin-level + principal you're using to create the ephemeral principal. Example: ```python -def test_run_as_lower_privilege_user(make_run_as, ws): - run_as = make_run_as(account_groups=['account.group.name']) - through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) - me = ws.current_user.me() - assert me.user_name != through_query.my_name +from pytest import fixture + +@fixture +def ws(make_run_as): + run_as = make_run_as(account_groups=['account.group.used.for.all.tests.in.this.file']) + return run_as.ws + +def test_creating_notebook_on_behalf_of_ephemeral_principal(make_notebook): + notebook = make_notebook() + assert notebook.exists() ``` See also [`acc`](#acc-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`env_or_skip`](#env_or_skip-fixture), [`log_account_link`](#log_account_link-fixture). diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index 56de1a5..35dee5c 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -250,6 +250,16 @@ def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_ski Use the `account_groups` argument to assign the service principal to account groups, which have the required permissions to perform a specific action. + Example: + + ```python + def test_run_as_lower_privilege_user(make_run_as, ws): + run_as = make_run_as(account_groups=['account.group.name']) + through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) + me = ws.current_user.me() + assert me.user_name != through_query.my_name + ``` + Returned object has the following properties: * `ws`: Workspace client that is authenticated as the ephemeral service principal. * `sql_backend`: SQL backend that is authenticated as the ephemeral service principal. @@ -257,17 +267,24 @@ def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_ski * `sql_fetch_all`: Function to fetch all rows from a SQL statement on behalf of the ephemeral service principal. * `display_name`: Display name of the ephemeral service principal. * `application_id`: Application ID of the ephemeral service principal. - * ... other fixtures are not currently available through the returned object yet, as it's quite complex to - implement, but there's a possibility to add generic support for them in the future. + * if you want to have other fixtures available in the context of the ephemeral service principal, you can override + the [`ws` fixture](#ws-fixture) on the file level, which would make all workspace fixtures provided by this + plugin to run as lower privilege ephemeral service principal. You cannot combine it with the account-admin-level + principal you're using to create the ephemeral principal. Example: ```python - def test_run_as_lower_privilege_user(make_run_as, ws): - run_as = make_run_as(account_groups=['account.group.name']) - through_query = next(run_as.sql_fetch_all("SELECT CURRENT_USER() AS my_name")) - me = ws.current_user.me() - assert me.user_name != through_query.my_name + from pytest import fixture + + @fixture + def ws(make_run_as): + run_as = make_run_as(account_groups=['account.group.used.for.all.tests.in.this.file']) + return run_as.ws + + def test_creating_notebook_on_behalf_of_ephemeral_principal(make_notebook): + notebook = make_notebook() + assert notebook.exists() ``` """ diff --git a/tests/integration/fixtures/test_run_as.py b/tests/integration/fixtures/test_run_as.py new file mode 100644 index 0000000..8583784 --- /dev/null +++ b/tests/integration/fixtures/test_run_as.py @@ -0,0 +1,12 @@ +from pytest import fixture + + +@fixture +def ws(make_run_as): + run_as = make_run_as(account_groups=['role.labs.lsql.write']) + return run_as.ws + + +def test_creating_notebook_on_behalf_of_ephemeral_principal(make_notebook): + notebook = make_notebook() + assert notebook.exists() From 17c2746865cd57f6854808c0c3f8dcfe696c0298 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 15 Nov 2024 14:10:11 +0100 Subject: [PATCH 6/7] ... --- src/databricks/labs/pytester/fixtures/iam.py | 36 ++++++++++++++++---- 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index 35dee5c..5995f82 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -3,6 +3,8 @@ from collections.abc import Callable, Generator, Iterable from datetime import timedelta +from databricks.sdk.credentials_provider import OAuthCredentialsProvider, OauthCredentialsStrategy +from databricks.sdk.oauth import ClientCredentials, Token from pytest import fixture from databricks.labs.lsql import Row from databricks.labs.lsql.backends import StatementExecutionBackend, SqlBackend @@ -313,17 +315,39 @@ def create(*, account_groups: list[str] | None = None): ) permissions = [WorkspacePermission.USER] acc.workspace_assignment.update(workspace_id, service_principal_id, permissions=permissions) - ws_as_spn = WorkspaceClient( - host=ws.config.host, - auth_type='oauth-m2m', - client_id=service_principal.application_id, - client_secret=created_secret.secret, - ) + ws_as_spn = _make_workspace_client(created_secret, service_principal) log_account_link('account service principal', f'users/serviceprincipals/{service_principal_id}') return RunAs(service_principal, ws_as_spn, env_or_skip) + def _make_workspace_client(created_secret, service_principal): + oidc = ws.config.oidc_endpoints + assert oidc is not None, 'OIDC is required' + application_id = service_principal.application_id + secret_value = created_secret.secret + assert application_id is not None + assert secret_value is not None + token_source = ClientCredentials( + client_id=application_id, + client_secret=secret_value, + token_url=oidc.token_endpoint, + scopes=["all-apis"], + use_header=True, + ) + + def inner() -> dict[str, str]: + inner_token = token_source.token() + return {'Authorization': f'{inner_token.token_type} {inner_token.access_token}'} + + def token() -> Token: + return token_source.token() + + credentials_provider = OAuthCredentialsProvider(inner, token) + credentials_strategy = OauthCredentialsStrategy('oauth-m2m', lambda _: credentials_provider) + ws_as_spn = WorkspaceClient(host=ws.config.host, credentials_strategy=credentials_strategy) + return ws_as_spn + def remove(run_as: RunAs): service_principal_id = run_as._service_principal.id # pylint: disable=protected-access assert service_principal_id is not None From 515b9644a7dd50f239994bf5f6f4a34ac9c8c6cc Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 15 Nov 2024 14:20:50 +0100 Subject: [PATCH 7/7] ... --- README.md | 6 +- src/databricks/labs/pytester/fixtures/iam.py | 78 ++++++++++++-------- 2 files changed, 52 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 95d9785..28b7564 100644 --- a/README.md +++ b/README.md @@ -338,7 +338,9 @@ def test_creating_notebook_on_behalf_of_ephemeral_principal(make_notebook): assert notebook.exists() ``` -See also [`acc`](#acc-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`env_or_skip`](#env_or_skip-fixture), [`log_account_link`](#log_account_link-fixture). +This fixture currently doesn't work with Databricks Metadata Service authentication on Azure Databricks. + +See also [`acc`](#acc-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`env_or_skip`](#env_or_skip-fixture), [`log_account_link`](#log_account_link-fixture), [`is_in_debug`](#is_in_debug-fixture). [[back to top](#python-testing-for-databricks)] @@ -1259,7 +1261,7 @@ Returns true if the test is running from a debugger in IDE, otherwise false. The following IDE are supported: IntelliJ IDEA (including Community Edition), PyCharm (including Community Edition), and Visual Studio Code. -See also [`debug_env`](#debug_env-fixture), [`env_or_skip`](#env_or_skip-fixture). +See also [`debug_env`](#debug_env-fixture), [`env_or_skip`](#env_or_skip-fixture), [`make_run_as`](#make_run_as-fixture). [[back to top](#python-testing-for-databricks)] diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index 5995f82..e86b14a 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -3,9 +3,11 @@ from collections.abc import Callable, Generator, Iterable from datetime import timedelta +import pytest +from pytest import fixture from databricks.sdk.credentials_provider import OAuthCredentialsProvider, OauthCredentialsStrategy from databricks.sdk.oauth import ClientCredentials, Token -from pytest import fixture +from databricks.sdk.service.oauth2 import CreateServicePrincipalSecretResponse from databricks.labs.lsql import Row from databricks.labs.lsql.backends import StatementExecutionBackend, SqlBackend from databricks.sdk import AccountGroupsAPI, GroupsAPI, WorkspaceClient, AccountClient @@ -242,8 +244,41 @@ def __repr__(self): return f'RunAs({self.display_name})' +def _make_workspace_client( + ws: WorkspaceClient, + created_secret: CreateServicePrincipalSecretResponse, + service_principal: ServicePrincipal, +) -> WorkspaceClient: + oidc = ws.config.oidc_endpoints + assert oidc is not None, 'OIDC is required' + application_id = service_principal.application_id + secret_value = created_secret.secret + assert application_id is not None + assert secret_value is not None + + token_source = ClientCredentials( + client_id=application_id, + client_secret=secret_value, + token_url=oidc.token_endpoint, + scopes=["all-apis"], + use_header=True, + ) + + def inner() -> dict[str, str]: + inner_token = token_source.token() + return {'Authorization': f'{inner_token.token_type} {inner_token.access_token}'} + + def token() -> Token: + return token_source.token() + + credentials_provider = OAuthCredentialsProvider(inner, token) + credentials_strategy = OauthCredentialsStrategy('oauth-m2m', lambda _: credentials_provider) + ws_as_spn = WorkspaceClient(host=ws.config.host, credentials_strategy=credentials_strategy) + return ws_as_spn + + @fixture -def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_skip, log_account_link): +def make_run_as(acc: AccountClient, ws: WorkspaceClient, make_random, env_or_skip, log_account_link, is_in_debug): """ This fixture provides a function to create an account service principal via [`acc` fixture](#acc-fixture) and assign it to a workspace. The service principal is removed after the test is complete. The service principal is @@ -288,8 +323,18 @@ def test_creating_notebook_on_behalf_of_ephemeral_principal(make_notebook): notebook = make_notebook() assert notebook.exists() ``` + + This fixture currently doesn't work with Databricks Metadata Service authentication on Azure Databricks. """ + if ws.config.auth_type == 'metadata-service' and ws.config.is_azure: + # TODO: fix `invalid_scope: AADSTS1002012: The provided value for scope all-apis is not valid.` error + # + # We're having issues with the Azure Metadata Service and service principals. The error message is: + # Client credential flows must have a scope value with /.default suffixed to the resource identifier + # (application ID URI) + pytest.skip('Azure Metadata Service does not support service principals') + def create(*, account_groups: list[str] | None = None): workspace_id = ws.get_workspace_id() service_principal = acc.service_principals.create(display_name=f'spn-{make_random()}') @@ -315,39 +360,12 @@ def create(*, account_groups: list[str] | None = None): ) permissions = [WorkspacePermission.USER] acc.workspace_assignment.update(workspace_id, service_principal_id, permissions=permissions) - ws_as_spn = _make_workspace_client(created_secret, service_principal) + ws_as_spn = _make_workspace_client(ws, created_secret, service_principal) log_account_link('account service principal', f'users/serviceprincipals/{service_principal_id}') return RunAs(service_principal, ws_as_spn, env_or_skip) - def _make_workspace_client(created_secret, service_principal): - oidc = ws.config.oidc_endpoints - assert oidc is not None, 'OIDC is required' - application_id = service_principal.application_id - secret_value = created_secret.secret - assert application_id is not None - assert secret_value is not None - token_source = ClientCredentials( - client_id=application_id, - client_secret=secret_value, - token_url=oidc.token_endpoint, - scopes=["all-apis"], - use_header=True, - ) - - def inner() -> dict[str, str]: - inner_token = token_source.token() - return {'Authorization': f'{inner_token.token_type} {inner_token.access_token}'} - - def token() -> Token: - return token_source.token() - - credentials_provider = OAuthCredentialsProvider(inner, token) - credentials_strategy = OauthCredentialsStrategy('oauth-m2m', lambda _: credentials_provider) - ws_as_spn = WorkspaceClient(host=ws.config.host, credentials_strategy=credentials_strategy) - return ws_as_spn - def remove(run_as: RunAs): service_principal_id = run_as._service_principal.id # pylint: disable=protected-access assert service_principal_id is not None