platform/server/mongo/src/storage.ts

//
// Copyright © 2020 Anticrm Platform Contributors.
//
// Licensed under the Eclipse Public License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License. You may
// obtain a copy of the License at https://www.eclipse.org/legal/epl-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//
// See the License for the specific language governing permissions and
// limitations under the License.
//

import core, {
  Class,
  Doc,
  DocumentQuery,
  DocumentUpdate,
  Domain,
  DOMAIN_MODEL,
  DOMAIN_TX,
  escapeLikeForRegexp,
  FindOptions,
  FindResult,
  Hierarchy,
  isOperator,
  Lookup,
  Mixin,
  ModelDb,
  Projection,
  QueryUpdate,
  Ref,
  ReverseLookups,
  SortingOrder,
  SortingQuery,
  StorageIterator,
  toFindResult,
  Tx,
  TxCreateDoc,
  TxMixin,
  TxProcessor,
  TxRemoveDoc,
  TxResult,
  TxUpdateDoc,
  WithLookup,
  WorkspaceId
} from '@hcengineering/core'
import type { DbAdapter, TxAdapter } from '@hcengineering/server-core'
import { Collection, Db, Document, Filter, MongoClient, Sort, UpdateFilter } from 'mongodb'
import { createHash } from 'node:crypto'
import { getMongoClient, getWorkspaceDB } from './utils'

function translateDoc (doc: Doc): Document {
  return doc as Document
}

function isLookupQuery<T extends Doc> (query: DocumentQuery<T>): boolean {
  for (const key in query) {
    if (key.includes('$lookup.')) return true
  }
  return false
}

function isLookupSort<T extends Doc> (sort: SortingQuery<T> | undefined): boolean {
  if (sort === undefined) return false
  for (const key in sort) {
    if (key.includes('$lookup.')) return true
  }
  return false
}

interface LookupStep {
  from: string
  localField?: string
  foreignField?: string
  as: string
  let?: any
  pipeline?: any
}

abstract class MongoAdapterBase extends TxProcessor {
  constructor (
    protected readonly db: Db,
    protected readonly hierarchy: Hierarchy,
    protected readonly modelDb: ModelDb,
    protected readonly client: MongoClient
  ) {
    super()
  }

  async init (): Promise<void> {}

  async close (): Promise<void> {
    await this.client.close()
  }

  private translateQuery<T extends Doc>(clazz: Ref<Class<T>>, query: DocumentQuery<T>): Filter<Document> {
    const translated: any = {}
    for (const key in query) {
      const value = (query as any)[key]

      const tkey = this.translateKey(key, clazz)
      if (value !== null && typeof value === 'object') {
        const keys = Object.keys(value)
        if (keys[0] === '$like') {
          const pattern = value.$like as string
          translated[tkey] = {
            $regex: `^${pattern
              .split('%')
              .map((it) => escapeLikeForRegexp(it))
              .join('.*')}$`,
            $options: 'i'
          }
          continue
        }
      }
      translated[tkey] = value
    }
    const baseClass = this.hierarchy.getBaseClass(clazz)
    const classes = this.hierarchy.getDescendants(baseClass)

    // Only replace if not specified
    if (translated._class === undefined) {
      translated._class = { $in: classes }
    }

    if (baseClass !== clazz) {
      // Add an mixin to be exists flag
      translated[clazz] = { $exists: true }
    }
    // return Object.assign({}, query, { _class: { $in: classes } })
    return translated
  }

