feat: option to disable docdb/s3/co sync (#135)

Author: helen-m-lin

src/aind_data_asset_indexer/aind_bucket_indexer.py

@@ -743,8 +743,8 @@ def _process_prefixes(self, prefixes: List[str]):
         )
         mapped_partitions.compute()
 
-    def run_job(self):
-        """Main method to run."""
+    def _run_docdb_sync(self):
+        """Sync changes in DocDB to S3"""
         with self._create_docdb_client() as iterator_docdb_client:
             filter = {
                 "location": {
@@ -768,6 +768,9 @@ def run_job(self):
                 if len(page) > 0:
                     self._process_records(records=page)
         logging.info("Finished scanning through DocDb.")
+
+    def _run_s3_sync(self):
+        """Sync changes in S3 to DocDB"""
         logging.info("Starting to scan through S3.")
         iterator_s3_client = boto3.client("s3")
         prefix_iterator = iterate_through_top_level(
@@ -779,6 +782,13 @@ def run_job(self):
         iterator_s3_client.close()
         logging.info("Finished scanning through S3.")
 
+    def run_job(self):
+        """Main method to run."""
+        if self.job_settings.run_docdb_sync is True:
+            self._run_docdb_sync()
+        if self.job_settings.run_s3_sync is True:
+            self._run_s3_sync()
+
 
 if __name__ == "__main__":
     sys_args = sys.argv[1:]

src/aind_data_asset_indexer/codeocean_bucket_indexer.py

@@ -401,6 +401,8 @@ def _delete_records_from_docdb(self, record_list: List[str]):
 
     def run_job(self):
         """Main method to run."""
+        if self.job_settings.run_co_sync is not True:
+            return
         logging.info("Starting to scan through CodeOcean.")
         retry = Retry(
             total=5,

src/aind_data_asset_indexer/models.py

@@ -53,6 +53,14 @@ class AindIndexBucketJobSettings(IndexJobSettings):
     doc_db_host: str
     doc_db_db_name: str
     doc_db_collection_name: str
+    run_docdb_sync: bool = Field(
+        default=True,
+        description="If true, then process DocDB records to sync to S3.",
+    )
+    run_s3_sync: bool = Field(
+        default=True,
+        description="If true, then process S3 prefixes to sync to DocDB.",
+    )
 
 
 class PopulateAindBucketsJobSettings(IndexJobSettings):
@@ -73,13 +81,19 @@ class AindIndexBucketsJobSettings(AindIndexBucketJobSettings):
 
 
 class CodeOceanIndexBucketJobSettings(IndexJobSettings):
-    """Aind Index Bucket Job Settings"""
+    """Code Ocean Index Bucket Job Settings"""
 
     doc_db_host: str
     doc_db_db_name: str
     doc_db_collection_name: str
     codeocean_domain: str
     codeocean_token: SecretStr
+    run_co_sync: bool = Field(
+        default=True,
+        description=(
+            "If true, then process Code Ocean results and external links."
+        ),
+    )
 
     @classmethod
     def from_param_store(cls, param_store_name: str):

tests/test_aind_bucket_indexer.py

@@ -1856,6 +1856,52 @@ def test_run_job_lookback_days(
             ]
         )
 
+    @patch(
+        "aind_data_asset_indexer.aind_bucket_indexer.AindIndexBucketJob."
+        "_run_s3_sync"
+    )
+    @patch(
+        "aind_data_asset_indexer.aind_bucket_indexer.AindIndexBucketJob."
+        "_run_docdb_sync"
+    )
+    def test_run_job_skip_docdb_sync(
+        self,
+        mock_run_docdb_sync: MagicMock,
+        mock_run_s3_sync: MagicMock,
+    ):
+        """Tests main run_job method when run_docdb_sync is False."""
+
+        job_configs_json = self.basic_job_configs.model_dump(mode="json")
+        job_configs_json["run_docdb_sync"] = False
+        job_configs = AindIndexBucketJobSettings(**job_configs_json)
+        job = AindIndexBucketJob(job_settings=job_configs)
+        job.run_job()
+        mock_run_docdb_sync.assert_not_called()
+        mock_run_s3_sync.assert_called_once()
+
+    @patch(
+        "aind_data_asset_indexer.aind_bucket_indexer.AindIndexBucketJob."
+        "_run_s3_sync"
+    )
+    @patch(
+        "aind_data_asset_indexer.aind_bucket_indexer.AindIndexBucketJob."
+        "_run_docdb_sync"
+    )
+    def test_run_job_skip_s3_sync(
+        self,
+        mock_run_docdb_sync: MagicMock,
+        mock_run_s3_sync: MagicMock,
+    ):
+        """Tests main run_job method when run_s3_sync is False."""
+
+        job_configs_json = self.basic_job_configs.model_dump(mode="json")
+        job_configs_json["run_s3_sync"] = False
+        job_configs = AindIndexBucketJobSettings(**job_configs_json)
+        job = AindIndexBucketJob(job_settings=job_configs)
+        job.run_job()
+        mock_run_docdb_sync.assert_called_once()
+        mock_run_s3_sync.assert_not_called()
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_codeocean_bucket_indexer.py

@@ -655,6 +655,41 @@ def test_run_job(
         )
         mock_docdb_client.return_value.__exit__.assert_called_once()
 
+    @patch(
+        "aind_data_asset_indexer.codeocean_bucket_indexer."
+        "CodeOceanIndexBucketJob._update_external_links_in_docdb"
+    )
+    @patch(
+        "aind_data_asset_indexer.codeocean_bucket_indexer."
+        "CodeOceanIndexBucketJob._delete_records_from_docdb"
+    )
+    @patch(
+        "aind_data_asset_indexer.codeocean_bucket_indexer."
+        "CodeOceanIndexBucketJob._process_codeocean_records"
+    )
+    @patch(
+        "aind_data_asset_indexer.codeocean_bucket_indexer."
+        "get_all_processed_codeocean_asset_records"
+    )
+    def test_run_job_skip(
+        self,
+        mock_get_all_co_records: MagicMock,
+        mock_process_codeocean_records: MagicMock,
+        mock_delete_records_from_docdb: MagicMock,
+        mock_update_external_links_in_docdb: MagicMock,
+    ):
+        """Tests run_job method. Given the example responses, should ignore
+        one record, add one record, and delete one record."""
+        job_configs_json = self.basic_job_configs.model_dump(mode="json")
+        job_configs_json["run_co_sync"] = False
+        job_configs = CodeOceanIndexBucketJobSettings(**job_configs_json)
+        job = CodeOceanIndexBucketJob(job_settings=job_configs)
+        job.run_job()
+        mock_get_all_co_records.assert_not_called()
+        mock_update_external_links_in_docdb.assert_not_called()
+        mock_process_codeocean_records.assert_not_called()
+        mock_delete_records_from_docdb.assert_not_called()
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_models.py

@@ -86,6 +86,8 @@ def test_defaults(self):
         self.assertEqual(
             "some_docdb_collection_name", job_settings.doc_db_collection_name
         )
+        self.assertTrue(job_settings.run_docdb_sync)
+        self.assertTrue(job_settings.run_s3_sync)
 
     @patch("boto3.client")
     def test_from_from_param_store(self, mock_boto3_client):
@@ -98,7 +100,8 @@ def test_from_from_param_store(self, mock_boto3_client):
                     '{"doc_db_host": "some_docdb_host",'
                     '"doc_db_db_name":"some_docdb_dbname",'
                     '"doc_db_collection_name":"some_docdb_collection_name",'
-                    '"s3_bucket":"some_bucket"}'
+                    '"s3_bucket":"some_bucket",'
+                    '"run_s3_sync":false}'
                 ),
                 "Version": 1,
                 "LastModifiedDate": datetime(
@@ -129,6 +132,7 @@ def test_from_from_param_store(self, mock_boto3_client):
             doc_db_host="some_docdb_host",
             doc_db_db_name="some_docdb_dbname",
             doc_db_collection_name="some_docdb_collection_name",
+            run_s3_sync=False,
         )
         self.assertEqual(expected_job_settings, job_settings)
 
@@ -167,6 +171,8 @@ def test_class_constructor(self):
         self.assertEqual(
             "some_docdb_collection_name", job_settings.doc_db_collection_name
         )
+        self.assertTrue(job_settings.run_docdb_sync)
+        self.assertTrue(job_settings.run_s3_sync)
 
 
 class TestCodeOceanIndexBucketJobSettings(unittest.TestCase):
@@ -194,6 +200,7 @@ def test_defaults(self):
         self.assertEqual(
             "some_co_token", job_settings.codeocean_token.get_secret_value()
         )
+        self.assertTrue(job_settings.run_co_sync)
 
     @patch("boto3.client")
     def test_from_from_param_store(self, mock_boto3_client):