From ec02ce76ad0f954ada6c846de8bc15519581f382 Mon Sep 17 00:00:00 2001
From: Adria Navarro <adria@budibase.com>
Date: Fri, 21 Mar 2025 09:48:08 +0100
Subject: [PATCH] Check usage across dev and prod app

---
 .../src/sdk/app/oauth2/tests/utils.spec.ts    | 54 ++++++++++++++-----
 packages/server/src/sdk/app/oauth2/utils.ts   | 20 +++++--
 2 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/packages/server/src/sdk/app/oauth2/tests/utils.spec.ts b/packages/server/src/sdk/app/oauth2/tests/utils.spec.ts
index d12509b9e0..b968f3138e 100644
--- a/packages/server/src/sdk/app/oauth2/tests/utils.spec.ts
+++ b/packages/server/src/sdk/app/oauth2/tests/utils.spec.ts
@@ -7,7 +7,7 @@ import path from "path"
 import { KEYCLOAK_IMAGE } from "../../../../integrations/tests/utils/images"
 import { startContainer } from "../../../../integrations/tests/utils"
 import { OAuth2CredentialsMethod } from "@budibase/types"
-import { cache, context, docIds } from "@budibase/backend-core"
+import { cache } from "@budibase/backend-core"
 import tk from "timekeeper"
 
 const config = new TestConfiguration()
@@ -140,16 +140,10 @@ describe("oauth2 utils", () => {
           )
 
           const usageLog = await config.doInContext(config.appId, () =>
-            context
-              .getAppDB()
-              .tryGet(docIds.generateOAuth2LogID(oauthConfig._id))
+            sdk.oauth2.getLastUsages([oauthConfig._id])
           )
 
-          expect(usageLog).toEqual(
-            expect.objectContaining({
-              lastUsage: Date.now(),
-            })
-          )
+          expect(usageLog[oauthConfig._id]).toEqual(Date.now())
         })
 
         it("does not track on failed usages", async () => {
@@ -171,12 +165,46 @@ describe("oauth2 utils", () => {
           )
 
           const usageLog = await config.doInContext(config.appId, () =>
-            context
-              .getAppDB()
-              .tryGet(docIds.generateOAuth2LogID(oauthConfig._id))
+            sdk.oauth2.getLastUsages([oauthConfig._id])
           )
 
-          expect(usageLog).toBeUndefined()
+          expect(usageLog[oauthConfig._id]).toBeUndefined()
+        })
+
+        it("tracks usages between prod and dev, keeping always the latest", async () => {
+          const oauthConfig = await config.doInContext(config.appId, () =>
+            sdk.oauth2.create({
+              name: generator.guid(),
+              url: `${keycloakUrl}/realms/myrealm/protocol/openid-connect/token`,
+              clientId: "my-client",
+              clientSecret: "my-secret",
+              method,
+            })
+          )
+
+          await config.doInContext(config.appId, () =>
+            getToken(oauthConfig._id)
+          )
+
+          await config.publish()
+
+          tk.travel(Date.now() + 100)
+          await config.doInContext(config.prodAppId, () =>
+            getToken(oauthConfig._id)
+          )
+          await testUtils.queue.processMessages(
+            cache.docWritethrough.DocWritethroughProcessor.queue
+          )
+
+          for (const appId of [config.appId, config.prodAppId]) {
+            const usageLog = await config.doInContext(appId, () =>
+              sdk.oauth2.getLastUsages([oauthConfig._id])
+            )
+
+            expect(usageLog).toEqual({
+              [oauthConfig._id]: Date.now(),
+            })
+          }
         })
       })
     }
diff --git a/packages/server/src/sdk/app/oauth2/utils.ts b/packages/server/src/sdk/app/oauth2/utils.ts
index 8a8952050d..349c62f9e7 100644
--- a/packages/server/src/sdk/app/oauth2/utils.ts
+++ b/packages/server/src/sdk/app/oauth2/utils.ts
@@ -99,15 +99,27 @@ export async function validateConfig(config: {
 }
 
 export async function getLastUsages(ids: string[]) {
-  const docs = await context
+  const devDocs = await context
     .getAppDB()
     .getMultiple<OAuth2LogDocument>(ids.map(docIds.generateOAuth2LogID), {
       allowMissing: true,
     })
+
+  const prodDocs = await context
+    .getProdAppDB()
+    .getMultiple<OAuth2LogDocument>(ids.map(docIds.generateOAuth2LogID), {
+      allowMissing: true,
+    })
+
   const result = ids.reduce<Record<string, number>>((acc, id) => {
-    const doc = docs.find(d => d._id === docIds.generateOAuth2LogID(id))
-    if (doc) {
-      acc[id] = doc.lastUsage
+    const devDoc = devDocs.find(d => d._id === docIds.generateOAuth2LogID(id))
+    if (devDoc) {
+      acc[id] = devDoc.lastUsage
+    }
+
+    const prodDoc = prodDocs.find(d => d._id === docIds.generateOAuth2LogID(id))
+    if (prodDoc && (!acc[id] || acc[id] < prodDoc.lastUsage)) {
+      acc[id] = prodDoc.lastUsage
     }
     return acc
   }, {})