Introduce ResourceWithVault layer

front/lib/api/assistant/configuration.ts

@@ -1405,6 +1405,21 @@ async function _createAgentDataSourcesConfigData(
     processConfigurationId: ModelId | null;
   }
 ): Promise<AgentDataSourceConfiguration[]> {
+  // 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 = [
+    ...new Set(dataSourceConfigurations.map((dsc) => dsc.workspaceId)),
+  ];
+  const hasUniqueAccessibleWorkspace =
+    allWorkspaceIds.length === 1 &&
+    auth.getNonNullableWorkspace().sId === allWorkspaceIds[0];
+  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 } },
@@ -1465,9 +1480,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
       );
 
@@ -1478,7 +1494,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
         );

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
   );
 

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(

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);
   }
 }
 

front/lib/models/data_source.ts

@@ -34,8 +34,9 @@ export class DataSource extends Model<
   declare workspaceId: ForeignKey<Workspace["id"]>;
   declare vaultId: ForeignKey<VaultModel["id"]>;
 
-  declare workspace: NonAttribute<Workspace>;
   declare editedByUser: NonAttribute<User>;
+  declare vault: NonAttribute<VaultModel>;
+  declare workspace: NonAttribute<Workspace>;
 }
 
 DataSource.init(

front/lib/resources/data_source_resource.ts

@@ -1,14 +1,13 @@
 import type {
   ConnectorProvider,
   DataSourceType,
-  LightWorkspaceType,
+  ModelId,
   Result,
 } from "@dust-tt/types";
 import { Err, formatUserFullName, Ok } from "@dust-tt/types";
 import type {
   Attributes,
   CreationAttributes,
-  FindOptions,
   ModelStatic,
   Transaction,
 } from "sequelize";
@@ -17,9 +16,11 @@ 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 type { VaultResource } from "@app/lib/resources/vault_resource";
 
 export type FetchDataSourceOptions = {
   includeEditedBy?: boolean;
@@ -33,31 +34,43 @@ export type FetchDataSourceOptions = {
 export interface DataSourceResource
   extends ReadonlyAttributesType<DataSource> {}
 // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging
-export class DataSourceResource extends BaseResource<DataSource> {
+export class DataSourceResource extends ResourceWithVault<DataSource> {
   static model: ModelStatic<DataSource> = DataSource;
 
-  editedByUser: Attributes<User> | undefined;
+  readonly editedByUser: Attributes<User> | undefined;
 
   constructor(
     model: ModelStatic<DataSource>,
     blob: Attributes<DataSource>,
-    editedByUser?: Attributes<User>
+    vault: VaultResource,
+    { editedByUser }: { editedByUser?: Attributes<User> } = {}
   ) {
-    super(DataSourceResource.model, blob);
+    super(DataSourceResource.model, blob, vault);
+
     this.editedByUser = editedByUser;
   }
 
-  static async makeNew(blob: CreationAttributes<DataSource>) {
-    const datasource = await DataSource.create(blob);
+  static async makeNew(
+    blob: Omit<CreationAttributes<DataSource>, "vaultId">,
+    vault: VaultResource
+  ) {
+    const datasource = await DataSource.create({
+      ...blob,
+      vaultId: vault.id,
+    });
 
-    return new this(DataSourceResource.model, datasource.get());
+    return new this(DataSourceResource.model, datasource.get(), vault);
   }
 
-  private static getOptions(options?: FetchDataSourceOptions) {
-    const result: FindOptions<DataSourceResource["model"]> = {};
+  // Fetching.
+
+  private static getOptions(
+    options?: FetchDataSourceOptions
+  ): ResourceFindOptions<DataSource> {
+    const result: ResourceFindOptions<DataSource> = {};
 
     if (options?.includeEditedBy) {
-      result.include = [
+      result.includes = [
         {
           model: User,
           as: "editedByUser",
@@ -81,82 +94,56 @@ export class DataSourceResource extends BaseResource<DataSource> {
     name: string,
     options?: Omit<FetchDataSourceOptions, "limit" | "order">
   ): Promise<DataSourceResource | null> {
-    const owner = await auth.getNonNullableWorkspace();
-    const datasource = await this.model.findOne({
+    const [dataSource] = await this.baseFetchWithAuthorization(auth, {
+      ...this.getOptions(options),
       where: {
-        workspaceId: owner.id,
         name,
       },
-      ...this.getOptions(options),
     });
 
-    if (!datasource) {
-      return null;
-    }
-    return new this(
-      DataSourceResource.model,
-      datasource.get(),
-      datasource.editedByUser?.get()
-    );
+    return dataSource ?? null;
   }
 
   static async listByWorkspace(
     auth: Authenticator,
     options?: FetchDataSourceOptions
   ): Promise<DataSourceResource[]> {
-    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.baseFetchWithAuthorization(auth, this.getOptions(options));
   }
 
   static async listByWorkspaceIdAndNames(
-    workspace: LightWorkspaceType,
+    auth: Authenticator,
     names: string[]
   ): Promise<DataSourceResource[]> {
-    const datasources = await this.model.findAll({
+    return this.baseFetchWithAuthorization(auth, {
       where: {
-        workspaceId: workspace.id,
         name: {
           [Op.in]: names,
         },
       },
     });
-
-    return datasources.map(
-      (datasource) => new this(DataSourceResource.model, datasource.get())
-    );
   }
 
   static async listByConnectorProvider(
     auth: Authenticator,
     connectorProvider: ConnectorProvider,
     options?: FetchDataSourceOptions
   ): Promise<DataSourceResource[]> {
-    const owner = await auth.getNonNullableWorkspace();
-    const datasources = await this.model.findAll({
+    return this.baseFetchWithAuthorization(auth, {
+      ...this.getOptions(options),
       where: {
-        workspaceId: owner.id,
         connectorProvider,
       },
-      ...this.getOptions(options),
+    });
+  }
+
+  // TODO(20240801 flav): Refactor this to make auth required on all fetchers.
+  static async fetchByModelIdWithAuth(auth: Authenticator, id: ModelId) {
+    const [dataSource] = await this.baseFetchWithAuthorization(auth, {
+      where: { id },
     });
 
-    return datasources.map(
-      (datasource) => new this(DataSourceResource.model, datasource.get())
-    );
+    return dataSource ?? null;
   }
 
   async delete(