Introducing User Resource + removing User.find callsites (#6222)

* Introducing User Resource * Lint * Lint * Add transaction to fecthByProvider * Lint * Removing a lot more renderUserTypes() * Use UserResource almost everywhere * cleanup * Remove transaction from delete() * Migrate the last few things to userResource * fetchAll* -> list* * Return resource * Rename userRes -> user * userRes ToCIOpayload * Fullname getter * getFullName * fecthbyid * Remove deleteUser() * Move fetchRevokedWorkspace to lib/api * UpdateInfo and updateAuth0Sub * UpdateName * Use toJSON * Migrate to usertype * Add fetchNonNullableByModelId * Remove try catch * fectNonNullable * Introduce unsafeDelete * move getFullname() to getAttribute() * Remove nonNullable * Remove userRes * lint * Updateinfo fix * fullName() * Change throw to error * Removing else * typo
dust-tt · Aug 28, 2024 · ef0a97f · ef0a97f
1 parent 9a0d1b8
commit ef0a97f
Show file tree

Hide file tree

Showing 28 changed files with 498 additions and 316 deletions.
diff --git a/front/admin/cli.ts b/front/admin/cli.ts
@@ -13,10 +13,8 @@ import { getConversation } from "@app/lib/api/assistant/conversation";
 import { renderConversationForModelMultiActions } from "@app/lib/api/assistant/generation";
 import config from "@app/lib/api/config";
 import { getDataSources } from "@app/lib/api/data_sources";
-import { renderUserType } from "@app/lib/api/user";
 import { Authenticator } from "@app/lib/auth";
 import { DataSource } from "@app/lib/models/data_source";
-import { User } from "@app/lib/models/user";
 import { Workspace } from "@app/lib/models/workspace";
 import { FREE_UPGRADED_PLAN_CODE } from "@app/lib/plans/plan_codes";
 import {
@@ -25,6 +23,7 @@ import {
 } from "@app/lib/plans/subscription";
 import { MembershipResource } from "@app/lib/resources/membership_resource";
 import { generateLegacyModelSId } from "@app/lib/resources/string_ids";
+import { UserResource } from "@app/lib/resources/user_resource";
 import logger from "@app/logger/logger";
 
 // `cli` takes an object type and a command as first two arguments and then a list of arguments.
@@ -183,11 +182,7 @@ const user = async (command: string, args: parseArgs.ParsedArgs) => {
         throw new Error("Missing --username argument");
       }
 
-      const users = await User.findAll({
-        where: {
-          username: args.username,
-        },
-      });
+      const users = await UserResource.listByUsername(args.username);
 
       users.forEach((u) => {
         console.log(
@@ -201,12 +196,7 @@ const user = async (command: string, args: parseArgs.ParsedArgs) => {
         throw new Error("Missing --userId argument");
       }
 
-      const u = await User.findOne({
-        where: {
-          id: args.userId,
-        },
-      });
-
+      const u = await UserResource.fetchByModelId(args.userId);
       if (!u) {
         throw new Error(`User not found: userId='${args.userId}'`);
       }
@@ -218,7 +208,7 @@ const user = async (command: string, args: parseArgs.ParsedArgs) => {
       console.log(`  email: ${u.email}`);
 
       const memberships = await MembershipResource.getLatestMemberships({
-        users: [renderUserType(u)],
+        users: [u],
       });
 
       const workspaces = await Workspace.findAll({

diff --git a/front/lib/api/assistant/conversation.ts b/front/lib/api/assistant/conversation.ts
@@ -70,12 +70,12 @@ import {
   Message,
   UserMessage,
 } from "@app/lib/models/assistant/conversation";
-import { User } from "@app/lib/models/user";
 import { countActiveSeatsInWorkspaceCached } from "@app/lib/plans/usage/seats";
 import { ContentFragmentResource } from "@app/lib/resources/content_fragment_resource";
 import { frontSequelize } from "@app/lib/resources/storage";
 import { ContentFragmentModel } from "@app/lib/resources/storage/models/content_fragment";
 import { generateLegacyModelSId } from "@app/lib/resources/string_ids";
+import { UserResource } from "@app/lib/resources/user_resource";
 import { ServerSideTracking } from "@app/lib/tracking/server";
 import logger from "@app/logger/logger";
 import { launchUpdateUsageWorkflow } from "@app/temporal/usage_queue/client";
@@ -751,10 +751,10 @@ export async function* postUserMessage(
         type: "user_message",
         visibility: "visible",
         version: 0,
-        user: user,
-        mentions: mentions,
+        user,
+        mentions,
         content,
-        context: context,
+        context,
         rank: m.rank,
       };
 
@@ -960,11 +960,7 @@ export async function* postUserMessage(
   async function logIfUserUnknown() {
     try {
       if (!user && context.email) {
-        const macthingUser = await User.findOne({
-          where: {
-            email: context.email,
-          },
-        });
+        const macthingUser = await UserResource.fetchByEmail(context.email);
 
         if (!macthingUser) {
           logger.warn(
@@ -1214,7 +1210,7 @@ export async function* editUserMessage(
         type: "user_message",
         visibility: m.visibility,
         version: m.version,
-        user: user,
+        user,
         mentions,
         content,
         context: message.context,

diff --git a/front/lib/api/assistant/messages.ts b/front/lib/api/assistant/messages.ts
@@ -32,9 +32,9 @@ import {
   Message,
   UserMessage,
 } from "@app/lib/models/assistant/conversation";
-import { User } from "@app/lib/models/user";
 import { ContentFragmentResource } from "@app/lib/resources/content_fragment_resource";
 import { ContentFragmentModel } from "@app/lib/resources/storage/models/content_fragment";
+import { UserResource } from "@app/lib/resources/user_resource";
 
 import { processActionTypesFromAgentMessageIds } from "./actions/process";
 import { retrievalActionTypesFromAgentMessageIds } from "./actions/retrieval";
@@ -58,11 +58,7 @@ export async function batchRenderUserMessages(
       if (userIds.length === 0) {
         return [];
       }
-      return User.findAll({
-        where: {
-          id: userIds,
-        },
-      });
+      return UserResource.listByModelIds(userIds);
     })(),
   ]);
 
@@ -83,21 +79,7 @@ export async function batchRenderUserMessages(
       visibility: message.visibility,
       version: message.version,
       created: message.createdAt.getTime(),
-      user: user
-        ? {
-            sId: user.sId,
-            id: user.id,
-            createdAt: user.createdAt.getTime(),
-            username: user.username,
-            email: user.email,
-            firstName: user.firstName,
-            lastName: user.lastName,
-            fullName:
-              user.firstName + (user.lastName ? ` ${user.lastName}` : ""),
-            provider: user.provider,
-            image: user.imageUrl,
-          }
-        : null,
+      user: user ? user.toJSON() : null,
       mentions: messageMentions.map((m) => {
         if (m.agentConfigurationId) {
           return {

diff --git a/front/lib/api/assistant/recent_authors.ts b/front/lib/api/assistant/recent_authors.ts
@@ -7,11 +7,10 @@ import { removeNulls } from "@dust-tt/types";
 import { Sequelize } from "sequelize";
 
 import { runOnRedis } from "@app/lib/api/redis";
-import { renderUserType } from "@app/lib/api/user";
 import { getGlobalAgentAuthorName } from "@app/lib/assistant";
 import type { Authenticator } from "@app/lib/auth";
 import { AgentConfiguration } from "@app/lib/models/assistant/agent";
-import { User } from "@app/lib/models/user";
+import { UserResource } from "@app/lib/resources/user_resource";
 
 // We keep the most recent authorIds for 3 days.
 const recentAuthorIdsKeyTTL = 60 * 60 * 24 * 3; // 3 days.
@@ -172,15 +171,13 @@ export async function getAgentsRecentAuthors({
   );
 
   const authorByUserId: Record<number, UserType> = (
-    await User.findAll({
-      where: {
-        id: removeNulls(
-          Array.from(new Set(Object.values(recentAuthorsIdsByAgentId).flat()))
-        ),
-      },
-    })
+    await UserResource.listByModelIds(
+      removeNulls(
+        Array.from(new Set(Object.values(recentAuthorsIdsByAgentId).flat()))
+      )
+    )
   ).reduce<Record<number, UserType>>((acc, user) => {
-    acc[user.id] = renderUserType(user);
+    acc[user.id] = user.toJSON();
     return acc;
   }, {});
 

diff --git a/front/lib/api/user.ts b/front/lib/api/user.ts
@@ -1,67 +1,49 @@
-import type { UserMetadataType, UserType } from "@dust-tt/types";
+import type {
+  Result,
+  UserMetadataType,
+  UserType,
+  UserTypeWithWorkspaces,
+} from "@dust-tt/types";
+import { Err, Ok } from "@dust-tt/types";
 
 import type { Authenticator } from "@app/lib/auth";
-import { User, UserMetadata } from "@app/lib/models/user";
+import { UserMetadata } from "@app/lib/models/user";
+import { Workspace } from "@app/lib/models/workspace";
+import { UserResource } from "@app/lib/resources/user_resource";
+import logger from "@app/logger/logger";
 
 import { MembershipResource } from "../resources/membership_resource";
 
-export function renderUserType(user: User): UserType {
-  return {
-    sId: user.sId,
-    id: user.id,
-    createdAt: user.createdAt.getTime(),
-    provider: user.provider,
-    username: user.username,
-    email: user.email,
-    firstName: user.firstName,
-    lastName: user.lastName,
-    fullName: user.firstName + (user.lastName ? ` ${user.lastName}` : ""),
-    image: user.imageUrl,
-  };
-}
-
 /**
  * This function checks that the user had at least one membership in the past for this workspace
  * otherwise returns null, preventing retrieving user information from their sId.
  */
 export async function getUserForWorkspace(
   auth: Authenticator,
   { userId }: { userId: string }
-): Promise<UserType | null> {
+): Promise<UserResource | null> {
   const owner = auth.workspace();
   if (!owner || !(auth.isAdmin() || auth.user()?.sId === userId)) {
     return null;
   }
 
-  const user = await User.findOne({
-    where: {
-      sId: userId,
-    },
-  });
+  const user = await UserResource.fetchById(userId);
 
   if (!user) {
     return null;
   }
 
   const membership =
     await MembershipResource.getLatestMembershipOfUserInWorkspace({
-      user: renderUserType(user),
+      user,
       workspace: owner,
     });
 
   if (!membership) {
     return null;
   }
 
-  return renderUserType(user);
-}
-
-export async function deleteUser(user: UserType): Promise<void> {
-  await User.destroy({
-    where: {
-      id: user.id,
-    },
-  });
+  return user;
 }
 
 /**
@@ -121,44 +103,37 @@ export async function setUserMetadata(
   await metadata.save();
 }
 
-export async function updateUserFullName({
-  user,
-  firstName,
-  lastName,
-}: {
-  user: UserType;
-  firstName: string;
-  lastName: string;
-}): Promise<boolean | null> {
-  const u = await User.findOne({
-    where: {
-      id: user.id,
-    },
-  });
+export async function fetchRevokedWorkspace(
+  user: UserTypeWithWorkspaces
+): Promise<Result<Workspace, Error>> {
+  // TODO(@fontanierh): this doesn't look very solid as it will start to behave
+  // weirdly if a user has multiple revoked memberships.
+  const u = await UserResource.fetchByModelId(user.id);
 
   if (!u) {
-    return null;
+    const message = "Unreachable: user not found.";
+    logger.error({ userId: user.id }, message);
+    return new Err(new Error(message));
   }
 
-  u.firstName = firstName;
-  u.lastName = lastName;
-  u.name = `${firstName} ${lastName}`;
-  await u.save();
+  const memberships = await MembershipResource.getLatestMemberships({
+    users: [u],
+  });
 
-  return true;
-}
+  if (!memberships.length) {
+    const message = "Unreachable: user has no memberships.";
+    logger.error({ userId: user.id }, message);
+    return new Err(new Error(message));
+  }
+
+  const revokedWorkspaceId = memberships[0].workspaceId;
+  const workspace = await Workspace.findByPk(revokedWorkspaceId);
 
-export async function unsafeGetUsersByModelId(
-  modelIds: number[]
-): Promise<UserType[]> {
-  if (modelIds.length === 0) {
-    return [];
+  if (!workspace) {
+    const message = "Unreachable: workspace not found.";
+    logger.error({ userId: user.id, workspaceId: revokedWorkspaceId }, message);
+    return new Err(new Error(message));
   }
-  const users = await User.findAll({
-    where: {
-      id: modelIds,
-    },
-  });
 
-  return users.map((u) => renderUserType(u));
+  return new Ok(workspace);
 }
diff --git a/front/lib/api/workspace.ts b/front/lib/api/workspace.ts
@@ -10,9 +10,9 @@ import type {
 } from "@dust-tt/types";
 
 import type { Authenticator } from "@app/lib/auth";
-import { User } from "@app/lib/models/user";
 import { Workspace, WorkspaceHasDomain } from "@app/lib/models/workspace";
 import { MembershipResource } from "@app/lib/resources/membership_resource";
+import { UserResource } from "@app/lib/resources/user_resource";
 import { renderLightWorkspaceType } from "@app/lib/workspace";
 
 export async function getWorkspaceInfos(
@@ -144,11 +144,9 @@ export async function getMembers(
         roles,
       });
 
-  const users = await User.findAll({
-    where: {
-      id: memberships.map((m) => m.userId),
-    },
-  });
+  const users = await UserResource.listByModelIds(
+    memberships.map((m) => m.userId)
+  );
 
   return users.map((u) => {
     const m = memberships.find((m) => m.userId === u.id);
@@ -166,16 +164,7 @@ export async function getMembers(
     }
 
     return {
-      sId: u.sId,
-      id: u.id,
-      createdAt: u.createdAt.getTime(),
-      provider: u.provider,
-      username: u.username,
-      email: u.email,
-      fullName: u.firstName + (u.lastName ? ` ${u.lastName}` : ""),
-      firstName: u.firstName,
-      lastName: u.lastName,
-      image: u.imageUrl,
+      ...u.toJSON(),
       workspaces: [{ ...owner, role, flags: null }],
     };
   });