From 7807b734bb4f9768803f95b53350fe3611be9034 Mon Sep 17 00:00:00 2001
From: Adria Navarro <adria@budibase.com>
Date: Thu, 16 May 2024 14:41:45 +0200
Subject: [PATCH] Persist googlesheet refs in context

---
 packages/backend-core/src/context/types.ts    |  6 ++
 .../server/src/integrations/googlesheets.ts   | 62 +++++++++++++------
 2 files changed, 49 insertions(+), 19 deletions(-)

diff --git a/packages/backend-core/src/context/types.ts b/packages/backend-core/src/context/types.ts
index f297d3089f..6694320978 100644
--- a/packages/backend-core/src/context/types.ts
+++ b/packages/backend-core/src/context/types.ts
@@ -1,4 +1,6 @@
 import { IdentityContext, Snippet, VM } from "@budibase/types"
+import { OAuth2Client } from "google-auth-library"
+import { GoogleSpreadsheet } from "google-spreadsheet"
 
 // keep this out of Budibase types, don't want to expose context info
 export type ContextMap = {
@@ -12,4 +14,8 @@ export type ContextMap = {
   vm?: VM
   cleanup?: (() => void | Promise<void>)[]
   snippets?: Snippet[]
+  googleSheets?: {
+    oauthClient: OAuth2Client
+    clients: Record<string, GoogleSpreadsheet>
+  }
 }
diff --git a/packages/server/src/integrations/googlesheets.ts b/packages/server/src/integrations/googlesheets.ts
index dc945b454a..28df55a981 100644
--- a/packages/server/src/integrations/googlesheets.ts
+++ b/packages/server/src/integrations/googlesheets.ts
@@ -228,32 +228,56 @@ class GoogleSheetsIntegration implements DatasourcePlus {
 
   private async connect() {
     try {
-      await setupCreationAuth(this.config)
+      const bbCtx = context.getCurrentContext()
+      let oauthClient = bbCtx?.googleSheets?.oauthClient
 
-      // Initialise oAuth client
-      const googleConfig = await configs.getGoogleDatasourceConfig()
-      if (!googleConfig) {
-        throw new HTTPError("Google config not found", 400)
-      }
+      if (!oauthClient) {
+        await setupCreationAuth(this.config)
 
-      const oauthClient = new OAuth2Client({
-        clientId: googleConfig.clientID,
-        clientSecret: googleConfig.clientSecret,
-      })
+        // Initialise oAuth client
+        const googleConfig = await configs.getGoogleDatasourceConfig()
+        if (!googleConfig) {
+          throw new HTTPError("Google config not found", 400)
+        }
 
-      const tokenResponse = await this.fetchAccessToken({
-        client_id: googleConfig.clientID,
-        client_secret: googleConfig.clientSecret,
-        refresh_token: this.config.auth.refreshToken,
-      })
+        oauthClient = new OAuth2Client({
+          clientId: googleConfig.clientID,
+          clientSecret: googleConfig.clientSecret,
+        })
 
-      oauthClient.setCredentials({
-        refresh_token: this.config.auth.refreshToken,
-        access_token: tokenResponse.access_token,
-      })
+        const tokenResponse = await this.fetchAccessToken({
+          client_id: googleConfig.clientID,
+          client_secret: googleConfig.clientSecret,
+          refresh_token: this.config.auth.refreshToken,
+        })
+
+        oauthClient.setCredentials({
+          refresh_token: this.config.auth.refreshToken,
+          access_token: tokenResponse.access_token,
+        })
+        if (bbCtx && !bbCtx.googleSheets) {
+          bbCtx.googleSheets = {
+            oauthClient,
+            clients: {},
+          }
+          bbCtx.cleanup = bbCtx.cleanup || []
 
       this.client = new GoogleSpreadsheet(this.spreadsheetId, oauthClient)
       await this.client.loadInfo()
+        }
+      }
+
+      let client = bbCtx?.googleSheets?.clients[this.spreadsheetId]
+      if (!client) {
+        client = new GoogleSpreadsheet(this.spreadsheetId, oauthClient)
+        await client.loadInfo()
+
+        if (bbCtx?.googleSheets?.clients) {
+          bbCtx.googleSheets.clients[this.spreadsheetId] = client
+        }
+      }
+
+      this.client = client
     } catch (err: any) {
       // this happens for xlsx imports
       if (err.message?.includes("operation is not supported")) {