feat: add schema to auth data

Author: pregnantboy

packages/backend/src/apps/databricks/actions/create-row.ts

@@ -2,12 +2,10 @@ import { IGlobalVariable, IRawAction } from '@plumber/types'
 
 import { z } from 'zod'
 
-import { databricksConfig } from '@/config/app-env-vars/databricks'
 import StepError from '@/errors/step'
 import logger from '@/helpers/logger'
 
-import { createClient } from '../auth/create-client'
-import { getSchemaName } from '../common/get-schema-name'
+import { createSession } from '../auth/create-client'
 
 const insertRowSchema = z.object({
   tableName: z.string().min(1),
@@ -108,12 +106,8 @@ const createRowAction: IRawAction = {
           $.app.name,
         )
       }
-      const client = await createClient($)
+      const { session, endSession } = await createSession($)
       const { tableName, rowData } = parametersParseResult.data
-      const session = await client.openSession({
-        initialSchema: getSchemaName($),
-        initialCatalog: databricksConfig.catalog,
-      })
       const columnNames = rowData.map((row) => row.columnName)
       const columnValues = rowData.map((row) => row.columnValue)
 
@@ -125,8 +119,7 @@ const createRowAction: IRawAction = {
         .join(',')})`
       const operation = await session.executeStatement(statement)
       await operation.fetchAll()
-      await session.close()
-      await client.close()
+      await endSession()
       $.setActionItem({
         raw: {
           success: true,

packages/backend/src/apps/databricks/auth/auth-data.ts

@@ -6,6 +6,14 @@ import { screenNameSchema } from '@/helpers/app-auth-schema'
 
 const databricksAuthDataSchema = z.object({
   screenName: screenNameSchema,
+  schema: z
+    .string()
+    .min(1, 'Empty Schema Name')
+    .regex(/^[a-z0-9_]+$/, {
+      message:
+        'Schema name can only contain lowercase letters, numbers and underscores',
+    })
+    .max(100),
   token: z.string().min(1, 'Empty Personal Access Token').max(100),
 })
 

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

@@ -5,7 +5,7 @@ import { DBSQLClient, LogLevel } from '@databricks/sql'
 import { databricksConfig } from '@/config/app-env-vars/databricks'
 import logger from '@/helpers/logger'
 
-export const createClient = async ($: IGlobalVariable) => {
+export const createSession = async ($: IGlobalVariable) => {
   const token = $.auth.data.token
   if (typeof token !== 'string') {
     throw new Error('Databricks personal access token is required')
@@ -14,16 +14,14 @@ export const createClient = async ($: IGlobalVariable) => {
   const client: DBSQLClient = new DBSQLClient({
     logger: {
       log(level: LogLevel, message: string) {
-        if (level <= LogLevel.info) {
-          logger[level]({
-            userId: $.user?.id,
-            stepId: $.step?.id,
-            flowId: $.flow?.id,
-            testRun: $.execution?.testRun,
-            event: 'databricks-client-log',
-            message,
-          })
-        }
+        logger[level]({
+          userId: $.user?.id,
+          stepId: $.step?.id,
+          flowId: $.flow?.id,
+          testRun: $.execution?.testRun,
+          event: 'databricks-client-log',
+          message,
+        })
       },
     },
   })
@@ -37,7 +35,16 @@ export const createClient = async ($: IGlobalVariable) => {
 
   try {
     const connectedClient = await client.connect(connectOptions)
-    return connectedClient
+    const session = await connectedClient.openSession({
+      initialSchema: $.auth.data.schema as string,
+      initialCatalog: databricksConfig.catalog,
+    })
+    const endSession = async () => {
+      await session.close()
+      await connectedClient.close()
+    }
+
+    return { session, endSession }
   } catch (error) {
     throw new Error('Failed to connect to Databricks')
   }

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

@@ -14,6 +14,13 @@ const auth: IUserAddedConnectionAuth = {
       required: true,
       readOnly: false,
     },
+    {
+      key: 'schema',
+      label: 'Schema',
+      type: 'string' as const,
+      required: true,
+      readOnly: false,
+    },
     {
       key: 'token',
       label: 'Personal Access Token',

packages/backend/src/apps/databricks/auth/verify-credentials.ts

@@ -3,21 +3,16 @@ import { IGlobalVariable } from '@plumber/types'
 import { ZodError } from 'zod'
 import { fromZodError } from 'zod-validation-error'
 
-import { databricksConfig } from '@/config/app-env-vars/databricks'
 import logger from '@/helpers/logger'
 
 import { validateAuthData } from './auth-data'
-import { createClient } from './create-client'
+import { createSession } from './create-client'
 
 const verifyCredentials = async ($: IGlobalVariable) => {
   try {
     validateAuthData($)
-    const client = await createClient($)
-    const session = await client.openSession({
-      initialCatalog: databricksConfig.catalog,
-    })
-    await session.close()
-    await client.close()
+    const { endSession } = await createSession($)
+    await endSession()
   } catch (error) {
     if (error instanceof ZodError) {
       // Auth data validation failed: throws message from first error

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

@@ -2,12 +2,10 @@ import { IDynamicAction, IGlobalVariable, IJSONObject } from '@plumber/types'
 
 import { z } from 'zod'
 
-import { databricksConfig } from '@/config/app-env-vars/databricks'
 import { BadUserInputError } from '@/errors/graphql-errors'
 import logger from '@/helpers/logger'
 
-import { createClient } from '../auth/create-client'
-import { getSchemaName } from '../common/get-schema-name'
+import { createSession } from '../auth/create-client'
 
 const createTableSchema = z.object({
   // table name must be lowercase and can only contain underscores
@@ -44,17 +42,12 @@ const dynamicData: IDynamicAction = {
 
       const { tableName, columnName } = parametersParseResult.data
 
-      const client = await createClient($)
-      const session = await client.openSession({
-        initialSchema: getSchemaName($),
-        initialCatalog: databricksConfig.catalog,
-      })
+      const { session, endSession } = await createSession($)
       // TODO: properly prepare this statement
       const statement = `ALTER TABLE \`${tableName}\` ADD COLUMN \`${columnName}\` STRING;`
       const operation = await session.executeStatement(statement)
       await operation.fetchAll()
