...

Author: nfx

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