improvement(credentials): report Databricks token failures in the shared service-account codes

Author: waleedlatif1

apps/sim/lib/credentials/databricks-service-account.server.test.ts

@@ -2,13 +2,19 @@
  * @vitest-environment node
  */
 
+import { queueTableRows, resetDbChainMock } from '@sim/testing'
 import { jsonResponse } from '@sim/testing/helpers/http'
 import { encryptionMock, encryptionMockFns } from '@sim/testing/mocks/encryption.mock'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 vi.mock('@/lib/core/security/encryption', () => encryptionMock)
 
-import { verifyAndEncryptDatabricksCredential } from '@/lib/credentials/databricks-service-account.server'
+import { credential } from '@sim/db/schema'
+import {
+  resolveDatabricksCredentialToken,
+  verifyAndEncryptDatabricksCredential,
+} from '@/lib/credentials/databricks-service-account.server'
+import { DATABRICKS_SERVICE_ACCOUNT_PROVIDER_ID } from '@/lib/oauth/types'
 
 const owner = { workspaceId: 'ws-1', actorUserId: 'user-1' }
 
@@ -49,3 +55,49 @@ describe('saving a Databricks connection', () => {
     expect(encryptionMockFns.mockEncryptSecret).not.toHaveBeenCalled()
   })
 })
+
+describe('getting a Databricks token for a run', () => {
+  const fetch = vi.fn()
+
+  beforeEach(() => {
+    resetDbChainMock()
+    queueTableRows(credential, [
+      {
+        workspaceId: 'ws-1',
+        organizationId: null,
+        type: 'service_account',
+        providerId: DATABRICKS_SERVICE_ACCOUNT_PROVIDER_ID,
+        revokedAt: null,
+        encryptedServiceAccountKey: 'ciphertext',
+      },
+    ])
+    encryptionMockFns.mockDecryptSecret.mockResolvedValueOnce({
+      decrypted: JSON.stringify({
+        type: 'databricks_credential_v1',
+        config: {
+          authType: 'service_principal',
+          host: 'dbc-1234abcd-5678.cloud.databricks.com',
+          clientId: '00000000-0000-0000-0000-000000000000',
+          clientSecret: 'dose-secret',
+        },
+      }),
+    })
+    vi.stubGlobal('fetch', fetch)
+  })
+
+  it('reports a refused client secret as a credential to reconnect', async () => {
+    fetch.mockResolvedValue(jsonResponse({ error: 'invalid_client' }, 401))
+
+    await expect(resolveDatabricksCredentialToken('cred-1')).rejects.toMatchObject({
+      code: 'invalid_credentials',
+    })
+  })
+
+  it('reports a Databricks outage as the provider being unavailable, not a bad credential', async () => {
+    fetch.mockResolvedValue(jsonResponse({ error: 'temporarily_unavailable' }, 503))
+
+    await expect(resolveDatabricksCredentialToken('cred-1')).rejects.toMatchObject({
+      code: 'provider_unavailable',
+    })
+  })
+})

apps/sim/lib/credentials/databricks-service-account.server.ts

@@ -6,6 +6,10 @@ import { z } from 'zod'
 import { decryptSecret, encryptSecret } from '@/lib/core/security/encryption'
 import { validateDatabricksWorkspaceHost } from '@/lib/core/security/input-validation'
 import { parseServiceAccountConfig } from '@/lib/credentials/service-account-config'
+import {
+  isTransientProviderStatus,
+  TokenServiceAccountValidationError,
+} from '@/lib/credentials/token-service-accounts/errors'
 import { WORKLOAD_IDENTITY_AUDIENCES } from '@/lib/credentials/workload-identity/audiences'
 import {
   mintWorkloadIdentityToken,
@@ -112,7 +116,11 @@ const storedSecretSchema = z
   .strict()
 
 export class DatabricksCredentialError extends Error {
-  constructor(message: string) {
+  /** @param status Databricks' HTTP status, when Databricks itself refused the request. */
+  constructor(
+    message: string,
+    readonly status?: number
+  ) {
     super(message)
     this.name = 'DatabricksCredentialError'
   }
@@ -145,7 +153,10 @@ async function requestWorkspaceToken(
         : typeof payload.error === 'string'
           ? payload.error
           : `HTTP ${response.status}`
-    throw new DatabricksCredentialError(`Databricks rejected the token request: ${reason}`)
+    throw new DatabricksCredentialError(
+      `Databricks rejected the token request: ${reason}`,
+      response.status
+    )
   }
   return payload.access_token
 }
@@ -299,8 +310,20 @@ export async function resolveDatabricksCredentialToken(
   const owner: WorkloadIdentityOwner = row.workspaceId
     ? { workspaceId: row.workspaceId }
     : { organizationId: row.organizationId ?? undefined }
-  return {
-    accessToken: await resolveDatabricksAccessToken(config, owner, signal),
-    domain: config.host,
+  try {
+    return {
+      accessToken: await resolveDatabricksAccessToken(config, owner, signal),
+      domain: config.host,
+    }
+  } catch (error) {
+    if (!(error instanceof DatabricksCredentialError)) throw error
+    // Report in the shared service-account codes: an outage is not a bad credential.
+    const unavailable =
+      error.status !== undefined && (error.status >= 500 || isTransientProviderStatus(error.status))
+    throw new TokenServiceAccountValidationError(
+      unavailable ? 'provider_unavailable' : 'invalid_credentials',
+      error.status ?? 400,
+      { step: 'token', reason: error.message }
+    )
   }
 }

apps/sim/lib/credentials/token-service-accounts/validators/token-checks.ts

@@ -667,8 +667,17 @@ export const validateYouComServiceAccount = tokenCheckValidator({
   }),
 })
 
-/** Google's 403 reasons that a valid key gets when it is only out of quota for now. */
-const YOUTUBE_QUOTA_REASONS = new Set(['quotaExceeded', 'dailyLimitExceeded', 'rateLimitExceeded'])
+/**
+ * Google's 403 reasons for a valid key that is only out of quota for now. Mirrors
+ * `isGoogleQuotaReason` in the Google connectors, whose module is too heavy for this graph.
+ */
+const YOUTUBE_QUOTA_REASONS = new Set([
+  'quotaExceeded',
+  'dailyLimitExceeded',
+  'rateLimitExceeded',
+  'userRateLimitExceeded',
+  'RATE_LIMIT_EXCEEDED',
+])
 
 /**
  * Listing video categories costs one quota unit. Google answers 400 for an invalid key and 403
@@ -684,7 +693,7 @@ export const validateYouTubeServiceAccount = tokenCheckValidator({
   forbiddenMeansScoped: true,
   rejectsKey: (errorBody) => {
     const error = toRecord(toRecord(errorBody).error)
-    if (error.code !== 400 && error.code !== 403) return false
+    if (error.code !== 403) return false
     return !toArray(error.errors).some((entry) =>
       YOUTUBE_QUOTA_REASONS.has(String(toRecord(entry).reason))
     )

apps/sim/lib/oauth/token-resolution.ts

@@ -435,11 +435,6 @@ export async function resolveCredentialToken(
             }
           }
         }
-        // Its messages name the trust, secret, or host problem and carry no secrets. Matched by
-        // name because the Databricks module is loaded lazily, off this module's graph.
-        if (error instanceof Error && error.name === 'DatabricksCredentialError') {
-          return { ok: false, status: 401, error: error.message }
-        }
         return { ok: false, status: 401, error: 'Failed to get service account token' }
       }
     }