CUMULUS-3757: Update granule upsert logic to allow updating collection info

CHANGELOG.md

@@ -30,6 +30,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- **CUMULUS-3757**
+  - Added a database helper function to assist with updating a granule and its files PG fields when moving the granule(s) across collections
 - **CUMULUS-3919**
   - Added terraform variables `disableSSL` and `rejectUnauthorized` to `tf-modules/cumulus-rds-tf` module.
 

packages/db/src/index.ts

@@ -134,6 +134,7 @@ export {
   getGranulesByApiPropertiesQuery,
   getGranulesByGranuleId,
   getGranuleAndCollection,
+  updateGranuleAndFiles,
 } from './lib/granule';
 
 export {

packages/db/src/lib/granule.ts

@@ -9,12 +9,15 @@ import {
 import { RecordDoesNotExist } from '@cumulus/errors';
 import Logger from '@cumulus/logger';
 
+import { ApiGranule } from '@cumulus/types/api/granules';
 import { CollectionPgModel } from '../models/collection';
 import { GranulePgModel } from '../models/granule';
 import { GranulesExecutionsPgModel } from '../models/granules-executions';
 import { PostgresGranule, PostgresGranuleRecord } from '../types/granule';
 import { GranuleWithProviderAndCollectionInfo } from '../types/query';
 import { UpdatedAtRange } from '../types/record';
+import { translateApiGranuleToPostgresGranule } from '../translate/granules';
+import { translateApiFiletoPostgresFile } from '../translate/file';
 const { deprecate } = require('@cumulus/common/util');
 
 const { TableNames } = require('../tables');
@@ -354,3 +357,41 @@ export const getGranulesByGranuleId = async (
     .where({ granule_id: granuleId });
   return records;
 };
+
+/**
+ * Change a granules' PG record and its files' PG record based on collection move
+ *
+ * @param {Knex | Knex.Transaction} knexOrTransaction - DB client or transaction
+ * @param {Object} [collectionPgModel] - Collection PG model class instance
+ * @param {string[]} granuleIds - list of granules by granuleIds to change
+ * @param {string} collectionId - collection ID
+ * @returns {Promise<void>}
+ */
+export const updateGranuleAndFiles = async (
+  knexOrTransaction: Knex | Knex.Transaction,
+  granules: Array<ApiGranule>
+):Promise<void> => {
+  const {
+    granules: granulesTable,
+    files: filesTable,
+  } = TableNames;
+  await Promise.all(granules.map(async (granule) => {
+    const pgGranule = await translateApiGranuleToPostgresGranule({
+      dynamoRecord: granule,
+      knexOrTransaction,
+    });
+    await knexOrTransaction(granulesTable).where('granule_id', '=', pgGranule.granule_id).update(
+      pgGranule
+    );
+
+    if (granule.files) {
+      await Promise.all(granule.files.map(async (file) => {
+        const pgFile = translateApiFiletoPostgresFile({ ...file, granuleId: pgGranule.granule_id });
+
+        await knexOrTransaction(filesTable).where('file_name', '=', String(pgFile.file_name)).update(
+          pgFile
+        );
+      }));
+    }
+  }));
+};

packages/db/tests/lib/test-updateGranuleAndFiles.js

@@ -0,0 +1,220 @@
+/* eslint-disable no-await-in-loop */
+const test = require('ava');
+const cryptoRandomString = require('crypto-random-string');
+const range = require('lodash/range');
+
+const { constructCollectionId } = require('@cumulus/message/Collections');
+const {
+  CollectionPgModel,
+  GranulePgModel,
+  FilePgModel,
+  generateLocalTestDb,
+  destroyLocalTestDb,
+  fakeCollectionRecordFactory,
+  fakeFileRecordFactory,
+  fakeGranuleRecordFactory,
+  getUniqueGranuleByGranuleId,
+  migrationDir,
+  updateGranuleAndFiles,
+  translatePostgresGranuleResultToApiGranule,
+  translatePostgresCollectionToApiCollection,
+} = require('../../dist');
+
+const testDbName = `granule_${cryptoRandomString({ length: 10 })}`;
+
+// this function is used to simulate granule records post-collection-move for database updates
+const simulateGranuleUpdate = async (knex, granules, collection, collectionId, collectionId2) => {
+  const movedGranules = [];
+  for (const granule of granules) {
+    const postMoveApiGranule = await translatePostgresGranuleResultToApiGranule(knex, {
+      ...granule,
+      collectionName: collection.name,
+      collectionVersion: collection.version,
+    });
+    postMoveApiGranule.collectionId = collectionId2;
+    postMoveApiGranule.updatedAt = Date.now();
+    postMoveApiGranule.lastUpdateDateTime = new Date().toISOString();
+    for (const apiFile of postMoveApiGranule.files) {
+      apiFile.bucket = apiFile.bucket.replace(collectionId, collectionId2);
+      apiFile.key = apiFile.key.replace(collectionId, collectionId2);
+      //apiFile.path = apiFile.path.replace(t.context.collectionId, t.context.collectionId2);
+      apiFile.updatedAt = Date.now();
+    }
+    movedGranules.push(postMoveApiGranule);
+  }
+  return movedGranules;
+};
+
+test.before(async (t) => {
+  const { knexAdmin, knex } = await generateLocalTestDb(
+    testDbName,
+    migrationDir
+  );
+  t.context.knexAdmin = knexAdmin;
+  t.context.knex = knex;
+
+  t.context.granulePgModel = new GranulePgModel();
+  t.context.collectionPgModel = new CollectionPgModel();
+  t.context.filePgModel = new FilePgModel();
+
+  // set up 2 collections
+  t.context.collection = fakeCollectionRecordFactory();
+  t.context.collection2 = fakeCollectionRecordFactory();
+  t.context.collectionId = constructCollectionId(
+    t.context.collection.name,
+    t.context.collection.version
+  );
+  t.context.collectionId2 = constructCollectionId(
+    t.context.collection2.name,
+    t.context.collection2.version
+  );
+  const collectionResponse = await t.context.collectionPgModel.create(
+    t.context.knex,
+    t.context.collection
+  );
+  const collectionResponse2 = await t.context.collectionPgModel.create(
+    t.context.knex,
+    t.context.collection2
+  );
+  t.context.collectionCumulusId = collectionResponse[0].cumulus_id;
+  t.context.collectionCumulusId2 = collectionResponse2[0].cumulus_id;
+  t.context.apiCollection1 = translatePostgresCollectionToApiCollection(collectionResponse[0]);
+  t.context.apiCollection2 = translatePostgresCollectionToApiCollection(collectionResponse2[0]);
+
+  // create 10 granules in one collection, 0 in the other
+  t.context.granuleIds = range(10).map((num) => 'granuleId___' + num);
+
+  t.context.granulePgModel = new GranulePgModel();
+  t.context.granules = range(10).map((num) => fakeGranuleRecordFactory({
+    granule_id: t.context.granuleIds[num],
+    collection_cumulus_id: t.context.collectionCumulusId,
+    cumulus_id: num,
+  }));
+  t.context.pgGranules = await t.context.granulePgModel.insert(
+    knex,
+    t.context.granules
+  );
+
+  t.context.movedGranules = [];
+
+  t.context.files = [];
+  // create fake files for each of the ten granules (3 per granule)
+  for (const pgGranule of t.context.granules) {
+    t.context.files.push(
+      fakeFileRecordFactory({
+        granule_cumulus_id: pgGranule.cumulus_id,
+        file_name: pgGranule.granule_id + '.hdf',
+        updated_at: new Date().toISOString(),
+        bucket: t.context.collectionId + '--bucket',
+        key: t.context.collectionId + pgGranule.granule_id + '/key-hdf.pem',
+        path: t.context.collectionId + '/' + pgGranule.granule_id,
+      }),
+      fakeFileRecordFactory({
+        granule_cumulus_id: pgGranule.cumulus_id,
+        file_name: pgGranule.granule_id + '.txt',
+        updated_at: new Date().toISOString(),
+        bucket: t.context.collectionId + '--bucket',
+        key: t.context.collectionId + pgGranule.granule_id + '/key-txt.pem',
+        path: t.context.collectionId + '/' + pgGranule.granule_id,
+      }),
+      fakeFileRecordFactory({
+        granule_cumulus_id: pgGranule.cumulus_id,
+        file_name: pgGranule.granule_id + '.cmr',
+        updated_at: new Date().toISOString(),
+        bucket: t.context.collectionId + '--bucket',
+        key: t.context.collectionId + pgGranule.granule_id + '/key-cmr.pem',
+        path: t.context.collectionId + '/' + pgGranule.granule_id,
+      })
+    );
+  }
+
+  t.context.pgFiles = await t.context.filePgModel.insert(knex, t.context.files);
+  // update 1/2 of the granules to be moved to the new collection
+  t.context.movedGranules.push(await simulateGranuleUpdate(knex, t.context.granules.slice(0, 5),
+    t.context.collection, t.context.collectionId, t.context.collectionId2));
+
+  // the other half will be unmoved but translated to an apiGranule
+  t.context.movedGranules.push(await simulateGranuleUpdate(knex, t.context.granules.slice(5),
+    t.context.collection, t.context.collectionId, t.context.collectionId));
+
+  t.context.movedGranules = t.context.movedGranules.flat();
+});
+
+test.after.always(async (t) => {
+  await destroyLocalTestDb({
+    ...t.context,
+    testDbName,
+  });
+});
+
+test.serial('updateGranuleAndFiles successfully updates a partial list of granules based on the collectionId change', async (t) => {
+  const {
+    granuleIds,
+    granulePgModel,
+    movedGranules,
+    collectionId2,
+    collectionId,
+    collection,
+    collection2,
+    knex,
+  } = t.context;
+  await updateGranuleAndFiles(knex, movedGranules);
+
+  const returnedGranules = await Promise.all(granuleIds.map((id) =>
+    getUniqueGranuleByGranuleId(knex, id, granulePgModel)));
+
+  for (const granule of returnedGranules) {
+    const testCollection = granule.cumulus_id >= 5 ? collection : collection2;
+    const testCollectionId = granule.cumulus_id >= 5 ? collectionId : collectionId2;
+    const apiGranule = await translatePostgresGranuleResultToApiGranule(knex, {
+      ...granule,
+      collectionName: testCollection.name,
+      collectionVersion: testCollection.version,
+    });
+    // the movedGranules param only has 1/2 of the granules to be moved to collection 2
+    // here we can check based on the granule's cumulus id which collection it should be a part of
+    t.true(apiGranule.collectionId === testCollectionId);
+    for (const file of apiGranule.files) {
+      t.true(file.key.includes(testCollectionId));
+      t.true(file.bucket.includes(testCollectionId));
+    }
+  }
+});
+
+test.serial('updateGranuleAndFiles successfully updates a complete list of granules, 1/2 of which have already been moved', async (t) => {
+  const {
+    granuleIds,
+    granulePgModel,
+    granules,
+    movedGranules,
+    collectionId2,
+    collectionId,
+    collection,
+    collection2,
+    knex,
+  } = t.context;
+  // the remaining granules of movedGranules in collection 1 will need to be updated to collection 2
+  movedGranules.splice(-5);
+  movedGranules.push(await simulateGranuleUpdate(knex, granules.slice(5), collection,
+    collectionId, collectionId2));
+
+  const testPostMoveApiGranules = movedGranules.flat();
+  await updateGranuleAndFiles(knex, testPostMoveApiGranules);
+
+  const returnedGranules = await Promise.all(granuleIds.map((id) =>
+    getUniqueGranuleByGranuleId(knex, id, granulePgModel)));
+
+  for (const granule of returnedGranules) {
+    const apiGranule = await translatePostgresGranuleResultToApiGranule(knex, {
+      ...granule,
+      collectionName: collection2.name,
+      collectionVersion: collection2.version,
+    });
+    // now every granule should be part of collection 2
+    t.true(apiGranule.collectionId === collectionId2);
+    for (const file of apiGranule.files) {
+      t.true(file.key.includes(collectionId2));
+      t.true(file.bucket.includes(collectionId2));
+    }
+  }
+});