From 739c993588e1d510362d53a4bf7130991a9d6ca0 Mon Sep 17 00:00:00 2001 From: Sam Rose Date: Tue, 1 Apr 2025 11:01:32 +0100 Subject: [PATCH 1/2] Create new global DB for self-host quota. --- .../backend-core/src/cache/writethrough.ts | 25 +++ packages/backend-core/src/constants/db.ts | 5 + .../backend-core/src/context/mainContext.ts | 31 +++ packages/backend-core/src/context/types.ts | 1 + packages/pro | 2 +- .../server/src/api/routes/tests/ai.spec.ts | 186 ++++++++++-------- 6 files changed, 164 insertions(+), 86 deletions(-) diff --git a/packages/backend-core/src/cache/writethrough.ts b/packages/backend-core/src/cache/writethrough.ts index cd7409ca15..5a1d9f6a14 100644 --- a/packages/backend-core/src/cache/writethrough.ts +++ b/packages/backend-core/src/cache/writethrough.ts @@ -96,6 +96,24 @@ async function get(db: Database, id: string): Promise { return cacheItem.doc } +async function tryGet( + db: Database, + id: string +): Promise { + const cache = await getCache() + const cacheKey = makeCacheKey(db, id) + let cacheItem: CacheItem | null = await cache.get(cacheKey) + if (!cacheItem) { + const doc = await db.tryGet(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 { 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(id: string) { return get(this.db, id) } + async tryGet(id: string) { + return tryGet(this.db, id) + } + async remove(docOrId: any, rev?: any) { return remove(this.db, docOrId, rev) } diff --git a/packages/backend-core/src/constants/db.ts b/packages/backend-core/src/constants/db.ts index 3085b91ef1..28d389e6ba 100644 --- a/packages/backend-core/src/constants/db.ts +++ b/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) diff --git a/packages/backend-core/src/context/mainContext.ts b/packages/backend-core/src/context/mainContext.ts index 8e0c71ff18..ed0c56daaf 100644 --- a/packages/backend-core/src/context/mainContext.ts +++ b/packages/backend-core/src/context/mainContext.ts @@ -157,6 +157,27 @@ export async function doInTenant( return newContext(updates, task) } +export async function doInSelfHostTenantUsingCloud( + tenantId: string, + task: () => T +): Promise { + 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( appId: string, task: () => T @@ -325,6 +346,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 +370,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) } diff --git a/packages/backend-core/src/context/types.ts b/packages/backend-core/src/context/types.ts index 23598b951e..adee495e60 100644 --- a/packages/backend-core/src/context/types.ts +++ b/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 diff --git a/packages/pro b/packages/pro index 4417bceb24..0f46b458f3 160000 --- a/packages/pro +++ b/packages/pro @@ -1 +1 @@ -Subproject commit 4417bceb24eabdd9a8c1615fb83c4e6fe8c0c914 +Subproject commit 0f46b458f3fd1edd15fa7ff1d16eeb92caee27e1 diff --git a/packages/server/src/api/routes/tests/ai.spec.ts b/packages/server/src/api/routes/tests/ai.spec.ts index 288ab888fd..ad2ae7dc50 100644 --- a/packages/server/src/api/routes/tests/ai.spec.ts +++ b/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): 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 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 + 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, + } + ) + }) + }) +}) From 853c00242f08db2d44232fca31f2eba3a1590f4f Mon Sep 17 00:00:00 2001 From: Sam Rose Date: Tue, 1 Apr 2025 11:28:13 +0100 Subject: [PATCH 2/2] Respond to PR comments. --- packages/backend-core/src/context/mainContext.ts | 6 ++++++ packages/pro | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/packages/backend-core/src/context/mainContext.ts b/packages/backend-core/src/context/mainContext.ts index ed0c56daaf..e701f111aa 100644 --- a/packages/backend-core/src/context/mainContext.ts +++ b/packages/backend-core/src/context/mainContext.ts @@ -157,6 +157,12 @@ export async function doInTenant( 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( tenantId: string, task: () => T diff --git a/packages/pro b/packages/pro index 0f46b458f3..8eb981cf01 160000 --- a/packages/pro +++ b/packages/pro @@ -1 +1 @@ -Subproject commit 0f46b458f3fd1edd15fa7ff1d16eeb92caee27e1 +Subproject commit 8eb981cf01151261697a8f26c08c4c28f66b8e15