Merge pull request #33 from datakind/sftp-ingestion

fix: test script

Author: Mesh-ach

src/worker/main.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import logging
-from typing import Any, Annotated, Dict
+from typing import Any, Annotated, Dict, Optional
 
 from fastapi import FastAPI, Depends, HTTPException, status
 from fastapi.responses import FileResponse
@@ -103,13 +103,17 @@ async def login_for_access_token(
 
 
 async def process_file(
-    storage_control: StorageControl, blob: str, env_vars: Dict[str, Any]
+    storage_control: StorageControl,
+    blob: str,
+    env_vars: Dict[str, Any],
+    pdp_id: Optional[str] = None,
 ) -> Dict[str, Any]:
     """Process a single file: generate URLs, transfer, and validate."""
     logger.debug(f">>>> Splitting {blob} to extract institution data")
     signed_urls = split_csv_and_generate_signed_urls(
         bucket_name=get_sftp_bucket_name(env_vars["BUCKET_ENV"]),
         source_blob_name=blob,
+        pdp_id=pdp_id,
     )
     logger.info(f">>>> Signed URLs, File names generated for {blob}")
 
@@ -197,6 +201,7 @@ async def execute_pdp_pull(
     sftp_source_filename: str,
     current_username: Annotated[str, Depends(get_current_username)],
     storage_control: Annotated[StorageControl, Depends(StorageControl)],
+    pdp_id: Optional[str] = None,
 ) -> Any:
     """Performs the PDP pull of the file."""
 
@@ -205,7 +210,7 @@ async def execute_pdp_pull(
     )
 
     gcs_blob = sftp_file_to_gcs_helper(storage_control, sftp_source_filename)
-    result = await process_file(storage_control, gcs_blob, env_vars)
+    result = await process_file(storage_control, gcs_blob, env_vars, pdp_id=pdp_id)
 
     # Aggregate results to return
     return {

src/worker/utilities.py

@@ -8,7 +8,7 @@
 from datetime import datetime, timedelta
 import io
 import logging
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 import requests
 import pandas as pd
 import re
@@ -383,7 +383,7 @@ def generate_signed_url(
 
 
 def split_csv_and_generate_signed_urls(
-    bucket_name: str, source_blob_name: str
+    bucket_name: str, source_blob_name: str, pdp_id: Optional[str] = None
 ) -> Dict[str, Dict[str, str]]:
     """
     Fetches a CSV from Google Cloud Storage, splits it by a specified column, uploads the results,
@@ -440,7 +440,11 @@ def split_csv_and_generate_signed_urls(
     all_data = {}
 
     # Processing the DataFrame
-    unique_inst_ids = df[institution_column].unique()
+    # unique_inst_ids = df[institution_column].unique()
+    unique_inst_ids = (
+        list(df[institution_column].unique()) if pdp_id is None else [pdp_id]
+    )
+
     for inst_id in unique_inst_ids:
         group = df[df[institution_column] == inst_id]
         output = io.StringIO()