<commit_before>import { Readable } from 'stream'

import { Observable } from 'rx'
import * as _ from 'lodash'
import { Timestamp } from 'mongodb'

import { Task, Controls, CheckPoint } from './config'
import { IR, MongoDoc, ESDoc, OpLog } from './types'
import Elasticsearch from './elasticsearch'
import MongoDB from './mongodb'

export default class Processor {
  static provisionedReadCapacity: number
  static consumedReadCapacity: number
  task: Task
  controls: Controls
  mongodb: MongoDB
  elasticsearch: Elasticsearch
  queue: OpLog[][] = []
  running: boolean = false

  constructor(task: Task, controls: Controls, mongodb: MongoDB, elasticsearch: Elasticsearch) {
    this.task = task
    this.controls = controls
    this.mongodb = mongodb
    this.elasticsearch = elasticsearch
    Processor.provisionedReadCapacity = controls.mongodbReadCapacity
    Processor.consumedReadCapacity = 0
  }

  static controlReadCapacity(stream: Readable): Readable {
    if (!Processor.provisionedReadCapacity) {
      return stream
    }
    const interval = setInterval(() => {
      Processor.consumedReadCapacity = 0
      stream.resume()
    }, 1000)
    stream.addListener('data', () => {
      Processor.consumedReadCapacity++
      if (Processor.consumedReadCapacity >= Processor.provisionedReadCapacity) {
        stream.pause()
      }
    })
    stream.addListener('end', () => {
      clearInterval(interval)
    })
    return stream
  }

  transformer(
    action: 'upsert' | 'delete',
    doc: MongoDoc | ESDoc,
    timestamp?: Timestamp,
    isESDoc: boolean = false,
  ): IR | null {
    if (action === 'delete') {
      return {
        action: 'delete',
        id: doc._id.toString(),
        parent: this.task.transform.parent && _.get(doc, this.task.transform.parent),
        timestamp: timestamp ? timestamp.getHighBits() : 0,
      }
    }

    const data = _.reduce(
      this.task.transform.mapping,
      (obj, value, key) => {
        if (isESDoc) {
          key = value
        }
        if (_.has(doc, key)) {
          _.set(obj, value, _.get(doc, key))
        }
        return obj
      },
      { ...this.task.transform.static } || {},
    )
    if (_.isEmpty(data)) {
      return null
    }
    return {
      action: 'upsert',
      id: doc._id.toString(),
      data,
      parent: this.task.transform.parent && _.get(doc, this.task.transform.parent),
      timestamp: timestamp ? timestamp.getHighBits() : 0,
    }
  }

  applyUpdateMongoDoc(
    doc: MongoDoc,
    set: { [key: string]: any } = {},
    unset: { [key: string]: any } = {},
  ): MongoDoc {
    _.forEach(this.task.transform.mapping, (ignored, key) => {
      if (_.get(unset, key)) {
        _.unset(doc, key)
      }
      if (_.has(set, key)) {
        _.set(doc, key, _.get(set, key))
      }
    })
    return doc
  }

  applyUpdateESDoc(
    doc: ESDoc,
    set: { [key: string]: any } = {},
    unset: { [key: string]: any } = {},
  ): ESDoc {
    _.forEach(this.task.transform.mapping, (value, key) => {
      if (_.get(unset, key)) {
        _.unset(doc, value)
      }
      if (_.has(set, key)) {
        _.set(doc, value, _.get(set, key))
      }
    })
    return doc
  }

  ignoreUpdate(oplog: OpLog): boolean {
    let ignore = true
    if (oplog.op === 'u') {
      _.forEach(this.task.transform.mapping, (value, key) => {
        ignore =
          ignore && !(_.has(oplog.o, key) || _.has(oplog.o.$set, key) || _.get(oplog.o.$unset, key))
      })
    }
    return ignore
  }

  scan(): Observable<MongoDoc> {
    return Observable.create<MongoDoc>((observer) => {
      try {
        const stream = Processor.controlReadCapacity(this.mongodb.getCollection())
        stream.addListener('data', (doc: MongoDoc) => {
          observer.onNext(doc)
        })
        stream.addListener('error', (err: Error) => {
          observer.onError(err)
        })
        stream.addListener('end', () => {
          observer.onCompleted()
        })
      } catch (err) {
        observer.onError(err)
      }
    })
  }

  tail(): Observable<OpLog> {
    return Observable.create<OpLog>((observer) => {
      try {
        const cursor = this.mongodb.getOplog()
        cursor.forEach(
          (log: OpLog) => {
            observer.onNext(log)
          },
          () => {
            observer.onCompleted()
          },
        )
      } catch (err) {
        observer.onError(err)
      }
    })
  }