  private async getLookupValue<T extends Doc>(
    clazz: Ref<Class<T>>,
    lookup: Lookup<T>,
    result: LookupStep[],
    parent?: string
  ): Promise<void> {
    for (const key in lookup) {
      if (key === '_id') {
        await this.getReverseLookupValue(lookup, result, parent)
        continue
      }
      const value = (lookup as any)[key]
      if (Array.isArray(value)) {
        const [_class, nested] = value
        const tkey = this.checkMixinKey(key, clazz)
        const fullKey = parent !== undefined ? parent + '.' + tkey : tkey
        const domain = this.hierarchy.getDomain(_class)
        if (domain !== DOMAIN_MODEL) {
          result.push({
            from: domain,
            localField: fullKey,
            foreignField: '_id',
            as: fullKey.split('.').join('') + '_lookup'
          })
        }
        await this.getLookupValue(_class, nested, result, fullKey + '_lookup')
      } else {
        const _class = value as Ref<Class<Doc>>
        const tkey = this.checkMixinKey(key, clazz)
        const fullKey = parent !== undefined ? parent + '.' + tkey : tkey
        const domain = this.hierarchy.getDomain(_class)
        if (domain !== DOMAIN_MODEL) {
          result.push({
            from: domain,
            localField: fullKey,
            foreignField: '_id',
            as: fullKey.split('.').join('') + '_lookup'
          })
        }
      }
    }
  }

  private async getReverseLookupValue (
    lookup: ReverseLookups,
    result: LookupStep[],
    parent?: string
  ): Promise<any | undefined> {
    const fullKey = parent !== undefined ? parent + '.' + '_id' : '_id'
    const lid = lookup?._id ?? {}
    for (const key in lid) {
      const as = parent !== undefined ? parent + key : key
      const value = lid[key]

      let _class: Ref<Class<Doc>>
      let attr = 'attachedTo'

      if (Array.isArray(value)) {
        _class = value[0]
        attr = value[1]
      } else {
        _class = value
      }
      const domain = this.hierarchy.getDomain(_class)
      const desc = this.hierarchy.getDescendants(_class)
      if (domain !== DOMAIN_MODEL) {
        const asVal = as.split('.').join('') + '_lookup'
        const step: LookupStep = {
          from: domain,
          localField: fullKey,
          foreignField: attr,
          pipeline: [
            {
              $match: {
                _class: { $in: desc }
              }
            }
          ],
          as: asVal
        }
        result.push(step)
      }
    }
  }

  private async getLookups<T extends Doc>(
    _class: Ref<Class<T>>,
    lookup: Lookup<T> | undefined,
    parent?: string
  ): Promise<LookupStep[]> {
    if (lookup === undefined) return []
    const result: [] = []
    await this.getLookupValue(_class, lookup, result, parent)
    return result
  }

  private async fillLookup<T extends Doc>(
    _class: Ref<Class<T>>,
    object: any,
    key: string,
    fullKey: string,
    targetObject: any
  ): Promise<void> {
    if (targetObject.$lookup === undefined) {
      targetObject.$lookup = {}
    }
    const domain = this.hierarchy.getDomain(_class)
    if (domain !== DOMAIN_MODEL) {
      const arr = object[fullKey]
      if (arr !== undefined && Array.isArray(arr)) {
        if (arr.length === 1) {
          targetObject.$lookup[key] = arr[0]
        } else if (arr.length > 1) {
          targetObject.$lookup[key] = arr
        }
      }
    } else {
      targetObject.$lookup[key] = (await this.modelDb.findAll(_class, { _id: targetObject[key] }))[0]
    }
  }

  private async fillLookupValue<T extends Doc>(
    clazz: Ref<Class<T>>,
    lookup: Lookup<T> | undefined,
    object: any,
    parent?: string,
    parentObject?: any
  ): Promise<void> {
    if (lookup === undefined) return
    for (const key in lookup) {
      if (key === '_id') {
        await this.fillReverseLookup(clazz, lookup, object, parent, parentObject)
        continue
      }
      const value = (lookup as any)[key]
      const tkey = this.checkMixinKey(key, clazz).split('.').join('')
      const fullKey = parent !== undefined ? parent + tkey + '_lookup' : tkey + '_lookup'
      const targetObject = parentObject ?? object
      if (Array.isArray(value)) {
        const [_class, nested] = value
        await this.fillLookup(_class, object, key, fullKey, targetObject)
        await this.fillLookupValue(_class, nested, object, fullKey, targetObject.$lookup[key])
      } else {
        await this.fillLookup(value, object, key, fullKey, targetObject)
      }
    }
  }

