backfill_library_references_search.mjs 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. // @ts-check
  2. import minimist from 'minimist'
  3. import logger from '@overleaf/logger'
  4. import { db } from '../app/src/infrastructure/mongodb.mjs'
  5. import {
  6. buildSearchTokens,
  7. buildMatchTokens,
  8. docSchema,
  9. } from '../modules/library/app/src/LibraryReferenceRepository.mts'
  10. import { tokenize } from '../modules/library/app/src/bibtex-search-tokens.mts'
  11. import { scriptRunner } from './lib/ScriptRunner.mjs'
  12. /** @typedef {import('mongodb').AnyBulkWriteOperation} AnyBulkWriteOperation */
  13. const argv = minimist(process.argv.slice(2), {
  14. boolean: ['commit', 'all', 'help'],
  15. default: { 'batch-size': 1000 },
  16. })
  17. function usage() {
  18. logger.info(
  19. {},
  20. `Usage: node backfill_library_references_search.mjs [options]
  21. Populates searchKey, searchTokens and matchTokens on libraryReferences so
  22. the account-level library search and import duplicate-detection can index
  23. them. Also unsets the obsolete fields.$[].searchValue. Safe to rerun; picks
  24. up only un-indexed rows by default.
  25. Options:
  26. --commit Apply changes. Without this, runs as a dry run.
  27. --all Re-index every row, not just rows where searchTokens
  28. is null. Use when the tokenization format has changed.
  29. --batch-size <n> bulkWrite batch size (default 1000).
  30. `
  31. )
  32. }
  33. if (argv.help) {
  34. usage()
  35. process.exit(0)
  36. }
  37. const BATCH_SIZE = Number(argv['batch-size'])
  38. /** @param {(message: string) => Promise<void>} trackProgress */
  39. async function backfill(trackProgress) {
  40. const filter = argv.all ? {} : { searchTokens: null }
  41. const cursor = db.libraryReferences
  42. .find(filter)
  43. .project({ key: 1, type: 1, fields: 1, updatedAt: 1 })
  44. let processed = 0
  45. /** @type {AnyBulkWriteOperation[]} */
  46. let ops = []
  47. const flush = async () => {
  48. if (ops.length === 0) return
  49. if (argv.commit) {
  50. await db.libraryReferences.bulkWrite(ops, { ordered: false })
  51. }
  52. processed += ops.length
  53. await trackProgress(
  54. `${argv.commit ? 'wrote' : '[dry-run]'} ${processed} docs`
  55. )
  56. ops = []
  57. }
  58. for await (const doc of cursor) {
  59. const entry = docSchema.parse({
  60. ...doc,
  61. type: doc.type ?? 'misc',
  62. updatedAt: doc.updatedAt ?? new Date(0),
  63. })
  64. const searchKey = tokenize(doc.key)
  65. const searchTokens = buildSearchTokens(entry)
  66. const matchTokens = buildMatchTokens(entry)
  67. ops.push({
  68. updateOne: {
  69. filter: { _id: doc._id },
  70. update: {
  71. $set: { searchKey, searchTokens, matchTokens },
  72. $unset: { 'fields.$[].searchValue': 1 },
  73. },
  74. },
  75. })
  76. if (ops.length >= BATCH_SIZE) {
  77. await flush()
  78. }
  79. }
  80. await flush()
  81. await trackProgress(`done; processed ${processed} docs`)
  82. }
  83. /** @param {(message: string) => Promise<void>} trackProgress */
  84. async function main(trackProgress) {
  85. if (!argv.commit) {
  86. await trackProgress('DRY RUN. Pass --commit to apply changes.')
  87. }
  88. await backfill(trackProgress)
  89. }
  90. try {
  91. await scriptRunner(main, {
  92. commit: Boolean(argv.commit),
  93. all: Boolean(argv.all),
  94. batchSize: BATCH_SIZE,
  95. })
  96. process.exit(0)
  97. } catch (err) {
  98. logger.error({ err }, 'backfill failed')
  99. process.exit(1)
  100. }