Fix Azure AD OAuth: discovery URL, scope resource, staging resource ID

The AzureOAuthManager had three coupled bugs that prevented Azure AD
OAuth from working for tenant-specific Entra apps and for any single-
tenant Entra app, and produced tokens with the wrong audience on
staging workspaces.

1) OIDC discovery URL was hardcoded to /organizations/, ignoring
   azureTenantId. Single-tenant Entra apps can't be resolved via
   /organizations/ and return AADSTS50059. Now:

     https://login.microsoftonline.com/${azureTenantId ?? 'organizations'}/v2.0/
                                       .well-known/openid-configuration

   When azureTenantId is unset, the URL is byte-identical to the
   previous behavior (/organizations/) — no regression for the
   multi-tenant default path.

2) OAuth scope was built from azureTenantId when provided:

     const tenantId = this.options.azureTenantId ?? datatricksAzureApp;
     azureScopes.push(`${tenantId}/.default`);

   The Azure v2.0 scope must be <resource-app-id>/.default — a tenant
   GUID isn't a resource, and Azure rejects with AADSTS500011.
   The variable is renamed to `resourceId` and always resolves to the
   Databricks Azure Login App ID (not the tenant).

3) The Databricks Azure Login App has a different ID in staging
   (4a67d088-db5c-48f1-9ff2-0aace800ae68) from prod
   (2ff814a6-3304-4ab8-85cb-cd0e6f879c1d). Using the prod resource ID
   on staging hosts mints a token with the wrong audience and the
   staging workspace rejects it. A new helper getAzureResourceId()
   picks the correct resource based on whether the host ends in
   .staging.azuredatabricks.net.

Empirical verification (matrix: baseline vs. patched):

  Prod Legacy, PAT:                                  PASS / PASS (identical)
  Prod Legacy, DB-M2M via DatabricksOAuthManager:    PASS / PASS (identical)
  Prod Legacy, AzureOAuthManager no azureTenantId:   AADSTS50059 / AADSTS50059
                                                     (identical; single-tenant
                                                      test cred)
  Prod Legacy, AzureOAuthManager + azureTenantId:    AADSTS50059 / AADSTS7000215
                                                     (patched reaches Azure AD;
                                                      both fail because test
                                                      env has a Databricks-side
                                                      secret, not an Azure-
                                                      Portal secret)
  Stg Legacy,  AzureOAuthManager + azureTenantId:    AADSTS50059 / PASS
                                                     (patch fixes staging)

No prod path that worked on baseline fails on patched. The
multi-tenant-app + no-azureTenantId path is byte-identical.

Unit tests cover:
  - getOIDCConfigUrl fallback (no tenant) and tenant-specific
  - getAzureResourceId for prod / prod SPOG / staging /
    staging SPOG / case-insensitive hosts
  - getScopes: M2M+U2M scope uses resource ID not tenant GUID,
    staging host uses staging resource ID, offline_access preserved

Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>

Author: msrathore-db

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Release History
 
+## Unreleased
+
+- Fix Azure AD OAuth for tenant-specific and single-tenant Entra apps, and correct the scope resource: use the Databricks Azure Login App ID (not the tenant GUID) as the OAuth scope; route OIDC discovery to `login.microsoftonline.com/${azureTenantId}/` when `azureTenantId` is provided (fallback `/organizations/` preserved); select the staging Azure resource ID on `.staging.azuredatabricks.net` hosts.
+
 ## 1.13.0
 
 - Add token federation support with custom token providers (databricks/databricks-sql-nodejs#318, databricks/databricks-sql-nodejs#319, databricks/databricks-sql-nodejs#320 by @madhav-db)

lib/connection/auth/DatabricksOAuth/OAuthManager.ts

@@ -276,8 +276,19 @@ export class AzureOAuthManager extends OAuthManager {
 
   public static datatricksAzureApp = '2ff814a6-3304-4ab8-85cb-cd0e6f879c1d';
 
+  public static datatricksAzureAppStaging = '4a67d088-db5c-48f1-9ff2-0aace800ae68';
+
+  private getAzureResourceId(): string {
+    const host = this.options.host.toLowerCase();
+    if (host.includes('.staging.azuredatabricks.net')) {
+      return AzureOAuthManager.datatricksAzureAppStaging;
+    }
+    return AzureOAuthManager.datatricksAzureApp;
+  }
+
   protected getOIDCConfigUrl(): string {
-    return 'https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration';
+    const tenantPath = this.options.azureTenantId ?? 'organizations';
+    return `https://login.microsoftonline.com/${tenantPath}/v2.0/.well-known/openid-configuration`;
   }
 
   protected getAuthorizationUrl(): string {
@@ -293,17 +304,18 @@ export class AzureOAuthManager extends OAuthManager {
   }
 
   protected getScopes(requestedScopes: OAuthScopes): OAuthScopes {
-    // There is no corresponding scopes in Azure, instead, access control will be delegated to Databricks
-    const tenantId = this.options.azureTenantId ?? AzureOAuthManager.datatricksAzureApp;
+    // There is no corresponding scopes in Azure, instead, access control will be delegated to Databricks.
+    // Scope must be the Azure *resource* ID (the Databricks Azure Login App), NOT the tenant ID.
+    const resourceId = this.getAzureResourceId();
 
     const azureScopes = [];
 
     switch (this.options.flow) {
       case OAuthFlow.U2M:
-        azureScopes.push(`${tenantId}/user_impersonation`);
+        azureScopes.push(`${resourceId}/user_impersonation`);
         break;
       case OAuthFlow.M2M:
-        azureScopes.push(`${tenantId}/.default`);
+        azureScopes.push(`${resourceId}/.default`);
         break;
       // no default
     }

tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.ts

@@ -519,3 +519,103 @@ class OpenIDClientStub implements BaseClient {
     });
   });
 });