  private async fillReverseLookup<T extends Doc>(
    clazz: Ref<Class<T>>,
    lookup: ReverseLookups,
    object: any,
    parent?: string,
    parentObject?: any
  ): Promise<void> {
    const targetObject = parentObject ?? object
    if (targetObject.$lookup === undefined) {
      targetObject.$lookup = {}
    }
    for (const key in lookup._id) {
      const value = lookup._id[key]
      let _class: Ref<Class<Doc>>
      let attr = 'attachedTo'

      if (Array.isArray(value)) {
        _class = value[0]
        attr = value[1]
      } else {
        _class = value
      }
      const domain = this.hierarchy.getDomain(_class)
      const tkey = this.checkMixinKey(key, clazz).split('.').join('')
      const fullKey = parent !== undefined ? parent + tkey + '_lookup' : tkey + '_lookup'
      if (domain !== DOMAIN_MODEL) {
        const arr = object[fullKey]
        targetObject.$lookup[key] = arr
      } else {
        const arr = await this.modelDb.findAll(_class, { [attr]: targetObject._id })
        targetObject.$lookup[key] = arr
      }
    }
  }

  private async lookup<T extends Doc>(
    clazz: Ref<Class<T>>,
    query: DocumentQuery<T>,
    options: FindOptions<T>
  ): Promise<FindResult<T>> {
    const pipeline = []
    const match = { $match: this.translateQuery(clazz, query) }
    const slowPipeline = isLookupQuery(query) || isLookupSort(options.sort)
    const steps = await this.getLookups(clazz, options.lookup)
    if (slowPipeline) {
      for (const step of steps) {
        pipeline.push({ $lookup: step })
      }
    }
    pipeline.push(match)
    const resultPipeline: any[] = []
    if (options.sort !== undefined) {
      const sort = {} as any
      for (const _key in options.sort) {
        const key: string = this.translateKey(_key, clazz)
        sort[key] = options.sort[_key] === SortingOrder.Ascending ? 1 : -1
      }
      pipeline.push({ $sort: sort })
    }
    if (options.limit !== undefined) {
      resultPipeline.push({ $limit: options.limit })
    }
    if (!slowPipeline) {
      for (const step of steps) {
        resultPipeline.push({ $lookup: step })
      }
    }
    if (options?.projection !== undefined) {
      const projection: Projection<T> = {}
      for (const key in options.projection) {
        const ckey = this.checkMixinKey<T>(key, clazz) as keyof T
        projection[ckey] = options.projection[key]
      }
      resultPipeline.push({ $project: projection })
    }
    pipeline.push({
      $facet: {
        results: resultPipeline,
        totalCount: [
          {
            $count: 'count'
          }
        ]
      }
    })
    const domain = this.hierarchy.getDomain(clazz)
    const cursor = this.db.collection(domain).aggregate(pipeline)
    const res = (await cursor.toArray())[0]
    const result = res.results as WithLookup<T>[]
    const total = res.totalCount?.shift()?.count
    for (const row of result) {
      row.$lookup = {}
      await this.fillLookupValue(clazz, options.lookup, row)
      this.clearExtraLookups(row)
    }
    return toFindResult(result, total)
  }

  private translateKey<T extends Doc>(key: string, clazz: Ref<Class<T>>): string {
    const arr = key.split('.').filter((p) => p)
    let tKey = ''

    for (let i = 0; i < arr.length; i++) {
      const element = arr[i]
      if (element === '$lookup') {
        tKey += arr[++i] + '_lookup'
      } else {
        if (!tKey.endsWith('.') && i > 0) {
          tKey += '.'
        }
        tKey += arr[i]
        if (i !== arr.length - 1) {
          tKey += '.'
        }
      }
      // Check if key is belong to mixin class, we need to add prefix.
      tKey = this.checkMixinKey<T>(tKey, clazz)
    }

    return tKey
  }

  private clearExtraLookups (row: any): void {
    for (const key in row) {
      if (key.endsWith('_lookup')) {
        // eslint-disable-next-line
        delete row[key]
      }
    }
  }

  private checkMixinKey<T extends Doc>(key: string, clazz: Ref<Class<T>>): string {
    if (!key.includes('.')) {
      try {
        const attr = this.hierarchy.findAttribute(clazz, key)
        if (attr !== undefined && this.hierarchy.isMixin(attr.attributeOf)) {
          // It is mixin
          key = attr.attributeOf + '.' + key
        }
      } catch (err: any) {
        // ignore, if
      }
    }
    return key
  }

