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..0a69a0fbd 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 = @@ -39,17 +58,20 @@ const loaderFor = Promise.resolve(profile); describe('DefaultCredentials chain', () => { + const selectedError = new Error('selected provider failed'); const resolutionCases: { name: string; strategies: readonly Strategy[]; profile: Profile; - wantHeaders: Header[]; + want: {headers: Header[]} | {error: Error}; }[] = [ { name: 'returns the first configured strategy', strategies: [patStrategy, configuredStrategy('oauth-m2m')], profile: {host: HOST, token: new Secret('dapi-abc')}, - wantHeaders: [{key: 'Authorization', value: 'Bearer dapi-abc'}], + want: { + headers: [{key: 'Authorization', value: 'Bearer dapi-abc'}], + }, }, { name: 'falls through to the next strategy when earlier ones are unconfigured', @@ -58,7 +80,9 @@ describe('DefaultCredentials chain', () => { configuredStrategy('oauth-m2m'), ], profile: {host: HOST}, - wantHeaders: [{key: 'X-Test-Strategy', value: 'oauth-m2m'}], + want: { + headers: [{key: 'X-Test-Strategy', value: 'oauth-m2m'}], + }, }, { // PAT is configured and comes first, but authType pins oauth-m2m, so @@ -70,23 +94,64 @@ describe('DefaultCredentials chain', () => { token: new Secret('dapi-abc'), authType: 'oauth-m2m', }, - wantHeaders: [{key: 'X-Test-Strategy', value: 'oauth-m2m'}], + want: { + headers: [{key: 'X-Test-Strategy', value: 'oauth-m2m'}], + }, + }, + { + name: 'skips unsupported strategies when a group is configured', + strategies: [ + configuredStrategy('pat', false), + configuredStrategy('oauth-m2m'), + ], + profile: {host: HOST, groupId: 'group-123'}, + want: { + headers: [{key: 'X-Test-Strategy', value: 'oauth-m2m'}], + }, + }, + { + name: 'preserves normal strategy ordering when the group is empty', + strategies: [ + configuredStrategy('pat', false), + configuredStrategy('oauth-m2m'), + ], + profile: {host: HOST, groupId: ''}, + want: {headers: [{key: 'X-Test-Strategy', value: 'pat'}]}, + }, + { + name: 'does not configure a fallback after the selected strategy fails', + strategies: [ + { + name: 'oauth-m2m', + supportsGroupAssumption: true, + configure: () => ({ + name: () => 'oauth-m2m', + authHeaders: () => Promise.reject(selectedError), + }), + }, + configuredStrategy('fallback', true, () => { + expect.fail('fallback strategy should not be configured'); + }), + ], + profile: {host: HOST, groupId: 'group-123'}, + want: {error: selectedError}, }, ]; - it.each(resolutionCases)( - '$name', - async ({strategies, profile, wantHeaders}) => { - const creds = new DefaultCredentials(strategies, loaderFor(profile)); - const headers = await creds.authHeaders(); - expect(headers).toEqual(wantHeaders); + it.each(resolutionCases)('$name', async ({strategies, profile, want}) => { + const creds = new DefaultCredentials(strategies, loaderFor(profile)); + if ('error' in want) { + await expect(creds.authHeaders()).rejects.toBe(want.error); + } else { + await expect(creds.authHeaders()).resolves.toEqual(want.headers); } - ); + }); it('caches the resolved strategy across calls', async () => { let buildCount = 0; const strategy: Strategy = { name: 'counting', + supportsGroupAssumption: true, configure: () => { buildCount += 1; return { @@ -141,6 +206,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..f4c52bbb4 100644 --- a/packages/auth/tests/credentials/u2m.test.ts +++ b/packages/auth/tests/credentials/u2m.test.ts @@ -2,8 +2,16 @@ 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 type { + DefaultCredentialsErrorCode, + U2mCredentialsErrorCode, +} from '../../src/credentials'; +import { + DefaultCredentialsError, + U2mCredentialsError, + defaultCredentials, + newU2mCredentials, +} from '../../src/credentials'; type ExecFileCallback = ( err: Error | null, @@ -154,21 +162,52 @@ describe('newU2mCredentials', () => { expect(args).toEqual(['auth', 'token', '--profile', profile]); }); - const errorCases: { + type ErrorCase = { name: string; setup?: () => void; - profile: string; - wantCode: U2mCredentialsErrorCode; - wantMessage: RegExp; - }[] = [ + verify?: () => void; + } & ( + | { + kind: 'u2m'; + profile: string; + wantCode: U2mCredentialsErrorCode; + wantMessage: RegExp; + } + | { + kind: 'default'; + options: Parameters[0]; + wantCode: DefaultCredentialsErrorCode; + } + ); + + const errorCases: ErrorCase[] = [ + { + name: 'grouped explicit CLI auth before invoking the CLI', + kind: 'default', + options: { + profile: { + name: DEFAULT_PROFILE, + host: 'https://workspace.example', + authType: 'databricks-cli', + groupId: 'group-123', + }, + }, + wantCode: 'GROUP_ROLE_UNSUPPORTED', + verify: (): void => { + expect(statMock).not.toHaveBeenCalled(); + expect(execFileMock).not.toHaveBeenCalled(); + }, + }, { name: 'empty profile', + kind: 'u2m', profile: '', wantCode: 'PROFILE_REQUIRED', wantMessage: /profile is required/, }, { name: 'binary missing from PATH', + kind: 'u2m', setup: (): void => { statMock.mockRejectedValue( Object.assign(new Error('ENOENT'), {code: 'ENOENT'}) @@ -180,6 +219,7 @@ describe('newU2mCredentials', () => { }, { name: 'only legacy (undersized) binary available', + kind: 'u2m', setup: (): void => { statReturnsFile(LEGACY_CLI_SIZE); }, @@ -189,6 +229,7 @@ describe('newU2mCredentials', () => { }, { name: 'CLI invocation surfaces stderr', + kind: 'u2m', setup: (): void => { statReturnsModernFile(); stubCliRun({kind: 'err', stderr: 'not logged in'}); @@ -199,6 +240,7 @@ describe('newU2mCredentials', () => { }, { name: 'CLI output is not valid JSON', + kind: 'u2m', setup: (): void => { statReturnsModernFile(); stubCliRun({kind: 'ok', stdout: 'not json'}); @@ -209,6 +251,7 @@ describe('newU2mCredentials', () => { }, { name: 'CLI response is missing access_token', + kind: 'u2m', setup: (): void => { statReturnsModernFile(); stubCliRun({ @@ -225,6 +268,7 @@ describe('newU2mCredentials', () => { }, { name: 'expiry cannot be parsed as a date', + kind: 'u2m', setup: (): void => { statReturnsModernFile(); stubCliRun(okResponse({expiry: 'totally-not-a-date'})); @@ -235,23 +279,31 @@ describe('newU2mCredentials', () => { }, ]; - it.each(errorCases)( - 'rejects on $name', - async ({setup, profile, wantCode, wantMessage}) => { - setup?.(); + it.each(errorCases)('rejects on $name', async testCase => { + testCase.setup?.(); - let caught: unknown; - try { - const creds = newU2mCredentials({profile}); - await creds.token(); - } catch (e) { - caught = e; + let caught: unknown; + try { + if (testCase.kind === 'u2m') { + await newU2mCredentials({profile: testCase.profile}).token(); + } else { + await defaultCredentials(testCase.options).authHeaders(); } + } catch (e) { + caught = e; + } + if (testCase.kind === 'u2m') { if (!(caught instanceof U2mCredentialsError)) { expect.fail(`expected U2mCredentialsError, got ${String(caught)}`); } - expect(caught.code).toBe(wantCode); - expect(caught.message).toMatch(wantMessage); + expect(caught.code).toBe(testCase.wantCode); + expect(caught.message).toMatch(testCase.wantMessage); + } else { + if (!(caught instanceof DefaultCredentialsError)) { + expect.fail(`expected DefaultCredentialsError, got ${String(caught)}`); + } + expect(caught.code).toBe(testCase.wantCode); } - ); + testCase.verify?.(); + }); }); 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..1944c9f11 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', () => { @@ -22,6 +23,7 @@ describe('property set and get', () => { def: PropertyDef; raw: string; wantGet: string; + wantNames?: {envVar: string; iniKey: string}; }[] = [ // String properties. { @@ -36,6 +38,16 @@ 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', + wantNames: { + envVar: 'DATABRICKS_GROUP_ID', + iniKey: 'group_id', + }, + }, // Secret properties. { name: 'secret: plain value', @@ -51,11 +63,17 @@ describe('property set and get', () => { }, ]; - it.each(roundTripCases)('should round-trip: $name', ({def, raw, wantGet}) => { - const profile: Profile = {}; - def.set(profile, raw); - expect(def.get(profile)).toBe(wantGet); - }); + it.each(roundTripCases)( + 'should round-trip: $name', + ({def, raw, wantGet, wantNames}) => { + const profile: Profile = {}; + def.set(profile, raw); + expect(def.get(profile)).toBe(wantGet); + if (wantNames !== undefined) { + expect(def).toMatchObject(wantNames); + } + } + ); it('should wrap secret fields in Secret instances', () => { const profile: Profile = {}; diff --git a/packages/core/tests/profiles/resolve.test.ts b/packages/core/tests/profiles/resolve.test.ts index caf1f6fce..e43879deb 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'); @@ -65,6 +66,7 @@ describe('resolve', () => { const resolveCases: { name: string; options?: Parameters[0]; + config?: string; env?: Record; want: Profile; wantErr?: ProfileErrorCode; @@ -233,6 +235,49 @@ 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: 'absent group configuration does not set a group ID', + config: '[DEFAULT]\nhost = https://workspace.example\n', + want: { + name: 'DEFAULT', + host: 'https://workspace.example', + }, + }, + { + name: 'empty profile group preserves the empty group ID', + config: '[DEFAULT]\nhost = https://workspace.example\ngroup_id =\n', + want: { + name: 'DEFAULT', + host: 'https://workspace.example', + groupId: '', + }, + }, + { + name: 'empty environment group does not set a group ID', + config: '[DEFAULT]\nhost = https://workspace.example\n', + env: {DATABRICKS_GROUP_ID: ''}, + want: { + name: 'DEFAULT', + host: 'https://workspace.example', + }, + }, { name: 'extra keys', options: {configFile: CFG, profile: 'extra-keys'}, @@ -314,22 +359,35 @@ describe('resolve', () => { }, ]; - it.each(resolveCases)('$name', async ({options, env, want, wantErr}) => { - if (env !== undefined) { - for (const [key, value] of Object.entries(env)) { - vi.stubEnv(key, value); + it.each(resolveCases)( + '$name', + async ({options, config, env, want, wantErr}) => { + if (env !== undefined) { + for (const [key, value] of Object.entries(env)) { + vi.stubEnv(key, value); + } } - } - if (wantErr !== undefined) { - await expect(resolve(options)).rejects.toMatchObject({ - code: wantErr, - }); - } else { - const got = await resolve(options); - expectProfileEqual(got, want); + const configFile = + config === undefined + ? undefined + : join(mkdtempSync(join(tmpdir(), 'profile-test-')), 'databrickscfg'); + if (configFile !== undefined) { + writeFileSync(configFile, config); + } + const resolvedOptions = + configFile === undefined ? options : {...options, configFile}; + + if (wantErr !== undefined) { + await expect(resolve(resolvedOptions)).rejects.toMatchObject({ + code: wantErr, + }); + } else { + const got = await resolve(resolvedOptions); + expectProfileEqual(got, want); + } } - }); + ); }); 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