populate.js

/*
 * used to populate or re-populate a MongoDB-compatible document
 * database with static fixtures at `cve-services/datadump/pre-population`
 */

const express = require('express')
const app = express()
const mongoose = require('mongoose')

const dataUtils = require('../utils/data')
const dbUtils = require('../utils/db')
const errors = require('../utils/error')
const logger = require('../middleware/logger')
const CveIdRange = require('../model/cve-id-range')
const CveId = require('../model/cve-id')
const Cve = require('../model/cve')
const Org = require('../model/org')
const User = require('../model/user')
const BaseOrg = require('../model/baseorg')
const BaseUser = require('../model/baseuser')
const ReviewObject = require('../model/reviewobject')
const Conversation = require('../model/conversation')
const Audit = require('../model/audit')

const error = new errors.IDRError()

const populateTheseCollections = {
  Cve: Cve,
  'Cve-Id-Range': CveIdRange,
  'Cve-Id': CveId,
  User: User,
  Org: Org,
  BaseOrg: BaseOrg,
  BaseUser: BaseUser,
  ReviewObject: ReviewObject,
  Conversation: Conversation,
  Audit: Audit
}

const indexesToCreate = {
  Cve: [{ 'cve.cveMetadata.cveId': 1 }, { 'cve.cveMetadata.dateUpdated': 1 }],
  'Cve-Id': [{ cve_id: 1 }, { owning_cna: 1, state: 1 }, { reserved: 1 }],
  User: [{ UUID: 1 }],
  Org: [{ UUID: 1 }, { 'authority.active_roles': 1 }]
}

// Body Parser Middleware
app.use(express.json()) // Allows us to handle raw JSON data
app.use(express.urlencoded({ extended: false })) // Allows us to handle url encoded data
// Make mongoose connection available globally
global.mongoose = mongoose

// Connect to MongoDB database
const dbConnectionStr = dbUtils.getMongoConnectionString()
mongoose.connect(dbConnectionStr, {
  useNewUrlParser: true,
  useUnifiedTopology: false,
  autoIndex: false
})

console.log('About to test connection')
const db = mongoose.connection
db.on('error', () => {
  console.error.bind(console, 'Connection Error: Something went wrong!')
  logger.error(error.connectionError())
})

db.once('open', async () => {
  logger.info('Successfully connected to database!')

  let userInput
  if (process.argv.length > 2 && process.argv.slice(2)[0] === 'y') {
    userInput = process.argv.slice(2)[0]
  } else {
    // script runner (currently) needs to agree to an action that drops collections
    userInput = dataUtils.getUserPopulateInput(Object.keys(populateTheseCollections))
  }

  // drops and re-populates collections
  if (userInput.toLowerCase() === 'y') {
    const collections = await db.db.listCollections().toArray()

    for (const collection of collections) {
      if (!collection.name.startsWith('system.')) {
        logger.info(`Dropping ${collection.name} collection !!!`)
        await db.dropCollection(collection.name)
      }
    }

    // Org
    await dataUtils.populateCollection(
      './datadump/pre-population/orgs.json',
      Org, dataUtils.newOrgTransform
    )

    // User, depends on Org
    const hash = await dataUtils.preprocessUserSecrets()
    await dataUtils.populateCollection(
      './datadump/pre-population/users.json',
      User, dataUtils.newUserTransform, hash
    )

    const populatePromises = []

    // CVE ID Range
    populatePromises.push(dataUtils.populateCollection(
      './datadump/pre-population/cve-ids-range.json',
      CveIdRange
    ))

    // CVE
    if (process.env.NODE_ENV === 'development' || process.env.NODE_ENV === 'test') {
      populatePromises.push(dataUtils.populateCollection(
        './datadump/pre-population/cves.json',
        Cve, dataUtils.newCveTransform
      ))
    }

    // CVE ID, depends on User and Org
    populatePromises.push(dataUtils.populateCollection(
      './datadump/pre-population/cve-ids.json',
      CveId, dataUtils.newCveIdTransform
    ))

    // don't close database connection until all remaining populate
    // promises are resolved
    Promise.all(populatePromises).then(async function () {
      logger.info('Successfully populated the database!')

      const indexPromises = []
      Object.keys(indexesToCreate).forEach(col => {
        indexesToCreate[col].forEach(index => {
          indexPromises.push(db.collections[col].createIndex(index))
        })
      })

      try {
        await Promise.all(indexPromises)
        logger.info('Successfully created indexes!')

        // Explicitly create collections for models that are not pre-populated but require transactions.
        // Implicit collection creation inside Mongo transactions acquires heavy locks and leads to LockTimeout.
        await Audit.createCollection()
        await ReviewObject.createCollection()
        await Conversation.createCollection()
      } catch (err) {
        logger.error('Error creating indexes:', err)
      } finally {
        mongoose.connection.close()
      }
    })
  } else {
    mongoose.connection.close()
  }
})