Merge pull request #15857 from Budibase/budibase-ai-3

Create new global DB for self-host quota.

packages/backend-core/src/cache/writethrough.ts

@@ -96,6 +96,24 @@ async function get<T extends Document>(db: Database, id: string): Promise<T> {
   return cacheItem.doc
 }
 
+async function tryGet<T extends Document>(
+  db: Database,
+  id: string
+): Promise<T | null> {
+  const cache = await getCache()
+  const cacheKey = makeCacheKey(db, id)
+  let cacheItem: CacheItem<T> | null = await cache.get(cacheKey)
+  if (!cacheItem) {
+    const doc = await db.tryGet<T>(id)
+    if (!doc) {
+      return null
+    }
+    cacheItem = makeCacheItem(doc)
+    await cache.store(cacheKey, cacheItem)
+  }
+  return cacheItem.doc
+}
+
 async function remove(db: Database, docOrId: any, rev?: any): Promise<void> {
   const cache = await getCache()
   if (!docOrId) {
@@ -123,10 +141,17 @@ export class Writethrough {
     return put(this.db, doc, writeRateMs)
   }
 
+  /**
+   * @deprecated use `tryGet` instead
+   */
   async get<T extends Document>(id: string) {
     return get<T>(this.db, id)
   }
 
+  async tryGet<T extends Document>(id: string) {
+    return tryGet<T>(this.db, id)
+  }
+
   async remove(docOrId: any, rev?: any) {
     return remove(this.db, docOrId, rev)
   }

packages/backend-core/src/constants/db.ts

@@ -60,6 +60,11 @@ export const StaticDatabases = {
   SCIM_LOGS: {
     name: "scim-logs",
   },
+  // Used by self-host users making use of Budicloud resources. Introduced when
+  // we started letting self-host users use Budibase AI in the cloud.
+  SELF_HOST_CLOUD: {
+    name: "self-host-cloud",
+  },
 }
 
 export const APP_PREFIX = prefixed(DocumentType.APP)

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

@@ -157,6 +157,33 @@ export async function doInTenant<T>(
   return newContext(updates, task)
 }
 
+// We allow self-host licensed users to make use of some Budicloud services
+// (e.g. Budibase AI). When they do this, they use their license key as an API
+// key. We use that license key to identify the tenant ID, and we set the
+// context to be self-host using cloud. This affects things like where their
+// quota documents get stored (because we want to avoid creating a new global
+// DB for each self-host tenant).
+export async function doInSelfHostTenantUsingCloud<T>(
+  tenantId: string,
+  task: () => T
+): Promise<T> {
+  const updates = { tenantId, isSelfHostUsingCloud: true }
+  return newContext(updates, task)
+}
+
+export function isSelfHostUsingCloud() {
+  const context = Context.get()
+  return !!context?.isSelfHostUsingCloud
+}
+
+export function getSelfHostCloudDB() {
+  const context = Context.get()
+  if (!context || !context.isSelfHostUsingCloud) {
+    throw new Error("Self-host cloud DB not found")
+  }
+  return getDB(StaticDatabases.SELF_HOST_CLOUD.name)
+}
+
 export async function doInAppContext<T>(
   appId: string,
   task: () => T
@@ -325,6 +352,11 @@ export function getGlobalDB(): Database {
   if (!context || (env.MULTI_TENANCY && !context.tenantId)) {
     throw new Error("Global DB not found")
   }
+  if (context.isSelfHostUsingCloud) {
+    throw new Error(
+      "Global DB not found - self-host users using cloud don't have a global DB"
+    )
+  }
   return getDB(baseGlobalDBName(context?.tenantId))
 }
 
@@ -344,6 +376,11 @@ export function getAppDB(opts?: any): Database {
   if (!appId) {
     throw new Error("Unable to retrieve app DB - no app ID.")
   }
+  if (isSelfHostUsingCloud()) {
+    throw new Error(
+      "App DB not found - self-host users using cloud don't have app DBs"
+    )
+  }
   return getDB(appId, opts)
 }
 

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

@@ -5,6 +5,7 @@ import { GoogleSpreadsheet } from "google-spreadsheet"
 // keep this out of Budibase types, don't want to expose context info
 export type ContextMap = {
   tenantId?: string
+  isSelfHostUsingCloud?: boolean
   appId?: string
   identity?: IdentityContext
   environmentVariables?: Record<string, string>

packages/pro

@@ -1 +1 @@
-Subproject commit 4417bceb24eabdd9a8c1615fb83c4e6fe8c0c914
+Subproject commit 8eb981cf01151261697a8f26c08c4c28f66b8e15

packages/server/src/api/routes/tests/ai.spec.ts

@@ -29,7 +29,6 @@ interface TestSetup {
   name: string
   setup: SetupFn
   mockLLMResponse: MockLLMResponseFn
-  selfHostOnly?: boolean
 }
 
 function budibaseAI(): SetupFn {
@@ -80,7 +79,7 @@ function customAIConfig(providerConfig: Partial<ProviderConfig>): SetupFn {
   }
 }
 
-const providers: TestSetup[] = [
+const allProviders: TestSetup[] = [
   {
     name: "OpenAI API key",
     setup: async () => {
@@ -89,7 +88,6 @@ const providers: TestSetup[] = [
       })
     },
     mockLLMResponse: mockChatGPTResponse,
-    selfHostOnly: true,
   },
   {
     name: "OpenAI API key with custom config",
@@ -126,9 +124,9 @@ describe("AI", () => {
     nock.cleanAll()
   })
 
-  describe.each(providers)(
+  describe.each(allProviders)(
     "provider: $name",
-    ({ setup, mockLLMResponse, selfHostOnly }: TestSetup) => {
+    ({ setup, mockLLMResponse }: TestSetup) => {
       let cleanup: () => Promise<void> | void
       beforeAll(async () => {
         cleanup = await setup(config)
@@ -243,86 +241,104 @@ describe("AI", () => {
           )
         })
       })
-
-      !selfHostOnly &&
-        describe("POST /api/ai/chat", () => {
-          let envCleanup: () => void
-          let featureCleanup: () => void
-          beforeAll(() => {
-            envCleanup = setEnv({ SELF_HOSTED: false })
-            featureCleanup = features.testutils.setFeatureFlags("*", {
-              AI_JS_GENERATION: true,
-            })
-          })
-
-          afterAll(() => {
-            featureCleanup()
-            envCleanup()
-          })
-
-          beforeEach(() => {
-            const license: License = {
-              plan: {
-                type: PlanType.FREE,
-                model: PlanModel.PER_USER,
-                usesInvoicing: false,
-              },
-              features: [],
-              quotas: {} as any,
-              tenantId: config.tenantId,
-            }
-            nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(200, license)
-          })
-
-          it("handles correct chat response", async () => {
-            mockLLMResponse("Hi there!")
-            const { message } = await config.api.ai.chat({
-              messages: [{ role: "user", content: "Hello!" }],
-              licenseKey: "test-key",
-            })
-            expect(message).toBe("Hi there!")
-          })
-
-          it("handles chat response error", async () => {
-            mockLLMResponse(() => {
-              throw new Error("LLM error")
-            })
-            await config.api.ai.chat(
-              {
-                messages: [{ role: "user", content: "Hello!" }],
-                licenseKey: "test-key",
-              },
-              { status: 500 }
-            )
-          })
-
-          it("handles no license", async () => {
-            nock.cleanAll()
-            nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(404)
-            await config.api.ai.chat(
-              {
-                messages: [{ role: "user", content: "Hello!" }],
-                licenseKey: "test-key",
-              },
-              {
-                status: 403,
-              }
-            )
-          })
-
-          it("handles no license key", async () => {
-            await config.api.ai.chat(
-              {
-                messages: [{ role: "user", content: "Hello!" }],
-                // @ts-expect-error - intentionally wrong
-                licenseKey: undefined,
-              },
-              {
-                status: 403,
-              }
-            )
-          })
-        })
     }
   )
 })
+
+describe("BudibaseAI", () => {
+  const config = new TestConfiguration()
+  let cleanup: () => void | Promise<void>
+  beforeAll(async () => {
+    await config.init()
+    cleanup = await budibaseAI()(config)
+  })
+
+  afterAll(async () => {
+    if ("then" in cleanup) {
+      await cleanup()
+    } else {
+      cleanup()
+    }
+    config.end()
+  })
+
+  describe("POST /api/ai/chat", () => {
+    let envCleanup: () => void
+    let featureCleanup: () => void
+    beforeAll(() => {
+      envCleanup = setEnv({ SELF_HOSTED: false })
+      featureCleanup = features.testutils.setFeatureFlags("*", {
+        AI_JS_GENERATION: true,
+      })
+    })
+
+    afterAll(() => {
+      featureCleanup()
+      envCleanup()
+    })
+
+    beforeEach(() => {
+      nock.cleanAll()
+      const license: License = {
+        plan: {
+          type: PlanType.FREE,
+          model: PlanModel.PER_USER,
+          usesInvoicing: false,
+        },
+        features: [],
+        quotas: {} as any,
+        tenantId: config.tenantId,
+      }
+      nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(200, license)
+    })
+
+    it("handles correct chat response", async () => {
+      mockChatGPTResponse("Hi there!")
+      const { message } = await config.api.ai.chat({
+        messages: [{ role: "user", content: "Hello!" }],
+        licenseKey: "test-key",
+      })
+      expect(message).toBe("Hi there!")
+    })
+
+    it("handles chat response error", async () => {
+      mockChatGPTResponse(() => {
+        throw new Error("LLM error")
+      })
+      await config.api.ai.chat(
+        {
+          messages: [{ role: "user", content: "Hello!" }],
+          licenseKey: "test-key",
+        },
+        { status: 500 }
+      )
+    })
+
+    it("handles no license", async () => {
+      nock.cleanAll()
+      nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(404)
+      await config.api.ai.chat(
+        {
+          messages: [{ role: "user", content: "Hello!" }],
+          licenseKey: "test-key",
+        },
+        {
+          status: 403,
+        }
+      )
+    })
+
+    it("handles no license key", async () => {
+      await config.api.ai.chat(
+        {
+          messages: [{ role: "user", content: "Hello!" }],
+          // @ts-expect-error - intentionally wrong
+          licenseKey: undefined,
+        },
+        {
+          status: 403,
+        }
+      )
+    })
+  })
+})