add_user_count_to_csv.mjs 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. // Usage: node scripts/add_user_count_to_csv.mjs [OPTS] [INPUT-FILE]
  2. // Looks up the number of users for each domain in the input csv file and adds
  3. // columns for the number of users in the domain, subdomains, and total.
  4. import fs from 'node:fs'
  5. // https://github.com/import-js/eslint-plugin-import/issues/1810
  6. // eslint-disable-next-line import/no-unresolved
  7. import * as csv from 'csv/sync'
  8. import minimist from 'minimist'
  9. import UserGetter from '../app/src/Features/User/UserGetter.mjs'
  10. import { db } from '../app/src/infrastructure/mongodb.js'
  11. import _ from 'lodash'
  12. import { scriptRunner } from './lib/ScriptRunner.mjs'
  13. const argv = minimist(process.argv.slice(2), {
  14. string: ['domain', 'output'],
  15. boolean: ['help'],
  16. alias: {
  17. domain: 'd',
  18. output: 'o',
  19. },
  20. default: {
  21. domain: 'Email domain',
  22. output: '/dev/stdout',
  23. },
  24. })
  25. if (argv.help || argv._.length > 1) {
  26. console.error(`Usage: node scripts/add_user_count_to_csv.mjs [OPTS] [INPUT-FILE]
  27. Looks up the number of users for each domain in the input file and adds
  28. columns for the number of users in the domain, subdomains, and total.
  29. Options:
  30. --domain name of the csv column containing the email domain (default: "Email domain")
  31. --output output file (default: /dev/stdout)
  32. `)
  33. process.exit(1)
  34. }
  35. const input = fs.readFileSync(argv._[0], 'utf8')
  36. const records = csv.parse(input, { columns: true })
  37. if (records.length === 0) {
  38. console.error('No records in input file')
  39. process.exit(1)
  40. }
  41. async function main() {
  42. for (const record of records) {
  43. const domain = record[argv.domain]
  44. const { domainUserCount, subdomainUserCount } = await getUserCount(domain, {
  45. _id: 1,
  46. })
  47. record['Domain Users'] = domainUserCount
  48. record['Subdomain Users'] = subdomainUserCount
  49. record['Total Users'] = domainUserCount + subdomainUserCount
  50. }
  51. const output = csv.stringify(records, { header: true })
  52. fs.writeFileSync(argv.output, output)
  53. }
  54. async function getUserCount(domain) {
  55. const domainUsers = await UserGetter.promises.getUsersByHostname(domain, {
  56. _id: 1,
  57. })
  58. const subdomainUsers = await getUsersByHostnameWithSubdomain(domain, {
  59. _id: 1,
  60. })
  61. return {
  62. domainUserCount: domainUsers.length,
  63. subdomainUserCount: subdomainUsers.length,
  64. }
  65. }
  66. async function getUsersByHostnameWithSubdomain(domain, projection) {
  67. const reversedDomain = domain.trim().split('').reverse().join('')
  68. const reversedDomainRegex = _.escapeRegExp(reversedDomain)
  69. const query = {
  70. emails: { $exists: true },
  71. // look for users in subdomains of a domain, but not the domain itself
  72. // e.g. for domain 'foo.edu', match 'cs.foo.edu' but not 'foo.edu'
  73. // we use the reversed hostname index to do this efficiently
  74. // we need to escape the domain name to prevent '.' from matching any character
  75. 'emails.reversedHostname': { $regex: '^' + reversedDomainRegex + '\\.' },
  76. }
  77. return await db.users.find(query, { projection }).toArray()
  78. }
  79. try {
  80. await scriptRunner(main)
  81. console.log('Done')
  82. process.exit(0)
  83. } catch (error) {
  84. console.error(error)
  85. process.exit(1)
  86. }