🗄️ refactor: Resource Migration Scripts for DocumentDB Compatibility (#9249)

* refactor: Resource Migration Scripts for DocumentDB compatibility

* fix: Correct type annotation for `db` parameter in ensureCollectionExists function
This commit is contained in:
Danny Avila 2025-08-25 03:01:50 -04:00 committed by GitHub
parent 1915d7b195
commit ac641e7cba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 168 additions and 28 deletions

View file

@ -1,3 +1,4 @@
const mongoose = require('mongoose');
const { logger } = require('@librechat/data-schemas');
const {
logAgentMigrationWarning,
@ -16,7 +17,8 @@ const { findRoleByIdentifier } = require('~/models');
async function checkMigrations() {
try {
const agentMigrationResult = await checkAgentPermissionsMigration({
db: {
mongoose,
methods: {
findRoleByIdentifier,
getProjectByName,
},
@ -28,7 +30,8 @@ async function checkMigrations() {
}
try {
const promptMigrationResult = await checkPromptPermissionsMigration({
db: {
mongoose,
methods: {
findRoleByIdentifier,
getProjectByName,
},

View file

@ -16,6 +16,38 @@ 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');
}
// Verify required roles exist
const ownerRole = await findRoleByIdentifier(AccessRoleIds.AGENT_OWNER);
const viewerRole = await findRoleByIdentifier(AccessRoleIds.AGENT_VIEWER);
@ -100,12 +132,19 @@ async function migrateAgentPermissionsEnhanced({ dryRun = true, batchSize = 100
}
});
logger.info('Agent categorization:', {
logger.info(
'Agent categorization:\n' +
JSON.stringify(
{
globalEditAccess: categories.globalEditAccess.length,
globalViewAccess: categories.globalViewAccess.length,
privateAgents: categories.privateAgents.length,
total: agentsToMigrate.length,
});
},
null,
2,
),
);
if (dryRun) {
return {

View file

@ -16,6 +16,38 @@ 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');
}
// Verify required roles exist
const ownerRole = await findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_OWNER);
const viewerRole = await findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_VIEWER);
@ -91,11 +123,18 @@ async function migrateToPromptGroupPermissions({ dryRun = true, batchSize = 100
}
});
logger.info('PromptGroup categorization:', {
logger.info(
'PromptGroup categorization:\n' +
JSON.stringify(
{
globalViewAccess: categories.globalViewAccess.length,
privateGroups: categories.privateGroups.length,
total: promptGroupsToMigrate.length,
});
},
null,
2,
),
);
if (dryRun) {
return {

View file

@ -1,7 +1,7 @@
import { logger } from '@librechat/data-schemas';
import { AccessRoleIds, ResourceType, PrincipalType, Constants } from 'librechat-data-provider';
import type { AccessRoleMethods, IAgent } from '@librechat/data-schemas';
import type { Model } from 'mongoose';
import type { Model, Mongoose, mongo } from 'mongoose';
const { GLOBAL_PROJECT_NAME } = Constants;
@ -17,7 +17,8 @@ export interface MigrationCheckDbMethods {
}
export interface MigrationCheckParams {
db: MigrationCheckDbMethods;
mongoose: Mongoose;
methods: MigrationCheckDbMethods;
AgentModel: Model<IAgent>;
}
@ -46,16 +47,44 @@ export interface MigrationCheckResult {
* This performs a dry-run check similar to the migration script
*/
export async function checkAgentPermissionsMigration({
db,
methods,
mongoose,
AgentModel,
}: MigrationCheckParams): Promise<MigrationCheckResult> {
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');
}
// Verify required roles exist
const ownerRole = await db.findRoleByIdentifier(AccessRoleIds.AGENT_OWNER);
const viewerRole = await db.findRoleByIdentifier(AccessRoleIds.AGENT_VIEWER);
const editorRole = await db.findRoleByIdentifier(AccessRoleIds.AGENT_EDITOR);
const ownerRole = await methods.findRoleByIdentifier(AccessRoleIds.AGENT_OWNER);
const viewerRole = await methods.findRoleByIdentifier(AccessRoleIds.AGENT_VIEWER);
const editorRole = await methods.findRoleByIdentifier(AccessRoleIds.AGENT_EDITOR);
if (!ownerRole || !viewerRole || !editorRole) {
logger.warn(
@ -70,7 +99,7 @@ export async function checkAgentPermissionsMigration({
}
// Get global project agent IDs
const globalProject = await db.getProjectByName(GLOBAL_PROJECT_NAME, ['agentIds']);
const globalProject = await methods.getProjectByName(GLOBAL_PROJECT_NAME, ['agentIds']);
const globalAgentIds = new Set(globalProject?.agentIds || []);
// Find agents without ACL entries (no batching for efficiency on startup)

View file

@ -1,7 +1,7 @@
import { logger } from '@librechat/data-schemas';
import { AccessRoleIds, ResourceType, PrincipalType, Constants } from 'librechat-data-provider';
import type { AccessRoleMethods, IPromptGroupDocument } from '@librechat/data-schemas';
import type { Model } from 'mongoose';
import type { Model, Mongoose, mongo } from 'mongoose';
const { GLOBAL_PROJECT_NAME } = Constants;
@ -17,7 +17,8 @@ export interface PromptMigrationCheckDbMethods {
}
export interface PromptMigrationCheckParams {
db: PromptMigrationCheckDbMethods;
mongoose: Mongoose;
methods: PromptMigrationCheckDbMethods;
PromptGroupModel: Model<IPromptGroupDocument>;
}
@ -44,16 +45,45 @@ export interface PromptMigrationCheckResult {
* This performs a dry-run check similar to the migration script
*/
export async function checkPromptPermissionsMigration({
db,
methods,
mongoose,
PromptGroupModel,
}: PromptMigrationCheckParams): Promise<PromptMigrationCheckResult> {
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');
}
// Verify required roles exist
const ownerRole = await db.findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_OWNER);
const viewerRole = await db.findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_VIEWER);
const editorRole = await db.findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_EDITOR);
const ownerRole = await methods.findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_OWNER);
const viewerRole = await methods.findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_VIEWER);
const editorRole = await methods.findRoleByIdentifier(AccessRoleIds.PROMPTGROUP_EDITOR);
if (!ownerRole || !viewerRole || !editorRole) {
logger.warn(
@ -66,8 +96,8 @@ export async function checkPromptPermissionsMigration({
};
}
// Get global project prompt group IDs
const globalProject = await db.getProjectByName(GLOBAL_PROJECT_NAME, ['promptGroupIds']);
/** Global project prompt group IDs */
const globalProject = await methods.getProjectByName(GLOBAL_PROJECT_NAME, ['promptGroupIds']);
const globalPromptGroupIds = new Set(
(globalProject?.promptGroupIds || []).map((id) => id.toString()),
);