+
+describe('AzureOAuthManager (host/tenant awareness)', () => {
+  function makeAzure(overrides: Partial<OAuthManagerOptions> = {}) {
+    return new AzureOAuthManager({
+      host: 'adb-1234567890123456.1.azuredatabricks.net',
+      flow: OAuthFlow.M2M,
+      context: new ClientContextStub(),
+      ...overrides,
+    });
+  }
+
+  // Access protected methods for unit inspection.
+  const call = <T>(mgr: AzureOAuthManager, name: string, ...args: unknown[]): T =>
+    (mgr as unknown as Record<string, (...a: unknown[]) => T>)[name](...args);
+
+  describe('getOIDCConfigUrl', () => {
+    it('falls back to /organizations/ when azureTenantId is not set (baseline-compatible)', () => {
+      const mgr = makeAzure();
+      const url = call<string>(mgr, 'getOIDCConfigUrl');
+      expect(url).to.equal('https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration');
+    });
+
+    it('uses the caller-supplied tenant in the discovery URL when provided', () => {
+      const tenant = '9f37a392-f0ae-4280-9796-f1864a10effc';
+      const mgr = makeAzure({ azureTenantId: tenant });
+      const url = call<string>(mgr, 'getOIDCConfigUrl');
+      expect(url).to.equal(`https://login.microsoftonline.com/${tenant}/v2.0/.well-known/openid-configuration`);
+    });
+  });
+
+  describe('getAzureResourceId', () => {
+    it('returns prod resource ID for prod azuredatabricks.net hosts', () => {
+      const mgr = makeAzure({ host: 'adb-6436897454825492.12.azuredatabricks.net' });
+      expect(call<string>(mgr, 'getAzureResourceId')).to.equal(AzureOAuthManager.datatricksAzureApp);
+    });
+
+    it('returns prod resource ID for SPOG (custom URL) prod hosts', () => {
+      const mgr = makeAzure({ host: 'peco.azuredatabricks.net' });
+      expect(call<string>(mgr, 'getAzureResourceId')).to.equal(AzureOAuthManager.datatricksAzureApp);
+    });
+
+    it('returns staging resource ID for .staging.azuredatabricks.net hosts', () => {
+      const mgr = makeAzure({
+        host: 'adb-7064161269814046.2.staging.azuredatabricks.net',
+      });
+      expect(call<string>(mgr, 'getAzureResourceId')).to.equal(AzureOAuthManager.datatricksAzureAppStaging);
+    });
+
+    it('returns staging resource ID for staging SPOG hosts', () => {
+      const mgr = makeAzure({ host: 'dogfood-spog.staging.azuredatabricks.net' });
+      expect(call<string>(mgr, 'getAzureResourceId')).to.equal(AzureOAuthManager.datatricksAzureAppStaging);
+    });
+
+    it('matches case-insensitively', () => {
+      const mgr = makeAzure({ host: 'ADB-X.STAGING.AZUREDATABRICKS.NET' });
+      expect(call<string>(mgr, 'getAzureResourceId')).to.equal(AzureOAuthManager.datatricksAzureAppStaging);
+    });
+  });
+
+  describe('getScopes — resource ID is always the Azure Login App, never a tenant GUID', () => {
+    it('M2M scope uses resource ID (not tenant GUID) even when azureTenantId is set — prod host', () => {
+      const tenant = '9f37a392-f0ae-4280-9796-f1864a10effc';
+      const mgr = makeAzure({
+        host: 'adb-6436897454825492.12.azuredatabricks.net',
+        azureTenantId: tenant,
+        flow: OAuthFlow.M2M,
+      });
+      const scopes = call<string[]>(mgr, 'getScopes', []);
+      expect(scopes).to.include(`${AzureOAuthManager.datatricksAzureApp}/.default`);
+      expect(scopes).to.not.include(`${tenant}/.default`);
+    });
+
+    it('U2M scope uses resource ID (not tenant GUID) even when azureTenantId is set — prod host', () => {
+      const tenant = '9f37a392-f0ae-4280-9796-f1864a10effc';
+      const mgr = makeAzure({
+        host: 'adb-6436897454825492.12.azuredatabricks.net',
+        azureTenantId: tenant,
+        flow: OAuthFlow.U2M,
+      });
+      const scopes = call<string[]>(mgr, 'getScopes', []);
+      expect(scopes).to.include(`${AzureOAuthManager.datatricksAzureApp}/user_impersonation`);
+      expect(scopes).to.not.include(`${tenant}/user_impersonation`);
+    });
+
+    it('M2M scope uses staging resource ID on staging hosts', () => {
+      const mgr = makeAzure({
+        host: 'adb-7064161269814046.2.staging.azuredatabricks.net',
+        flow: OAuthFlow.M2M,
+      });
+      const scopes = call<string[]>(mgr, 'getScopes', []);
+      expect(scopes).to.include(`${AzureOAuthManager.datatricksAzureAppStaging}/.default`);
+    });
+
+    it('preserves offline_access when requested alongside M2M', () => {
+      const mgr = makeAzure({ flow: OAuthFlow.M2M });
+      const scopes = call<string[]>(mgr, 'getScopes', [OAuthScope.offlineAccess]);
+      expect(scopes).to.include(OAuthScope.offlineAccess);
+    });
+  });
+});