Adding the ability to store environment variables to context, to make them more easily available/re-usable, as well as exposing them to queries.

packages/backend-core/src/context/Context.ts

@@ -1,17 +1,14 @@
 import { AsyncLocalStorage } from "async_hooks"
+import { ContextMap } from "./mainContext"
 
 export default class Context {
-  static storage = new AsyncLocalStorage<Record<string, any>>()
+  static storage = new AsyncLocalStorage<ContextMap>()
 
-  static run(context: Record<string, any>, func: any) {
+  static run(context: ContextMap, func: any) {
     return Context.storage.run(context, () => func())
   }
 
-  static get(): Record<string, any> {
-    return Context.storage.getStore() as Record<string, any>
-  }
-
-  static set(context: Record<string, any>) {
-    Context.storage.enterWith(context)
+  static get(): ContextMap {
+    return Context.storage.getStore() as ContextMap
   }
 }

packages/backend-core/src/context/mainContext.ts

@@ -16,6 +16,7 @@ export type ContextMap = {
   tenantId?: string
   appId?: string
   identity?: IdentityContext
+  environmentVariables?: Record<string, string>
 }
 
 let TEST_APP_ID: string | null = null
@@ -75,7 +76,7 @@ export function getTenantIDFromAppID(appId: string) {
   }
 }
 
-function updateContext(updates: ContextMap) {
+function updateContext(updates: ContextMap): ContextMap {
   let context: ContextMap
   try {
     context = Context.get()
@@ -120,15 +121,23 @@ export async function doInTenant(
   return newContext(updates, task)
 }
 
-export async function doInAppContext(appId: string, task: any): Promise<any> {
-  if (!appId) {
+export async function doInAppContext(
+  appId: string | null,
+  task: any
+): Promise<any> {
+  if (!appId && !env.isTest()) {
     throw new Error("appId is required")
   }
 
-  const tenantId = getTenantIDFromAppID(appId)
-  const updates: ContextMap = { appId }
-  if (tenantId) {
-    updates.tenantId = tenantId
+  let updates: ContextMap
+  if (!appId) {
+    updates = { appId: "" }
+  } else {
+    const tenantId = getTenantIDFromAppID(appId)
+    updates = { appId }
+    if (tenantId) {
+      updates.tenantId = tenantId
+    }
   }
   return newContext(updates, task)
 }
@@ -189,25 +198,25 @@ export const getProdAppId = () => {
   return conversions.getProdAppID(appId)
 }
 
-export function updateTenantId(tenantId?: string) {
-  let context: ContextMap = updateContext({
-    tenantId,
-  })
-  Context.set(context)
+export function doInEnvironmentContext(
+  values: Record<string, string>,
+  task: any
+) {
+  if (!values) {
+    throw new Error("Must supply environment variables.")
+  }
+  const updates = {
+    environmentVariables: values,
+  }
+  return newContext(updates, task)
 }
 
-export function updateAppId(appId: string) {
-  let context: ContextMap = updateContext({
-    appId,
-  })
-  try {
-    Context.set(context)
-  } catch (err) {
-    if (env.isTest()) {
-      TEST_APP_ID = appId
-    } else {
-      throw err
-    }
+export function getEnvironmentVariables() {
+  const context = Context.get()
+  if (!context.environmentVariables) {
+    return null
+  } else {
+    return context.environmentVariables
   }
 }
 

packages/server/src/api/controllers/application.ts

@@ -116,42 +116,42 @@ async function createInstance(template: any, includeSampleData: boolean) {
   const tenantId = tenancy.isMultiTenant() ? tenancy.getTenantId() : null
   const baseAppId = generateAppID(tenantId)
   const appId = generateDevAppID(baseAppId)
-  await context.updateAppId(appId)
+  return await context.doInAppContext(appId, async () => {
+    const db = context.getAppDB()
+    await db.put({
+      _id: "_design/database",
+      // view collation information, read before writing any complex views:
+      // https://docs.couchdb.org/en/master/ddocs/views/collation.html#collation-specification
+      views: {},
+    })
 
-  const db = context.getAppDB()
-  await db.put({
-    _id: "_design/database",
-    // view collation information, read before writing any complex views:
-    // https://docs.couchdb.org/en/master/ddocs/views/collation.html#collation-specification
-    views: {},
+    // NOTE: indexes need to be created before any tables/templates
+    // add view for linked rows
+    await createLinkView()
+    await createRoutingView()
+    await createAllSearchIndex()
+
+    // replicate the template data to the instance DB
+    // this is currently very hard to test, downloading and importing template files
+    if (template && template.templateString) {
+      const { ok } = await db.load(stringToReadStream(template.templateString))
+      if (!ok) {
+        throw "Error loading database dump from memory."
+      }
+    } else if (template && template.useTemplate === "true") {
+      await sdk.backups.importApp(appId, db, template)
+    } else {
+      // create the users table
+      await db.put(USERS_TABLE_SCHEMA)
+
+      if (includeSampleData) {
+        // create ootb stock db
+        await addDefaultTables(db)
+      }
+    }
+
+    return { _id: appId }
   })
-
-  // NOTE: indexes need to be created before any tables/templates
-  // add view for linked rows
-  await createLinkView()
-  await createRoutingView()
-  await createAllSearchIndex()
-
-  // replicate the template data to the instance DB
-  // this is currently very hard to test, downloading and importing template files
-  if (template && template.templateString) {
-    const { ok } = await db.load(stringToReadStream(template.templateString))
-    if (!ok) {
-      throw "Error loading database dump from memory."
-    }
-  } else if (template && template.useTemplate === "true") {
-    await sdk.backups.importApp(appId, db, template)
-  } else {
-    // create the users table
-    await db.put(USERS_TABLE_SCHEMA)
-
-    if (includeSampleData) {
-      // create ootb stock db
-      await addDefaultTables(db)
-    }
-  }
-
-  return { _id: appId }
 }
 
 async function addDefaultTables(db: Database) {

packages/server/src/api/controllers/datasource.ts

@@ -15,6 +15,7 @@ import { db as dbCore, context, events } from "@budibase/backend-core"
 import { BBContext, Datasource, Row } from "@budibase/types"
 import sdk from "../../sdk"
 import { cloneDeep } from "lodash/fp"
+import { enrich } from "../../sdk/app/datasources/datasources"
 
 export async function fetch(ctx: BBContext) {
   // Get internal tables
@@ -315,8 +316,7 @@ function updateError(error: any, newError: any, tables: string[]) {
 
 async function buildSchemaHelper(datasource: Datasource) {
   const Connector = await getIntegration(datasource.source)
-  datasource = await sdk.datasources.enrichDatasourceWithValues(datasource)
-
+  datasource = await sdk.datasources.enrich(datasource)
   // Connect to the DB and build the schema
   const connector = new Connector(datasource.config)
   await connector.buildSchema(datasource._id, datasource.entities)

packages/server/src/api/controllers/query/index.ts

@@ -8,6 +8,7 @@ import env from "../../../environment"
 import { quotas } from "@budibase/pro"
 import { events, context, utils, constants } from "@budibase/backend-core"
 import sdk from "../../../sdk"
+import { QueryEvent } from "../../../threads/definitions"
 
 const Runner = new Thread(ThreadType.QUERY, {
   timeoutMs: env.QUERY_THREAD_TIMEOUT || 10000,
@@ -127,9 +128,9 @@ function getAuthConfig(ctx: any) {
 }
 
 export async function preview(ctx: any) {
-  const datasource = await sdk.datasources.get(ctx.request.body.datasourceId, {
-    withEnvVars: true,
-  })
+  const { datasource, envVars } = await sdk.datasources.getWithEnvVars(
+    ctx.request.body.datasourceId
+  )
   const query = ctx.request.body
   // preview may not have a queryId as it hasn't been saved, but if it does
   // this stops dynamic variables from calling the same query
@@ -138,20 +139,21 @@ export async function preview(ctx: any) {
   const authConfigCtx: any = getAuthConfig(ctx)
 
   try {
-    const runFn = () =>
-      Runner.run({
-        appId: ctx.appId,
-        datasource,
-        queryVerb,
-        fields,
-        parameters,
-        transformer,
-        queryId,
-        ctx: {
-          user: ctx.user,
-          auth: { ...authConfigCtx },
-        },
-      })
+    const inputs: QueryEvent = {
+      appId: ctx.appId,
+      datasource,
+      queryVerb,
+      fields,
+      parameters,
+      transformer,
+      queryId,
+      environmentVariables: envVars,
+      ctx: {
+        user: ctx.user,
+        auth: { ...authConfigCtx },
+      },
+    }
+    const runFn = () => Runner.run(inputs)
 
     const { rows, keys, info, extra } = await quotas.addQuery(runFn, {
       datasourceId: datasource._id,
@@ -202,9 +204,9 @@ async function execute(
   const db = context.getAppDB()
 
   const query = await db.get(ctx.params.queryId)
-  const datasource = await sdk.datasources.get(query.datasourceId, {
-    withEnvVars: true,
-  })
+  const { datasource, envVars } = await sdk.datasources.getWithEnvVars(
+    query.datasourceId
+  )
 
   let authConfigCtx: any = {}
   if (!opts.isAutomation) {
@@ -222,21 +224,22 @@ async function execute(
 
   // call the relevant CRUD method on the integration class
   try {
-    const runFn = () =>
-      Runner.run({
-        appId: ctx.appId,
-        datasource,
-        queryVerb: query.queryVerb,
-        fields: query.fields,
-        pagination: ctx.request.body.pagination,
-        parameters: enrichedParameters,
-        transformer: query.transformer,
-        queryId: ctx.params.queryId,
-        ctx: {
-          user: ctx.user,
-          auth: { ...authConfigCtx },
-        },
-      })
+    const inputs: QueryEvent = {
+      appId: ctx.appId,
+      datasource,
+      queryVerb: query.queryVerb,
+      fields: query.fields,
+      pagination: ctx.request.body.pagination,
+      parameters: enrichedParameters,
+      transformer: query.transformer,
+      queryId: ctx.params.queryId,
+      environmentVariables: envVars,
+      ctx: {
+        user: ctx.user,
+        auth: { ...authConfigCtx },
+      },
+    }
+    const runFn = () => Runner.run(inputs)
 
     const { rows, pagination, extra } = await quotas.addQuery(runFn, {
       datasourceId: datasource._id,

packages/server/src/api/controllers/static/index.ts

@@ -155,7 +155,7 @@ export const getSignedUploadURL = async function (ctx: any) {
   let datasource
   try {
     const { datasourceId } = ctx.params
-    datasource = await sdk.datasources.get(datasourceId, { withEnvVars: true })
+    datasource = await sdk.datasources.get(datasourceId, { enriched: true })
     if (!datasource) {
       ctx.throw(400, "The specified datasource could not be found")
     }

packages/server/src/api/controllers/webhook.ts

@@ -39,60 +39,62 @@ export async function destroy(ctx: BBContext) {
 }
 
 export async function buildSchema(ctx: BBContext) {
-  await context.updateAppId(ctx.params.instance)
-  const db = context.getAppDB()
-  const webhook = (await db.get(ctx.params.id)) as Webhook
-  webhook.bodySchema = toJsonSchema(ctx.request.body)
-  // update the automation outputs
-  if (webhook.action.type === WebhookActionType.AUTOMATION) {
-    let automation = (await db.get(webhook.action.target)) as Automation
-    const autoOutputs = automation.definition.trigger.schema.outputs
-    let properties = webhook.bodySchema.properties
-    // reset webhook outputs
-    autoOutputs.properties = {
-      body: autoOutputs.properties.body,
-    }
-    for (let prop of Object.keys(properties)) {
-      autoOutputs.properties[prop] = {
-        type: properties[prop].type,
-        description: AUTOMATION_DESCRIPTION,
+  await context.doInAppContext(ctx.params.instance, async () => {
+    const db = context.getAppDB()
+    const webhook = (await db.get(ctx.params.id)) as Webhook
+    webhook.bodySchema = toJsonSchema(ctx.request.body)
+    // update the automation outputs
+    if (webhook.action.type === WebhookActionType.AUTOMATION) {
+      let automation = (await db.get(webhook.action.target)) as Automation
+      const autoOutputs = automation.definition.trigger.schema.outputs
+      let properties = webhook.bodySchema.properties
+      // reset webhook outputs
+      autoOutputs.properties = {
+        body: autoOutputs.properties.body,
       }
+      for (let prop of Object.keys(properties)) {
+        autoOutputs.properties[prop] = {
+          type: properties[prop].type,
+          description: AUTOMATION_DESCRIPTION,
+        }
+      }
+      await db.put(automation)
     }
-    await db.put(automation)
-  }
-  ctx.body = await db.put(webhook)
+    ctx.body = await db.put(webhook)
+  })
 }
 
 export async function trigger(ctx: BBContext) {
   const prodAppId = dbCore.getProdAppID(ctx.params.instance)
-  await context.updateAppId(prodAppId)
-  try {
-    const db = context.getAppDB()
-    const webhook = (await db.get(ctx.params.id)) as Webhook
-    // validate against the schema
-    if (webhook.bodySchema) {
-      validate(ctx.request.body, webhook.bodySchema)
-    }
-    const target = await db.get(webhook.action.target)
-    if (webhook.action.type === WebhookActionType.AUTOMATION) {
-      // trigger with both the pure request and then expand it
-      // incase the user has produced a schema to bind to
-      await triggers.externalTrigger(target, {
-        body: ctx.request.body,
-        ...ctx.request.body,
-        appId: prodAppId,
-      })
-    }
-    ctx.status = 200
-    ctx.body = {
-      message: "Webhook trigger fired successfully",
-    }
-  } catch (err: any) {
-    if (err.status === 404) {
+  await context.doInAppContext(prodAppId, async () => {
+    try {
+      const db = context.getAppDB()
+      const webhook = (await db.get(ctx.params.id)) as Webhook
+      // validate against the schema
+      if (webhook.bodySchema) {
+        validate(ctx.request.body, webhook.bodySchema)
+      }
+      const target = await db.get(webhook.action.target)
+      if (webhook.action.type === WebhookActionType.AUTOMATION) {
+        // trigger with both the pure request and then expand it
+        // incase the user has produced a schema to bind to
+        await triggers.externalTrigger(target, {
+          body: ctx.request.body,
+          ...ctx.request.body,
+          appId: prodAppId,
+        })
+      }
       ctx.status = 200
       ctx.body = {
-        message: "Application not deployed yet.",
+        message: "Webhook trigger fired successfully",
+      }
+    } catch (err: any) {
+      if (err.status === 404) {
+        ctx.status = 200
+        ctx.body = {
+          message: "Application not deployed yet.",
+        }
       }
     }
-  }
+  })
 }

packages/server/src/integrations/base/query.ts

@@ -6,7 +6,7 @@ export async function makeExternalQuery(
   datasource: Datasource,
   json: QueryJson
 ) {
-  datasource = await sdk.datasources.enrichDatasourceWithValues(datasource)
+  datasource = await sdk.datasources.enrich(datasource)
   const Integration = await getIntegration(datasource.source)
   // query is the opinionated function
   if (Integration.prototype.query) {

packages/server/src/integrations/queries/sql.ts

@@ -1,55 +1,10 @@
-import { findHBSBlocks, processStringSync } from "@budibase/string-templates"
+import { findHBSBlocks } from "@budibase/string-templates"
 import { DatasourcePlus } from "@budibase/types"
+import sdk from "../../sdk"
 
 const CONST_CHAR_REGEX = new RegExp("'[^']*'", "g")
 
-export function enrichQueryFields(
-  fields: { [key: string]: any },
-  parameters = {}
-) {
-  const enrichedQuery: { [key: string]: any } = Array.isArray(fields) ? [] : {}
-  if (!fields || !parameters) {
-    return enrichedQuery
-  }
-  // enrich the fields with dynamic parameters
-  for (let key of Object.keys(fields)) {
-    if (fields[key] == null) {
-      continue
-    }
-    if (typeof fields[key] === "object") {
-      // enrich nested fields object
-      enrichedQuery[key] = enrichQueryFields(fields[key], parameters)
-    } else if (typeof fields[key] === "string") {
-      // enrich string value as normal
-      enrichedQuery[key] = processStringSync(fields[key], parameters, {
-        noEscaping: true,
-        noHelpers: true,
-        escapeNewlines: true,
-      })
-    } else {
-      enrichedQuery[key] = fields[key]
-    }
-  }
-  if (
-    enrichedQuery.json ||
-    enrichedQuery.customData ||
-    enrichedQuery.requestBody
-  ) {
-    try {
-      enrichedQuery.json = JSON.parse(
-        enrichedQuery.json ||
-          enrichedQuery.customData ||
-          enrichedQuery.requestBody
-      )
-    } catch (err) {
-      // no json found, ignore
-    }
-    delete enrichedQuery.customData
-  }
-  return enrichedQuery
-}
-
-export function interpolateSQL(
+export async function interpolateSQL(
   fields: { [key: string]: any },
   parameters: { [key: string]: any },
   integration: DatasourcePlus
@@ -90,7 +45,7 @@ export function interpolateSQL(
     else if (listRegexMatch) {
       arrays.push(binding)
       // determine the length of the array
-      const value = enrichQueryFields([binding], parameters)[0]
+      const value = (await sdk.queries.enrichContext([binding], parameters))[0]
         .split(",")
         .map((val: string) => val.trim())
       // build a string like ($1, $2, $3)
@@ -109,7 +64,7 @@ export function interpolateSQL(
   }
   // replicate the knex structure
   fields.sql = sql
-  fields.bindings = enrichQueryFields(variables, parameters)
+  fields.bindings = await sdk.queries.enrichContext(variables, parameters)
   // check for arrays in the data
   let updated: string[] = []
   for (let i = 0; i < variables.length; i++) {

packages/server/src/sdk/app/datasources/datasources.ts

@@ -1,29 +1,39 @@
-import { environmentVariables } from "@budibase/pro"
-import { context, db as dbCore } from "@budibase/backend-core"
+import { context } from "@budibase/backend-core"
 import { processObjectSync } from "@budibase/string-templates"
-import { AppEnvironment, Datasource } from "@budibase/types"
+import { Datasource } from "@budibase/types"
 import { cloneDeep } from "lodash/fp"
+import { getEnvironmentVariables } from "../../utils"
 
-export async function enrichDatasourceWithValues(datasource: Datasource) {
-  const appId = context.getAppId()
-  const appEnv = dbCore.isDevAppID(appId)
-    ? AppEnvironment.DEVELOPMENT
-    : AppEnvironment.PRODUCTION
+async function enrichDatasourceWithValues(datasource: Datasource) {
   const cloned = cloneDeep(datasource)
-  const envVars = await environmentVariables.fetchValues(appEnv)
-  const processed = processObjectSync(cloned, { env: envVars })
-  return processed as Datasource
+  const env = await getEnvironmentVariables()
+  const processed = processObjectSync(cloned, env)
+  return {
+    datasource: processed as Datasource,
+    envVars: env.env as Record<string, string>,
+  }
+}
+
+export async function enrich(datasource: Datasource) {
+  const { datasource: response } = await enrichDatasourceWithValues(datasource)
+  return response
 }
 
 export async function get(
   datasourceId: string,
-  opts?: { withEnvVars: boolean }
+  opts?: { enriched: boolean }
 ): Promise<Datasource> {
   const appDb = context.getAppDB()
   const datasource = await appDb.get(datasourceId)
-  if (opts?.withEnvVars) {
-    return await enrichDatasourceWithValues(datasource)
+  if (opts?.enriched) {
+    return (await enrichDatasourceWithValues(datasource)).datasource
   } else {
     return datasource
   }
 }
+
+export async function getWithEnvVars(datasourceId: string) {
+  const appDb = context.getAppDB()
+  const datasource = await appDb.get(datasourceId)
+  return enrichDatasourceWithValues(datasource)
+}

packages/server/src/sdk/app/queries/index.ts

@@ -0,0 +1,5 @@
+import * as queries from "./queries"
+
+export default {
+  ...queries,
+}

packages/server/src/sdk/app/queries/queries.ts

@@ -0,0 +1,50 @@
+import { getEnvironmentVariables } from "../../utils"
+import { processStringSync } from "@budibase/string-templates"
+
+export async function enrichContext(
+  fields: Record<string, any>,
+  inputs = {}
+): Promise<Record<string, any>> {
+  const enrichedQuery: Record<string, any> = Array.isArray(fields) ? [] : {}
+  if (!fields || !inputs) {
+    return enrichedQuery
+  }
+  const env = await getEnvironmentVariables()
+  const parameters = { ...inputs, ...env }
+  // enrich the fields with dynamic parameters
+  for (let key of Object.keys(fields)) {
+    if (fields[key] == null) {
+      continue
+    }
+    if (typeof fields[key] === "object") {
+      // enrich nested fields object
+      enrichedQuery[key] = await enrichContext(fields[key], parameters)
+    } else if (typeof fields[key] === "string") {
+      // enrich string value as normal
+      enrichedQuery[key] = processStringSync(fields[key], parameters, {
+        noEscaping: true,
+        noHelpers: true,
+        escapeNewlines: true,
+      })
+    } else {
+      enrichedQuery[key] = fields[key]
+    }
+  }
+  if (
+    enrichedQuery.json ||
+    enrichedQuery.customData ||
+    enrichedQuery.requestBody
+  ) {
+    try {
+      enrichedQuery.json = JSON.parse(
+        enrichedQuery.json ||
+          enrichedQuery.customData ||
+          enrichedQuery.requestBody
+      )
+    } catch (err) {
+      // no json found, ignore
+    }
+    delete enrichedQuery.customData
+  }
+  return enrichedQuery
+}

packages/server/src/sdk/app/tables/index.ts

@@ -28,7 +28,7 @@ async function getAllExternalTables(
   datasourceId: any
 ): Promise<Record<string, Table>> {
   const db = context.getAppDB()
-  const datasource = await datasources.get(datasourceId, { withEnvVars: true })
+  const datasource = await datasources.get(datasourceId, { enriched: true })
   if (!datasource || !datasource.entities) {
     throw "Datasource is not configured fully."
   }

packages/server/src/sdk/index.ts

@@ -3,6 +3,7 @@ import { default as tables } from "./app/tables"
 import { default as automations } from "./app/automations"
 import { default as applications } from "./app/applications"
 import { default as datasources } from "./app/datasources"
+import { default as queries } from "./app/queries"
 import { default as rows } from "./app/rows"
 import { default as users } from "./users"
 
@@ -14,6 +15,7 @@ const sdk = {
   rows,
   users,
   datasources,
+  queries,
 }
 
 // default export for TS

packages/server/src/sdk/utils/index.ts

@@ -0,0 +1,16 @@
+import { environmentVariables } from "@budibase/pro"
+import { context, db as dbCore } from "@budibase/backend-core"
+import { AppEnvironment } from "@budibase/types"
+
+export async function getEnvironmentVariables() {
+  let envVars = context.getEnvironmentVariables()
+  if (!envVars) {
+    const appId = context.getAppId()
+    const appEnv = dbCore.isDevAppID(appId)
+      ? AppEnvironment.DEVELOPMENT
+      : AppEnvironment.PRODUCTION
+
+    envVars = await environmentVariables.fetchValues(appEnv)
+  }
+  return { env: envVars }
+}

packages/server/src/tests/utilities/TestConfiguration.ts

@@ -364,20 +364,23 @@ class TestConfiguration {
     // create dev app
     // clear any old app
     this.appId = null
-    // @ts-ignore
-    await context.updateAppId(null)
-    this.app = await this._req({ name: appName }, null, controllers.app.create)
-    this.appId = this.app.appId
-    // @ts-ignore
-    await context.updateAppId(this.appId)
+    await context.doInAppContext(null, async () => {
+      this.app = await this._req(
+        { name: appName },
+        null,
+        controllers.app.create
+      )
+      this.appId = this.app.appId
+    })
+    return await context.doInAppContext(this.appId, async () => {
+      // create production app
+      this.prodApp = await this.publish()
 
-    // create production app
-    this.prodApp = await this.publish()
+      this.allApps.push(this.prodApp)
+      this.allApps.push(this.app)
 
-    this.allApps.push(this.prodApp)
-    this.allApps.push(this.app)
-
-    return this.app
+      return this.app
+    })
   }
 
   async publish() {

packages/server/src/threads/definitions.ts

@@ -1,3 +1,5 @@
+import { EnvironmentVariablesDecrypted } from "@budibase/types"
+
 export type WorkerCallback = (error: any, response?: any) => void
 
 export interface QueryEvent {
@@ -9,6 +11,7 @@ export interface QueryEvent {
   pagination?: any
   transformer: any
   queryId: string
+  environmentVariables?: Record<string, string>
   ctx?: any
 }
 

packages/server/src/threads/query.ts

@@ -10,7 +10,7 @@ import sdk from "../sdk"
 import { cloneDeep } from "lodash/fp"
 
 import { isSQL } from "../integrations/utils"
-import { enrichQueryFields, interpolateSQL } from "../integrations/queries/sql"
+import { interpolateSQL } from "../integrations/queries/sql"
 
 class QueryRunner {
   datasource: any
@@ -60,10 +60,11 @@ class QueryRunner {
     }
 
     if (datasourceClone.config.authConfigs) {
-      datasourceClone.config.authConfigs =
-        datasourceClone.config.authConfigs.map((config: any) => {
-          return enrichQueryFields(config, this.ctx)
-        })
+      const updatedConfigs = []
+      for (let config of datasourceClone.config.authConfigs) {
+        updatedConfigs.push(await sdk.queries.enrichContext(config, this.ctx))
+      }
+      datasourceClone.config.authConfigs = updatedConfigs
     }
 
     const integration = new Integration(datasourceClone.config)
@@ -73,12 +74,15 @@ class QueryRunner {
 
     // Enrich the parameters with the addition context items.
     // 'user' is now a reserved variable key in mapping parameters
-    const enrichedParameters = enrichQueryFields(parameters, this.ctx)
+    const enrichedParameters = await sdk.queries.enrichContext(
+      parameters,
+      this.ctx
+    )
     const enrichedContext = { ...enrichedParameters, ...this.ctx }
 
     // Parse global headers
     if (datasourceClone.config.defaultHeaders) {
-      datasourceClone.config.defaultHeaders = enrichQueryFields(
+      datasourceClone.config.defaultHeaders = await sdk.queries.enrichContext(
         datasourceClone.config.defaultHeaders,
         enrichedContext
       )
@@ -87,9 +91,9 @@ class QueryRunner {
     let query
     // handle SQL injections by interpolating the variables
     if (isSQL(datasourceClone)) {
-      query = interpolateSQL(fieldsClone, enrichedParameters, integration)
+      query = await interpolateSQL(fieldsClone, enrichedParameters, integration)
     } else {
-      query = enrichQueryFields(fieldsClone, enrichedContext)
+      query = await sdk.queries.enrichContext(fieldsClone, enrichedContext)
     }
 
     // Add pagination values for REST queries
@@ -165,7 +169,7 @@ class QueryRunner {
     const db = context.getAppDB()
     const query = await db.get(queryId)
     const datasource = await sdk.datasources.get(query.datasourceId, {
-      withEnvVars: true,
+      enriched: true,
     })
     return new QueryRunner(
       {
@@ -280,7 +284,7 @@ class QueryRunner {
 }
 
 export function execute(input: QueryEvent, callback: WorkerCallback) {
-  context.doInAppContext(input.appId!, async () => {
+  const run = async () => {
     const Runner = new QueryRunner(input)
     try {
       const response = await Runner.execute()
@@ -288,5 +292,14 @@ export function execute(input: QueryEvent, callback: WorkerCallback) {
     } catch (err) {
       callback(err)
     }
+  }
+  context.doInAppContext(input.appId!, async () => {
+    if (input.environmentVariables) {
+      return context.doInEnvironmentContext(input.environmentVariables, () => {
+        return run()
+      })
+    } else {
+      return run()
+    }
   })
 }