oauth: Gated Github support & dual mode (#6314)

* setup working

* oauth: fix error_response logging

* fix

* oauth: add logging when going to provider

* flagged test_oauth_setup

* oauth: github connector dual-mode

* nit

* add migration

* clean-up

* add connectionId to ConnectorType

* fix expiry logic

* dry out setupConnection for managed and udpate

* fix infra oauth

* add github-app-private-key to oauth service

* add github-app-private-key to oauth service

connectors/migrations/db/migration_06.sql

@@ -0,0 +1,2 @@
+-- Migration created on Jul 18, 2024
+CREATE INDEX github_connector_states_installation_id ON github_connector_states ("installationId");

connectors/src/api/webhooks/webhook_github.ts

@@ -1,9 +1,8 @@
-import type { WithConnectorsAPIErrorReponse } from "@dust-tt/types";
+import type { ModelId, WithConnectorsAPIErrorReponse } from "@dust-tt/types";
 import { assertNever } from "@dust-tt/types";
 import type { Request, Response } from "express";
 import { isLeft } from "fp-ts/lib/Either";
 import * as reporter from "io-ts-reporters";
-import { Op } from "sequelize";
 
 import {
   GithubWebhookPayloadSchema,
@@ -110,59 +109,49 @@ const _webhookGithubAPIHandler = async (
 
   const installationId = payload.installation.id.toString();
 
-  // TODO(spolu): GITHUB_MIGRATION find by ConnectorState.installationId
-  const connectors = await ConnectorResource.listByType("github", {
-    connectionId: installationId,
+  const githubConnectorStates = await GithubConnectorState.findAll({
+    where: {
+      installationId,
+    },
   });
 
-  if (!connectors.length) {
-    logger.error(
-      {
-        installationId,
-      },
-      "No GitHub connectors found for installation."
-    );
-    // return 200 to avoid github retrying
-    return res.status(200);
-  }
-
-  const githubConnectorStates = (
-    await GithubConnectorState.findAll({
-      where: {
-        connectorId: {
-          [Op.in]: connectors.map((c) => c.id),
-        },
-      },
-    })
+  const connectors = (
+    await ConnectorResource.fetchByIds(
+      "github",
+      githubConnectorStates.map((s) => s.connectorId)
+    )
   ).reduce(
-    (acc, curr) => Object.assign(acc, { [curr.connectorId]: curr }),
-    {} as Record<string, GithubConnectorState>
+    (acc, curr) => Object.assign(acc, { [curr.id]: curr }),
+    {} as Record<ModelId, ConnectorResource>
   );
 
   const enabledConnectors: ConnectorResource[] = [];
-  for (const connector of connectors) {
-    if (connector.isPaused()) {
-      logger.info(
+
+  for (const connectorState of githubConnectorStates) {
+    const connector = connectors[connectorState.connectorId];
+
+    if (!connector) {
+      logger.error(
         {
-          connectorId: connector.id,
+          connectorId: connectorState.connectorId,
           installationId,
         },
-        "Skipping webhook for Github connector because it is paused."
+        "Connector unexpectedly not found"
       );
       continue;
     }
-    const connectorState = githubConnectorStates[connector.id];
-    if (!connectorState) {
-      logger.error(
+
+    if (connector.isPaused()) {
+      logger.info(
         {
           connectorId: connector.id,
           installationId,
         },
-        "Connector state not found"
+        "Skipping webhook for Github connector because it is paused."
       );
-      // return 200 to avoid github retrying
       continue;
     }
+
     if (
       !connectorState.webhooksEnabledAt ||
       connectorState.webhooksEnabledAt.getTime() > Date.now()

connectors/src/connectors/github/index.ts

@@ -85,11 +85,12 @@ export class GithubConnectorManager extends BaseConnectorManager<null> {
       });
     }
 
-    // TOOD(spolu): GITHUB_MIGRATION we will have to retrieve installationId from connector state
-    // for existing and pull from oauth scrubbed_raw_json to get the new installationId.
     if (connectionId) {
-      const oldGithubInstallationId = c.connectionId;
-      const newGithubInstallationId = connectionId;
+      const [oldGithubInstallationId, newGithubInstallationId] =
+        await Promise.all([
+          installationIdFromConnectionId(c.connectionId),
+          installationIdFromConnectionId(connectionId),
+        ]);
 
       if (oldGithubInstallationId !== newGithubInstallationId) {
         return new Err({

connectors/src/connectors/github/lib/github_api.ts

@@ -1,5 +1,5 @@
 import type { Result } from "@dust-tt/types";
-import { Err, Ok } from "@dust-tt/types";
+import { Err, getOAuthConnectionAccessToken, Ok } from "@dust-tt/types";
 import { createAppAuth } from "@octokit/auth-app";
 import { hash as blake3 } from "blake3";
 import { isLeft } from "fp-ts/lib/Either";
@@ -29,7 +29,9 @@ import {
   GetDiscussionPayloadSchema,
   GetRepoDiscussionsPayloadSchema,
 } from "@connectors/connectors/github/lib/github_graphql";
+import { apiConfig } from "@connectors/lib/api/config";
 import { ExternalOauthTokenError } from "@connectors/lib/error";
+import { isDualUseOAuthConnectionId } from "@connectors/lib/oauth";
 import logger from "@connectors/logger/logger";
 
 const API_PAGE_SIZE = 100;
@@ -97,18 +99,37 @@ export async function getGithubAppPrivateKey(): Promise<string> {
 export async function installationIdFromConnectionId(
   connectionId: string
 ): Promise<string | null> {
-  // TODO(spolu): GITHUB_MIGRATION dual mode will pull the refresh token to get the installation id
-  // from the scrubbed_raw_json.
-  const octokit = await getOctokit(connectionId);
+  if (isDualUseOAuthConnectionId(connectionId)) {
+    const tokRes = await getOAuthConnectionAccessToken({
+      config: apiConfig.getOAuthAPIConfig(),
+      logger,
+      provider: "github",
+      connectionId,
+    });
+    if (tokRes.isErr()) {
+      logger.error(
+        { connectionId, error: tokRes.error },
+        "Error retrieving Github access token"
+      );
+      return null;
+    }
 
-  try {
-    await octokit.rest.apps.getAuthenticated();
-  } catch (e) {
-    logger.error({ error: e }, "Error validating github installation id");
-    return null;
-  }
+    return (tokRes.value.scrubbed_raw_json as { installation_id: string })[
+      "installation_id"
+    ];
+  } else {
+    // TODO(spolu) remove once fully migrated to oauth.
+    const octokit = await getOctokit(connectionId);
+
+    try {
+      await octokit.rest.apps.getAuthenticated();
+    } catch (e) {
+      logger.error({ error: e }, "Error validating github installation id");
+      return null;
+    }
 
-  return connectionId;
+    return connectionId;
+  }
 }
 
 export async function getReposPage(
@@ -556,19 +577,39 @@ export async function getDiscussion(
 }
 
 export async function getOctokit(connectionId: string): Promise<Octokit> {
-  if (!GITHUB_APP_ID) {
-    throw new Error("GITHUB_APP_ID not set");
+  if (isDualUseOAuthConnectionId(connectionId)) {
+    const tokRes = await getOAuthConnectionAccessToken({
+      config: apiConfig.getOAuthAPIConfig(),
+      logger,
+      provider: "github",
+      connectionId,
+    });
+
+    if (tokRes.isErr()) {
+      logger.error(
+        { connectionId, error: tokRes.error },
+        "Error retrieving Github access token"
+      );
+      throw new Error("Error retrieving Github access token");
+    }
+
+    return new Octokit({ auth: tokRes.value.access_token });
+  } else {
+    // TODO(spolu) remove once fully migrated to oauth.
+    if (!GITHUB_APP_ID) {
+      throw new Error("GITHUB_APP_ID not set");
+    }
+    const privateKey = await getGithubAppPrivateKey();
+
+    return new Octokit({
+      authStrategy: createAppAuth,
+      auth: {
+        appId: GITHUB_APP_ID,
+        privateKey: privateKey,
+        installationId: connectionId,
+      },
+    });
   }
-  const privateKey = await getGithubAppPrivateKey();
-
-  return new Octokit({
-    authStrategy: createAppAuth,
-    auth: {
-      appId: GITHUB_APP_ID,
-      privateKey: privateKey,
-      installationId: connectionId,
-    },
-  });
 }
 
 // Repository processing

connectors/src/lib/models/github.ts

@@ -57,7 +57,10 @@ GithubConnectorState.init(
   {
     sequelize: sequelizeConnection,
     modelName: "github_connector_states",
-    indexes: [{ fields: ["connectorId"], unique: true }],
+    indexes: [
+      { fields: ["connectorId"], unique: true },
+      { fields: ["installationId"] },
+    ],
   }
 );
 ConnectorModel.hasOne(GithubConnectorState);

connectors/src/lib/oauth.ts

@@ -0,0 +1,7 @@
+// This function is used to discreminate between a new OAuth connection and an old Nango/Github
+// connection. It is used to support dual-use while migrating and should be unused by a connector
+// once fully migrated
+export function isDualUseOAuthConnectionId(connectionId: string): boolean {
+  // TODO(spolu): make sure this function is removed once fully migrated.
+  return connectionId.startsWith("con_");
+}

connectors/src/resources/connector_resource.ts

@@ -211,6 +211,7 @@ export class ConnectorResource extends BaseResource<ConnectorModel> {
     return {
       id: this.id.toString(),
       type: this.type,
+      connectionId: this.connectionId,
       workspaceId: this.workspaceId,
       dataSourceName: this.dataSourceName,
       lastSyncStatus: this.lastSyncStatus,