From 77733aab0932b161dc9ccfa230f488276adda711 Mon Sep 17 00:00:00 2001 From: Aric Lasry Date: Fri, 5 Apr 2024 10:02:01 +0200 Subject: [PATCH] Address flvdnd and spolu comments. --- connectors/src/connectors/slack/bot.ts | 36 ----- connectors/src/connectors/slack/index.ts | 22 ++- .../connectors/slack/lib/workspace_limits.ts | 14 +- connectors/src/lib/cli.ts | 19 ++- .../resources/slack_configuration_resource.ts | 134 +++++++++++++++--- 5 files changed, 151 insertions(+), 74 deletions(-) diff --git a/connectors/src/connectors/slack/bot.ts b/connectors/src/connectors/slack/bot.ts index d3dfadc0b7d68..75b5d3a68184b 100644 --- a/connectors/src/connectors/slack/bot.ts +++ b/connectors/src/connectors/slack/bot.ts @@ -626,42 +626,6 @@ export async function getBotEnabled( return new Ok(slackConfig.botEnabled); } -export async function toggleSlackbot( - connectorId: ModelId, - botEnabled: boolean -): Promise> { - const slackConfig = await SlackConfigurationResource.fetchByConnectorId( - connectorId - ); - - if (!slackConfig) { - return new Err( - new Error( - `Failed to find a Slack configuration for connector ${connectorId}` - ) - ); - } - - if (botEnabled) { - const otherSlackConfigWithBotEnabled = - await SlackConfigurationResource.fetchByActiveBot( - slackConfig.slackTeamId - ); - - if (otherSlackConfigWithBotEnabled) { - return new Err( - new Error( - "Another Dust workspace has already enabled the slack bot for your Slack workspace." - ) - ); - } - } - - await slackConfig.update({ botEnabled: botEnabled }); - - return new Ok(void 0); -} - async function makeContentFragment( slackClient: WebClient, channelId: string, diff --git a/connectors/src/connectors/slack/index.ts b/connectors/src/connectors/slack/index.ts index f380a62262e3a..53846f03e0a56 100644 --- a/connectors/src/connectors/slack/index.ts +++ b/connectors/src/connectors/slack/index.ts @@ -14,10 +14,7 @@ import type { ConnectorPermissionRetriever, } from "@connectors/connectors/interface"; import { getChannels } from "@connectors/connectors/slack//temporal/activities"; -import { - getBotEnabled, - toggleSlackbot, -} from "@connectors/connectors/slack/bot"; +import { getBotEnabled } from "@connectors/connectors/slack/bot"; import { joinChannel } from "@connectors/connectors/slack/lib/channels"; import { getSlackAccessToken, @@ -620,8 +617,21 @@ export async function setSlackConfig( switch (configKey) { case "botEnabled": { - const res = await toggleSlackbot(connectorId, configValue === "true"); - return res; + const slackConfig = await SlackConfigurationResource.fetchByConnectorId( + connectorId + ); + if (!slackConfig) { + return new Err( + new Error( + `Slack configuration not found for connector ${connectorId}` + ) + ); + } + if (configValue === "true") { + return slackConfig.enableBot(); + } else { + return slackConfig.disableBot(); + } } default: { diff --git a/connectors/src/connectors/slack/lib/workspace_limits.ts b/connectors/src/connectors/slack/lib/workspace_limits.ts index 57872a2482589..3739b050808f2 100644 --- a/connectors/src/connectors/slack/lib/workspace_limits.ts +++ b/connectors/src/connectors/slack/lib/workspace_limits.ts @@ -6,9 +6,9 @@ import type {} from "@slack/web-api/dist/response/UsersInfoResponse"; import type { SlackUserInfo } from "@connectors/connectors/slack/lib/slack_client"; import { getSlackConversationInfo } from "@connectors/connectors/slack/lib/slack_client"; import { dataSourceConfigFromConnector } from "@connectors/lib/api/data_source_config"; -import { SlackBotWhitelistModel } from "@connectors/lib/models/slack"; import logger from "@connectors/logger/logger"; import type { ConnectorResource } from "@connectors/resources/connector_resource"; +import { SlackConfigurationResource } from "@connectors/resources/slack_configuration_resource"; async function getActiveMemberEmails( connector: ConnectorResource @@ -213,13 +213,11 @@ async function isBotAllowed( // This means that even a non verified user of a given Slack workspace who can trigger a bot // that talks to our bot (@dust) will be able to use the Dust bot. // Make sure to be explicit about this with users as you whitelist a new bot. - // Example: non-verfied-user -> @AnyWhitelistedBot -> @dust -> Dust answers with potentially private information. - const whitelist = await SlackBotWhitelistModel.findOne({ - where: { - connectorId: connector.id, - botName: names, - }, - }); + // Example: non-verified-user -> @AnyWhitelistedBot -> @dust -> Dust answers with potentially private information. + const slackConfig = await SlackConfigurationResource.fetchByConnectorId( + connector.id + ); + const whitelist = await slackConfig?.isBotWhitelisted(names); if (!whitelist) { logger.info( { user: slackUserInfo, connectorId: connector.id }, diff --git a/connectors/src/lib/cli.ts b/connectors/src/lib/cli.ts index 44e463e761ee2..f864c5778d0dc 100644 --- a/connectors/src/lib/cli.ts +++ b/connectors/src/lib/cli.ts @@ -70,7 +70,6 @@ import { upsertPageWorkflow, } from "@connectors/connectors/notion/temporal/workflows"; import { uninstallSlack } from "@connectors/connectors/slack"; -import { toggleSlackbot } from "@connectors/connectors/slack/bot"; import { maybeLaunchSlackSyncWorkflowForChannelId } from "@connectors/connectors/slack/lib/cli"; import { launchSlackSyncOneThreadWorkflow } from "@connectors/connectors/slack/temporal/client"; import { launchCrawlWebsiteSchedulerWorkflow } from "@connectors/connectors/webcrawler/temporal/client"; @@ -917,7 +916,19 @@ export const slack = async ({ if (!connector) { throw new Error(`Could not find connector for workspace ${args.wId}`); } - await throwOnError(toggleSlackbot(connector.id, true)); + const slackConfig = await SlackConfigurationResource.fetchByConnectorId( + connector.id + ); + if (!slackConfig) { + throw new Error( + `Could not find slack configuration for connector ${connector.id}` + ); + } + + const res = await slackConfig.enableBot(); + if (res.isErr()) { + throw res.error; + } return { success: true }; } @@ -1075,9 +1086,7 @@ export const slack = async ({ connector.id ); if (slackConfig) { - await slackConfig.update({ - whitelistedDomains: whitelistedDomainsArray, - }); + await slackConfig.setWhitelistedDomains(whitelistedDomainsArray); } return { success: true }; diff --git a/connectors/src/resources/slack_configuration_resource.ts b/connectors/src/resources/slack_configuration_resource.ts index 3bd3bbfb707c0..c549b19293e46 100644 --- a/connectors/src/resources/slack_configuration_resource.ts +++ b/connectors/src/resources/slack_configuration_resource.ts @@ -2,9 +2,14 @@ import type { ModelId, Result } from "@dust-tt/types"; import { Err, Ok } from "@dust-tt/types"; import type { Attributes, ModelStatic, Transaction } from "sequelize"; -import { SlackConfigurationModel } from "@connectors/lib/models/slack"; +import { + SlackBotWhitelistModel, + SlackChannel, + SlackConfigurationModel, + SlackMessages, +} from "@connectors/lib/models/slack"; +import logger from "@connectors/logger/logger"; import { BaseResource } from "@connectors/resources/base_resource"; -import { sequelizeConnection } from "@connectors/resources/storage"; import type { ReadonlyAttributesType } from "@connectors/resources/storage/types"; // Attributes are marked as read-only to reflect the stateless nature of our Resource. @@ -50,10 +55,10 @@ export class SlackConfigurationResource extends BaseResource { + return !!(await SlackBotWhitelistModel.findOne({ + where: { + connectorId: this.id, + botName: botName, + }, + })); + } + static async listAll() { const blobs = await SlackConfigurationResource.model.findAll({}); @@ -88,7 +102,9 @@ export class SlackConfigurationResource extends BaseResource { const blobs = await this.model.findAll({ where: { slackTeamId, @@ -101,20 +117,100 @@ export class SlackConfigurationResource extends BaseResource> { - return sequelizeConnection.transaction(async (transaction) => { - try { - await this.model.destroy({ - where: { - id: this.id, - }, - transaction, - }); - - return new Ok(undefined); - } catch (err) { - return new Err(err as Error); + async enableBot(): Promise> { + const otherSlackConfigurationWithBotEnabled = + await SlackConfigurationModel.findOne({ + where: { + slackTeamId: this.slackTeamId, + botEnabled: true, + }, + }); + if ( + otherSlackConfigurationWithBotEnabled && + otherSlackConfigurationWithBotEnabled.id !== this.id + ) { + logger.error( + { + slackTeamId: this.slackTeamId, + }, + "Another Dust workspace has already enabled the slack bot for your Slack workspace." + ); + return new Err( + new Error( + "Another Dust workspace has already enabled the slack bot for your Slack workspace." + ) + ); + } + await this.model.update( + { botEnabled: true }, + { + where: { + id: this.id, + }, } - }); + ); + + return new Ok(undefined); + } + + async disableBot(): Promise> { + await this.model.update( + { botEnabled: false }, + { + where: { + id: this.id, + }, + } + ); + + return new Ok(undefined); + } + + async setWhitelistedDomains(domain: string[]) { + await this.model.update( + { whitelistedDomains: domain }, + { + where: { + id: this.id, + }, + } + ); + + return new Ok(undefined); + } + async delete(transaction: Transaction): Promise> { + try { + await SlackChannel.destroy({ + where: { + connectorId: this.id, + }, + transaction, + }); + + await SlackMessages.destroy({ + where: { + connectorId: this.id, + }, + transaction, + }); + + await SlackBotWhitelistModel.destroy({ + where: { + connectorId: this.id, + }, + transaction, + }); + + await this.model.destroy({ + where: { + id: this.id, + }, + transaction, + }); + + return new Ok(undefined); + } catch (err) { + return new Err(err as Error); + } } }