From 85aa3e7d9cfe539687d83a06f5651a63180d6b46 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 10 Sep 2025 20:40:58 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20refactor:=20Centralize=20Collect?= =?UTF-8?q?ion=20Checks=20for=20Permissions=20Migration=20(#9565)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🔧 refactor: Centralize Collection Existence Checks for Permissions Migration * Replace individual collection existence checks with a unified function `ensureRequiredCollectionsExist` in the database utility module. * Update migration scripts for agents and prompts to utilize the new function, ensuring all required collections are verified for existence in a single call. * Remove redundant collection existence logic from migration files, improving code maintainability and clarity. * chore: import order in migration scripts * 🔧 test: Update Token Test Cases for Realistic Scenarios * Changed email in test data to 'user1-alt@example.com' for a more realistic scenario. * Clarified expectation comment for token retrieval to indicate it finds the only matching token based on criteria. --- config/migrate-agent-permissions.js | 28 +--------- config/migrate-prompt-permissions.js | 28 +--------- packages/api/src/agents/migration.ts | 27 ++-------- packages/api/src/db/utils.ts | 54 +++++++++++++++++++ packages/api/src/index.ts | 1 + packages/api/src/prompts/migration.ts | 27 ++-------- .../data-schemas/src/methods/token.spec.ts | 4 +- 7 files changed, 67 insertions(+), 102 deletions(-) create mode 100644 packages/api/src/db/utils.ts diff --git a/config/migrate-agent-permissions.js b/config/migrate-agent-permissions.js index ac7ca5b76..b206c648c 100644 --- a/config/migrate-agent-permissions.js +++ b/config/migrate-agent-permissions.js @@ -1,5 +1,6 @@ const path = require('path'); const { logger } = require('@librechat/data-schemas'); +const { ensureRequiredCollectionsExist } = require('@librechat/api'); const { AccessRoleIds, ResourceType, PrincipalType } = require('librechat-data-provider'); const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants; @@ -16,36 +17,11 @@ async function migrateAgentPermissionsEnhanced({ dryRun = true, batchSize = 100 logger.info('Starting Enhanced Agent Permissions Migration', { dryRun, batchSize }); - /** Ensurse `aclentries` collection exists for DocumentDB compatibility - * @param {import('mongoose').mongo.Db} db - * @param {string} collectionName - */ - async function ensureCollectionExists(db, collectionName) { - try { - const collections = await db.listCollections({ name: collectionName }).toArray(); - if (collections.length === 0) { - await db.createCollection(collectionName); - logger.info(`Created collection: ${collectionName}`); - } else { - logger.info(`Collection already exists: ${collectionName}`); - } - } catch (error) { - logger.error(`'Failed to check/create "${collectionName}" collection:`, error); - // If listCollections fails, try alternative approach - try { - // Try to access the collection directly - this will create it in MongoDB if it doesn't exist - await db.collection(collectionName).findOne({}, { projection: { _id: 1 } }); - } catch (createError) { - logger.error(`Could not ensure collection ${collectionName} exists:`, createError); - } - } - } - const mongoose = require('mongoose'); /** @type {import('mongoose').mongo.Db | undefined} */ const db = mongoose.connection.db; if (db) { - await ensureCollectionExists(db, 'aclentries'); + await ensureRequiredCollectionsExist(db); } // Verify required roles exist diff --git a/config/migrate-prompt-permissions.js b/config/migrate-prompt-permissions.js index e1df481cd..6018b1663 100644 --- a/config/migrate-prompt-permissions.js +++ b/config/migrate-prompt-permissions.js @@ -1,5 +1,6 @@ const path = require('path'); const { logger } = require('@librechat/data-schemas'); +const { ensureRequiredCollectionsExist } = require('@librechat/api'); const { AccessRoleIds, ResourceType, PrincipalType } = require('librechat-data-provider'); const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants; @@ -16,36 +17,11 @@ async function migrateToPromptGroupPermissions({ dryRun = true, batchSize = 100 logger.info('Starting PromptGroup Permissions Migration', { dryRun, batchSize }); - /** Ensurse `aclentries` collection exists for DocumentDB compatibility - * @param {import('mongoose').mongo.Db} db - * @param {string} collectionName - */ - async function ensureCollectionExists(db, collectionName) { - try { - const collections = await db.listCollections({ name: collectionName }).toArray(); - if (collections.length === 0) { - await db.createCollection(collectionName); - logger.info(`Created collection: ${collectionName}`); - } else { - logger.info(`Collection already exists: ${collectionName}`); - } - } catch (error) { - logger.error(`'Failed to check/create "${collectionName}" collection:`, error); - // If listCollections fails, try alternative approach - try { - // Try to access the collection directly - this will create it in MongoDB if it doesn't exist - await db.collection(collectionName).findOne({}, { projection: { _id: 1 } }); - } catch (createError) { - logger.error(`Could not ensure collection ${collectionName} exists:`, createError); - } - } - } - const mongoose = require('mongoose'); /** @type {import('mongoose').mongo.Db | undefined} */ const db = mongoose.connection.db; if (db) { - await ensureCollectionExists(db, 'aclentries'); + await ensureRequiredCollectionsExist(db); } // Verify required roles exist diff --git a/packages/api/src/agents/migration.ts b/packages/api/src/agents/migration.ts index e9b76c3b9..f8cad88b6 100644 --- a/packages/api/src/agents/migration.ts +++ b/packages/api/src/agents/migration.ts @@ -1,7 +1,8 @@ import { logger } from '@librechat/data-schemas'; import { AccessRoleIds, ResourceType, PrincipalType, Constants } from 'librechat-data-provider'; +import { ensureRequiredCollectionsExist } from '../db/utils'; import type { AccessRoleMethods, IAgent } from '@librechat/data-schemas'; -import type { Model, Mongoose, mongo } from 'mongoose'; +import type { Model, Mongoose } from 'mongoose'; const { GLOBAL_PROJECT_NAME } = Constants; @@ -54,31 +55,9 @@ export async function checkAgentPermissionsMigration({ logger.debug('Checking if agent permissions migration is needed'); try { - /** Ensurse `aclentries` collection exists for DocumentDB compatibility */ - async function ensureCollectionExists(db: mongo.Db, collectionName: string) { - try { - const collections = await db.listCollections({ name: collectionName }).toArray(); - if (collections.length === 0) { - await db.createCollection(collectionName); - logger.info(`Created collection: ${collectionName}`); - } else { - logger.debug(`Collection already exists: ${collectionName}`); - } - } catch (error) { - logger.error(`'Failed to check/create "${collectionName}" collection:`, error); - // If listCollections fails, try alternative approach - try { - // Try to access the collection directly - this will create it in MongoDB if it doesn't exist - await db.collection(collectionName).findOne({}, { projection: { _id: 1 } }); - } catch (createError) { - logger.error(`Could not ensure collection ${collectionName} exists:`, createError); - } - } - } - const db = mongoose.connection.db; if (db) { - await ensureCollectionExists(db, 'aclentries'); + await ensureRequiredCollectionsExist(db); } // Verify required roles exist diff --git a/packages/api/src/db/utils.ts b/packages/api/src/db/utils.ts new file mode 100644 index 000000000..b53478d04 --- /dev/null +++ b/packages/api/src/db/utils.ts @@ -0,0 +1,54 @@ +import { logger } from '@librechat/data-schemas'; +import type { mongo } from 'mongoose'; + +/** + * Ensures that a collection exists in the database. + * For DocumentDB compatibility, it tries multiple approaches. + * @param db - The MongoDB database instance + * @param collectionName - The name of the collection to ensure exists + */ +export async function ensureCollectionExists(db: mongo.Db, collectionName: string): Promise { + try { + const collections = await db.listCollections({ name: collectionName }).toArray(); + if (collections.length === 0) { + await db.createCollection(collectionName); + logger.info(`Created collection: ${collectionName}`); + } else { + logger.debug(`Collection already exists: ${collectionName}`); + } + } catch (error) { + logger.error(`Failed to check/create "${collectionName}" collection:`, error); + // If listCollections fails, try alternative approach + try { + // Try to access the collection directly - this will create it in MongoDB if it doesn't exist + await db.collection(collectionName).findOne({}, { projection: { _id: 1 } }); + } catch (createError) { + logger.error(`Could not ensure collection ${collectionName} exists:`, createError); + } + } +} + +/** + * Ensures that all required collections exist for the permission system. + * This includes aclentries, groups, accessroles, and any other collections + * needed for migrations and permission checks. + * @param db - The MongoDB database instance + */ +export async function ensureRequiredCollectionsExist(db: mongo.Db): Promise { + const requiredCollections = [ + 'aclentries', // ACL permission entries + 'groups', // User groups + 'accessroles', // Access roles for permissions + 'agents', // Agents collection + 'promptgroups', // Prompt groups collection + 'projects', // Projects collection + ]; + + logger.debug('Ensuring required collections exist for permission system'); + + for (const collectionName of requiredCollections) { + await ensureCollectionExists(db, collectionName); + } + + logger.debug('All required collections have been checked/created'); +} diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index d0c797aba..e9cbfd1ff 100644 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -11,6 +11,7 @@ export * from './mcp/zod'; export * from './format'; export * from './mcp/utils'; export * from './utils'; +export * from './db/utils'; /* OAuth */ export * from './oauth'; /* Crypto */ diff --git a/packages/api/src/prompts/migration.ts b/packages/api/src/prompts/migration.ts index 938d8ca8f..40f0a585d 100644 --- a/packages/api/src/prompts/migration.ts +++ b/packages/api/src/prompts/migration.ts @@ -1,7 +1,8 @@ import { logger } from '@librechat/data-schemas'; import { AccessRoleIds, ResourceType, PrincipalType, Constants } from 'librechat-data-provider'; +import { ensureRequiredCollectionsExist } from '../db/utils'; import type { AccessRoleMethods, IPromptGroupDocument } from '@librechat/data-schemas'; -import type { Model, Mongoose, mongo } from 'mongoose'; +import type { Model, Mongoose } from 'mongoose'; const { GLOBAL_PROJECT_NAME } = Constants; @@ -52,32 +53,10 @@ export async function checkPromptPermissionsMigration({ logger.debug('Checking if prompt permissions migration is needed'); try { - /** Ensurse `aclentries` collection exists for DocumentDB compatibility */ - async function ensureCollectionExists(db: mongo.Db, collectionName: string) { - try { - const collections = await db.listCollections({ name: collectionName }).toArray(); - if (collections.length === 0) { - await db.createCollection(collectionName); - logger.info(`Created collection: ${collectionName}`); - } else { - logger.debug(`Collection already exists: ${collectionName}`); - } - } catch (error) { - logger.error(`'Failed to check/create "${collectionName}" collection:`, error); - // If listCollections fails, try alternative approach - try { - // Try to access the collection directly - this will create it in MongoDB if it doesn't exist - await db.collection(collectionName).findOne({}, { projection: { _id: 1 } }); - } catch (createError) { - logger.error(`Could not ensure collection ${collectionName} exists:`, createError); - } - } - } - /** Native MongoDB database instance */ const db = mongoose.connection.db; if (db) { - await ensureCollectionExists(db, 'aclentries'); + await ensureRequiredCollectionsExist(db); } // Verify required roles exist diff --git a/packages/data-schemas/src/methods/token.spec.ts b/packages/data-schemas/src/methods/token.spec.ts index bb9824187..248061998 100644 --- a/packages/data-schemas/src/methods/token.spec.ts +++ b/packages/data-schemas/src/methods/token.spec.ts @@ -118,7 +118,7 @@ describe('Token Methods - Detailed Tests', () => { { token: 'token-3', userId: user1Id, - email: 'user1@example.com', + email: 'user1-alt@example.com', // Different email for realistic scenario createdAt: new Date(), expiresAt: new Date(Date.now() + 3600000), }, @@ -164,7 +164,7 @@ describe('Token Methods - Detailed Tests', () => { }); expect(found).toBeDefined(); - expect(found?.token).toBe('token-1'); // Should find first matching + expect(found?.token).toBe('token-1'); // Should find the only token matching both criteria }); test('should return null for non-existent token', async () => {