  async oplog(oplog: OpLog): Promise<IR | null> {
    try {
      switch (oplog.op) {
        case 'i': {
          return this.transformer('upsert', oplog.o, oplog.ts)
        }
        case 'u': {
          if (_.size(oplog.o2) !== 1 || !oplog.o2._id) {
            console.warn('oplog', 'cannot transform', oplog)
            return null
          }<commit_msg>放开oplog.o2的大小限制 #37 (#39)<commit_after>import { Readable } from 'stream'

import { Observable } from 'rx'
import * as _ from 'lodash'
import { Timestamp } from 'mongodb'

import { Task, Controls, CheckPoint } from './config'
import { IR, MongoDoc, ESDoc, OpLog } from './types'
import Elasticsearch from './elasticsearch'
import MongoDB from './mongodb'

export default class Processor {
  static provisionedReadCapacity: number
  static consumedReadCapacity: number
  task: Task
  controls: Controls
  mongodb: MongoDB
  elasticsearch: Elasticsearch
  queue: OpLog[][] = []
  running: boolean = false

  constructor(task: Task, controls: Controls, mongodb: MongoDB, elasticsearch: Elasticsearch) {
    this.task = task
    this.controls = controls
    this.mongodb = mongodb
    this.elasticsearch = elasticsearch
    Processor.provisionedReadCapacity = controls.mongodbReadCapacity
    Processor.consumedReadCapacity = 0
  }

  static controlReadCapacity(stream: Readable): Readable {
    if (!Processor.provisionedReadCapacity) {
      return stream
    }
    const interval = setInterval(() => {
      Processor.consumedReadCapacity = 0
      stream.resume()
    }, 1000)
    stream.addListener('data', () => {
      Processor.consumedReadCapacity++
      if (Processor.consumedReadCapacity >= Processor.provisionedReadCapacity) {
        stream.pause()
      }
    })
    stream.addListener('end', () => {
      clearInterval(interval)
    })
    return stream
  }

  transformer(
    action: 'upsert' | 'delete',
    doc: MongoDoc | ESDoc,
    timestamp?: Timestamp,
    isESDoc: boolean = false,
  ): IR | null {
    if (action === 'delete') {
      return {
        action: 'delete',
        id: doc._id.toString(),
        parent: this.task.transform.parent && _.get(doc, this.task.transform.parent),
        timestamp: timestamp ? timestamp.getHighBits() : 0,
      }
    }

    const data = _.reduce(
      this.task.transform.mapping,
      (obj, value, key) => {
        if (isESDoc) {
          key = value
        }
        if (_.has(doc, key)) {
          _.set(obj, value, _.get(doc, key))
        }
        return obj
      },
      { ...this.task.transform.static } || {},
    )
    if (_.isEmpty(data)) {
      return null
    }
    return {
      action: 'upsert',
      id: doc._id.toString(),
      data,
      parent: this.task.transform.parent && _.get(doc, this.task.transform.parent),
      timestamp: timestamp ? timestamp.getHighBits() : 0,
    }
  }

  applyUpdateMongoDoc(
    doc: MongoDoc,
    set: { [key: string]: any } = {},
    unset: { [key: string]: any } = {},
  ): MongoDoc {
    _.forEach(this.task.transform.mapping, (ignored, key) => {
      if (_.get(unset, key)) {
        _.unset(doc, key)
      }
      if (_.has(set, key)) {
        _.set(doc, key, _.get(set, key))
      }
    })
    return doc
  }

  applyUpdateESDoc(
    doc: ESDoc,
    set: { [key: string]: any } = {},
    unset: { [key: string]: any } = {},
  ): ESDoc {
    _.forEach(this.task.transform.mapping, (value, key) => {
      if (_.get(unset, key)) {
        _.unset(doc, value)
      }
      if (_.has(set, key)) {
        _.set(doc, value, _.get(set, key))
      }
    })
    return doc
  }

  ignoreUpdate(oplog: OpLog): boolean {
    let ignore = true
    if (oplog.op === 'u') {
      _.forEach(this.task.transform.mapping, (value, key) => {
        ignore =
          ignore && !(_.has(oplog.o, key) || _.has(oplog.o.$set, key) || _.get(oplog.o.$unset, key))
      })
    }
    return ignore
  }

  scan(): Observable<MongoDoc> {
    return Observable.create<MongoDoc>((observer) => {
      try {
        const stream = Processor.controlReadCapacity(this.mongodb.getCollection())
        stream.addListener('data', (doc: MongoDoc) => {
          observer.onNext(doc)
        })
        stream.addListener('error', (err: Error) => {
          observer.onError(err)
        })
        stream.addListener('end', () => {
          observer.onCompleted()
        })
      } catch (err) {
        observer.onError(err)
      }
    })
  }

  tail(): Observable<OpLog> {
    return Observable.create<OpLog>((observer) => {
      try {
        const cursor = this.mongodb.getOplog()
        cursor.forEach(
          (log: OpLog) => {
            observer.onNext(log)
          },
          () => {
            observer.onCompleted()
          },
        )
      } catch (err) {
        observer.onError(err)
      }
    })
  }

  async oplog(oplog: OpLog): Promise<IR | null> {
    try {
      switch (oplog.op) {
        case 'i': {
          return this.transformer('upsert', oplog.o, oplog.ts)
        }
        case 'u': {
          if (!oplog.o2._id) {
            console.warn('oplog', 'cannot transform', oplog)
            return null
          }
{'fraction_non_alphanumeric': 0.14737282813481264, 'fraction_numerical': 0.002616705045007327, 'mean_word_length': 2.176529255319149, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 43, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4963670', 'n_tokens_mistral': 3098, 'n_tokens_neox': 3017, 'n_words': 636}