Added make_run_as fixture

Author: nfx

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)]

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)

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

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.

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)

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:

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',

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

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