From 411510b722eea8283e3dc9e1742a318ee05a2d0a Mon Sep 17 00:00:00 2001 From: Aric Lasry Date: Thu, 14 Dec 2023 18:55:22 +0100 Subject: [PATCH] Cache Nango tokens in Gdrive (#2881) * Cache Nango token in Gdrive * Clean up after self review * Use cache in Gdrive --- .../google_drive/temporal/activities.ts | 17 +-- connectors/src/connectors/notion/index.ts | 2 - connectors/src/lib/nango_client.ts | 1 - connectors/src/lib/nango_helpers.ts | 103 ++++++++++++------ 4 files changed, 79 insertions(+), 44 deletions(-) diff --git a/connectors/src/connectors/google_drive/temporal/activities.ts b/connectors/src/connectors/google_drive/temporal/activities.ts index 5cd8ca2f8004..e2822e01bfd9 100644 --- a/connectors/src/connectors/google_drive/temporal/activities.ts +++ b/connectors/src/connectors/google_drive/temporal/activities.ts @@ -33,6 +33,7 @@ import { GoogleDriveSyncToken, GoogleDriveWebhook, } from "@connectors/lib/models/google_drive"; +import { getConnectionFromNango } from "@connectors/lib/nango_helpers"; import logger from "@connectors/logger/logger"; import { registerWebhook } from "../lib"; @@ -93,8 +94,7 @@ export async function getGoogleCredentials( return await nango_client().getConnection( NANGO_GOOGLE_DRIVE_CONNECTOR_ID, nangoConnectionId, - false, - true + false ); } @@ -104,12 +104,13 @@ export async function getAuthObject( if (!NANGO_GOOGLE_DRIVE_CONNECTOR_ID) { throw new Error("NANGO_GOOGLE_DRIVE_CONNECTOR_ID is not defined"); } - const res: NangoGetConnectionRes = await nango_client().getConnection( - NANGO_GOOGLE_DRIVE_CONNECTOR_ID, - nangoConnectionId, - false, - true - ); + const res: NangoGetConnectionRes = await getConnectionFromNango({ + connectionId: nangoConnectionId, + integrationId: NANGO_GOOGLE_DRIVE_CONNECTOR_ID, + refreshToken: false, + useCache: true, + }); + const oauth2Client = new google.auth.OAuth2(); oauth2Client.setCredentials({ access_token: res.credentials.access_token, diff --git a/connectors/src/connectors/notion/index.ts b/connectors/src/connectors/notion/index.ts index 7abc2936d6a7..0ed93fc49392 100644 --- a/connectors/src/connectors/notion/index.ts +++ b/connectors/src/connectors/notion/index.ts @@ -133,13 +133,11 @@ export async function updateNotionConnector( const connectionRes = await nango_client().getConnection( NANGO_NOTION_CONNECTOR_ID, oldConnectionId, - false, false ); const newConnectionRes = await nango_client().getConnection( NANGO_NOTION_CONNECTOR_ID, connectionId, - false, false ); diff --git a/connectors/src/lib/nango_client.ts b/connectors/src/lib/nango_client.ts index 6da66d279d6c..260e284e7bc2 100644 --- a/connectors/src/lib/nango_client.ts +++ b/connectors/src/lib/nango_client.ts @@ -11,7 +11,6 @@ class CustomNango extends Nango { async getConnection( providerConfigKey: string, connectionId: string, - forceRefresh?: boolean, refreshToken?: boolean ) { try { diff --git a/connectors/src/lib/nango_helpers.ts b/connectors/src/lib/nango_helpers.ts index 9f989fd3c990..562722f74fa2 100644 --- a/connectors/src/lib/nango_helpers.ts +++ b/connectors/src/lib/nango_helpers.ts @@ -1,10 +1,33 @@ -import { redisClient } from "@connectors/lib/redis"; +import { cacheWithRedis } from "@dust-tt/types"; + import { NangoConnectionId } from "@connectors/types/nango_connection_id"; import { nango_client } from "./nango_client"; const NANGO_ACCESS_TOKEN_TTL_SECONDS = 60 * 5; // 5 minutes +async function _getAccessTokenFromNango({ + connectionId, + integrationId, +}: { + connectionId: NangoConnectionId; + integrationId: string; +}) { + const accessToken = await nango_client().getToken( + integrationId, + connectionId + ); + return accessToken; +} + +const _cachedGetAccessTokenFromNango = cacheWithRedis( + _getAccessTokenFromNango, + ({ connectionId, integrationId }) => { + return `${integrationId}-${connectionId}`; + }, + NANGO_ACCESS_TOKEN_TTL_SECONDS * 1000 +); + export async function getAccessTokenFromNango({ connectionId, integrationId, @@ -14,49 +37,63 @@ export async function getAccessTokenFromNango({ integrationId: string; useCache?: boolean; }) { - const cacheKey = `nango_access_token:${integrationId}/${connectionId}`; - const redis = await redisClient(); - - try { - const _setCache = (token: string) => - redis.set(cacheKey, token, { - EX: NANGO_ACCESS_TOKEN_TTL_SECONDS, - }); - - if (!useCache) { - const accessToken = await _getAccessTokenFromNango({ - connectionId, - integrationId, - }); - await _setCache(accessToken); - return accessToken; - } - - const maybeAccessToken = await redis.get(cacheKey); - if (maybeAccessToken) { - return maybeAccessToken; - } - const accessToken = await nango_client().getToken( + if (useCache) { + return await _cachedGetAccessTokenFromNango({ + connectionId, integrationId, - connectionId - ); - await _setCache(accessToken); - return accessToken; - } finally { - await redis.quit(); + }); + } else { + return await _getAccessTokenFromNango({ connectionId, integrationId }); } } -async function _getAccessTokenFromNango({ +async function _getConnectionFromNango({ connectionId, integrationId, + refreshToken, }: { connectionId: NangoConnectionId; integrationId: string; + refreshToken?: boolean; }) { - const accessToken = await nango_client().getToken( + const accessToken = await nango_client().getConnection( integrationId, - connectionId + connectionId, + refreshToken ); return accessToken; } + +const _getCachedConnectionFromNango = cacheWithRedis( + _getConnectionFromNango, + ({ connectionId, integrationId, refreshToken }) => { + return `${integrationId}-${connectionId}-${refreshToken}`; + }, + NANGO_ACCESS_TOKEN_TTL_SECONDS * 1000 +); + +export async function getConnectionFromNango({ + connectionId, + integrationId, + refreshToken = false, + useCache = false, +}: { + connectionId: NangoConnectionId; + integrationId: string; + refreshToken?: boolean; + useCache?: boolean; +}) { + if (useCache) { + return await _getCachedConnectionFromNango({ + connectionId, + integrationId, + refreshToken, + }); + } else { + return await _getConnectionFromNango({ + connectionId, + integrationId, + refreshToken, + }); + } +}