From fd8f94ef3dadd34f11edb308e46f37f022970c9f Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 13:30:56 +0200 Subject: [PATCH 1/9] Implement ACL on data source and data source views. --- front/lib/api/assistant/configuration.ts | 21 ++- front/lib/auth.ts | 3 +- front/lib/models/data_source.ts | 3 +- front/lib/resources/data_source_resource.ts | 129 ++++++++++-------- .../resources/data_source_view_resource.ts | 124 +++++++++++------ .../storage/models/data_source_view.ts | 4 + front/lib/resources/string_ids.ts | 7 + ...ews_in_agent_data_source_configurations.ts | 2 +- front/pages/api/w/[wId]/data_sources/index.ts | 20 +-- .../pages/api/w/[wId]/data_sources/managed.ts | 22 +-- types/src/front/auth.ts | 11 ++ types/src/front/groups.ts | 2 +- types/src/index.ts | 1 + 13 files changed, 226 insertions(+), 123 deletions(-) create mode 100644 types/src/front/auth.ts diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index 0d7743fa75b8..adf2349e6fd0 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -1442,6 +1442,19 @@ async function _createAgentDataSourcesConfigData( processConfigurationId: ModelId | null; } ): Promise { + // Although we have the capability to support multiple workspaces, + // currently, we only support one workspace, which is the one the user is in. + // This allows us to use the current authenticator to fetch resources. + const allWorkspaceIds = _.uniqBy(dataSourceConfigurations, "workspaceSId"); + const hasUniqueAccessibleWorkspace = + allWorkspaceIds.length === 1 && + auth.getNonNullableWorkspace().sId === allWorkspaceIds[0].workspaceId; + if (!hasUniqueAccessibleWorkspace) { + throw new Error( + "Can't create AgentDataSourcesConfig for retrieval: Multiple workspaces." + ); + } + // dsConfig contains this format: // [ // { workspaceSId: s1o1u1p, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, @@ -1502,9 +1515,10 @@ async function _createAgentDataSourcesConfigData( // Then we get to do one findAllQuery per workspaceId, in a Promise.all. const getDataSourcesQueries = dsNamesPerWorkspaceId.map( - async ({ workspace, dataSourceNames }) => { + async ({ dataSourceNames }) => { const dataSources = await DataSourceResource.listByWorkspaceIdAndNames( - workspace, + // We can use `auth` because we limit to one workspace. + auth, dataSourceNames ); @@ -1515,7 +1529,8 @@ async function _createAgentDataSourcesConfigData( // and assign it to the agent data source configuration. const dataSourceViews = await DataSourceViewResource.listForDataSourcesInVault( - workspace, + // We can use `auth` because we limit to one workspace. + auth, uniqueDataSources.filter((ds) => ds.isManaged()), globalVault ); diff --git a/front/lib/auth.ts b/front/lib/auth.ts index 8af4c6297cf5..caec9ead77bf 100644 --- a/front/lib/auth.ts +++ b/front/lib/auth.ts @@ -13,6 +13,7 @@ import type { PlanType, SubscriptionType } from "@dust-tt/types"; import type { DustAPICredentials } from "@dust-tt/types"; import type { Result } from "@dust-tt/types"; import type { APIErrorWithStatusCode } from "@dust-tt/types"; +import type { BaseAuthenticator } from "@dust-tt/types"; import { Err, groupHasPermission, @@ -60,7 +61,7 @@ const DUST_INTERNAL_EMAIL_REGEXP = /^[^@]+@dust\.tt$/; * It explicitely does not store a reference to the current user to make sure our permissions are * workspace oriented. Use `getUserFromSession` if needed. */ -export class Authenticator { +export class Authenticator implements BaseAuthenticator { _flags: WhitelistableFeature[]; _key?: KeyAuthType; _role: RoleType; diff --git a/front/lib/models/data_source.ts b/front/lib/models/data_source.ts index 2554dcdfec09..1d9be47281a9 100644 --- a/front/lib/models/data_source.ts +++ b/front/lib/models/data_source.ts @@ -34,8 +34,9 @@ export class DataSource extends Model< declare workspaceId: ForeignKey; declare vaultId: ForeignKey; - declare workspace: NonAttribute; declare editedByUser: NonAttribute; + declare vault: NonAttribute; + declare workspace: NonAttribute; } DataSource.init( diff --git a/front/lib/resources/data_source_resource.ts b/front/lib/resources/data_source_resource.ts index 7f1923095733..363e22dbab67 100644 --- a/front/lib/resources/data_source_resource.ts +++ b/front/lib/resources/data_source_resource.ts @@ -1,16 +1,18 @@ import type { + ACLType, ConnectorProvider, DataSourceType, - LightWorkspaceType, Result, } from "@dust-tt/types"; -import { Err, formatUserFullName, Ok } from "@dust-tt/types"; +import { Err, formatUserFullName, Ok, removeNulls } from "@dust-tt/types"; import type { Attributes, CreationAttributes, FindOptions, + Includeable, ModelStatic, Transaction, + WhereOptions, } from "sequelize"; import { Op } from "sequelize"; @@ -20,6 +22,7 @@ import { User } from "@app/lib/models/user"; import { BaseResource } from "@app/lib/resources/base_resource"; import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource"; import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types"; +import { VaultResource } from "@app/lib/resources/vault_resource"; export type FetchDataSourceOptions = { includeEditedBy?: boolean; @@ -36,23 +39,31 @@ export interface DataSourceResource export class DataSourceResource extends BaseResource { static model: ModelStatic = DataSource; - editedByUser: Attributes | undefined; + readonly editedByUser: Attributes | undefined; + readonly vault: VaultResource; constructor( model: ModelStatic, blob: Attributes, + vault: VaultResource, editedByUser?: Attributes ) { super(DataSourceResource.model, blob); this.editedByUser = editedByUser; + this.vault = vault; } - static async makeNew(blob: CreationAttributes) { + static async makeNew( + blob: Omit, "vaultId">, + vault: VaultResource + ) { const datasource = await DataSource.create(blob); - return new this(DataSourceResource.model, datasource.get()); + return new this(DataSourceResource.model, datasource.get(), vault); } + // Fetching. + private static getOptions(options?: FetchDataSourceOptions) { const result: FindOptions = {}; @@ -76,68 +87,75 @@ export class DataSourceResource extends BaseResource { return result; } + private static async baseFetch( + auth: Authenticator, + where: WhereOptions, + options?: FindOptions + ): Promise { + const includeClauses: Includeable[] = [ + { + model: VaultResource.model, + as: "vault", + }, + ]; + + if (options?.include) { + if (Array.isArray(options.include)) { + includeClauses.push(...options.include); + } else { + includeClauses.push(options.include); + } + } + + const blobs = await this.model.findAll({ + where: { + ...where, + workspaceId: auth.getNonNullableWorkspace().id, + }, + include: includeClauses, + }); + + const dataSources = blobs.map((b) => { + const vault = new VaultResource(VaultResource.model, b.vault.get()); + + return new this(this.model, b.get(), vault, b.editedByUser?.get()); + }); + + return removeNulls(dataSources); + } + static async fetchByName( auth: Authenticator, name: string, options?: Omit ): Promise { - const owner = await auth.getNonNullableWorkspace(); - const datasource = await this.model.findOne({ - where: { - workspaceId: owner.id, + const [dataSource] = await this.baseFetch( + auth, + { name, }, - ...this.getOptions(options), - }); - - if (!datasource) { - return null; - } - return new this( - DataSourceResource.model, - datasource.get(), - datasource.editedByUser?.get() + this.getOptions(options) ); + + return dataSource ?? null; } static async listByWorkspace( auth: Authenticator, options?: FetchDataSourceOptions ): Promise { - const owner = await auth.getNonNullableWorkspace(); - const datasources = await this.model.findAll({ - where: { - workspaceId: owner.id, - }, - ...this.getOptions(options), - }); - - return datasources.map( - (datasource) => - new this( - DataSourceResource.model, - datasource.get(), - datasource.editedByUser?.get() - ) - ); + return this.baseFetch(auth, {}, this.getOptions(options)); } static async listByWorkspaceIdAndNames( - workspace: LightWorkspaceType, + auth: Authenticator, names: string[] ): Promise { - const datasources = await this.model.findAll({ - where: { - workspaceId: workspace.id, - name: { - [Op.in]: names, - }, + return this.baseFetch(auth, { + name: { + [Op.in]: names, }, }); - - return datasources.map( - (datasource) => new this(DataSourceResource.model, datasource.get()) - ); } static async listByConnectorProvider( @@ -145,17 +163,12 @@ export class DataSourceResource extends BaseResource { connectorProvider: ConnectorProvider, options?: FetchDataSourceOptions ): Promise { - const owner = await auth.getNonNullableWorkspace(); - const datasources = await this.model.findAll({ - where: { - workspaceId: owner.id, + return this.baseFetch( + auth, + { connectorProvider, }, - ...this.getOptions(options), - }); - - return datasources.map( - (datasource) => new this(DataSourceResource.model, datasource.get()) + this.getOptions(options) ); } @@ -222,6 +235,12 @@ export class DataSourceResource extends BaseResource { ); } + // Permissions. + + acl(): ACLType { + return this.vault.acl(); + } + // Serialization. toJSON(): DataSourceType { diff --git a/front/lib/resources/data_source_view_resource.ts b/front/lib/resources/data_source_view_resource.ts index 9d2f64dd2b83..d63f5ff7d661 100644 --- a/front/lib/resources/data_source_view_resource.ts +++ b/front/lib/resources/data_source_view_resource.ts @@ -2,9 +2,9 @@ // This design will be moved up to BaseResource once we transition away from Sequelize. // eslint-disable-next-line @typescript-eslint/no-empty-interface import type { + ACLType, DataSourceType, DataSourceViewType, - LightWorkspaceType, ModelId, Result, } from "@dust-tt/types"; @@ -14,6 +14,7 @@ import type { CreationAttributes, ModelStatic, Transaction, + WhereOptions, } from "sequelize"; import type { Authenticator } from "@app/lib/auth"; @@ -21,8 +22,12 @@ import { BaseResource } from "@app/lib/resources/base_resource"; import type { DataSourceResource } from "@app/lib/resources/data_source_resource"; import { DataSourceViewModel } from "@app/lib/resources/storage/models/data_source_view"; import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types"; -import { getResourceIdFromSId, makeSId } from "@app/lib/resources/string_ids"; -import type { VaultResource } from "@app/lib/resources/vault_resource"; +import { + getResourceIdFromSId, + isResourceSId, + makeSId, +} from "@app/lib/resources/string_ids"; +import { VaultResource } from "@app/lib/resources/vault_resource"; // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging export interface DataSourceViewResource @@ -31,19 +36,27 @@ export interface DataSourceViewResource export class DataSourceViewResource extends BaseResource { static model: ModelStatic = DataSourceViewModel; + readonly vault: VaultResource; + constructor( model: ModelStatic, - blob: Attributes + blob: Attributes, + vault: VaultResource ) { super(DataSourceViewModel, blob); + + this.vault = vault; } // Creation. - private static async makeNew(blob: CreationAttributes) { + private static async makeNew( + blob: Omit, "vaultId">, + vault: VaultResource + ) { const key = await DataSourceViewResource.model.create(blob); - return new this(DataSourceViewResource.model, key.get()); + return new this(DataSourceViewResource.model, key.get(), vault); } static async createViewInVaultFromDataSource( @@ -51,11 +64,14 @@ export class DataSourceViewResource extends BaseResource { dataSource: DataSourceType, parentsIn: string[] ) { - return this.makeNew({ - dataSourceId: dataSource.id, - parentsIn, - workspaceId: vault.workspaceId, - }); + return this.makeNew( + { + dataSourceId: dataSource.id, + parentsIn, + workspaceId: vault.workspaceId, + }, + vault + ); } // For now, we create a default view for all data sources in the global vault. @@ -64,42 +80,55 @@ export class DataSourceViewResource extends BaseResource { vault: VaultResource, dataSource: DataSourceResource ) { - return this.makeNew({ - dataSourceId: dataSource.id, - parentsIn: null, - vaultId: vault.id, - workspaceId: vault.workspaceId, - }); + return this.makeNew( + { + dataSourceId: dataSource.id, + parentsIn: null, + workspaceId: vault.workspaceId, + }, + vault + ); } // Fetching. - static async listByWorkspace(auth: Authenticator) { + private static async baseFetch( + auth: Authenticator, + where: Omit, "workspaceId"> = {} + ) { const blobs = await this.model.findAll({ where: { + ...where, workspaceId: auth.getNonNullableWorkspace().id, }, + include: [ + { + model: VaultResource.model, + as: "vault", + }, + ], + }); + + return blobs.map((b) => { + const vault = new VaultResource(VaultResource.model, b.vault.get()); + + return new this(this.model, b.get(), vault); }); + } - return blobs.map((b) => new this(this.model, b.get())); + static async listByWorkspace(auth: Authenticator) { + return this.baseFetch(auth); } static async listForDataSourcesInVault( - owner: LightWorkspaceType, + auth: Authenticator, dataSources: DataSourceResource[], - vault: VaultResource, - { transaction }: { transaction?: Transaction } = {} + vault: VaultResource ) { - const blobs = await this.model.findAll({ - where: { - workspaceId: owner.id, - dataSourceId: dataSources.map((ds) => ds.id), - vaultId: vault.id, - }, - transaction, + return this.baseFetch(auth, { + dataSourceId: dataSources.map((ds) => ds.id), + vaultId: vault.id, }); - - return blobs.map((b) => new this(this.model, b.get())); } static async fetchById(auth: Authenticator, id: string) { @@ -108,20 +137,19 @@ export class DataSourceViewResource extends BaseResource { return null; } - const blob = await this.model.findOne({ - where: { - workspaceId: auth.getNonNullableWorkspace().id, - id: fileModelId, - }, - }); - if (!blob) { - return null; - } + const [dataSource] = await this.baseFetch(auth, { id: fileModelId }); - // Use `.get` to extract model attributes, omitting Sequelize instance metadata. - return new this(this.model, blob.get()); + return dataSource ?? null; } + // Peer fetching. + + // async fetchDataSource( + // auth: Authenticator + // ): Promise { + // return DataSourceResource.fetchByModelId(auth, this.dataSourceId); + // } + // Deletion. async delete( @@ -167,6 +195,8 @@ export class DataSourceViewResource extends BaseResource { }); } + // sId logic. + get sId(): string { return DataSourceViewResource.modelIdToSId({ id: this.id, @@ -187,6 +217,16 @@ export class DataSourceViewResource extends BaseResource { }); } + static isDataSourceViewSId(sId: string): boolean { + return isResourceSId("data_source_view", sId); + } + + // Permissions. + + acl(): ACLType { + return this.vault.acl(); + } + // Serialization logic. toJSON(): DataSourceViewType { diff --git a/front/lib/resources/storage/models/data_source_view.ts b/front/lib/resources/storage/models/data_source_view.ts index 543958fd2fb8..24fd30604d59 100644 --- a/front/lib/resources/storage/models/data_source_view.ts +++ b/front/lib/resources/storage/models/data_source_view.ts @@ -3,6 +3,7 @@ import type { ForeignKey, InferAttributes, InferCreationAttributes, + NonAttribute, } from "sequelize"; import { DataTypes, Model } from "sequelize"; @@ -24,6 +25,9 @@ export class DataSourceViewModel extends Model< declare dataSourceId: ForeignKey; declare vaultId: ForeignKey; declare workspaceId: ForeignKey; + + declare dataSource: NonAttribute; + declare vault: NonAttribute; } DataSourceViewModel.init( { diff --git a/front/lib/resources/string_ids.ts b/front/lib/resources/string_ids.ts index d8d7f3f2a44f..ca500d4e6bc1 100644 --- a/front/lib/resources/string_ids.ts +++ b/front/lib/resources/string_ids.ts @@ -104,6 +104,13 @@ export function getResourceIdFromSId(sId: string): ModelId | null { return sIdsRes.value.resourceId; } +export function isResourceSId( + resourceName: ResourceNameType, + sId: string +): boolean { + return sId.startsWith(RESOURCES_PREFIX[resourceName]); +} + // Legacy behavior. /** diff --git a/front/migrations/20240731_backfill_views_in_agent_data_source_configurations.ts b/front/migrations/20240731_backfill_views_in_agent_data_source_configurations.ts index 8317199b0d1d..66db11c83044 100644 --- a/front/migrations/20240731_backfill_views_in_agent_data_source_configurations.ts +++ b/front/migrations/20240731_backfill_views_in_agent_data_source_configurations.ts @@ -31,7 +31,7 @@ async function backfillViewsInAgentDataSourceConfigurationForWorkspace( // Retrieve data source views for managed data sources. const dataSourceViews = await DataSourceViewResource.listForDataSourcesInVault( - auth.getNonNullableWorkspace(), + auth, managedDataSources, globalVault ); diff --git a/front/pages/api/w/[wId]/data_sources/index.ts b/front/pages/api/w/[wId]/data_sources/index.ts index 3c8656428c54..e186bfc0f8b2 100644 --- a/front/pages/api/w/[wId]/data_sources/index.ts +++ b/front/pages/api/w/[wId]/data_sources/index.ts @@ -172,15 +172,17 @@ async function handler( } const globalVault = await VaultResource.fetchWorkspaceGlobalVault(auth); - const ds = await DataSourceResource.makeNew({ - name: req.body.name, - description: description, - dustAPIProjectId: dustProject.value.project.project_id.toString(), - workspaceId: owner.id, - assistantDefaultSelected: req.body.assistantDefaultSelected, - editedByUserId: user.id, - vaultId: globalVault.id, - }); + const ds = await DataSourceResource.makeNew( + { + name: req.body.name, + description: description, + dustAPIProjectId: dustProject.value.project.project_id.toString(), + workspaceId: owner.id, + assistantDefaultSelected: req.body.assistantDefaultSelected, + editedByUserId: user.id, + }, + globalVault + ); const dataSourceType = await getDataSource(auth, ds.name); if (dataSourceType) { diff --git a/front/pages/api/w/[wId]/data_sources/managed.ts b/front/pages/api/w/[wId]/data_sources/managed.ts index e4b1e80c2ff2..a265cb8dfa4e 100644 --- a/front/pages/api/w/[wId]/data_sources/managed.ts +++ b/front/pages/api/w/[wId]/data_sources/managed.ts @@ -331,16 +331,18 @@ async function handler( const vault = await (provider === "webcrawler" ? VaultResource.fetchWorkspaceGlobalVault(auth) : VaultResource.fetchWorkspaceSystemVault(auth)); - const dataSource = await DataSourceResource.makeNew({ - assistantDefaultSelected, - connectorProvider: provider, - description: dataSourceDescription, - dustAPIProjectId: dustProject.value.project.project_id.toString(), - editedByUserId: user.id, - name: dataSourceName, - vaultId: vault.id, - workspaceId: owner.id, - }); + const dataSource = await DataSourceResource.makeNew( + { + assistantDefaultSelected, + connectorProvider: provider, + description: dataSourceDescription, + dustAPIProjectId: dustProject.value.project.project_id.toString(), + editedByUserId: user.id, + name: dataSourceName, + workspaceId: owner.id, + }, + vault + ); // For managed data source, we create a default view in the workspace vault. if (dataSource.isManaged()) { diff --git a/types/src/front/auth.ts b/types/src/front/auth.ts new file mode 100644 index 000000000000..51643d90ca0e --- /dev/null +++ b/types/src/front/auth.ts @@ -0,0 +1,11 @@ +import { GroupType } from "./groups"; +import { LightWorkspaceType } from "./user"; + +// Authenticator is a concept that lives in front, +// but it's used when doing api calls to other services. +// This interface is a cheap way to represent the concept of an authenticator within types. +export interface BaseAuthenticator { + groups: () => GroupType[]; + + getNonNullableWorkspace: () => LightWorkspaceType; +} diff --git a/types/src/front/groups.ts b/types/src/front/groups.ts index d431291a2aea..39aa3c688904 100644 --- a/types/src/front/groups.ts +++ b/types/src/front/groups.ts @@ -30,8 +30,8 @@ export function isGlobalGroupType(value: SupportedGroupType): boolean { export type GroupType = { id: ModelId; + name: string; sId: string; type: SupportedGroupType; - name: string; workspaceId: ModelId; }; diff --git a/types/src/index.ts b/types/src/index.ts index 0a89ac5906a7..bcef466992f9 100644 --- a/types/src/index.ts +++ b/types/src/index.ts @@ -32,6 +32,7 @@ export * from "./front/assistant/avatar"; export * from "./front/assistant/builder"; export * from "./front/assistant/conversation"; export * from "./front/assistant/templates"; +export * from "./front/auth"; export * from "./front/content_fragment"; export * from "./front/data_source"; export * from "./front/data_source_view"; From 8e0689d549eea6d5c33d187d85616e5b5a69c1f1 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 15:01:07 +0200 Subject: [PATCH 2/9] Introduce ResourceWithVault --- front/lib/resources/data_source_resource.ts | 103 ++++---------- .../resources/data_source_view_resource.ts | 56 ++------ front/lib/resources/resource_with_vault.ts | 128 ++++++++++++++++++ 3 files changed, 168 insertions(+), 119 deletions(-) create mode 100644 front/lib/resources/resource_with_vault.ts diff --git a/front/lib/resources/data_source_resource.ts b/front/lib/resources/data_source_resource.ts index 363e22dbab67..6d8dc1c38e3a 100644 --- a/front/lib/resources/data_source_resource.ts +++ b/front/lib/resources/data_source_resource.ts @@ -1,28 +1,21 @@ -import type { - ACLType, - ConnectorProvider, - DataSourceType, - Result, -} from "@dust-tt/types"; -import { Err, formatUserFullName, Ok, removeNulls } from "@dust-tt/types"; +import type { ConnectorProvider, DataSourceType, Result } from "@dust-tt/types"; +import { Err, formatUserFullName, Ok } from "@dust-tt/types"; import type { Attributes, CreationAttributes, - FindOptions, - Includeable, ModelStatic, Transaction, - WhereOptions, } from "sequelize"; import { Op } from "sequelize"; import type { Authenticator } from "@app/lib/auth"; import { DataSource } from "@app/lib/models/data_source"; import { User } from "@app/lib/models/user"; -import { BaseResource } from "@app/lib/resources/base_resource"; import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource"; +import type { ResourceFindOptions } from "@app/lib/resources/resource_with_vault"; +import { ResourceWithVault } from "@app/lib/resources/resource_with_vault"; import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types"; -import { VaultResource } from "@app/lib/resources/vault_resource"; +import type { VaultResource } from "@app/lib/resources/vault_resource"; export type FetchDataSourceOptions = { includeEditedBy?: boolean; @@ -36,21 +29,20 @@ export type FetchDataSourceOptions = { export interface DataSourceResource extends ReadonlyAttributesType {} // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging -export class DataSourceResource extends BaseResource { +export class DataSourceResource extends ResourceWithVault { static model: ModelStatic = DataSource; readonly editedByUser: Attributes | undefined; - readonly vault: VaultResource; constructor( model: ModelStatic, blob: Attributes, vault: VaultResource, - editedByUser?: Attributes + { editedByUser }: { editedByUser?: Attributes } = {} ) { - super(DataSourceResource.model, blob); + super(DataSourceResource.model, blob, vault); + this.editedByUser = editedByUser; - this.vault = vault; } static async makeNew( @@ -64,11 +56,13 @@ export class DataSourceResource extends BaseResource { // Fetching. - private static getOptions(options?: FetchDataSourceOptions) { - const result: FindOptions = {}; + private static getOptions( + options?: FetchDataSourceOptions + ): ResourceFindOptions { + const result: ResourceFindOptions = {}; if (options?.includeEditedBy) { - result.include = [ + result.includes = [ { model: User, as: "editedByUser", @@ -87,55 +81,17 @@ export class DataSourceResource extends BaseResource { return result; } - private static async baseFetch( - auth: Authenticator, - where: WhereOptions, - options?: FindOptions - ): Promise { - const includeClauses: Includeable[] = [ - { - model: VaultResource.model, - as: "vault", - }, - ]; - - if (options?.include) { - if (Array.isArray(options.include)) { - includeClauses.push(...options.include); - } else { - includeClauses.push(options.include); - } - } - - const blobs = await this.model.findAll({ - where: { - ...where, - workspaceId: auth.getNonNullableWorkspace().id, - }, - include: includeClauses, - }); - - const dataSources = blobs.map((b) => { - const vault = new VaultResource(VaultResource.model, b.vault.get()); - - return new this(this.model, b.get(), vault, b.editedByUser?.get()); - }); - - return removeNulls(dataSources); - } - static async fetchByName( auth: Authenticator, name: string, options?: Omit ): Promise { - const [dataSource] = await this.baseFetch( - auth, - { + const [dataSource] = await this.baseFetch(auth, { + ...this.getOptions(options), + where: { name, }, - this.getOptions(options) - ); + }); return dataSource ?? null; } @@ -144,7 +100,7 @@ export class DataSourceResource extends BaseResource { auth: Authenticator, options?: FetchDataSourceOptions ): Promise { - return this.baseFetch(auth, {}, this.getOptions(options)); + return this.baseFetch(auth, this.getOptions(options)); } static async listByWorkspaceIdAndNames( @@ -152,8 +108,10 @@ export class DataSourceResource extends BaseResource { names: string[] ): Promise { return this.baseFetch(auth, { - name: { - [Op.in]: names, + where: { + name: { + [Op.in]: names, + }, }, }); } @@ -163,13 +121,12 @@ export class DataSourceResource extends BaseResource { connectorProvider: ConnectorProvider, options?: FetchDataSourceOptions ): Promise { - return this.baseFetch( - auth, - { + return this.baseFetch(auth, { + ...this.getOptions(options), + where: { connectorProvider, }, - this.getOptions(options) - ); + }); } async delete( @@ -235,12 +192,6 @@ export class DataSourceResource extends BaseResource { ); } - // Permissions. - - acl(): ACLType { - return this.vault.acl(); - } - // Serialization. toJSON(): DataSourceType { diff --git a/front/lib/resources/data_source_view_resource.ts b/front/lib/resources/data_source_view_resource.ts index d63f5ff7d661..a02356fa07ab 100644 --- a/front/lib/resources/data_source_view_resource.ts +++ b/front/lib/resources/data_source_view_resource.ts @@ -2,7 +2,6 @@ // This design will be moved up to BaseResource once we transition away from Sequelize. // eslint-disable-next-line @typescript-eslint/no-empty-interface import type { - ACLType, DataSourceType, DataSourceViewType, ModelId, @@ -14,12 +13,11 @@ import type { CreationAttributes, ModelStatic, Transaction, - WhereOptions, } from "sequelize"; import type { Authenticator } from "@app/lib/auth"; -import { BaseResource } from "@app/lib/resources/base_resource"; import type { DataSourceResource } from "@app/lib/resources/data_source_resource"; +import { ResourceWithVault } from "@app/lib/resources/resource_with_vault"; import { DataSourceViewModel } from "@app/lib/resources/storage/models/data_source_view"; import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types"; import { @@ -27,25 +25,21 @@ import { isResourceSId, makeSId, } from "@app/lib/resources/string_ids"; -import { VaultResource } from "@app/lib/resources/vault_resource"; +import type { VaultResource } from "@app/lib/resources/vault_resource"; // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging export interface DataSourceViewResource extends ReadonlyAttributesType {} // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging -export class DataSourceViewResource extends BaseResource { +export class DataSourceViewResource extends ResourceWithVault { static model: ModelStatic = DataSourceViewModel; - readonly vault: VaultResource; - constructor( model: ModelStatic, blob: Attributes, vault: VaultResource ) { - super(DataSourceViewModel, blob); - - this.vault = vault; + super(DataSourceViewModel, blob, vault); } // Creation. @@ -92,30 +86,6 @@ export class DataSourceViewResource extends BaseResource { // Fetching. - private static async baseFetch( - auth: Authenticator, - where: Omit, "workspaceId"> = {} - ) { - const blobs = await this.model.findAll({ - where: { - ...where, - workspaceId: auth.getNonNullableWorkspace().id, - }, - include: [ - { - model: VaultResource.model, - as: "vault", - }, - ], - }); - - return blobs.map((b) => { - const vault = new VaultResource(VaultResource.model, b.vault.get()); - - return new this(this.model, b.get(), vault); - }); - } - static async listByWorkspace(auth: Authenticator) { return this.baseFetch(auth); } @@ -126,8 +96,10 @@ export class DataSourceViewResource extends BaseResource { vault: VaultResource ) { return this.baseFetch(auth, { - dataSourceId: dataSources.map((ds) => ds.id), - vaultId: vault.id, + where: { + dataSourceId: dataSources.map((ds) => ds.id), + vaultId: vault.id, + }, }); } @@ -137,7 +109,11 @@ export class DataSourceViewResource extends BaseResource { return null; } - const [dataSource] = await this.baseFetch(auth, { id: fileModelId }); + const [dataSource] = await this.baseFetch(auth, { + where: { + id: fileModelId, + }, + }); return dataSource ?? null; } @@ -221,12 +197,6 @@ export class DataSourceViewResource extends BaseResource { return isResourceSId("data_source_view", sId); } - // Permissions. - - acl(): ACLType { - return this.vault.acl(); - } - // Serialization logic. toJSON(): DataSourceViewType { diff --git a/front/lib/resources/resource_with_vault.ts b/front/lib/resources/resource_with_vault.ts new file mode 100644 index 000000000000..7fd65d22ee4c --- /dev/null +++ b/front/lib/resources/resource_with_vault.ts @@ -0,0 +1,128 @@ +import type { + Attributes, + FindOptions, + ForeignKey, + Includeable, + ModelStatic, + NonAttribute, + WhereOptions, +} from "sequelize"; +import { Model } from "sequelize"; + +import type { Authenticator } from "@app/lib/auth"; +import type { Workspace } from "@app/lib/models/workspace"; +import { BaseResource } from "@app/lib/resources/base_resource"; +import type { VaultModel } from "@app/lib/resources/storage/models/vaults"; +import { VaultResource } from "@app/lib/resources/vault_resource"; + +// Interface to enforce workspaceId and vaultId. +interface ModelWithVault { + workspaceId: ForeignKey; + vaultId: ForeignKey; + vault: NonAttribute; +} + +export type NonAttributeKeys = { + [K in keyof M]: M[K] extends NonAttribute> ? K : never; +}[keyof M] & + string; + +type InferIncludeType = { + [K in NonAttributeKeys]: M[K] extends NonAttribute + ? T extends Model + ? T + : never + : never; +}; + +export type TypedIncludeable = { + [K in NonAttributeKeys]: { + model: ModelStatic[K]>; + as: K; + }; +}[NonAttributeKeys]; + +export interface ResourceFindOptions { + includes?: TypedIncludeable[]; + limit?: number; + order?: FindOptions["order"]; + where?: WhereOptions; +} + +export abstract class ResourceWithVault< + M extends Model & ModelWithVault, +> extends BaseResource { + protected constructor( + model: ModelStatic, + blob: Attributes, + public readonly vault: VaultResource + ) { + super(model, blob); + } + + protected static async baseFetch< + T extends ResourceWithVault, + M extends Model & ModelWithVault, + IncludeType extends Partial>, + >( + this: { + new ( + model: ModelStatic, + blob: Attributes, + vault: VaultResource, + includes?: IncludeType + ): T; + } & { model: ModelStatic }, + auth: Authenticator, + { includes, limit, order, where }: ResourceFindOptions = {} + ): Promise { + const includeClauses: Includeable[] = [ + { + model: VaultResource.model, + as: "vault", + }, + ...(includes || []), + ]; + + const blobs = await this.model.findAll({ + where: { + ...where, + workspaceId: auth.getNonNullableWorkspace().id, + } as WhereOptions, + include: includeClauses, + limit, + order, + }); + + return blobs.map((b) => { + const vault = new VaultResource(VaultResource.model, b.vault.get()); + + const includedResults = (includes || []).reduce( + (acc, current) => { + if ( + typeof current === "object" && + "as" in current && + typeof current.as === "string" + ) { + const key = current.as as keyof IncludeType; + // Only handle other includes if they are not vault. + if (key !== "vault") { + const includedModel = b[key as keyof typeof b]; + if (includedModel instanceof Model) { + acc[key] = includedModel.get(); + } + } + } + return acc; + }, + {} as IncludeType + ); + + return new this(this.model, b.get(), vault, includedResults); + }); + } + + acl() { + this.vault.acl(); + } +} From 8abb791234b7ddf1debe9a6892ee720f7b834714 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 17:34:33 +0200 Subject: [PATCH 3/9] :sparkles: --- front/lib/api/data_sources.ts | 3 ++- front/lib/document_tracker.ts | 13 ++++++++----- .../update_tracked_documents/lib.ts | 2 +- front/lib/resources/data_source_resource.ts | 14 +++++++++++++- 4 files changed, 24 insertions(+), 8 deletions(-) diff --git a/front/lib/api/data_sources.ts b/front/lib/api/data_sources.ts index 59b19a013cfa..d028d81f99c6 100644 --- a/front/lib/api/data_sources.ts +++ b/front/lib/api/data_sources.ts @@ -88,7 +88,8 @@ export async function updateDataSourceEditedBy( }); } - const dataSourceResource = await DataSourceResource.fetchByModelId( + const dataSourceResource = await DataSourceResource.fetchByModelIdWithAuth( + auth, dataSource.id ); diff --git a/front/lib/document_tracker.ts b/front/lib/document_tracker.ts index 455e6ab363a0..b734e70f8fb8 100644 --- a/front/lib/document_tracker.ts +++ b/front/lib/document_tracker.ts @@ -1,7 +1,7 @@ import { CoreAPI } from "@dust-tt/types"; import { literal, Op } from "sequelize"; -import { Authenticator } from "@app/lib/auth"; +import type { Authenticator } from "@app/lib/auth"; import { TrackedDocument } from "@app/lib/models/doc_tracker"; import { User } from "@app/lib/models/user"; import { Workspace } from "@app/lib/models/workspace"; @@ -12,14 +12,19 @@ import logger from "@app/logger/logger"; import config from "./api/config"; export async function updateTrackedDocuments( + auth: Authenticator, dataSourceId: number, documentId: string, documentContent: string ) { - const dataSource = await DataSourceResource.fetchByModelId(dataSourceId); + const dataSource = await DataSourceResource.fetchByModelIdWithAuth( + auth, + dataSourceId + ); if (!dataSource) { throw new Error(`Could not find data source with id ${dataSourceId}`); } + if (!dataSource.workspaceId) { throw new Error( `Data source with id ${dataSourceId} has no workspace id set` @@ -31,9 +36,7 @@ export async function updateTrackedDocuments( `Could not find workspace with id ${dataSource.workspaceId}` ); } - const auth = await Authenticator.internalBuilderForWorkspace( - workspaceModel.sId - ); + const owner = auth.workspace(); if (!owner) { throw new Error( diff --git a/front/lib/documents_post_process_hooks/hooks/document_tracker/update_tracked_documents/lib.ts b/front/lib/documents_post_process_hooks/hooks/document_tracker/update_tracked_documents/lib.ts index 721124b18232..3c2aa520f11c 100644 --- a/front/lib/documents_post_process_hooks/hooks/document_tracker/update_tracked_documents/lib.ts +++ b/front/lib/documents_post_process_hooks/hooks/document_tracker/update_tracked_documents/lib.ts @@ -99,7 +99,7 @@ export async function documentTrackerUpdateTrackedDocumentsOnUpsert({ ) ) { logger.info("Updating tracked documents."); - await updateTrackedDocuments(dataSource.id, documentId, documentText); + await updateTrackedDocuments(auth, dataSource.id, documentId, documentText); } } diff --git a/front/lib/resources/data_source_resource.ts b/front/lib/resources/data_source_resource.ts index 6d8dc1c38e3a..15fa1eefec68 100644 --- a/front/lib/resources/data_source_resource.ts +++ b/front/lib/resources/data_source_resource.ts @@ -1,4 +1,9 @@ -import type { ConnectorProvider, DataSourceType, Result } from "@dust-tt/types"; +import type { + ConnectorProvider, + DataSourceType, + ModelId, + Result, +} from "@dust-tt/types"; import { Err, formatUserFullName, Ok } from "@dust-tt/types"; import type { Attributes, @@ -129,6 +134,13 @@ export class DataSourceResource extends ResourceWithVault { }); } + // TODO(20240801 flav): Refactor this to make auth required on all fetchers. + static async fetchByModelIdWithAuth(auth: Authenticator, id: ModelId) { + const [dataSource] = await this.baseFetch(auth, { where: { id } }); + + return dataSource ?? null; + } + async delete( auth: Authenticator, transaction?: Transaction From df17b8413b2e5f299233e9d68204928a0a721b48 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 17:35:37 +0200 Subject: [PATCH 4/9] Add fetchDataSource on DataSourceView --- front/lib/resources/data_source_view_resource.ts | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/front/lib/resources/data_source_view_resource.ts b/front/lib/resources/data_source_view_resource.ts index a02356fa07ab..5c304d2781d9 100644 --- a/front/lib/resources/data_source_view_resource.ts +++ b/front/lib/resources/data_source_view_resource.ts @@ -16,7 +16,7 @@ import type { } from "sequelize"; import type { Authenticator } from "@app/lib/auth"; -import type { DataSourceResource } from "@app/lib/resources/data_source_resource"; +import { DataSourceResource } from "@app/lib/resources/data_source_resource"; import { ResourceWithVault } from "@app/lib/resources/resource_with_vault"; import { DataSourceViewModel } from "@app/lib/resources/storage/models/data_source_view"; import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types"; @@ -120,11 +120,11 @@ export class DataSourceViewResource extends ResourceWithVault { - // return DataSourceResource.fetchByModelId(auth, this.dataSourceId); - // } + async fetchDataSource( + auth: Authenticator + ): Promise { + return DataSourceResource.fetchByModelIdWithAuth(auth, this.dataSourceId); + } // Deletion. From a6c699a4b4f601da15b0071fd58ab13eac4336d6 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 17:49:52 +0200 Subject: [PATCH 5/9] :see_no_evil: --- front/lib/resources/resource_with_vault.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/front/lib/resources/resource_with_vault.ts b/front/lib/resources/resource_with_vault.ts index 7fd65d22ee4c..99824fabf9a3 100644 --- a/front/lib/resources/resource_with_vault.ts +++ b/front/lib/resources/resource_with_vault.ts @@ -123,6 +123,6 @@ export abstract class ResourceWithVault< } acl() { - this.vault.acl(); + return this.vault.acl(); } } From 12957494809f8561f2898b54ad5917fcdd70ba7a Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 18:56:57 +0200 Subject: [PATCH 6/9] Address comments from review --- front/lib/api/assistant/configuration.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index b2fc3955a673..78d71e4312a4 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -1408,10 +1408,12 @@ async function _createAgentDataSourcesConfigData( // Although we have the capability to support multiple workspaces, // currently, we only support one workspace, which is the one the user is in. // This allows us to use the current authenticator to fetch resources. - const allWorkspaceIds = _.uniqBy(dataSourceConfigurations, "workspaceSId"); + const allWorkspaceIds = [ + ...new Set(dataSourceConfigurations.map((dsc) => dsc.workspaceId)), + ]; const hasUniqueAccessibleWorkspace = allWorkspaceIds.length === 1 && - auth.getNonNullableWorkspace().sId === allWorkspaceIds[0].workspaceId; + auth.getNonNullableWorkspace().sId === allWorkspaceIds[0]; if (!hasUniqueAccessibleWorkspace) { throw new Error( "Can't create AgentDataSourcesConfig for retrieval: Multiple workspaces." From 4453eb3524997c5c9c49b236d56967af300c5ba3 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 19:01:21 +0200 Subject: [PATCH 7/9] :see_no_evil: --- front/lib/resources/data_source_resource.ts | 5 ++++- front/lib/resources/data_source_view_resource.ts | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/front/lib/resources/data_source_resource.ts b/front/lib/resources/data_source_resource.ts index 15fa1eefec68..f8cc3577287a 100644 --- a/front/lib/resources/data_source_resource.ts +++ b/front/lib/resources/data_source_resource.ts @@ -54,7 +54,10 @@ export class DataSourceResource extends ResourceWithVault { blob: Omit, "vaultId">, vault: VaultResource ) { - const datasource = await DataSource.create(blob); + const datasource = await DataSource.create({ + ...blob, + vaultId: vault.id, + }); return new this(DataSourceResource.model, datasource.get(), vault); } diff --git a/front/lib/resources/data_source_view_resource.ts b/front/lib/resources/data_source_view_resource.ts index 5c304d2781d9..2819c16ffe15 100644 --- a/front/lib/resources/data_source_view_resource.ts +++ b/front/lib/resources/data_source_view_resource.ts @@ -48,7 +48,10 @@ export class DataSourceViewResource extends ResourceWithVault, "vaultId">, vault: VaultResource ) { - const key = await DataSourceViewResource.model.create(blob); + const key = await DataSourceViewResource.model.create({ + ...blob, + vaultId: vault.id, + }); return new this(DataSourceViewResource.model, key.get(), vault); } From ce8b72540900836b31daf2aded7411d4bf7e24cd Mon Sep 17 00:00:00 2001 From: Flavien David Date: Thu, 1 Aug 2024 19:28:42 +0200 Subject: [PATCH 8/9] :scissors: --- front/lib/auth.ts | 3 +-- types/src/front/auth.ts | 11 ----------- types/src/index.ts | 1 - 3 files changed, 1 insertion(+), 14 deletions(-) delete mode 100644 types/src/front/auth.ts diff --git a/front/lib/auth.ts b/front/lib/auth.ts index caec9ead77bf..8af4c6297cf5 100644 --- a/front/lib/auth.ts +++ b/front/lib/auth.ts @@ -13,7 +13,6 @@ import type { PlanType, SubscriptionType } from "@dust-tt/types"; import type { DustAPICredentials } from "@dust-tt/types"; import type { Result } from "@dust-tt/types"; import type { APIErrorWithStatusCode } from "@dust-tt/types"; -import type { BaseAuthenticator } from "@dust-tt/types"; import { Err, groupHasPermission, @@ -61,7 +60,7 @@ const DUST_INTERNAL_EMAIL_REGEXP = /^[^@]+@dust\.tt$/; * It explicitely does not store a reference to the current user to make sure our permissions are * workspace oriented. Use `getUserFromSession` if needed. */ -export class Authenticator implements BaseAuthenticator { +export class Authenticator { _flags: WhitelistableFeature[]; _key?: KeyAuthType; _role: RoleType; diff --git a/types/src/front/auth.ts b/types/src/front/auth.ts deleted file mode 100644 index 51643d90ca0e..000000000000 --- a/types/src/front/auth.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { GroupType } from "./groups"; -import { LightWorkspaceType } from "./user"; - -// Authenticator is a concept that lives in front, -// but it's used when doing api calls to other services. -// This interface is a cheap way to represent the concept of an authenticator within types. -export interface BaseAuthenticator { - groups: () => GroupType[]; - - getNonNullableWorkspace: () => LightWorkspaceType; -} diff --git a/types/src/index.ts b/types/src/index.ts index 8efb3e848098..1cc99b31acc7 100644 --- a/types/src/index.ts +++ b/types/src/index.ts @@ -32,7 +32,6 @@ export * from "./front/assistant/builder"; export * from "./front/assistant/conversation"; export * from "./front/assistant/templates"; export * from "./front/assistant/visualization"; -export * from "./front/auth"; export * from "./front/content_fragment"; export * from "./front/data_source"; export * from "./front/data_source_view"; From 968546ff4028b440c9c437ae2c70413c0e407059 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Fri, 2 Aug 2024 10:20:29 +0200 Subject: [PATCH 9/9] s/baseFetch/baseFetchWithAuthorization --- front/lib/resources/data_source_resource.ts | 12 +++++++----- front/lib/resources/data_source_view_resource.ts | 6 +++--- front/lib/resources/resource_with_vault.ts | 2 +- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/front/lib/resources/data_source_resource.ts b/front/lib/resources/data_source_resource.ts index f8cc3577287a..46a0b870e9f4 100644 --- a/front/lib/resources/data_source_resource.ts +++ b/front/lib/resources/data_source_resource.ts @@ -94,7 +94,7 @@ export class DataSourceResource extends ResourceWithVault { name: string, options?: Omit ): Promise { - const [dataSource] = await this.baseFetch(auth, { + const [dataSource] = await this.baseFetchWithAuthorization(auth, { ...this.getOptions(options), where: { name, @@ -108,14 +108,14 @@ export class DataSourceResource extends ResourceWithVault { auth: Authenticator, options?: FetchDataSourceOptions ): Promise { - return this.baseFetch(auth, this.getOptions(options)); + return this.baseFetchWithAuthorization(auth, this.getOptions(options)); } static async listByWorkspaceIdAndNames( auth: Authenticator, names: string[] ): Promise { - return this.baseFetch(auth, { + return this.baseFetchWithAuthorization(auth, { where: { name: { [Op.in]: names, @@ -129,7 +129,7 @@ export class DataSourceResource extends ResourceWithVault { connectorProvider: ConnectorProvider, options?: FetchDataSourceOptions ): Promise { - return this.baseFetch(auth, { + return this.baseFetchWithAuthorization(auth, { ...this.getOptions(options), where: { connectorProvider, @@ -139,7 +139,9 @@ export class DataSourceResource extends ResourceWithVault { // TODO(20240801 flav): Refactor this to make auth required on all fetchers. static async fetchByModelIdWithAuth(auth: Authenticator, id: ModelId) { - const [dataSource] = await this.baseFetch(auth, { where: { id } }); + const [dataSource] = await this.baseFetchWithAuthorization(auth, { + where: { id }, + }); return dataSource ?? null; } diff --git a/front/lib/resources/data_source_view_resource.ts b/front/lib/resources/data_source_view_resource.ts index 2819c16ffe15..c9f1a6531fc9 100644 --- a/front/lib/resources/data_source_view_resource.ts +++ b/front/lib/resources/data_source_view_resource.ts @@ -90,7 +90,7 @@ export class DataSourceViewResource extends ResourceWithVault ds.id), vaultId: vault.id, @@ -112,7 +112,7 @@ export class DataSourceViewResource extends ResourceWithVault, M extends Model & ModelWithVault, IncludeType extends Partial>,