import type { Collection , Db , MongoClientOptions , ClientSession , BulkWriteResult , Filter , UpdateFilter , OptionalUnlessRequiredId , UpdateResult , DeleteResult , InsertManyResult , InsertOneResult , } from 'mongodb' ; import { MongoClient } from 'mongodb' ; import { ObjectId } from 'bson' ; import { inspect } from 'util' ; import type { ConnectionConfig , QueryResult , Transaction , QueryOrderMap , FilterQuery , AnyEntity , EntityName , Dictionary , EntityData , TransactionEventBroadcaster , IsolationLevel , Configuration , ConnectionOptions , } from '@mikro-orm/core' ; import { Connection , Utils , QueryOrder , EventType , ValidationError } from '@mikro-orm/core' ; export class MongoConnection extends Connection { protected client ! : MongoClient ; protected db ! : Db ; private connected = false ; constructor ( config , options ? , type = 'write' ) { super ( config , options , type ) ; ObjectId . prototype [ inspect . custom ] = function ( ) { return `ObjectId('${ this . toHexString ( ) }')` ; } ; Date . prototype [ inspect . custom ] = function ( ) { return `ISODate('${ this . toISOString ( ) }')` ; } ; } async connect ( ) { this . client = new MongoClient ( this . config . getClientUrl ( ) , this . getConnectionOptions ( ) ) ; await this . client . connect ( ) ; this . db = this . client . db ( this . config . get ( 'dbName' ) ) ; this . connected = true ; } async close ( force ? ) { await this . client ?. close ( ! ! force ) ; this . connected = false ; } async isConnected ( ) { return this . connected ; } getCollection < T > ( name ) { return this . db . collection < T > ( this . getCollectionName ( name ) ) ; } async createCollection < T > ( name ) { return this . db . createCollection ( this . getCollectionName ( name ) ) ; } async listCollections ( ) { const collections = await this . db . listCollections ( { } , { nameOnly : true } ) . toArray ( ) ; return collections . map ( c => c . name ) ; } async dropCollection ( name ) { return this . db . dropCollection ( this . getCollectionName ( name ) ) ; } getDefaultClientUrl ( ) { return 'mongodb://127.0.0.1:27017' ; } getConnectionOptions ( ) { const ret = { } ; const pool = this . config . get ( 'pool' ) ! ; const username = this . config . get ( 'user' ) ; const password = this . config . get ( 'password' ) as string ; if ( this . config . get ( 'host' ) ) { throw new ValidationError ( 'Mongo driver does not support `host` options, use `clientUrl` instead!' ) ; } if ( username && password ) { ret . auth = { username , password } ; } if ( pool . min ) { ret . minPoolSize = pool . min ; } if ( pool . max ) { ret . maxPoolSize = pool . max ; } return Utils . merge ( ret , this . config . get ( 'driverOptions' ) ) ; } getClientUrl ( ) { const options = this . getConnectionOptions ( ) ; const clientUrl = this . config . getClientUrl ( true ) ; const match = clientUrl . match ( /^(\w+):\/\/((.*@.+)|.+)$/ ) ; return match ? `${ match [ 1 ] }://${ options . auth ? options . auth . username + ':*****@' : '' }${ match [ 2 ] }` : clientUrl ; } getDb ( ) { return this . db ; } async execute ( query ) { throw new Error ( `${ this . constructor . name } does not support generic execute method` ) ; } async find < T extends AnyEntity < T > > ( collection , where , orderBy ? , limit ? , offset ? , fields ? , ctx ? ) { collection = this . getCollectionName ( collection ) ; const options = { session : ctx } ; if ( fields ) { options . projection = fields . reduce ( ( o , k ) => ( { ... o , [ k ] : 1 } ) , { } ) ; } const resultSet = this . getCollection < T > ( collection ) . find ( where as Filter < T > , options ) ; let query = `db.getCollection('${ collection }').find(${ this . logObject ( where ) }, ${ this . logObject ( options ) })` ; orderBy = Utils . asArray ( orderBy ) ; if ( orderBy . length > 0 ) { const orderByTuples = [ ] ; orderBy . forEach ( o => { Object . keys ( o ) . forEach ( k => { const direction = o [ k ] ; orderByTuples . push ( [ k , Utils . isString ( direction ) ? direction . toUpperCase ( ) === QueryOrder . ASC ? 1 : - 1 : direction ] ) ; } ) ; } ) ; if ( orderByTuples . length > 0 ) { query += `.sort(${ this . logObject ( orderByTuples ) })` ; resultSet . sort ( orderByTuples ) ; } } if ( limit !== undefined ) { query += `.limit(${ limit })` ; resultSet . limit ( limit ) ; } if ( offset !== undefined ) { query += `.skip(${ offset })` ; resultSet . skip ( offset ) ; } const now = Date . now ( ) ; const res = await resultSet . toArray ( ) ; this . logQuery ( `${ query }.toArray();` , { took : Date . now ( ) - now } ) ; return res as EntityData < T > [ ] ; } async insertOne < T > ( collection , data , ctx ? ) { return this . runQuery < T > ( 'insertOne' , collection , data , undefined , ctx ) ; } async insertMany < T > ( collection , data , ctx ? ) { return this . runQuery < T > ( 'insertMany' , collection , data , undefined , ctx ) ; } async updateMany < T > ( collection , where , data , ctx ? ) { return this . runQuery < T > ( 'updateMany' , collection , data , where , ctx ) ; } async bulkUpdateMany < T > ( collection , where , data , ctx ? ) { return this . runQuery < T > ( 'bulkUpdateMany' , collection , data , where , ctx ) ; } async deleteMany < T > ( collection , where , ctx ? ) { return this . runQuery < T > ( 'deleteMany' , collection , undefined , where , ctx ) ; } async aggregate < T = any > ( collection , pipeline , ctx ? ) { collection = this . getCollectionName ( collection ) ; const options = { session : ctx } ; const query = `db.getCollection('${ collection }').aggregate(${ this . logObject ( pipeline ) }, ${ this . logObject ( options ) }).toArray();` ; const now = Date . now ( ) ; const res = this . getCollection ( collection ) . aggregate < T > ( pipeline , options ) . toArray ( ) ; this . logQuery ( query , { took : Date . now ( ) - now } ) ; return res ; } async countDocuments < T > ( collection , where , ctx ? ) { return this . runQuery < T , number > ( 'countDocuments' , collection , undefined , where , ctx ) ; } async transactional < T > ( cb , options = { } ) { const session = await this . begin ( options ) ; try { const ret = await cb ( session ) ; await this . commit ( session , options . eventBroadcaster ) ; return ret ; } catch ( error ) { await this . rollback ( session , options . eventBroadcaster ) ; throw error ; } finally { await session . endSession ( ) ; } } async begin ( options = { } ) { if ( ! options . ctx ) { await options . eventBroadcaster ?. dispatchEvent ( EventType . beforeTransactionStart ) ; } const session = options . ctx || this . client . startSession ( ) ; session . startTransaction ( ) ; this . logQuery ( 'db.begin();' ) ; await options . eventBroadcaster ?. dispatchEvent ( EventType . afterTransactionStart , session ) ; return session ; } async commit ( ctx , eventBroadcaster ? ) { await eventBroadcaster ?. dispatchEvent ( EventType . beforeTransactionCommit , ctx ) ; await ctx . commitTransaction ( ) ; this . logQuery ( 'db.commit();' ) ; await eventBroadcaster ?. dispatchEvent ( EventType . afterTransactionCommit , ctx ) ; } async rollback ( ctx , eventBroadcaster ? ) { await eventBroadcaster ?. dispatchEvent ( EventType . beforeTransactionRollback , ctx ) ; await ctx . abortTransaction ( ) ; this . logQuery ( 'db.rollback();' ) ; await eventBroadcaster ?. dispatchEvent ( EventType . afterTransactionRollback , ctx ) ; } private async runQuery < T , U extends QueryResult < T > | number = QueryResult < T > > ( method , collection , data ? , where ? , ctx ? ) { collection = this . getCollectionName ( collection ) ; const logger = this . config . getLogger ( ) ; const options = { session : ctx } ; const now = Date . now ( ) ; let res ; let query ; const log = ( msg ) => logger . isEnabled ( 'query' ) ? msg ( ) : '' ; switch ( method ) { case 'insertOne' : query = log ( ( ) => `db.getCollection('${ collection }').insertOne(${ this . logObject ( data ) }, ${ this . logObject ( options ) });` ) ; res = await this . getCollection < T > ( collection ) . insertOne ( data as OptionalUnlessRequiredId < T > , options ) ; break ; case 'insertMany' : query = log ( ( ) => `db.getCollection('${ collection }').insertMany(${ this . logObject ( data ) }, ${ this . logObject ( options ) });` ) ; res = await this . getCollection < T > ( collection ) . insertMany ( data as OptionalUnlessRequiredId < T > [ ] , options ) ; break ; case 'updateMany' : { const payload = Object . keys ( data ! ) . some ( k => k . startsWith ( '$' ) ) ? data : { $set : data } ; query = log ( ( ) => `db.getCollection('${ collection }').updateMany(${ this . logObject ( where ) }, ${ this . logObject ( payload ) }, ${ this . logObject ( options ) });` ) ; res = await this . getCollection < T > ( collection ) . updateMany ( where as Filter < T > , payload as UpdateFilter < T > , options ) as UpdateResult ; break ; } case 'bulkUpdateMany' : { query = log ( ( ) => `bulk = db.getCollection('${ collection }').initializeUnorderedBulkOp(${ this . logObject ( options ) });\n` ) ; const bulk = this . getCollection < T > ( collection ) . initializeUnorderedBulkOp ( options ) ; ( data as T [ ] ) . forEach ( ( row , idx ) => { const cond = { _id : ( where as Dictionary [ ] ) [ idx ] } ; const doc = { $set : row } ; query += log ( ( ) => `bulk.find(${ this . logObject ( cond ) }).update(${ this . logObject ( doc ) });\n` ) ; bulk . find ( cond ) . update ( doc ) ; } ) ; query += log ( ( ) => `bulk.execute()` ) ; res = await bulk . execute ( ) ! ; break ; } case 'deleteMany' : case 'countDocuments' : query = log ( ( ) => `db.getCollection('${ collection }').${ method }(${ this . logObject ( where ) }, ${ this . logObject ( options ) });` ) ; res = await this . getCollection < T > ( collection ) [ method as 'deleteMany' ] ( where as Filter < T > , options ) ; break ; } this . logQuery ( query ! , { took : Date . now ( ) - now } ) ; if ( method === 'countDocuments' ) { return res ! as unknown as U ; } return this . transformResult < T > ( res ! ) as U ; } private transformResult < T > ( res ) { return { affectedRows : res . modifiedCount || res . deletedCount || res . insertedCount || 0 , insertId : res . insertedId ?? res . insertedIds ?. [ 0 ] , insertedIds : res . insertedIds ? Object . values ( res . insertedIds ) : undefined , } ; } private getCollectionName < T > ( name ) { name = Utils . className ( name ) ; const meta = this . metadata . find ( name ) ; return meta ? meta . collection : name ; } private logObject ( o ) { if ( o . session ) { o = { ... o , session : `[ClientSession]` } ; } return inspect ( o , { depth : 5 , compact : true , breakLength : 300 } ) ; } }