[Issue #3520] Add the ability to only process a subset of opportunity attachments

api/src/data_migration/transformation/subtask/abstract_transform_subtask.py

@@ -119,6 +119,7 @@ def fetch_with_opportunity(
         source_model: Type[transform_constants.S],
         destination_model: Type[transform_constants.D],
         join_clause: Sequence,
+        batch_size: int = 5000,
     ) -> list[Tuple[transform_constants.S, transform_constants.D | None, Opportunity | None]]:
         # Similar to the above fetch function, but also grabs an opportunity record
         # Note that this requires your source_model to have an opportunity_id field defined.
@@ -134,7 +135,7 @@ def fetch_with_opportunity(
                     isouter=True,
                 )
                 .where(source_model.transformed_at.is_(None))
-                .execution_options(yield_per=5000)
+                .execution_options(yield_per=batch_size)
             ),
         )
 

api/src/data_migration/transformation/subtask/transform_opportunity_attachment.py

@@ -13,10 +13,17 @@
 from src.services.opportunity_attachments import attachment_util
 from src.task.task import Task
 from src.util import file_util
+from src.util.env_config import PydanticBaseEnvConfig
 
 logger = logging.getLogger(__name__)
 
 
+class TransformOpportunityAttachmentConfig(PydanticBaseEnvConfig):
+    # This is just for testing, we want to be able to only
+    # import a few attachments when manually testing.
+    total_attachments_to_process: int | None = None
+
+
 class TransformOpportunityAttachment(AbstractTransformSubTask):
 
     def __init__(self, task: Task, s3_config: S3Config | None = None):
@@ -27,13 +34,20 @@ def __init__(self, task: Task, s3_config: S3Config | None = None):
 
         self.s3_config = s3_config
 
+        self.attachment_config = TransformOpportunityAttachmentConfig()
+
     def transform_records(self) -> None:
 
         # Fetch staging attachment / our attachment / opportunity groups
         records = self.fetch_with_opportunity(
             TsynopsisAttachment,
             OpportunityAttachment,
             [TsynopsisAttachment.syn_att_id == OpportunityAttachment.attachment_id],
+            batch_size=(
+                self.attachment_config.total_attachments_to_process
+                if self.attachment_config.total_attachments_to_process
+                else 1000
+            ),
         )
 
         self.process_opportunity_attachment_group(records)
@@ -44,8 +58,19 @@ def process_opportunity_attachment_group(
             tuple[TsynopsisAttachment, OpportunityAttachment | None, Opportunity | None]
         ],
     ) -> None:
+
+        records_processed = 0
         for source_attachment, target_attachment, opportunity in records:
             try:
+                # Note we increment first in case there are errors, want it to always increment
+                records_processed += 1
+                if (
+                    self.attachment_config.total_attachments_to_process is not None
+                    and self.attachment_config.total_attachments_to_process <= records_processed
+                ):
+                    logger.info("Ending processing early due to configuration")
+                    break
+
                 self.process_opportunity_attachment(
                     source_attachment, target_attachment, opportunity
                 )