  async findAll<T extends Doc>(
    _class: Ref<Class<T>>,
    query: DocumentQuery<T>,
    options?: FindOptions<T>
  ): Promise<FindResult<T>> {
    // TODO: rework this
    if (options !== null && options !== undefined) {
      if (options.lookup !== undefined) {
        return await this.lookup(_class, query, options)
      }
    }
    const domain = this.hierarchy.getDomain(_class)
    const coll = this.db.collection(domain)
    const mongoQuery = this.translateQuery(_class, query)
    let cursor = coll.find<T>(mongoQuery)

    if (options?.projection !== undefined) {
      const projection: Projection<T> = {}
      for (const key in options.projection) {
        const ckey = this.checkMixinKey<T>(key, _class) as keyof T
        projection[ckey] = options.projection[key]
      }
      cursor = cursor.project(projection)
    }
    let total: number | undefined
    if (options !== null && options !== undefined) {
      if (options.sort !== undefined) {
        const sort: Sort = {}
        for (const key in options.sort) {
          const ckey = this.checkMixinKey<T>(key, _class)
          const order = options.sort[key] === SortingOrder.Ascending ? 1 : -1
          sort[ckey] = order
        }
        cursor = cursor.sort(sort)
      }
      if (options.limit !== undefined) {
        total = await coll.countDocuments(mongoQuery)
        cursor = cursor.limit(options.limit)
      }
    }
    const res = await cursor.toArray()
    return toFindResult(res, total)
  }

  find (domain: Domain): StorageIterator {
    const coll = this.db.collection<Doc>(domain)
    const iterator = coll.find({}, {})

    return {
      next: async () => {
        const d = await iterator.next()
        if (d === null) {
          return undefined
        }
        const doc = JSON.stringify(d)
        const hash = createHash('sha256')
        hash.update(doc)
        const digest = hash.digest('base64')
        return {
          id: d._id,
          hash: digest,
          size: doc.length // Some approx size for document.
        }
      },
      close: async () => {
        await iterator.close()
      }
    }
  }

  async load (domain: Domain, docs: Ref<Doc>[]): Promise<Doc[]> {
    return await this.db
      .collection(domain)
      .find<Doc>({ _id: { $in: docs } })
      .toArray()
  }

  async upload (domain: Domain, docs: Doc[]): Promise<void> {
    const coll = this.db.collection(domain)

    const docMap = new Map(docs.map((it) => [it._id, it]))

    // remove old and insert new ones
    const keys = Array.from(docMap.keys())
    if (keys.length > 0) {
      await coll.bulkWrite(
        docs.map((it) => ({
          replaceOne: {
            filter: { _id: it._id },
            replacement: it,
            upsert: true
          }
        }))
      )

      // await coll.deleteMany({ _id: { $in: keys } })
      // await coll.insertMany(Array.from(docMap.values()) as Document[])
    }
  }

  async update (domain: Domain, operations: Map<Ref<Doc>, DocumentUpdate<Doc>>): Promise<void> {
    const coll = this.db.collection(domain)

    try {
      // remove old and insert new ones
      if (operations.size > 0) {
        await coll.bulkWrite(
          Array.from(operations.entries()).map((it) => ({
            updateOne: {
              filter: { _id: it[0] },
              update: {
                $set: it[1]
              }
            }
          }))
        )
      }
    } catch (err: any) {
      console.error(err)
    }
  }

  async clean (domain: Domain, docs: Ref<Doc>[]): Promise<void> {
    await this.db.collection(domain).deleteMany({ _id: { $in: docs } })
  }
}

class MongoAdapter extends MongoAdapterBase {
  protected override async txRemoveDoc (tx: TxRemoveDoc<Doc>): Promise<TxResult> {
    const domain = this.hierarchy.getDomain(tx.objectClass)
    await this.db.collection(domain).deleteOne({ _id: tx.objectId })
    return {}
  }

