[DATABRICKS-2] PLU-600: add databricks app skeleton, auth, and client setup

packages/backend/src/apps/app-rankings.ts

@@ -8,6 +8,7 @@
 import aisayApp from './aisay'
 import calculatorApp from './calculator'
 import customApiApp from './custom-api'
+import databricksApp from './databricks'
 import delayApp from './delay'
 import formatterApp from './formatter'
 import formsgApp from './formsg'
@@ -37,6 +38,7 @@ export const ACTION_APPS_RANKING = [
   postmanApp.key,
   tilesApp.key,
   m365ExcelApp.key,
+  databricksApp.key,
   toolboxApp.key,
   formatterApp.key,
   calculatorApp.key,

packages/backend/src/apps/databricks/actions/index.ts

@@ -0,0 +1 @@
+export default []

packages/backend/src/apps/databricks/auth/create-client.ts

@@ -0,0 +1,65 @@
+import { IGlobalVariable } from '@plumber/types'
+
+import { DBSQLClient, LogLevel } from '@databricks/sql'
+import IDBSQLClient, {
+  ConnectionOptions,
+} from '@databricks/sql/dist/contracts/IDBSQLClient'
+
+import { databricksConfig } from '@/config/app-env-vars/databricks'
+import logger from '@/helpers/logger'
+
+import { constructSchemaName } from '../common/construct-schema-name'
+
+import { getDatabricksToken } from './token-persistence'
+
+export const createSession = async ($: IGlobalVariable) => {
+  const client: DBSQLClient = new DBSQLClient({
+    logger: {
+      log(level: LogLevel, message: string) {
+        logger[level]({
+          userId: $.user?.id,
+          stepId: $.step?.id,
+          flowId: $.flow?.id,
+          testRun: $.execution?.testRun,
+          event: 'databricks-client-log',
+          message,
+        })
+      },
+    },
+  })
+
+  const token = await getDatabricksToken()
+  const connectOptions = {
+    authType: 'access-token',
+    host: databricksConfig.serverHostname,
+    path: databricksConfig.httpPath,
+    token,
+  } satisfies ConnectionOptions
+
+  const schemaName = constructSchemaName($)
+
+  let connectedClient: IDBSQLClient
+  try {
+    connectedClient = await client.connect(connectOptions)
+    const session = await connectedClient.openSession({
+      initialSchema: schemaName,
+      initialCatalog: databricksConfig.catalog,
+    })
+    const endSession = async () => {
+      await session.close()
+      await connectedClient.close()
+    }
+
+    return { session, endSession }
+  } catch (error) {
+    // Clean up the connected client if it was created
+    if (connectedClient) {
+      await connectedClient.close().catch(() => {})
+    }
+    logger.error('Failed to connect to Databricks', {
+      event: 'databricks-connect-error',
+      error,
+    })
+    throw new Error('Failed to connect to Databricks')
+  }
+}

packages/backend/src/apps/databricks/auth/token-persistence.ts

@@ -0,0 +1,62 @@
+import axios from 'axios'
+
+import { databricksConfig } from '@/config/app-env-vars/databricks'
+import { createRedisClient, REDIS_DB_INDEX } from '@/config/redis'
+import logger from '@/helpers/logger'
+
+const redisClient = createRedisClient(REDIS_DB_INDEX.APP_DATA)
+
+const DATABRICKS_AUTH_TOKEN_REDIS_PREFIX = 'databricks:authToken:'
+
+export async function getDatabricksToken(): Promise<string> {
+  const redisKey =
+    DATABRICKS_AUTH_TOKEN_REDIS_PREFIX + databricksConfig.serverHostname
+  try {
+    const cachedToken = await redisClient.get(redisKey)
+    if (cachedToken) {
+      logger.info('Databricks OAuth token read', {
+        event: 'databricks-oauth-token-read',
+        redisKey,
+      })
+      return cachedToken
+    }
+
+    const response = await axios.post(
+      '/oidc/v1/token',
+      new URLSearchParams({
+        grant_type: 'client_credentials',
+        scope: 'all-apis',
+      }),
+      {
+        baseURL: `https://${databricksConfig.serverHostname}`,
+        auth: {
+          username: databricksConfig.clientId,
+          password: databricksConfig.clientSecret,
+        },
+        headers: {
+          'Content-Type': 'application/x-www-form-urlencoded',
+        },
+      },
+    )
+    const accessToken = response.data.access_token
+    logger.info('Databricks OAuth token response', {
+      event: 'databricks-oauth-token-response',
+      redisKey,
+      accessToken: accessToken.slice(0, 10) + '...',
+    })
+
+    // expires_in is in seconds, minus 1 minute of buffer
+    const expiresIn = response.data.expires_in - 60
+    // Write it into Redis
+    await redisClient.set(redisKey, accessToken, 'EX', expiresIn)
+
+    return accessToken
+  } catch (e) {
+    logger.error('Databricks OAuth token read error', {
+      event: 'databricks-oauth-token-read-error',
+      error: e,
+      redisKey,
+    })
+    throw new Error('Failed to get Databricks OAuth token')
+  }
+}

packages/backend/src/apps/databricks/common/construct-schema-name.ts

@@ -0,0 +1,10 @@
+import { IGlobalVariable } from '@plumber/types'
+
+export const constructSchemaName = ($: IGlobalVariable) => {
+  const userEmail = $.user?.email
+  if (!userEmail) {
+    throw new Error('User email is required')
+  }
+  // replace non-alphanumeric characters with underscore
+  return userEmail.replace(/[^a-zA-Z0-9]/g, '_').toLowerCase()
+}

packages/backend/src/apps/databricks/common/types.ts

@@ -0,0 +1,18 @@
+export interface DatabrickColumnRes {
+  TABLE_CAT: string
+  TABLE_SCHEM: string
+  TABLE_NAME: string
+  COLUMN_NAME: string
+  DATA_TYPE: number
+  TYPE_NAME: string
+  NULLABLE: number
+  IS_NULLABLE: 'YES' | 'NO'
+}
+
+export interface DatabrickTableRes {
+  TABLE_CAT: string
+  TABLE_SCHEM: string
+  TABLE_NAME: string
+  TABLE_TYPE: string
+  REMARKS: string
+}

packages/backend/src/apps/databricks/dynamic-data/index.ts

@@ -0,0 +1 @@
+export default []

packages/backend/src/apps/databricks/index.ts

@@ -0,0 +1,21 @@
+import { IApp } from '@plumber/types'
+
+import actions from './actions'
+import dynamicData from './dynamic-data'
+
+const app: IApp = {
+  name: 'Databricks',
+  key: 'databricks',
+  description: 'Store data for analytics and machine learning',
+  iconUrl: '{BASE_URL}/apps/databricks/assets/favicon.svg',
+  authDocUrl: '',
+  beforeRequest: [],
+  baseUrl: '',
+  apiBaseUrl: '',
+  primaryColor: '0059F7',
+  actions,
+  dynamicData,
+  category: 'data',
+}
+
+export default app

packages/backend/src/apps/index.ts

@@ -3,6 +3,7 @@ import type { IApp } from '@plumber/types'
 import aisayApp from './aisay'
 import calculatorApp from './calculator'
 import customApiApp from './custom-api'
+import databricksApp from './databricks'
 import delayApp from './delay'
 import formatterApp from './formatter'
 import formsgApp from './formsg'
@@ -42,6 +43,7 @@ const apps: Record<string, IApp> = {
   [webhookApp.key]: webhookApp,
   [aisayApp.key]: aisayApp,
   [gathersgApp.key]: gathersgApp,
+  [databricksApp.key]: databricksApp,
 }
 
 export default apps