diff --git a/README.md b/README.md index d209141..28b7564 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 @@ -255,7 +265,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 +288,59 @@ 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)] + +### `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. + +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. +* `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. +* 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 +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() +``` + +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)] @@ -305,7 +367,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 [`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)] @@ -372,7 +434,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)] @@ -1059,6 +1121,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 @@ -1191,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)] @@ -1200,11 +1270,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/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..e86b14a 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -1,15 +1,30 @@ import logging import warnings -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Iterable from datetime import timedelta +import pytest from pytest import fixture -from databricks.sdk import AccountGroupsAPI, GroupsAPI, WorkspaceClient +from databricks.sdk.credentials_provider import OAuthCredentialsProvider, OauthCredentialsStrategy +from databricks.sdk.oauth import ClientCredentials, Token +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 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 +198,177 @@ 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, 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__: + return self.__dict__[item] + fixture_value = self._request.getfixturevalue(item) + return fixture_value + + @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): + 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, 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 + 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. + + 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. + * `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. + * 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 + 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() + ``` + + 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()}') + 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: + 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=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) + 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 remove(run_as: RunAs): + 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/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..86a84f8 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, @@ -60,6 +61,7 @@ 'debug_env', 'env_or_skip', 'ws', + 'make_run_as', 'acc', 'spark', 'sql_backend', @@ -105,6 +107,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_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}" 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 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() 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