feature: databricks sync

backend/src/services/app-connection/app-connection-enums.ts

@@ -1,6 +1,7 @@
 export enum AppConnection {
   GitHub = "github",
   AWS = "aws",
+  Databricks = "databricks",
   GCP = "gcp",
   AzureKeyVault = "azure-key-vault",
   AzureAppConfiguration = "azure-app-configuration"

backend/src/services/app-connection/app-connection-fns.ts

@@ -5,12 +5,17 @@ import { TAppConnectionServiceFactoryDep } from "@app/services/app-connection/ap
 import { TAppConnection, TAppConnectionConfig } from "@app/services/app-connection/app-connection-types";
 import {
   AwsConnectionMethod,
-  getAwsAppConnectionListItem,
+  getAwsConnectionListItem,
   validateAwsConnectionCredentials
 } from "@app/services/app-connection/aws";
+import {
+  DatabricksConnectionMethod,
+  getDatabricksConnectionListItem,
+  validateDatabricksConnectionCredentials
+} from "@app/services/app-connection/databricks";
 import {
   GcpConnectionMethod,
-  getGcpAppConnectionListItem,
+  getGcpConnectionListItem,
   validateGcpConnectionCredentials
 } from "@app/services/app-connection/gcp";
 import {
@@ -33,11 +38,12 @@ import {
 
 export const listAppConnectionOptions = () => {
   return [
-    getAwsAppConnectionListItem(),
+    getAwsConnectionListItem(),
     getGitHubConnectionListItem(),
-    getGcpAppConnectionListItem(),
+    getGcpConnectionListItem(),
     getAzureKeyVaultConnectionListItem(),
-    getAzureAppConfigurationConnectionListItem()
+    getAzureAppConfigurationConnectionListItem(),
+    getDatabricksConnectionListItem()
   ].sort((a, b) => a.name.localeCompare(b.name));
 };
 
@@ -90,6 +96,8 @@ export const validateAppConnectionCredentials = async (
   switch (app) {
     case AppConnection.AWS:
       return validateAwsConnectionCredentials(appConnection);
+    case AppConnection.Databricks:
+      return validateDatabricksConnectionCredentials(appConnection);
     case AppConnection.GitHub:
       return validateGitHubConnectionCredentials(appConnection);
     case AppConnection.GCP:
@@ -118,6 +126,8 @@ export const getAppConnectionMethodName = (method: TAppConnection["method"]) =>
       return "Assume Role";
     case GcpConnectionMethod.ServiceAccountImpersonation:
       return "Service Account Impersonation";
+    case DatabricksConnectionMethod.ServicePrincipal:
+      return "Service Principal";
     default:
       // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
       throw new Error(`Unhandled App Connection Method: ${method}`);

backend/src/services/app-connection/app-connection-maps.ts

@@ -5,5 +5,6 @@ export const APP_CONNECTION_NAME_MAP: Record<AppConnection, string> = {
   [AppConnection.GitHub]: "GitHub",
   [AppConnection.GCP]: "GCP",
   [AppConnection.AzureKeyVault]: "Azure Key Vault",
-  [AppConnection.AzureAppConfiguration]: "Azure App Configuration"
+  [AppConnection.AzureAppConfiguration]: "Azure App Configuration",
+  [AppConnection.Databricks]: "Databricks"
 };

backend/src/services/app-connection/app-connection-service.ts

@@ -23,6 +23,8 @@ import {
   TValidateAppConnectionCredentials
 } from "@app/services/app-connection/app-connection-types";
 import { ValidateAwsConnectionCredentialsSchema } from "@app/services/app-connection/aws";
+import { ValidateDatabricksConnectionCredentialsSchema } from "@app/services/app-connection/databricks";
+import { databricksConnectionService } from "@app/services/app-connection/databricks/databricks-connection-service";
 import { ValidateGitHubConnectionCredentialsSchema } from "@app/services/app-connection/github";
 import { githubConnectionService } from "@app/services/app-connection/github/github-connection-service";
 import { TKmsServiceFactory } from "@app/services/kms/kms-service";
@@ -46,7 +48,8 @@ const VALIDATE_APP_CONNECTION_CREDENTIALS_MAP: Record<AppConnection, TValidateAp
   [AppConnection.GitHub]: ValidateGitHubConnectionCredentialsSchema,
   [AppConnection.GCP]: ValidateGcpConnectionCredentialsSchema,
   [AppConnection.AzureKeyVault]: ValidateAzureKeyVaultConnectionCredentialsSchema,
-  [AppConnection.AzureAppConfiguration]: ValidateAzureAppConfigurationConnectionCredentialsSchema
+  [AppConnection.AzureAppConfiguration]: ValidateAzureAppConfigurationConnectionCredentialsSchema,
+  [AppConnection.Databricks]: ValidateDatabricksConnectionCredentialsSchema
 };
 
 export const appConnectionServiceFactory = ({
@@ -365,6 +368,7 @@ export const appConnectionServiceFactory = ({
     connectAppConnectionById,
     listAvailableAppConnectionsForUser,
     github: githubConnectionService(connectAppConnectionById),
-    gcp: gcpConnectionService(connectAppConnectionById)
+    gcp: gcpConnectionService(connectAppConnectionById),
+    databricks: databricksConnectionService(connectAppConnectionById, appConnectionDAL, kmsService)
   };
 };

backend/src/services/app-connection/app-connection-types.ts

@@ -4,6 +4,12 @@ import {
   TAwsConnectionInput,
   TValidateAwsConnectionCredentials
 } from "@app/services/app-connection/aws";
+import {
+  TDatabricksConnection,
+  TDatabricksConnectionConfig,
+  TDatabricksConnectionInput,
+  TValidateDatabricksConnectionCredentials
+} from "@app/services/app-connection/databricks";
 import {
   TGitHubConnection,
   TGitHubConnectionConfig,
@@ -31,6 +37,7 @@ export type TAppConnection = { id: string } & (
   | TGcpConnection
   | TAzureKeyVaultConnection
   | TAzureAppConfigurationConnection
+  | TDatabricksConnection
 );
 
 export type TAppConnectionInput = { id: string } & (
@@ -39,6 +46,7 @@ export type TAppConnectionInput = { id: string } & (
   | TGcpConnectionInput
   | TAzureKeyVaultConnectionInput
   | TAzureAppConfigurationConnectionInput
+  | TDatabricksConnectionInput
 );
 
 export type TCreateAppConnectionDTO = Pick<
@@ -55,11 +63,13 @@ export type TAppConnectionConfig =
   | TGitHubConnectionConfig
   | TGcpConnectionConfig
   | TAzureKeyVaultConnectionConfig
-  | TAzureAppConfigurationConnectionConfig;
+  | TAzureAppConfigurationConnectionConfig
+  | TDatabricksConnectionConfig;
 
 export type TValidateAppConnectionCredentials =
   | TValidateAwsConnectionCredentials
   | TValidateGitHubConnectionCredentials
   | TValidateGcpConnectionCredentials
   | TValidateAzureKeyVaultConnectionCredentials
-  | TValidateAzureAppConfigurationConnectionCredentials;
+  | TValidateAzureAppConfigurationConnectionCredentials
+  | TValidateDatabricksConnectionCredentials;

backend/src/services/app-connection/aws/aws-connection-fns.ts

@@ -9,7 +9,7 @@ import { AppConnection, AWSRegion } from "@app/services/app-connection/app-conne
 import { AwsConnectionMethod } from "./aws-connection-enums";
 import { TAwsConnectionConfig } from "./aws-connection-types";
 
-export const getAwsAppConnectionListItem = () => {
+export const getAwsConnectionListItem = () => {
   const { INF_APP_CONNECTION_AWS_ACCESS_KEY_ID } = getConfig();
 
   return {

backend/src/services/app-connection/databricks/databricks-connection-enums.ts

@@ -0,0 +1,3 @@
+export enum DatabricksConnectionMethod {
+  ServicePrincipal = "service-principal"
+}

backend/src/services/app-connection/databricks/databricks-connection-fns.ts

@@ -0,0 +1,92 @@
+import { request } from "@app/lib/config/request";
+import { BadRequestError } from "@app/lib/errors";
+import { removeTrailingSlash } from "@app/lib/fn";
+import { TAppConnectionDALFactory } from "@app/services/app-connection/app-connection-dal";
+import { AppConnection } from "@app/services/app-connection/app-connection-enums";
+import { encryptAppConnectionCredentials } from "@app/services/app-connection/app-connection-fns";
+import { TKmsServiceFactory } from "@app/services/kms/kms-service";
+
+import { DatabricksConnectionMethod } from "./databricks-connection-enums";
+import {
+  TAuthorizeDatabricksConnection,
+  TDatabricksConnection,
+  TDatabricksConnectionConfig
+} from "./databricks-connection-types";
+
+export const getDatabricksConnectionListItem = () => {
+  return {
+    name: "Databricks" as const,
+    app: AppConnection.Databricks as const,
+    methods: Object.values(DatabricksConnectionMethod) as [DatabricksConnectionMethod.ServicePrincipal]
+  };
+};
+
+const authorizeDatabricksConnection = async ({
+  clientId,
+  clientSecret,
+  workspaceUrl
+}: Pick<TDatabricksConnection["credentials"], "workspaceUrl" | "clientId" | "clientSecret">) => {
+  const { data } = await request.post<TAuthorizeDatabricksConnection>(
+    `${removeTrailingSlash(workspaceUrl)}/oidc/v1/token`,
+    "grant_type=client_credentials&scope=all-apis",
+    {
+      auth: {
+        username: clientId,
+        password: clientSecret
+      },
+      headers: {
+        "Content-Type": "application/x-www-form-urlencoded"
+      }
+    }
+  );
+
+  return { accessToken: data.access_token, expiresAt: data.expires_in * 1000 + Date.now() };
+};
+
+export const getDatabricksConnectionAccessToken = async (
+  { id, orgId, credentials }: TDatabricksConnection,
+  appConnectionDAL: Pick<TAppConnectionDALFactory, "updateById">,
+  kmsService: Pick<TKmsServiceFactory, "createCipherPairWithDataKey">
+) => {
+  const { clientSecret, clientId, workspaceUrl, accessToken, expiresAt } = credentials;
+
+  // get new token if less than 10 minutes from expiry
+  if (Date.now() < expiresAt - 10_000) {
+    return accessToken;
+  }
+
+  const authData = await authorizeDatabricksConnection({ clientId, clientSecret, workspaceUrl });
+
+  const updatedCredentials: TDatabricksConnection["credentials"] = {
+    ...credentials,
+    ...authData
+  };
+
+  const encryptedCredentials = await encryptAppConnectionCredentials({
+    credentials: updatedCredentials,
+    orgId,
+    kmsService
+  });
+
+  await appConnectionDAL.updateById(id, { encryptedCredentials });
+
+  return authData.accessToken;
+};
+
+export const validateDatabricksConnectionCredentials = async (appConnection: TDatabricksConnectionConfig) => {
+  const { credentials } = appConnection;
+
+  try {
+    const { accessToken, expiresAt } = await authorizeDatabricksConnection(appConnection.credentials);
+
+    return {
+      ...credentials,
+      accessToken,
+      expiresAt
+    };
+  } catch (e: unknown) {
+    throw new BadRequestError({
+      message: `Unable to validate connection - verify credentials`
+    });
+  }
+};

backend/src/services/app-connection/databricks/databricks-connection-schemas.ts

@@ -0,0 +1,77 @@
+import { z } from "zod";
+
+import { AppConnections } from "@app/lib/api-docs";
+import { AppConnection } from "@app/services/app-connection/app-connection-enums";
+import {
+  BaseAppConnectionSchema,
+  GenericCreateAppConnectionFieldsSchema,
+  GenericUpdateAppConnectionFieldsSchema
+} from "@app/services/app-connection/app-connection-schemas";
+
+import { DatabricksConnectionMethod } from "./databricks-connection-enums";
+
+export const DatabricksConnectionServicePrincipalInputCredentialsSchema = z.object({
+  clientId: z.string().trim().min(1, "Client ID required"),
+  clientSecret: z.string().trim().min(1, "Client Secret required"),
+  workspaceUrl: z.string().trim().url().min(1, "Workspace URL required")
+});
+
+export const DatabricksConnectionServicePrincipalOutputCredentialsSchema = z
+  .object({
+    accessToken: z.string(),
+    expiresAt: z.number()
+  })
+  .merge(DatabricksConnectionServicePrincipalInputCredentialsSchema);
+
+const BaseDatabricksConnectionSchema = BaseAppConnectionSchema.extend({ app: z.literal(AppConnection.Databricks) });
+
+export const DatabricksConnectionSchema = z.intersection(
+  BaseDatabricksConnectionSchema,
+  z.discriminatedUnion("method", [
+    z.object({
+      method: z.literal(DatabricksConnectionMethod.ServicePrincipal),
+      credentials: DatabricksConnectionServicePrincipalOutputCredentialsSchema
+    })
+  ])
+);
+
+export const SanitizedDatabricksConnectionSchema = z.discriminatedUnion("method", [
+  BaseDatabricksConnectionSchema.extend({
+    method: z.literal(DatabricksConnectionMethod.ServicePrincipal),
+    credentials: DatabricksConnectionServicePrincipalOutputCredentialsSchema.pick({
+      clientId: true,
+      workspaceUrl: true
+    })
+  })
+]);
+
+export const ValidateDatabricksConnectionCredentialsSchema = z.discriminatedUnion("method", [
+  z.object({
+    method: z
+      .literal(DatabricksConnectionMethod.ServicePrincipal)
+      .describe(AppConnections?.CREATE(AppConnection.Databricks).method),
+    credentials: DatabricksConnectionServicePrincipalInputCredentialsSchema.describe(
+      AppConnections.CREATE(AppConnection.Databricks).credentials
+    )
+  })
+]);
+
+export const CreateDatabricksConnectionSchema = ValidateDatabricksConnectionCredentialsSchema.and(
+  GenericCreateAppConnectionFieldsSchema(AppConnection.Databricks)
+);
+
+export const UpdateDatabricksConnectionSchema = z
+  .object({
+    credentials: DatabricksConnectionServicePrincipalInputCredentialsSchema.optional().describe(
+      AppConnections.UPDATE(AppConnection.Databricks).credentials
+    )
+  })
+  .and(GenericUpdateAppConnectionFieldsSchema(AppConnection.Databricks));
+
+export const DatabricksConnectionListItemSchema = z.object({
+  name: z.literal("Databricks"),
+  app: z.literal(AppConnection.Databricks),
+  // the below is preferable but currently breaks with our zod to json schema parser
+  // methods: z.tuple([z.literal(AwsConnectionMethod.ServicePrincipal), z.literal(AwsConnectionMethod.AccessKey)]),
+  methods: z.nativeEnum(DatabricksConnectionMethod).array()
+});

backend/src/services/app-connection/databricks/databricks-connection-service.ts

@@ -0,0 +1,60 @@
+import { request } from "@app/lib/config/request";
+import { removeTrailingSlash } from "@app/lib/fn";
+import { OrgServiceActor } from "@app/lib/types";
+import { TAppConnectionDALFactory } from "@app/services/app-connection/app-connection-dal";
+import { AppConnection } from "@app/services/app-connection/app-connection-enums";
+import { getDatabricksConnectionAccessToken } from "@app/services/app-connection/databricks/databricks-connection-fns";
+import {
+  TDatabricksConnection,
+  TDatabricksListSecretScopesResponse
+} from "@app/services/app-connection/databricks/databricks-connection-types";
+import { TKmsServiceFactory } from "@app/services/kms/kms-service";
+
+type TGetAppConnectionFunc = (
+  app: AppConnection,
+  connectionId: string,
+  actor: OrgServiceActor
+) => Promise<TDatabricksConnection>;
+
+const listDatabricksSecretScopes = async (
+  appConnection: TDatabricksConnection,
+  appConnectionDAL: Pick<TAppConnectionDALFactory, "updateById">,
+  kmsService: Pick<TKmsServiceFactory, "createCipherPairWithDataKey">
+) => {
+  const {
+    credentials: { workspaceUrl }
+  } = appConnection;
+
+  const accessToken = await getDatabricksConnectionAccessToken(appConnection, appConnectionDAL, kmsService);
+
+  const { data } = await request.get<TDatabricksListSecretScopesResponse>(
+    `${removeTrailingSlash(workspaceUrl)}/api/2.0/secrets/scopes/list`,
+    {
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Accept-Encoding": "application/json"
+      }
+    }
+  );
+
+  // not present in response if no scopes exists
+  return data.scopes ?? [];
+};
+
+export const databricksConnectionService = (
+  getAppConnection: TGetAppConnectionFunc,
+  appConnectionDAL: Pick<TAppConnectionDALFactory, "updateById">,
+  kmsService: Pick<TKmsServiceFactory, "createCipherPairWithDataKey">
+) => {
+  const listSecretScopes = async (connectionId: string, actor: OrgServiceActor) => {
+    const appConnection = await getAppConnection(AppConnection.Databricks, connectionId, actor);
+
+    const secretScopes = await listDatabricksSecretScopes(appConnection, appConnectionDAL, kmsService);
+
+    return secretScopes;
+  };
+
+  return {
+    listSecretScopes
+  };
+};

backend/src/services/app-connection/databricks/databricks-connection-types.ts

@@ -0,0 +1,36 @@
+import { z } from "zod";
+
+import { DiscriminativePick } from "@app/lib/types";
+import { AppConnection } from "@app/services/app-connection/app-connection-enums";
+
+import {
+  CreateDatabricksConnectionSchema,
+  DatabricksConnectionSchema,
+  ValidateDatabricksConnectionCredentialsSchema
+} from "./databricks-connection-schemas";
+
+export type TDatabricksConnection = z.infer<typeof DatabricksConnectionSchema>;
+
+export type TDatabricksConnectionInput = z.infer<typeof CreateDatabricksConnectionSchema> & {
+  app: AppConnection.Databricks;
+};
+
+export type TValidateDatabricksConnectionCredentials = typeof ValidateDatabricksConnectionCredentialsSchema;
+
+export type TDatabricksConnectionConfig = DiscriminativePick<
+  TDatabricksConnection,
+  "method" | "app" | "credentials"
+> & {
+  orgId: string;
+};
+
+export type TAuthorizeDatabricksConnection = {
+  access_token: string;
+  scope: string;
+  token_type: string;
+  expires_in: number;
+};
+
+export type TDatabricksListSecretScopesResponse = {
+  scopes?: { name: string; backend_type: string; keyvault_metadata: { resource_id: string; dns_name: string } }[];
+};

backend/src/services/app-connection/databricks/index.ts

@@ -0,0 +1,4 @@
+export * from "./databricks-connection-enums";
+export * from "./databricks-connection-fns";
+export * from "./databricks-connection-schemas";
+export * from "./databricks-connection-types";

backend/src/services/app-connection/gcp/gcp-connection-fns.ts

@@ -17,7 +17,7 @@ import {
   TGcpConnectionConfig
 } from "./gcp-connection-types";
 
-export const getGcpAppConnectionListItem = () => {
+export const getGcpConnectionListItem = () => {
   return {
     name: "GCP" as const,
     app: AppConnection.GCP as const,