  protected async txMixin (tx: TxMixin<Doc, Doc>): Promise<TxResult> {
    const domain = this.hierarchy.getDomain(tx.objectClass)

    if (isOperator(tx.attributes)) {
      const operator = Object.keys(tx.attributes)[0]
      if (operator === '$move') {
        const keyval = (tx.attributes as any).$move
        const arr = tx.mixin + '.' + Object.keys(keyval)[0]
        const desc = keyval[arr]
        const ops = [
          {
            updateOne: {
              filter: { _id: tx.objectId },
              update: {
                $pull: {
                  [arr]: desc.$value
                }
              }
            }
          },
          {
            updateOne: {
              filter: { _id: tx.objectId },
              update: {
                $set: {
                  modifiedBy: tx.modifiedBy,
                  modifiedOn: tx.modifiedOn
                },
                $push: {
                  [arr]: {
                    $each: [desc.$value],
                    $position: desc.$position
                  }
                }
              }
            }
          }
        ]
        return await this.db.collection(domain).bulkWrite(ops as any)
      } else {
        return await this.db.collection(domain).updateOne(
          { _id: tx.objectId },
          {
            ...this.translateMixinAttrs(tx.mixin, tx.attributes),
            $set: {
              modifiedBy: tx.modifiedBy,
              modifiedOn: tx.modifiedOn
            }
          }
        )
      }
    } else {
      return await this.db.collection(domain).updateOne(
        { _id: tx.objectId },
        {
          $set: {
            ...this.translateMixinAttrs(tx.mixin, tx.attributes),
            modifiedBy: tx.modifiedBy,
            modifiedOn: tx.modifiedOn
          }
        }
      )
    }
  }

  private translateMixinAttrs (mixin: Ref<Mixin<Doc>>, attributes: Record<string, any>): Record<string, any> {
    const attrs: Record<string, any> = {}
    let count = 0
    for (const [k, v] of Object.entries(attributes)) {
      if (k.startsWith('$')) {
        attrs[k] = this.translateMixinAttrs(mixin, v)
      } else {
        attrs[mixin + '.' + k] = v
      }
      count++
    }

    if (count === 0) {
      // We need at least one attribute, to be inside for first time,
      // for mongo to create embedded object, if we don't want to get object first.
      attrs[mixin + '.' + '__mixin'] = 'true'
    }
    return attrs
  }

  protected override async txCreateDoc (tx: TxCreateDoc<Doc>): Promise<TxResult> {
    const doc = TxProcessor.createDoc2Doc(tx)
    const domain = this.hierarchy.getDomain(doc._class)
    await this.db.collection(domain).insertOne(translateDoc(doc))
    return {}
  }

