backfill_library_references_search.mjs 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  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 { buildSearchFields } from '../modules/library/app/src/LibraryReferenceRepository.mts'
  6. import { scriptRunner } from './lib/ScriptRunner.mjs'
  7. /** @typedef {import('mongodb').AnyBulkWriteOperation} AnyBulkWriteOperation */
  8. const argv = minimist(process.argv.slice(2), {
  9. boolean: ['commit', 'rollback', 'all', 'help'],
  10. default: { 'batch-size': 1000 },
  11. })
  12. function usage() {
  13. logger.info(
  14. {},
  15. `Usage: node backfill_library_references_search.mjs [options]
  16. Populates searchKey and fields.searchValue on libraryReferences so the
  17. account-level library search can index them. Safe to rerun; picks up only
  18. un-indexed rows by default.
  19. Options:
  20. --commit Apply changes. Without this, runs as a dry run.
  21. --rollback Unset searchKey and fields.searchValue on all rows that
  22. have them. Mirrors the original migration's rollback.
  23. --all Re-index every row, not just rows where searchKey is null.
  24. Use when the tokenization format has changed.
  25. --batch-size <n> bulkWrite batch size (default 1000).
  26. `
  27. )
  28. }
  29. if (argv.help) {
  30. usage()
  31. process.exit(0)
  32. }
  33. const BATCH_SIZE = Number(argv['batch-size'])
  34. /** @param {(message: string) => Promise<void>} trackProgress */
  35. async function backfill(trackProgress) {
  36. const filter = argv.all ? {} : { searchKey: null }
  37. const cursor = db.libraryReferences
  38. .find(filter)
  39. .hint({ userId: 1, searchKey: 1 })
  40. .project({ key: 1, fields: 1 })
  41. let processed = 0
  42. /** @type {AnyBulkWriteOperation[]} */
  43. let ops = []
  44. const flush = async () => {
  45. if (ops.length === 0) return
  46. if (argv.commit) {
  47. await db.libraryReferences.bulkWrite(ops, { ordered: false })
  48. }
  49. processed += ops.length
  50. await trackProgress(
  51. `${argv.commit ? 'wrote' : '[dry-run]'} ${processed} docs`
  52. )
  53. ops = []
  54. }
  55. for await (const doc of cursor) {
  56. const { searchKey, fields } = buildSearchFields({
  57. key: doc.key,
  58. fields: (doc.fields ?? []).map(
  59. (/** @type {{ name: string; editableValue?: string }} */ f) => ({
  60. name: f.name,
  61. editableValue: f.editableValue ?? '',
  62. })
  63. ),
  64. })
  65. ops.push({
  66. updateOne: {
  67. filter: { _id: doc._id },
  68. update: { $set: { searchKey, fields } },
  69. },
  70. })
  71. if (ops.length >= BATCH_SIZE) {
  72. await flush()
  73. }
  74. }
  75. await flush()
  76. await trackProgress(`done; processed ${processed} docs`)
  77. }
  78. /** @param {(message: string) => Promise<void>} trackProgress */
  79. async function rollback(trackProgress) {
  80. if (!argv.commit) {
  81. const count = await db.libraryReferences.countDocuments({
  82. searchKey: { $ne: null },
  83. })
  84. await trackProgress(`[dry-run] would unset search fields on ${count} docs`)
  85. return
  86. }
  87. const result = await db.libraryReferences.updateMany(
  88. { searchKey: { $ne: null } },
  89. { $unset: { searchKey: 1, 'fields.$[].searchValue': 1 } },
  90. { hint: { userId: 1, searchKey: 1 } }
  91. )
  92. await trackProgress(`unset search fields on ${result.modifiedCount} docs`)
  93. }
  94. /** @param {(message: string) => Promise<void>} trackProgress */
  95. async function main(trackProgress) {
  96. if (!argv.commit) {
  97. await trackProgress('DRY RUN. Pass --commit to apply changes.')
  98. }
  99. if (argv.rollback) {
  100. await rollback(trackProgress)
  101. } else {
  102. await backfill(trackProgress)
  103. }
  104. }
  105. try {
  106. await scriptRunner(main, {
  107. commit: Boolean(argv.commit),
  108. rollback: Boolean(argv.rollback),
  109. all: Boolean(argv.all),
  110. batchSize: BATCH_SIZE,
  111. })
  112. process.exit(0)
  113. } catch (err) {
  114. logger.error({ err }, 'backfill failed')
  115. process.exit(1)
  116. }