Merge pull request #31 from datakind/sftp-ingestion

feat: added new sftp-files endpoint

Author: Mesh-ach

src/worker/main.py

@@ -49,7 +49,6 @@ class PdpPullRequest(BaseModel):
 class PdpPullResponse(BaseModel):
     """Fields for the PDP pull response."""
 
-    sftp_files: list[dict]
     pdp_inst_generated: list[Any]
     pdp_inst_not_found: list[Any]
     upload_status: dict
@@ -58,6 +57,12 @@ class Config:
         json_encoders = {np.int64: lambda v: int(v)}
 
 
+class PdpListFiles(BaseModel):
+    """Fields for the PDP pull response."""
+
+    sftp_files: list[dict]
+
+
 @app.on_event("startup")
 def on_startup():
     print("Starting up app...")
@@ -162,8 +167,8 @@ async def process_file(
     }
 
 
-@app.post("/execute-pdp-pull", response_model=PdpPullResponse)
-async def execute_pdp_pull(
+@app.get("/sftp_files", response_model=PdpListFiles)
+def sftp_files(
     req: PdpPullRequest,
     current_username: Annotated[str, Depends(get_current_username)],
     storage_control: Annotated[StorageControl, Depends(StorageControl)],
@@ -181,15 +186,30 @@ async def execute_pdp_pull(
         remote_path="./receive",
     )
 
-    results = []
-    for file in files:
-        gcs_blob = sftp_file_to_gcs_helper(storage_control, file)
-        result = await process_file(storage_control, gcs_blob, env_vars)
-        results.append(result)
-
     # Aggregate results to return
     return {
         "sftp_files": files,
+    }
+
+
+@app.post("/execute-pdp-pull", response_model=PdpPullResponse)
+async def execute_pdp_pull(
+    req: PdpPullRequest,
+    sftp_source_filename: str,
+    current_username: Annotated[str, Depends(get_current_username)],
+    storage_control: Annotated[StorageControl, Depends(StorageControl)],
+) -> Any:
+    """Performs the PDP pull of the file."""
+
+    storage_control.create_bucket_if_not_exists(
+        get_sftp_bucket_name(env_vars["BUCKET_ENV"])
+    )
+
+    gcs_blob = sftp_file_to_gcs_helper(storage_control, sftp_source_filename)
+    result = await process_file(storage_control, gcs_blob, env_vars)
+
+    # Aggregate results to return
+    return {
         "pdp_inst_generated": list(result["valid_inst_ids"]),
         "pdp_inst_not_found": list(result["invalid_ids"]),
         "upload_status": dict(result["uploads"]),

src/worker/main_test.py

@@ -47,6 +47,13 @@ def test_retrieve_token(client: TestClient) -> Any:
     assert response.status_code == 200
 
 
+def sftp_files(client: TestClient) -> Any:
+    """Test GET /sftp-files endpoint"""
+    response = client.get("/sftp-files")
+    assert response.status_code == 200
+    assert response.json() == {"sftp_files": {}}
+
+
 @patch("google.auth.default")
 def test_execute_pdp_pull(
     mock_auth_default: Any, client: TestClient, monkeypatch: Any
@@ -65,19 +72,17 @@ def test_execute_pdp_pull(
         lambda filename: f"processed_{filename}"
     )
     MOCK_STORAGE.create_bucket_if_not_exists.return_value = None
-    MOCK_STORAGE.list_sftp_files.return_value = [
-        {"path": "file1.csv"},
-        {"path": "file2.csv"},
-    ]
+
     # Optionally, if there's a process_file or similar function, you can mock it too.
     # For this test, we're focusing on the overall endpoint behavior.
 
-    response = client.post("/execute-pdp-pull", json={"placeholder": "val"})
+    response = client.post(
+        "/execute-pdp-pull?sftp_source_filename=file1.csv", json={"placeholder": "val"}
+    )
 
     # Verify the response status and content.
     assert response.status_code == 200
     assert response.json() == {
-        "sftp_files": [{"path": "file1.csv"}, {"path": "file2.csv"}],
         "pdp_inst_generated": [],
         "pdp_inst_not_found": [],
         "upload_status": {},

src/worker/utilities.py

@@ -54,7 +54,7 @@ def copy_from_sftp_to_gcs(
             if client is None:
                 raise RuntimeError("Failed to create SFTP client.")
             # Open the file in binary read mode.
-            with client.open(sftp_file, "rb") as f:
+            with client.open(f"./receive/{sftp_file}", "rb") as f:
                 blob.upload_from_file(f)
 
     def list_sftp_files(
@@ -485,7 +485,7 @@ def split_csv_and_generate_signed_urls(
 
 
 def sftp_file_to_gcs_helper(
-    storage_control: StorageControl, sftp_source_filename: dict
+    storage_control: StorageControl, sftp_source_filename: str
 ) -> str:
     """
     For each source file in sftp_source_filenames, copies the file from the SFTP
@@ -502,7 +502,7 @@ def sftp_file_to_gcs_helper(
     )
     # for sftp_source_filename in sftp_source_filenames:
     # Extract the base filename and prepare the destination filename
-    source_filename = sftp_source_filename["path"]
+    source_filename = sftp_source_filename
     # Extract the base filename.
     base_filename = os.path.basename(source_filename)
     dest_filename = f"{base_filename}"
@@ -521,7 +521,7 @@ def sftp_file_to_gcs_helper(
             22,
             sftp_vars["SFTP_USER"],
             sftp_vars["SFTP_PASSWORD"],
-            source_filename,
+            dest_filename,
             get_sftp_bucket_name(env_vars["ENV"]),
             dest_filename,
         )