-      await session.close()
-      await client.close()
+      await endSession()
       return {
         newValue: columnName,
       }

packages/backend/src/apps/databricks/dynamic-data/create-column.ts

@@ -2,12 +2,10 @@ import { IDynamicAction, IGlobalVariable, IJSONObject } from '@plumber/types'
 
 import { z } from 'zod'
 
-import { databricksConfig } from '@/config/app-env-vars/databricks'
 import { BadUserInputError } from '@/errors/graphql-errors'
 import logger from '@/helpers/logger'
 
-import { createClient } from '../auth/create-client'
-import { getSchemaName } from '../common/get-schema-name'
+import { createSession } from '../auth/create-client'
 
 const createTableSchema = z.object({
   // table name must be lowercase and can only contain underscores
@@ -32,17 +30,12 @@ const dynamicData: IDynamicAction = {
     try {
       const { tableName } = parametersParseResult.data
 
-      const client = await createClient($)
-      const session = await client.openSession({
-        initialSchema: getSchemaName($),
-        initialCatalog: databricksConfig.catalog,
-      })
+      const { session, endSession } = await createSession($)
       // TODO: properly prepare this statement
       const statement = `CREATE TABLE \`${tableName}\`;`
       const operation = await session.executeStatement(statement)
       await operation.fetchAll()
-      await session.close()
-      await client.close()
+      await endSession()
       return {
         newValue: tableName,
       }

packages/backend/src/apps/databricks/dynamic-data/create-table.ts

@@ -4,11 +4,9 @@ import {
   IGlobalVariable,
 } from '@plumber/types'
 
-import { databricksConfig } from '@/config/app-env-vars/databricks'
 import logger from '@/helpers/logger'
 
-import { createClient } from '../auth/create-client'
-import { getSchemaName } from '../common/get-schema-name'
+import { createSession } from '../auth/create-client'
 import { DatabrickColumnRes } from '../common/types'
 
 const dynamicData: IDynamicData = {
@@ -26,19 +24,14 @@ const dynamicData: IDynamicData = {
           },
         }
       }
-      const client = await createClient($)
-      const session = await client.openSession({
-        initialSchema: getSchemaName($),
-        initialCatalog: databricksConfig.catalog,
-      })
+      const { session, endSession } = await createSession($)
       const operation = await session.getColumns({
         tableName: $.step.parameters.tableName as string,
       })
       const columns = (await operation.fetchAll({
         maxRows: 1000,
       })) as DatabrickColumnRes[]
-      await session.close()
-      await client.close()
+      await endSession()
       return {
         data: columns.map((column) => ({
           name: column.COLUMN_NAME,

packages/backend/src/apps/databricks/dynamic-data/list-table-columns.ts

@@ -7,8 +7,7 @@ import {
 import { databricksConfig } from '@/config/app-env-vars/databricks'
 import logger from '@/helpers/logger'
 
-import { createClient } from '../auth/create-client'
-import { getSchemaName } from '../common/get-schema-name'
+import { createSession } from '../auth/create-client'
 import { DatabrickTableRes } from '../common/types'
 
 const dynamicData: IDynamicData = {
@@ -17,16 +16,16 @@ const dynamicData: IDynamicData = {
 
   async run($: IGlobalVariable): Promise<DynamicDataOutput> {
     try {
-      const client = await createClient($)
-      const session = await client.openSession()
+      const { session, endSession } = await createSession($)
       const operation = await session.getTables({
         catalogName: databricksConfig.catalog,
-        schemaName: getSchemaName($),
+        schemaName: $.auth.data.schema as string,
         tableTypes: ['TABLE'],
       })
       const tables = (await operation.fetchAll({
         maxRows: 1000,
       })) as DatabrickTableRes[]
+      await endSession()
       return {
         data: tables.map((row) => ({
           name: row.TABLE_NAME,