add_user_count_to_csv.mjs 3.0 KB

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