refactor(data-schemas): enhance method organization and add librechat-data-provider dependency

package-lock.json

@@ -45642,6 +45642,7 @@
       },
       "peerDependencies": {
         "keyv": "^5.3.2",
+        "librechat-data-provider": "*",
         "mongoose": "^8.12.1"
       }
     },

packages/data-schemas/package.json

@@ -70,7 +70,8 @@
   },
   "peerDependencies": {
     "keyv": "^5.3.2",
-    "mongoose": "^8.12.1"
+    "mongoose": "^8.12.1",
+    "librechat-data-provider": "*"
   },
   "publishConfig": {
     "registry": "https://registry.npmjs.org/",

packages/data-schemas/src/methods/index.ts

@@ -1,8 +1,24 @@
-// User methods
+import { createUserMethods, type UserMethods } from './user';
-export * from './user';
+import { createSessionMethods, type SessionMethods } from './session';
+import { createTokenMethods, type TokenMethods } from './token';
+import { createRoleMethods, type RoleMethods } from './role';
 
-// Session methods
+/**
-export * from './session';
+ * Creates all database methods for all collections
+ */
+export function createAllMethods(mongoose: typeof import('mongoose')) {
+  return {
+    ...createUserMethods(mongoose),
+    ...createSessionMethods(mongoose),
+    ...createTokenMethods(mongoose),
+    ...createRoleMethods(mongoose),
+  };
+}
 
-// Token methods
+export type AllMethods = UserMethods & SessionMethods & TokenMethods & RoleMethods;
-export * from './token';
+
+// Also export individual factory functions for granular usage if needed
+export { createUserMethods, type UserMethods } from './user';
+export { createSessionMethods, type SessionMethods } from './session';
+export { createTokenMethods, type TokenMethods } from './token';
+export { createRoleMethods, type RoleMethods } from './role';

packages/data-schemas/src/methods/session.ts

@@ -1,5 +1,3 @@
-import mongoose from 'mongoose';
-import { Session } from '~/models';
 import {
   ISession,
   CreateSessionOptions,
@@ -16,187 +14,208 @@ import logger from '~/config/winston';
 const { REFRESH_TOKEN_EXPIRY } = process.env ?? {};
 const expires = eval(REFRESH_TOKEN_EXPIRY ?? '0') ?? 1000 * 60 * 60 * 24 * 7; // 7 days default
 
-/**
+// Factory function that takes mongoose instance and returns the methods
- * Creates a new session for a user
+export function createSessionMethods(mongoose: typeof import('mongoose')) {
- */
+  /**
-export async function createSession(
+   * Creates a new session for a user
-  userId: string,
+   */
-  options: CreateSessionOptions = {},
+  async function createSession(
-): Promise<SessionResult> {
+    userId: string,
-  if (!userId) {
+    options: CreateSessionOptions = {},
-    throw new SessionError('User ID is required', 'INVALID_USER_ID');
+  ): Promise<SessionResult> {
-  }
-
-  try {
-    const session = {
-      _id: new mongoose.Types.ObjectId(),
-      user: new mongoose.Types.ObjectId(userId),
-      expiration: options.expiration || new Date(Date.now() + expires),
-    };
-    const refreshToken = await generateRefreshToken(session);
-
-    return { session, refreshToken };
-  } catch (error) {
-    logger.error('[createSession] Error creating session:', error);
-    throw new SessionError('Failed to create session', 'CREATE_SESSION_FAILED');
-  }
-}
-
-/**
- * Finds a session by various parameters
- */
-export async function findSession(
-  params: SessionSearchParams,
-  options: SessionQueryOptions = { lean: true },
-): Promise<ISession | null> {
-  try {
-    const query: Record<string, unknown> = {};
-
-    if (!params.refreshToken && !params.userId && !params.sessionId) {
-      throw new SessionError('At least one search parameter is required', 'INVALID_SEARCH_PARAMS');
-    }
-
-    if (params.refreshToken) {
-      const tokenHash = await hashToken(params.refreshToken);
-      query.refreshTokenHash = tokenHash;
-    }
-
-    if (params.userId) {
-      query.user = params.userId;
-    }
-
-    if (params.sessionId) {
-      const sessionId =
-        typeof params.sessionId === 'object' && 'sessionId' in params.sessionId
-          ? params.sessionId.sessionId
-          : params.sessionId;
-      if (!mongoose.Types.ObjectId.isValid(sessionId)) {
-        throw new SessionError('Invalid session ID format', 'INVALID_SESSION_ID');
-      }
-      query._id = sessionId;
-    }
-
-    // Add expiration check to only return valid sessions
-    query.expiration = { $gt: new Date() };
-
-    const sessionQuery = Session.findOne(query);
-
-    if (options.lean) {
-      return await sessionQuery.lean();
-    }
-
-    return await sessionQuery.exec();
-  } catch (error) {
-    logger.error('[findSession] Error finding session:', error);
-    throw new SessionError('Failed to find session', 'FIND_SESSION_FAILED');
-  }
-}
-
-/**
- * Deletes a session by refresh token or session ID
- */
-export async function deleteSession(
-  params: DeleteSessionParams,
-): Promise<{ deletedCount?: number }> {
-  try {
-    if (!params.refreshToken && !params.sessionId) {
-      throw new SessionError(
-        'Either refreshToken or sessionId is required',
-        'INVALID_DELETE_PARAMS',
-      );
-    }
-
-    const query: Record<string, unknown> = {};
-
-    if (params.refreshToken) {
-      query.refreshTokenHash = await hashToken(params.refreshToken);
-    }
-
-    if (params.sessionId) {
-      query._id = params.sessionId;
-    }
-
-    const result = await Session.deleteOne(query);
-
-    if (result.deletedCount === 0) {
-      logger.warn('[deleteSession] No session found to delete');
-    }
-
-    return result;
-  } catch (error) {
-    logger.error('[deleteSession] Error deleting session:', error);
-    throw new SessionError('Failed to delete session', 'DELETE_SESSION_FAILED');
-  }
-}
-
-/**
- * Generates a refresh token for a session
- */
-export async function generateRefreshToken(session: Partial<ISession>): Promise<string> {
-  if (!session || !session.user) {
-    throw new SessionError('Invalid session object', 'INVALID_SESSION');
-  }
-
-  try {
-    const expiresIn = session.expiration ? session.expiration.getTime() : Date.now() + expires;
-    if (!session.expiration) {
-      session.expiration = new Date(expiresIn);
-    }
-
-    const refreshToken = await signPayload({
-      payload: {
-        id: session.user,
-        sessionId: session._id,
-      },
-      secret: process.env.JWT_REFRESH_SECRET,
-      expirationTime: Math.floor((expiresIn - Date.now()) / 1000),
-    });
-
-    session.refreshTokenHash = await hashToken(refreshToken);
-    await Session.create(session);
-    return refreshToken;
-  } catch (error) {
-    logger.error('[generateRefreshToken] Error generating refresh token:', error);
-    throw new SessionError('Failed to generate refresh token', 'GENERATE_TOKEN_FAILED');
-  }
-}
-
-/**
- * Deletes all sessions for a user
- */
-export async function deleteAllUserSessions(
-  userId: string | { userId: string },
-  options: DeleteAllSessionsOptions = {},
-): Promise<{ deletedCount?: number }> {
-  try {
     if (!userId) {
       throw new SessionError('User ID is required', 'INVALID_USER_ID');
     }
 
-    // Extract userId if it's passed as an object
+    try {
-    const userIdString = typeof userId === 'object' ? userId.userId : userId;
+      const session = {
+        _id: new mongoose.Types.ObjectId(),
+        user: new mongoose.Types.ObjectId(userId),
+        expiration: options.expiration || new Date(Date.now() + expires),
+      };
+      const refreshToken = await generateRefreshToken(session);
 
-    if (!mongoose.Types.ObjectId.isValid(userIdString)) {
+      return { session, refreshToken };
-      throw new SessionError('Invalid user ID format', 'INVALID_USER_ID_FORMAT');
+    } catch (error) {
+      logger.error('[createSession] Error creating session:', error);
+      throw new SessionError('Failed to create session', 'CREATE_SESSION_FAILED');
     }
-
-    const query: Record<string, unknown> = { user: userIdString };
-
-    if (options.excludeCurrentSession && options.currentSessionId) {
-      query._id = { $ne: options.currentSessionId };
-    }
-
-    const result = await Session.deleteMany(query);
-
-    if (result.deletedCount && result.deletedCount > 0) {
-      logger.debug(
-        `[deleteAllUserSessions] Deleted ${result.deletedCount} sessions for user ${userIdString}.`,
-      );
-    }
-
-    return result;
-  } catch (error) {
-    logger.error('[deleteAllUserSessions] Error deleting user sessions:', error);
-    throw new SessionError('Failed to delete user sessions', 'DELETE_ALL_SESSIONS_FAILED');
   }
+
+  /**
+   * Finds a session by various parameters
+   */
+  async function findSession(
+    params: SessionSearchParams,
+    options: SessionQueryOptions = { lean: true },
+  ): Promise<ISession | null> {
+    try {
+      const Session = mongoose.models.Session;
+      const query: Record<string, unknown> = {};
+
+      if (!params.refreshToken && !params.userId && !params.sessionId) {
+        throw new SessionError(
+          'At least one search parameter is required',
+          'INVALID_SEARCH_PARAMS',
+        );
+      }
+
+      if (params.refreshToken) {
+        const tokenHash = await hashToken(params.refreshToken);
+        query.refreshTokenHash = tokenHash;
+      }
+
+      if (params.userId) {
+        query.user = params.userId;
+      }
+
+      if (params.sessionId) {
+        const sessionId =
+          typeof params.sessionId === 'object' && 'sessionId' in params.sessionId
+            ? params.sessionId.sessionId
+            : params.sessionId;
+        if (!mongoose.Types.ObjectId.isValid(sessionId)) {
+          throw new SessionError('Invalid session ID format', 'INVALID_SESSION_ID');
+        }
+        query._id = sessionId;
+      }
+
+      // Add expiration check to only return valid sessions
+      query.expiration = { $gt: new Date() };
+
+      const sessionQuery = Session.findOne(query);
+
+      if (options.lean) {
+        return (await sessionQuery.lean()) as ISession | null;
+      }
+
+      return await sessionQuery.exec();
+    } catch (error) {
+      logger.error('[findSession] Error finding session:', error);
+      throw new SessionError('Failed to find session', 'FIND_SESSION_FAILED');
+    }
+  }
+
+  /**
+   * Deletes a session by refresh token or session ID
+   */
+  async function deleteSession(params: DeleteSessionParams): Promise<{ deletedCount?: number }> {
+    try {
+      const Session = mongoose.models.Session;
+
+      if (!params.refreshToken && !params.sessionId) {
+        throw new SessionError(
+          'Either refreshToken or sessionId is required',
+          'INVALID_DELETE_PARAMS',
+        );
+      }
+
+      const query: Record<string, unknown> = {};
+
+      if (params.refreshToken) {
+        query.refreshTokenHash = await hashToken(params.refreshToken);
+      }
+
+      if (params.sessionId) {
+        query._id = params.sessionId;
+      }
+
+      const result = await Session.deleteOne(query);
+
+      if (result.deletedCount === 0) {
+        logger.warn('[deleteSession] No session found to delete');
+      }
+
+      return result;
+    } catch (error) {
+      logger.error('[deleteSession] Error deleting session:', error);
+      throw new SessionError('Failed to delete session', 'DELETE_SESSION_FAILED');
+    }
+  }
+
+  /**
+   * Generates a refresh token for a session
+   */
+  async function generateRefreshToken(session: Partial<ISession>): Promise<string> {
+    if (!session || !session.user) {
+      throw new SessionError('Invalid session object', 'INVALID_SESSION');
+    }
+
+    try {
+      const Session = mongoose.models.Session;
+      const expiresIn = session.expiration ? session.expiration.getTime() : Date.now() + expires;
+      if (!session.expiration) {
+        session.expiration = new Date(expiresIn);
+      }
+
+      const refreshToken = await signPayload({
+        payload: {
+          id: session.user,
+          sessionId: session._id,
+        },
+        secret: process.env.JWT_REFRESH_SECRET,
+        expirationTime: Math.floor((expiresIn - Date.now()) / 1000),
+      });
+
+      session.refreshTokenHash = await hashToken(refreshToken);
+      await Session.create(session);
+      return refreshToken;
+    } catch (error) {
+      logger.error('[generateRefreshToken] Error generating refresh token:', error);
+      throw new SessionError('Failed to generate refresh token', 'GENERATE_TOKEN_FAILED');
+    }
+  }
+
+  /**
+   * Deletes all sessions for a user
+   */
+  async function deleteAllUserSessions(
+    userId: string | { userId: string },
+    options: DeleteAllSessionsOptions = {},
+  ): Promise<{ deletedCount?: number }> {
+    try {
+      const Session = mongoose.models.Session;
+
+      if (!userId) {
+        throw new SessionError('User ID is required', 'INVALID_USER_ID');
+      }
+
+      // Extract userId if it's passed as an object
+      const userIdString = typeof userId === 'object' ? userId.userId : userId;
+
+      if (!mongoose.Types.ObjectId.isValid(userIdString)) {
+        throw new SessionError('Invalid user ID format', 'INVALID_USER_ID_FORMAT');
+      }
+
+      const query: Record<string, unknown> = { user: userIdString };
+
+      if (options.excludeCurrentSession && options.currentSessionId) {
+        query._id = { $ne: options.currentSessionId };
+      }
+
+      const result = await Session.deleteMany(query);
+
+      if (result.deletedCount && result.deletedCount > 0) {
+        logger.debug(
+          `[deleteAllUserSessions] Deleted ${result.deletedCount} sessions for user ${userIdString}.`,
+        );
+      }
+
+      return result;
+    } catch (error) {
+      logger.error('[deleteAllUserSessions] Error deleting user sessions:', error);
+      throw new SessionError('Failed to delete user sessions', 'DELETE_ALL_SESSIONS_FAILED');
+    }
+  }
+
+  // Return all methods
+  return {
+    createSession,
+    findSession,
+    deleteSession,
+    generateRefreshToken,
+    deleteAllUserSessions,
+  };
 }
+
+export type SessionMethods = ReturnType<typeof createSessionMethods>;

packages/data-schemas/src/methods/token.ts

@@ -1,105 +1,121 @@
-import { Token } from '~/models';
 import { IToken, TokenCreateData, TokenQuery, TokenUpdateData, TokenDeleteResult } from '~/types';
 import logger from '~/config/winston';
 
-/**
+// Factory function that takes mongoose instance and returns the methods
- * Creates a new Token instance.
+export function createTokenMethods(mongoose: typeof import('mongoose')) {
- */
+  /**
-export async function createToken(tokenData: TokenCreateData): Promise<IToken> {
+   * Creates a new Token instance.
-  try {
+   */
-    const currentTime = new Date();
+  async function createToken(tokenData: TokenCreateData): Promise<IToken> {
-    const expiresAt = new Date(currentTime.getTime() + tokenData.expiresIn * 1000);
+    try {
+      const Token = mongoose.models.Token;
+      const currentTime = new Date();
+      const expiresAt = new Date(currentTime.getTime() + tokenData.expiresIn * 1000);
 
-    const newTokenData = {
+      const newTokenData = {
-      ...tokenData,
+        ...tokenData,
-      createdAt: currentTime,
+        createdAt: currentTime,
-      expiresAt,
+        expiresAt,
-    };
+      };
 
-    return await Token.create(newTokenData);
+      return await Token.create(newTokenData);
-  } catch (error) {
+    } catch (error) {
-    logger.debug('An error occurred while creating token:', error);
+      logger.debug('An error occurred while creating token:', error);
-    throw error;
+      throw error;
+    }
   }
+
+  /**
+   * Updates a Token document that matches the provided query.
+   */
+  async function updateToken(
+    query: TokenQuery,
+    updateData: TokenUpdateData,
+  ): Promise<IToken | null> {
+    try {
+      const Token = mongoose.models.Token;
+      return await Token.findOneAndUpdate(query, updateData, { new: true });
+    } catch (error) {
+      logger.debug('An error occurred while updating token:', error);
+      throw error;
+    }
+  }
+
+  /**
+   * Deletes all Token documents that match the provided token, user ID, or email.
+   */
+  async function deleteTokens(query: TokenQuery): Promise<TokenDeleteResult> {
+    try {
+      const Token = mongoose.models.Token;
+      const conditions = [];
+
+      if (query.userId) {
+        conditions.push({ userId: query.userId });
+      }
+      if (query.token) {
+        conditions.push({ token: query.token });
+      }
+      if (query.email) {
+        conditions.push({ email: query.email });
+      }
+      if (query.identifier) {
+        conditions.push({ identifier: query.identifier });
+      }
+
+      if (conditions.length === 0) {
+        throw new Error('At least one query parameter must be provided');
+      }
+
+      return await Token.deleteMany({
+        $or: conditions,
+      });
+    } catch (error) {
+      logger.debug('An error occurred while deleting tokens:', error);
+      throw error;
+    }
+  }
+
+  /**
+   * Finds a Token document that matches the provided query.
+   */
+  async function findToken(query: TokenQuery): Promise<IToken | null> {
+    try {
+      const Token = mongoose.models.Token;
+      const conditions = [];
+
+      if (query.userId) {
+        conditions.push({ userId: query.userId });
+      }
+      if (query.token) {
+        conditions.push({ token: query.token });
+      }
+      if (query.email) {
+        conditions.push({ email: query.email });
+      }
+      if (query.identifier) {
+        conditions.push({ identifier: query.identifier });
+      }
+
+      if (conditions.length === 0) {
+        throw new Error('At least one query parameter must be provided');
+      }
+
+      return (await Token.findOne({
+        $and: conditions,
+      }).lean()) as IToken | null;
+    } catch (error) {
+      logger.debug('An error occurred while finding token:', error);
+      throw error;
+    }
+  }
+
+  // Return all methods
+  return {
+    createToken,
+    updateToken,
+    deleteTokens,
+    findToken,
+  };
 }
 
-/**
+export type TokenMethods = ReturnType<typeof createTokenMethods>;
- * Updates a Token document that matches the provided query.
- */
-export async function updateToken(
-  query: TokenQuery,
-  updateData: TokenUpdateData,
-): Promise<IToken | null> {
-  try {
-    return await Token.findOneAndUpdate(query, updateData, { new: true });
-  } catch (error) {
-    logger.debug('An error occurred while updating token:', error);
-    throw error;
-  }
-}
-
-/**
- * Deletes all Token documents that match the provided token, user ID, or email.
- */
-export async function deleteTokens(query: TokenQuery): Promise<TokenDeleteResult> {
-  try {
-    const conditions = [];
-
-    if (query.userId) {
-      conditions.push({ userId: query.userId });
-    }
-    if (query.token) {
-      conditions.push({ token: query.token });
-    }
-    if (query.email) {
-      conditions.push({ email: query.email });
-    }
-    if (query.identifier) {
-      conditions.push({ identifier: query.identifier });
-    }
-
-    if (conditions.length === 0) {
-      throw new Error('At least one query parameter must be provided');
-    }
-
-    return await Token.deleteMany({
-      $or: conditions,
-    });
-  } catch (error) {
-    logger.debug('An error occurred while deleting tokens:', error);
-    throw error;
-  }
-}
-
-/**
- * Finds a Token document that matches the provided query.
- */
-export async function findToken(query: TokenQuery): Promise<IToken | null> {
-  try {
-    const conditions = [];
-
-    if (query.userId) {
-      conditions.push({ userId: query.userId });
-    }
-    if (query.token) {
-      conditions.push({ token: query.token });
-    }
-    if (query.email) {
-      conditions.push({ email: query.email });
-    }
-    if (query.identifier) {
-      conditions.push({ identifier: query.identifier });
-    }
-
-    if (conditions.length === 0) {
-      throw new Error('At least one query parameter must be provided');
-    }
-
-    return (await Token.findOne({
-      $and: conditions,
-    }).lean()) as IToken | null;
-  } catch (error) {
-    logger.debug('An error occurred while finding token:', error);
-    throw error;
-  }
-}

packages/data-schemas/src/methods/user.ts

@@ -1,151 +1,174 @@
 import mongoose, { FilterQuery } from 'mongoose';
-import { User, Balance } from '~/models';
 import { IUser, BalanceConfig, UserCreateData, UserUpdateResult } from '~/types';
 import { signPayload } from '~/schema/session';
 
-/**
+/** Factory function that takes mongoose instance and returns the methods */
- * Search for a single user based on partial data and return matching user document as plain object.
+export function createUserMethods(mongoose: typeof import('mongoose')) {
- */
+  /**
-export async function findUser(
+   * Search for a single user based on partial data and return matching user document as plain object.
-  searchCriteria: FilterQuery<IUser>,
+   */
-  fieldsToSelect?: string | string[] | null,
+  async function findUser(
-): Promise<IUser | null> {
+    searchCriteria: FilterQuery<IUser>,
-  const query = User.findOne(searchCriteria);
+    fieldsToSelect?: string | string[] | null,
-  if (fieldsToSelect) {
+  ): Promise<IUser | null> {
-    query.select(fieldsToSelect);
+    const User = mongoose.models.User;
-  }
+    const query = User.findOne(searchCriteria);
-  return await query.lean();
+    if (fieldsToSelect) {
-}
+      query.select(fieldsToSelect);
-
+    }
-/**
+    return (await query.lean()) as IUser | null;
- * Count the number of user documents in the collection based on the provided filter.
- */
-export async function countUsers(filter: FilterQuery<IUser> = {}): Promise<number> {
-  return await User.countDocuments(filter);
-}
-
-/**
- * Creates a new user, optionally with a TTL of 1 week.
- */
-export async function createUser(
-  data: UserCreateData,
-  balanceConfig?: BalanceConfig,
-  disableTTL: boolean = true,
-  returnUser: boolean = false,
-): Promise<mongoose.Types.ObjectId | Partial<IUser>> {
-  const userData: Partial<IUser> = {
-    ...data,
-    expiresAt: disableTTL ? undefined : new Date(Date.now() + 604800 * 1000), // 1 week in milliseconds
-  };
-
-  if (disableTTL) {
-    delete userData.expiresAt;
   }
 
-  const user = await User.create(userData);
+  /**
+   * Count the number of user documents in the collection based on the provided filter.
+   */
+  async function countUsers(filter: FilterQuery<IUser> = {}): Promise<number> {
+    const User = mongoose.models.User;
+    return await User.countDocuments(filter);
+  }
 
-  // If balance is enabled, create or update a balance record for the user
+  /**
-  if (balanceConfig?.enabled && balanceConfig?.startBalance) {
+   * Creates a new user, optionally with a TTL of 1 week.
-    const update: {
+   */
-      $inc: { tokenCredits: number };
+  async function createUser(
-      $set?: {
+    data: UserCreateData,
-        autoRefillEnabled: boolean;
+    balanceConfig?: BalanceConfig,
-        refillIntervalValue: number;
+    disableTTL: boolean = true,
-        refillIntervalUnit: string;
+    returnUser: boolean = false,
-        refillAmount: number;
+  ): Promise<mongoose.Types.ObjectId | Partial<IUser>> {
-      };
+    const User = mongoose.models.User;
-    } = {
+    const Balance = mongoose.models.Balance;
-      $inc: { tokenCredits: balanceConfig.startBalance },
+
+    const userData: Partial<IUser> = {
+      ...data,
+      expiresAt: disableTTL ? undefined : new Date(Date.now() + 604800 * 1000), // 1 week in milliseconds
     };
 
-    if (
+    if (disableTTL) {
-      balanceConfig.autoRefillEnabled &&
+      delete userData.expiresAt;
-      balanceConfig.refillIntervalValue != null &&
+    }
-      balanceConfig.refillIntervalUnit != null &&
+
-      balanceConfig.refillAmount != null
+    const user = await User.create(userData);
-    ) {
+
-      update.$set = {
+    // If balance is enabled, create or update a balance record for the user
-        autoRefillEnabled: true,
+    if (balanceConfig?.enabled && balanceConfig?.startBalance) {
-        refillIntervalValue: balanceConfig.refillIntervalValue,
+      const update: {
-        refillIntervalUnit: balanceConfig.refillIntervalUnit,
+        $inc: { tokenCredits: number };
-        refillAmount: balanceConfig.refillAmount,
+        $set?: {
+          autoRefillEnabled: boolean;
+          refillIntervalValue: number;
+          refillIntervalUnit: string;
+          refillAmount: number;
+        };
+      } = {
+        $inc: { tokenCredits: balanceConfig.startBalance },
       };
+
+      if (
+        balanceConfig.autoRefillEnabled &&
+        balanceConfig.refillIntervalValue != null &&
+        balanceConfig.refillIntervalUnit != null &&
+        balanceConfig.refillAmount != null
+      ) {
+        update.$set = {
+          autoRefillEnabled: true,
+          refillIntervalValue: balanceConfig.refillIntervalValue,
+          refillIntervalUnit: balanceConfig.refillIntervalUnit,
+          refillAmount: balanceConfig.refillAmount,
+        };
+      }
+
+      await Balance.findOneAndUpdate({ user: user._id }, update, {
+        upsert: true,
+        new: true,
+      }).lean();
     }
 
-    await Balance.findOneAndUpdate({ user: user._id }, update, { upsert: true, new: true }).lean();
+    if (returnUser) {
+      return user.toObject() as Partial<IUser>;
+    }
+    return user._id as mongoose.Types.ObjectId;
   }
 
-  if (returnUser) {
+  /**
-    return user.toObject() as Partial<IUser>;
+   * Update a user with new data without overwriting existing properties.
+   */
+  async function updateUser(userId: string, updateData: Partial<IUser>): Promise<IUser | null> {
+    const User = mongoose.models.User;
+    const updateOperation = {
+      $set: updateData,
+      $unset: { expiresAt: '' }, // Remove the expiresAt field to prevent TTL
+    };
+    return (await User.findByIdAndUpdate(userId, updateOperation, {
+      new: true,
+      runValidators: true,
+    }).lean()) as IUser | null;
   }
-  return user._id as mongoose.Types.ObjectId;
-}
 
-/**
+  /**
- * Update a user with new data without overwriting existing properties.
+   * Retrieve a user by ID and convert the found user document to a plain object.
- */
+   */
-export async function updateUser(
+  async function getUserById(
-  userId: string,
+    userId: string,
-  updateData: Partial<IUser>,
+    fieldsToSelect?: string | string[] | null,
-): Promise<IUser | null> {
+  ): Promise<IUser | null> {
-  const updateOperation = {
+    const User = mongoose.models.User;
-    $set: updateData,
+    const query = User.findById(userId);
-    $unset: { expiresAt: '' }, // Remove the expiresAt field to prevent TTL
+    if (fieldsToSelect) {
+      query.select(fieldsToSelect);
+    }
+    return (await query.lean()) as IUser | null;
+  }
+
+  /**
+   * Delete a user by their unique ID.
+   */
+  async function deleteUserById(userId: string): Promise<UserUpdateResult> {
+    try {
+      const User = mongoose.models.User;
+      const result = await User.deleteOne({ _id: userId });
+      if (result.deletedCount === 0) {
+        return { deletedCount: 0, message: 'No user found with that ID.' };
+      }
+      return { deletedCount: result.deletedCount, message: 'User was deleted successfully.' };
+    } catch (error: unknown) {
+      const errorMessage = error instanceof Error ? error.message : 'Unknown error';
+      throw new Error('Error deleting user: ' + errorMessage);
+    }
+  }
+
+  /**
+   * Generates a JWT token for a given user.
+   */
+  async function generateToken(user: IUser): Promise<string> {
+    if (!user) {
+      throw new Error('No user provided');
+    }
+
+    const expires = eval(process.env.SESSION_EXPIRY ?? '0') ?? 1000 * 60 * 15;
+
+    return await signPayload({
+      payload: {
+        id: user._id,
+        username: user.username,
+        provider: user.provider,
+        email: user.email,
+      },
+      secret: process.env.JWT_SECRET,
+      expirationTime: expires / 1000,
+    });
+  }
+
+  // Return all methods
+  return {
+    findUser,
+    countUsers,
+    createUser,
+    updateUser,
+    getUserById,
+    deleteUserById,
+    generateToken,
   };
-  return await User.findByIdAndUpdate(userId, updateOperation, {
-    new: true,
-    runValidators: true,
-  }).lean();
 }
 
-/**
+export type UserMethods = ReturnType<typeof createUserMethods>;
- * Retrieve a user by ID and convert the found user document to a plain object.
- */
-export async function getUserById(
-  userId: string,
-  fieldsToSelect?: string | string[] | null,
-): Promise<IUser | null> {
-  const query = User.findById(userId);
-  if (fieldsToSelect) {
-    query.select(fieldsToSelect);
-  }
-  return await query.lean();
-}
-
-/**
- * Delete a user by their unique ID.
- */
-export async function deleteUserById(userId: string): Promise<UserUpdateResult> {
-  try {
-    const result = await User.deleteOne({ _id: userId });
-    if (result.deletedCount === 0) {
-      return { deletedCount: 0, message: 'No user found with that ID.' };
-    }
-    return { deletedCount: result.deletedCount, message: 'User was deleted successfully.' };
-  } catch (error: unknown) {
-    const errorMessage = error instanceof Error ? error.message : 'Unknown error';
-    throw new Error('Error deleting user: ' + errorMessage);
-  }
-}
-
-/**
- * Generates a JWT token for a given user.
- */
-export async function generateToken(user: IUser): Promise<string> {
-  if (!user) {
-    throw new Error('No user provided');
-  }
-
-  const expires = eval(process.env.SESSION_EXPIRY ?? '0') ?? 1000 * 60 * 15;
-
-  return await signPayload({
-    payload: {
-      id: user._id,
-      username: user.username,
-      provider: user.provider,
-      email: user.email,
-    },
-    secret: process.env.JWT_SECRET,
-    expirationTime: expires / 1000,
-  });
-}