diff --git a/packages/auth/src/credentials/default/chain.ts b/packages/auth/src/credentials/default/chain.ts index 6d2af7933..07e19eb56 100644 --- a/packages/auth/src/credentials/default/chain.ts +++ b/packages/auth/src/credentials/default/chain.ts @@ -15,6 +15,8 @@ import {DefaultCredentialsError} from './errors'; export interface Strategy { /** Short identifier, e.g. `pat`, `oauth-m2m`, or `databricks-cli`. */ readonly name: string; + /** Whether this strategy can request credentials for an assumed group. */ + readonly supportsGroupAssumption: boolean; readonly configure: (profile: Profile) => Credentials | undefined; } @@ -57,6 +59,13 @@ export class DefaultCredentials implements Credentials { return this.resolveByAuthType(profile, profile.authType); } for (const strategy of this.strategies) { + if ( + profile.groupId !== undefined && + profile.groupId !== '' && + !strategy.supportsGroupAssumption + ) { + continue; + } const built = strategy.configure(profile); if (built !== undefined) { return built; @@ -76,6 +85,16 @@ export class DefaultCredentials implements Credentials { `auth type "${authType}" not found, please check ${AUTH_DOC_URL} for a list of supported auth types` ); } + if ( + profile.groupId !== undefined && + profile.groupId !== '' && + !strategy.supportsGroupAssumption + ) { + throw new DefaultCredentialsError( + 'GROUP_ROLE_UNSUPPORTED', + `auth type "${authType}" does not support group role assumption. Use OAuth M2M or Workload Identity Federation` + ); + } const built = strategy.configure(profile); if (built === undefined) { throw new DefaultCredentialsError( @@ -90,6 +109,7 @@ export class DefaultCredentials implements Credentials { /** PAT strategy: configured when `token` is set in the profile. */ export const patStrategy: Strategy = { name: 'pat', + supportsGroupAssumption: false, configure: profile => { if (profile.host === undefined) return undefined; if (profile.token === undefined) return undefined; @@ -103,6 +123,7 @@ export const patStrategy: Strategy = { */ export const m2mStrategy: Strategy = { name: 'oauth-m2m', + supportsGroupAssumption: true, configure: profile => { if (profile.host === undefined) return undefined; if (profile.clientId === undefined) return undefined; @@ -112,6 +133,7 @@ export const m2mStrategy: Strategy = { clientId: profile.clientId, clientSecret: profile.clientSecret.value, ...(profile.accountId !== undefined && {accountId: profile.accountId}), + ...(profile.groupId !== undefined && {groupId: profile.groupId}), }); }, }; diff --git a/packages/auth/src/credentials/default/default-credentials.ts b/packages/auth/src/credentials/default/default-credentials.ts index e6030bc69..67ccf9a02 100644 --- a/packages/auth/src/credentials/default/default-credentials.ts +++ b/packages/auth/src/credentials/default/default-credentials.ts @@ -26,6 +26,10 @@ interface DefaultCredentialsOptions { * 2. OAuth M2M (`oauth-m2m`). * 3. Databricks CLI (`databricks-cli`). * + * When the resolved profile contains a non-empty group ID, strategies that + * cannot assume a group are skipped. Explicitly selecting such a strategy + * through `authType` returns an error. + * * When no profile is provided via `options.profile`, the profile is * resolved on first use from the default config file (~/.databrickscfg) * and environment variables. diff --git a/packages/auth/src/credentials/default/errors.ts b/packages/auth/src/credentials/default/errors.ts index b08400cf7..772da4c3a 100644 --- a/packages/auth/src/credentials/default/errors.ts +++ b/packages/auth/src/credentials/default/errors.ts @@ -1,7 +1,8 @@ /** Discriminant codes for {@link DefaultCredentialsError}. */ export type DefaultCredentialsErrorCode = | 'NO_AUTH_CONFIGURED' - | 'AUTH_TYPE_NOT_FOUND'; + | 'AUTH_TYPE_NOT_FOUND' + | 'GROUP_ROLE_UNSUPPORTED'; /** * Error thrown when the default credentials chain cannot resolve a diff --git a/packages/auth/src/credentials/default/u2m-strategy.ts b/packages/auth/src/credentials/default/u2m-strategy.ts index c19f6f5e5..ccdb73af9 100644 --- a/packages/auth/src/credentials/default/u2m-strategy.ts +++ b/packages/auth/src/credentials/default/u2m-strategy.ts @@ -10,6 +10,7 @@ import type {Strategy} from './chain'; */ export const u2mStrategy: Strategy = { name: 'databricks-cli', + supportsGroupAssumption: false, configure: profile => { if (profile.host === undefined) return undefined; if (profile.name === undefined) return undefined; diff --git a/packages/auth/src/credentials/m2m.ts b/packages/auth/src/credentials/m2m.ts index 3d7f4779b..f687c23e0 100644 --- a/packages/auth/src/credentials/m2m.ts +++ b/packages/auth/src/credentials/m2m.ts @@ -33,6 +33,12 @@ export interface M2mCredentialsOptions { */ accountId?: string; + /** + * ID of the group whose role is assumed by the issued token. When omitted + * or empty, no group role is assumed. + */ + groupId?: string; + /** * OAuth scopes to request. When omitted or empty, defaults to * `['all-apis']`. diff --git a/packages/auth/tests/credentials/default/chain.test.ts b/packages/auth/tests/credentials/default/chain.test.ts index ab328160d..eb163cbd0 100644 --- a/packages/auth/tests/credentials/default/chain.test.ts +++ b/packages/auth/tests/credentials/default/chain.test.ts @@ -6,7 +6,7 @@ import {describe, expect, it} from 'vitest'; import {Secret} from '@databricks/sdk-core/profiles/browser'; import type {Profile} from '@databricks/sdk-core/profiles/browser'; -import type {Header} from '../../../src/auth'; +import type {Credentials, Header} from '../../../src/auth'; import { DefaultCredentials, m2mStrategy, @@ -18,19 +18,38 @@ import type {DefaultCredentialsErrorCode} from '../../../src/credentials/default const HOST = 'https://workspace.example'; -function configuredStrategy(label: string): Strategy { +function configuredStrategy( + label: string, + supportsGroupAssumption = true, + onConfigure?: (profile: Profile) => void +): Strategy { return { name: label, - configure: () => ({ - name: () => label, - authHeaders: () => - Promise.resolve([{key: 'X-Test-Strategy', value: label}]), - }), + supportsGroupAssumption, + configure: (profile): Credentials => { + onConfigure?.(profile); + return { + name: () => label, + authHeaders: () => + Promise.resolve([{key: 'X-Test-Strategy', value: label}]), + }; + }, }; } -function unconfiguredStrategy(label: string): Strategy { - return {name: label, configure: () => undefined}; +function unconfiguredStrategy( + label: string, + supportsGroupAssumption = true, + onConfigure?: () => void +): Strategy { + return { + name: label, + supportsGroupAssumption, + configure: (): undefined => { + onConfigure?.(); + return undefined; + }, + }; } const loaderFor = @@ -87,6 +106,7 @@ describe('DefaultCredentials chain', () => { let buildCount = 0; const strategy: Strategy = { name: 'counting', + supportsGroupAssumption: true, configure: () => { buildCount += 1; return { @@ -101,6 +121,35 @@ describe('DefaultCredentials chain', () => { expect(buildCount).toBe(1); }); + it('skips unsupported strategies when a group is configured', async () => { + let unsupportedCalls = 0; + const creds = new DefaultCredentials( + [ + configuredStrategy('pat', false, () => { + unsupportedCalls += 1; + }), + configuredStrategy('oauth-m2m'), + ], + loaderFor({host: HOST, groupId: 'group-123'}) + ); + + expect(await creds.authHeaders()).toEqual([ + {key: 'X-Test-Strategy', value: 'oauth-m2m'}, + ]); + expect(unsupportedCalls).toBe(0); + }); + + it('preserves normal strategy ordering when the group is empty', async () => { + const creds = new DefaultCredentials( + [configuredStrategy('pat', false), configuredStrategy('oauth-m2m')], + loaderFor({host: HOST, groupId: ''}) + ); + + expect(await creds.authHeaders()).toEqual([ + {key: 'X-Test-Strategy', value: 'pat'}, + ]); + }); + it('invokes the profile loader exactly once', async () => { let loaderCalls = 0; const loader = (): Promise => { @@ -113,6 +162,29 @@ describe('DefaultCredentials chain', () => { expect(loaderCalls).toBe(1); }); + it('does not configure a fallback after the selected strategy fails', async () => { + const selectedError = new Error('selected provider failed'); + let fallbackCalls = 0; + const selected: Strategy = { + name: 'oauth-m2m', + supportsGroupAssumption: true, + configure: () => ({ + name: () => 'oauth-m2m', + authHeaders: () => Promise.reject(selectedError), + }), + }; + const fallback = configuredStrategy('fallback', true, () => { + fallbackCalls += 1; + }); + const creds = new DefaultCredentials( + [selected, fallback], + loaderFor({host: HOST, groupId: 'group-123'}) + ); + + await expect(creds.authHeaders()).rejects.toBe(selectedError); + expect(fallbackCalls).toBe(0); + }); + const errorCases: { name: string; strategies: readonly Strategy[]; @@ -141,6 +213,37 @@ describe('DefaultCredentials chain', () => { profile: {host: HOST, authType: 'pat'}, wantCode: 'NO_AUTH_CONFIGURED', }, + { + name: 'throws GROUP_ROLE_UNSUPPORTED for an explicitly selected PAT strategy', + strategies: [patStrategy, m2mStrategy], + profile: { + host: HOST, + token: new Secret('dapi-abc'), + groupId: 'group-123', + authType: 'pat', + }, + wantCode: 'GROUP_ROLE_UNSUPPORTED', + }, + { + name: 'throws GROUP_ROLE_UNSUPPORTED for an explicitly selected CLI strategy', + strategies: [configuredStrategy('databricks-cli', false)], + profile: { + host: HOST, + groupId: 'group-123', + authType: 'databricks-cli', + }, + wantCode: 'GROUP_ROLE_UNSUPPORTED', + }, + { + name: 'throws NO_AUTH_CONFIGURED when grouped strategies are exhausted', + strategies: [ + configuredStrategy('pat', false), + unconfiguredStrategy('oauth-m2m'), + configuredStrategy('databricks-cli', false), + ], + profile: {host: HOST, groupId: 'group-123'}, + wantCode: 'NO_AUTH_CONFIGURED', + }, ]; it.each(errorCases)('$name', async ({strategies, profile, wantCode}) => { diff --git a/packages/auth/tests/credentials/u2m.test.ts b/packages/auth/tests/credentials/u2m.test.ts index 31c0fde07..1ce21142f 100644 --- a/packages/auth/tests/credentials/u2m.test.ts +++ b/packages/auth/tests/credentials/u2m.test.ts @@ -3,7 +3,11 @@ import type {Stats} from 'node:fs'; import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; import type {U2mCredentialsErrorCode} from '../../src/credentials'; -import {U2mCredentialsError, newU2mCredentials} from '../../src/credentials'; +import { + U2mCredentialsError, + defaultCredentials, + newU2mCredentials, +} from '../../src/credentials'; type ExecFileCallback = ( err: Error | null, @@ -83,6 +87,23 @@ describe('newU2mCredentials', () => { vi.unstubAllEnvs(); }); + it('rejects grouped explicit CLI auth before invoking the CLI', async () => { + const credentials = defaultCredentials({ + profile: { + name: DEFAULT_PROFILE, + host: 'https://workspace.example', + authType: 'databricks-cli', + groupId: 'group-123', + }, + }); + + await expect(credentials.authHeaders()).rejects.toMatchObject({ + code: 'GROUP_ROLE_UNSUPPORTED', + }); + expect(statMock).not.toHaveBeenCalled(); + expect(execFileMock).not.toHaveBeenCalled(); + }); + const successCases: { name: string; profile: string; diff --git a/packages/core/NEXT_CHANGELOG.md b/packages/core/NEXT_CHANGELOG.md index 50df105cb..62736d234 100644 --- a/packages/core/NEXT_CHANGELOG.md +++ b/packages/core/NEXT_CHANGELOG.md @@ -4,6 +4,9 @@ ### New Features and Improvements +- Added `groupId` profile resolution from `DATABRICKS_GROUP_ID` and profile + `group_id`. + ### Bug Fixes ### Documentation diff --git a/packages/core/src/profiles/profile.ts b/packages/core/src/profiles/profile.ts index 59ef80721..6d7fd817c 100644 --- a/packages/core/src/profiles/profile.ts +++ b/packages/core/src/profiles/profile.ts @@ -28,6 +28,9 @@ export interface Profile { /** Databricks Account ID for Accounts API. */ accountId?: string; + /** ID of the group whose role is assumed when obtaining OAuth tokens. */ + groupId?: string; + /** Personal access token for PAT authentication. */ token?: Secret; @@ -142,6 +145,15 @@ export const PROPERTY_DEFS: readonly PropertyDef[] = [ }, get: (p: Profile): string | undefined => p.accountId, }, + { + field: 'groupId', + envVar: 'DATABRICKS_GROUP_ID', + iniKey: 'group_id', + set: (p: Profile, v: string): void => { + p.groupId = v; + }, + get: (p: Profile): string | undefined => p.groupId, + }, { field: 'token', envVar: 'DATABRICKS_TOKEN', diff --git a/packages/core/tests/profiles/profile.test.ts b/packages/core/tests/profiles/profile.test.ts index 658d298b4..7af02099c 100644 --- a/packages/core/tests/profiles/profile.test.ts +++ b/packages/core/tests/profiles/profile.test.ts @@ -14,6 +14,7 @@ function findDef(field: string): PropertyDef { } const STRING_DEF = findDef('host'); +const GROUP_ID_DEF = findDef('groupId'); const SECRET_DEF = findDef('token'); describe('property set and get', () => { @@ -36,6 +37,12 @@ describe('property set and get', () => { raw: 'https://x.com?a=1&b=2', wantGet: 'https://x.com?a=1&b=2', }, + { + name: 'group ID', + def: GROUP_ID_DEF, + raw: 'group-123', + wantGet: 'group-123', + }, // Secret properties. { name: 'secret: plain value', @@ -77,6 +84,13 @@ describe('property set and get', () => { }); describe('PROPERTY_DEFS', () => { + it('maps groupId to the Databricks environment and INI names', () => { + expect(GROUP_ID_DEF).toMatchObject({ + envVar: 'DATABRICKS_GROUP_ID', + iniKey: 'group_id', + }); + }); + it('should cover every Profile field except name and extra', () => { // Set every property to a sentinel value via PROPERTY_DEFS, then check // that no Profile field was missed. The source of truth is the Profile diff --git a/packages/core/tests/profiles/resolve.test.ts b/packages/core/tests/profiles/resolve.test.ts index caf1f6fce..2f86dae58 100644 --- a/packages/core/tests/profiles/resolve.test.ts +++ b/packages/core/tests/profiles/resolve.test.ts @@ -1,4 +1,4 @@ -import {mkdtempSync} from 'node:fs'; +import {mkdtempSync, writeFileSync} from 'node:fs'; import {tmpdir} from 'node:os'; import {dirname, join} from 'node:path'; import {fileURLToPath} from 'node:url'; @@ -15,6 +15,7 @@ import {PROPERTY_DEFS} from '../../src/profiles/profile'; const TEST_DIR = dirname(fileURLToPath(import.meta.url)); const TESTDATA = join(TEST_DIR, 'testdata'); const CFG = join(TESTDATA, 'databrickscfg'); +const CFG_GROUP_ID = join(TESTDATA, 'databrickscfg_group_id'); const CFG_NO_DEFAULT = join(TESTDATA, 'databrickscfg_no_default'); const CFG_SETTINGS = join(TESTDATA, 'databrickscfg_settings'); const CFG_SETTINGS_EMPTY = join(TESTDATA, 'databrickscfg_settings_empty'); @@ -233,6 +234,23 @@ describe('resolve', () => { clientSecret: new Secret('secret-xyz'), }, }, + { + name: 'group ID from profile', + options: {configFile: CFG_GROUP_ID, profile: 'workspace'}, + want: {name: 'workspace', groupId: 'profile-group'}, + }, + { + name: 'group ID from environment', + options: {noProfile: true}, + env: {DATABRICKS_GROUP_ID: 'env-group'}, + want: {groupId: 'env-group'}, + }, + { + name: 'environment group ID overrides profile group ID', + options: {configFile: CFG_GROUP_ID, profile: 'workspace'}, + env: {DATABRICKS_GROUP_ID: 'env-group'}, + want: {name: 'workspace', groupId: 'env-group'}, + }, { name: 'extra keys', options: {configFile: CFG, profile: 'extra-keys'}, @@ -330,6 +348,48 @@ describe('resolve', () => { expectProfileEqual(got, want); } }); + + const emptyGroupCases: { + name: string; + config: string; + env?: string; + wantGroupId: string | undefined; + }[] = [ + { + name: 'absent group configuration', + config: '[DEFAULT]\nhost = https://workspace.example\n', + wantGroupId: undefined, + }, + { + name: 'empty profile group', + config: '[DEFAULT]\nhost = https://workspace.example\ngroup_id =\n', + wantGroupId: '', + }, + { + name: 'empty environment group', + config: '[DEFAULT]\nhost = https://workspace.example\n', + env: '', + wantGroupId: undefined, + }, + ]; + + it.each(emptyGroupCases)( + 'treats $name as no group assumption', + async ({config, env, wantGroupId}) => { + const configFile = join( + mkdtempSync(join(tmpdir(), 'group-profile-test-')), + 'databrickscfg' + ); + writeFileSync(configFile, config); + if (env !== undefined) { + vi.stubEnv('DATABRICKS_GROUP_ID', env); + } + + const profile = await resolve({configFile}); + + expect(profile.groupId).toBe(wantGroupId); + } + ); }); describe('listProfiles', () => { diff --git a/packages/core/tests/profiles/testdata/databrickscfg_group_id b/packages/core/tests/profiles/testdata/databrickscfg_group_id new file mode 100644 index 000000000..4e285f3fe --- /dev/null +++ b/packages/core/tests/profiles/testdata/databrickscfg_group_id @@ -0,0 +1,2 @@ +[workspace] +group_id = profile-group