feature: databricks sync

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";