Add group assumption for OAuth authentication

packages/auth/NEXT_CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### New Features and Improvements
 
+- Added group role assumption for OAuth M2M and OIDC token exchange, including
+  default credential configuration through `DATABRICKS_GROUP_ID` and profile
+  `group_id`.
+
 ### Bug Fixes
 
 ### Documentation

packages/auth/src/credentials/m2m.ts

@@ -78,10 +78,14 @@ export function newM2mCredentials(
       ? options.scopes
       : DEFAULT_SCOPES;
 
-  const body = new URLSearchParams({
+  const params = new URLSearchParams({
     grant_type: 'client_credentials',
     scope: scopes.join(' '),
-  }).toString();
+  });
+  if (options.groupId !== undefined && options.groupId !== '') {
+    params.set('assume_group', options.groupId);
+  }
+  const body = params.toString();
 
   // Client ID and secret are URL-encoded before Basic auth encoding to
   // avoid ambiguity with special characters in either value, matching the

packages/auth/src/oidc/tokensource.ts

@@ -43,6 +43,12 @@ export interface DatabricksOidcTokenProviderConfig {
    */
   host: string;
 
+  /**
+   * ID of the group whose role is assumed by the exchanged token. When
+   * omitted or empty, no group role is assumed.
+   */
+  groupId?: string;
+
   /**
    * TokenEndpointProvider returns the token endpoint for the Databricks OIDC
    * application.
@@ -89,6 +95,9 @@ async function exchangeIdToken(
   params.set('subject_token_type', 'urn:ietf:params:oauth:token-type:jwt');
   params.set('subject_token', idToken.value);
   params.set('grant_type', 'urn:ietf:params:oauth:grant-type:token-exchange');
+  if (config.groupId !== undefined && config.groupId !== '') {
+    params.set('assume_group', config.groupId);
+  }
 
   const response = await fetch(endpoints.tokenEndpoint, {
     method: 'POST',

packages/auth/tests/credentials/default/chain.test.ts

@@ -1,4 +1,4 @@
-import {describe, expect, it} from 'vitest';
+import {afterEach, describe, expect, it, vi} from 'vitest';
 
 // Import Secret from the browser subpath so this test can run under both
 // the Node and browser runners (the default `/profiles` entry pulls in
@@ -59,6 +59,11 @@ const loaderFor =
 
 describe('DefaultCredentials chain', () => {
   const selectedError = new Error('selected provider failed');
+
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
   const resolutionCases: {
     name: string;
     strategies: readonly Strategy[];
@@ -166,6 +171,52 @@ describe('DefaultCredentials chain', () => {
     expect(buildCount).toBe(1);
   });
 
+  it('passes the configured group to the M2M token request', async () => {
+    const tokenEndpoint = `${HOST}/oidc/v1/token`;
+    const fetchMock = vi.fn<typeof fetch>((input, init) => {
+      const url = input instanceof Request ? input.url : input.toString();
+      if (url === `${HOST}/.well-known/databricks-config`) {
+        return Promise.resolve(
+          new Response(JSON.stringify({oidc_endpoint: `${HOST}/oidc`}), {
+            status: 200,
+          })
+        );
+      }
+      if (url === `${HOST}/oidc/.well-known/oauth-authorization-server`) {
+        return Promise.resolve(
+          new Response(JSON.stringify({token_endpoint: tokenEndpoint}), {
+            status: 200,
+          })
+        );
+      }
+      const body = init?.body;
+      if (typeof body !== 'string') {
+        expect.fail('expected body to be a string');
+      }
+      const params = new URLSearchParams(body);
+      expect(params.get('assume_group')).toBe('group-123');
+      return Promise.resolve(
+        new Response(JSON.stringify({access_token: 'token'}), {status: 200})
+      );
+    });
+    vi.stubGlobal('fetch', fetchMock);
+    const creds = new DefaultCredentials(
+      [patStrategy, m2mStrategy],
+      loaderFor({
+        host: HOST,
+        groupId: 'group-123',
+        token: new Secret('ignored-pat'),
+        clientId: 'client-id',
+        clientSecret: new Secret('client-secret'),
+      })
+    );
+
+    expect(await creds.authHeaders()).toEqual([
+      {key: 'Authorization', value: 'Bearer token'},
+    ]);
+    expect(fetchMock).toHaveBeenCalledTimes(3);
+  });
+
   it('invokes the profile loader exactly once', async () => {
     let loaderCalls = 0;
     const loader = (): Promise<Profile> => {