  protected override async txUpdateDoc (tx: TxUpdateDoc<Doc>): Promise<TxResult> {
    const domain = this.hierarchy.getDomain(tx.objectClass)
    if (isOperator(tx.operations)) {
      const operator = Object.keys(tx.operations)[0]
      if (operator === '$move') {
        const keyval = (tx.operations as any).$move
        const arr = Object.keys(keyval)[0]
        const desc = keyval[arr]
        const ops = [
          {
            updateOne: {
              filter: { _id: tx.objectId },
              update: {
                $pull: {
                  [arr]: desc.$value
                }
              }
            }
          },
          {
            updateOne: {
              filter: { _id: tx.objectId },
              update: {
                $set: {
                  modifiedBy: tx.modifiedBy,
                  modifiedOn: tx.modifiedOn
                },
                $push: {
                  [arr]: {
                    $each: [desc.$value],
                    $position: desc.$position
                  }
                }
              }
            }
          }
        ]
        return await this.db.collection(domain).bulkWrite(ops as any)
      } else if (operator === '$update') {
        const keyval = (tx.operations as any).$update
        const arr = Object.keys(keyval)[0]
        const desc = keyval[arr] as QueryUpdate<any>
        const ops = [
          {
            updateOne: {
              filter: {
                _id: tx.objectId,
                ...Object.fromEntries(Object.entries(desc.$query).map((it) => [arr + '.' + it[0], it[1]]))
              },
              update: {
                $set: {
                  ...Object.fromEntries(Object.entries(desc.$update).map((it) => [arr + '.$.' + it[0], it[1]]))
                }
              }
            }
          },
          {
            updateOne: {
              filter: { _id: tx.objectId },
              update: {
                $set: {
                  modifiedBy: tx.modifiedBy,
                  modifiedOn: tx.modifiedOn
                }
              }
            }
          }
        ]
        return await this.db.collection(domain).bulkWrite(ops as any)
      } else {
        if (tx.retrieve === true) {
          const result = await this.db.collection(domain).findOneAndUpdate(
            { _id: tx.objectId },
            {
              ...tx.operations,
              $set: {
                modifiedBy: tx.modifiedBy,
                modifiedOn: tx.modifiedOn
              }
            } as unknown as UpdateFilter<Document>,
            { returnDocument: 'after' }
          )
          return { object: result.value }
        } else {
          return await this.db.collection(domain).updateOne(
            { _id: tx.objectId },
            {
              ...tx.operations,
              $set: {
                modifiedBy: tx.modifiedBy,
                modifiedOn: tx.modifiedOn
              }
            }
          )
        }
      }
    } else {
      if (tx.retrieve === true) {
        const result = await this.db.collection(domain).findOneAndUpdate(
          { _id: tx.objectId },
          {
            $set: { ...tx.operations, modifiedBy: tx.modifiedBy, modifiedOn: tx.modifiedOn }
          } as unknown as UpdateFilter<Document>,
          { returnDocument: 'after' }
        )
        return { object: result.value }
      } else {
        return await this.db
          .collection(domain)
          .updateOne(
            { _id: tx.objectId },
            { $set: { ...tx.operations, modifiedBy: tx.modifiedBy, modifiedOn: tx.modifiedOn } }
          )
      }
    }
  }
}

class MongoTxAdapter extends MongoAdapterBase implements TxAdapter {
  txColl: Collection | undefined
  protected txCreateDoc (tx: TxCreateDoc<Doc>): Promise<TxResult> {
    throw new Error('Method not implemented.')
  }

  protected txUpdateDoc (tx: TxUpdateDoc<Doc>): Promise<TxResult> {
    throw new Error('Method not implemented.')
  }

  protected txRemoveDoc (tx: TxRemoveDoc<Doc>): Promise<TxResult> {
    throw new Error('Method not implemented.')
  }

  protected txMixin (tx: TxMixin<Doc, Doc>): Promise<TxResult> {
    throw new Error('Method not implemented.')
  }

  override async tx (tx: Tx, user: string): Promise<TxResult>
  override async tx (tx: Tx): Promise<TxResult>

  override async tx (tx: Tx, user?: string): Promise<TxResult> {
    await this.txCollection().insertOne(translateDoc(tx))
    return {}
  }

  private txCollection (): Collection {
    if (this.txColl !== undefined) {
      return this.txColl
    }
    this.txColl = this.db.collection(DOMAIN_TX)
    return this.txColl
  }

  async getModel (): Promise<Tx[]> {
    const model = await this.db
      .collection(DOMAIN_TX)
      .find<Tx>({ objectSpace: core.space.Model })
      .sort({ _id: 1 })
      .toArray()
    // We need to put all core.account.System transactions first
    const systemTr: Tx[] = []
    const userTx: Tx[] = []

    model.forEach((tx) => (tx.modifiedBy === core.account.System ? systemTr : userTx).push(tx))

    return systemTr.concat(userTx)
  }
}

/**
 * @public
 */
export async function createMongoAdapter (
  hierarchy: Hierarchy,
  url: string,
  workspaceId: WorkspaceId,
  modelDb: ModelDb
): Promise<DbAdapter> {
  const client = await getMongoClient(url)
  const db = getWorkspaceDB(client, workspaceId)
  return new MongoAdapter(db, hierarchy, modelDb, client)
}

/**
 * @public
 */
export async function createMongoTxAdapter (
  hierarchy: Hierarchy,
  url: string,
  workspaceId: WorkspaceId,
  modelDb: ModelDb
): Promise<TxAdapter> {
  const client = await getMongoClient(url)
  const db = getWorkspaceDB(client, workspaceId)
  return new MongoTxAdapter(db, hierarchy, modelDb, client)
}