Merge pull request #3 from datakind/sftp-integration

fix: linting errors

Author: Mesh-ach

.github/workflows/type-check.yml

@@ -15,7 +15,7 @@ jobs:
           python-version: "3.10"
       - name: Get changed files
         id: changed-files
-        uses: step-security/changed-files@45
+        uses: step-security/changed-files@v45
         with:
           files: |
             src/**/*.py

pyproject.toml

@@ -21,6 +21,10 @@ dependencies = [
     "strenum>=0.4.15",
     "tomli~=2.0; python_version<'3.11'",
     "jsonpickle>=4.0.1",
+    "requests>=2.0.0",
+    "types-requests",
+    "types-paramiko",
+    "pandas"
 ]
 
 [project.urls]

src/worker/authn.py

@@ -2,21 +2,28 @@
 
 from fastapi.security import (
     OAuth2PasswordBearer,
-    OAuth2PasswordRequestForm,
     APIKeyHeader,
-    APIKeyQuery,
 )
 from pydantic import BaseModel
 from datetime import timedelta, datetime, timezone
 from .config import env_vars
 from typing import Annotated
-from fastapi import Depends, HTTPException, status
+from fastapi import Depends, HTTPException, status, Security
 from jwt.exceptions import InvalidTokenError
 
 oauth2_scheme = OAuth2PasswordBearer(
     tokenUrl="token",
 )
 
+api_key_header = APIKeyHeader(name="X-API-KEY", scheme_name="api-key", auto_error=False)
+api_key_inst_header = APIKeyHeader(
+    name="INST", scheme_name="api-inst", auto_error=False
+)
+# The following is for use by the frontend enduser only.
+api_key_enduser_header = APIKeyHeader(
+    name="ENDUSER", scheme_name="api-enduser", auto_error=False
+)
+
 
 class Token(BaseModel):
     access_token: str
@@ -27,7 +34,30 @@ class TokenData(BaseModel):
     username: str | None = None
 
 
-def check_creds(username: str, password: str):
+def get_api_key(
+    api_key_header: str = Security(api_key_header),
+    api_key_inst_header: str = Security(api_key_inst_header),
+    api_key_enduser_header: str = Security(api_key_enduser_header),
+) -> tuple:
+    """Retrieve the api key and enduser header key if present.
+
+    Args:
+        api_key_header: The API key passed in the HTTP header.
+
+    Returns:
+        A tuple with the api key and enduser header if present. Authentication happens elsewhere.
+    Raises:
+        HTTPException: If the API key is invalid or missing.
+    """
+    if api_key_header:
+        return (api_key_header, api_key_inst_header, api_key_enduser_header)
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid or missing API Key",
+    )
+
+
+def check_creds(username: str, password: str) -> bool:
     if username == env_vars["USERNAME"] and password == env_vars["PASSWORD"]:
         return True
     raise HTTPException(
@@ -36,13 +66,13 @@ def check_creds(username: str, password: str):
     )
 
 
-def create_access_token(data: dict, expires_delta: timedelta | None = None):
+def create_access_token(data: dict, expires_delta: timedelta | None = None) -> str:
     to_encode = data.copy()
     if expires_delta:
         expire = datetime.now(timezone.utc) + expires_delta
     else:
         expire = datetime.now(timezone.utc) + timedelta(
-            minutes=env_vars["ACCESS_TOKEN_EXPIRE_MINUTES"]
+            minutes=float(env_vars["ACCESS_TOKEN_EXPIRE_MINUTES"])
         )
     to_encode.update({"exp": expire})
     encoded_jwt = jwt.encode(

src/worker/main.py

@@ -2,18 +2,27 @@
 
 import logging
 from typing import Any, Annotated
-from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi import FastAPI, Depends, HTTPException, status, Security
 from fastapi.responses import FileResponse
 
 from pydantic import BaseModel
-from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
+from fastapi.security import OAuth2PasswordRequestForm
 from .utilities import (
     get_sftp_bucket_name,
     StorageControl,
+    split_csv_and_generate_signed_urls,
+    fetch_institution_ids,
 )
 from .config import sftp_vars, env_vars, startup_env_vars
-from .authn import Token, get_current_username, check_creds, create_access_token
-from datetime import timedelta, datetime, timezone
+from .authn import (
+    Token,
+    get_current_username,
+    check_creds,
+    create_access_token,
+    get_api_key,
+)
+from datetime import timedelta
+import os
 
 # Set the logging
 logging.basicConfig(format="%(asctime)s [%(levelname)s]: %(message)s")
@@ -41,8 +50,9 @@ class PdpPullRequest(BaseModel):
 class PdpPullResponse(BaseModel):
     """Fields for the PDP pull response."""
 
-    pdp_inst_generated: list[int]
-    pdp_inst_not_found: list[int]
+    sftp_files: list[dict]
+    pdp_inst_generated: list[str]
+    pdp_inst_not_found: list[str]
 
 
 @app.on_event("startup")
@@ -84,30 +94,86 @@ async def login_for_access_token(
     return Token(access_token=access_token, token_type="bearer")
 
 
-def sftp_helper(
-    storage_control: StorageControl, sftp_source_filename: str, dest_filename: str
-):
-    storage_control.copy_from_sftp_to_gcs(
-        sftp_vars["SFTP_HOST"],
-        sftp_vars["SFTP_PORT"],
-        sftp_vars["SFTP_USER"],
-        sftp_vars["SFTP_PASSWORD"],
-        sftp_source_filename,
-        get_sftp_bucket_name(env_vars["ENV"]),
-        dest_filename,
-    )
+def sftp_helper(storage_control: StorageControl, sftp_source_filenames: list) -> list:
+    """
+    For each source file in sftp_source_filenames, copies the file from the SFTP
+    server to GCS. The destination filename is automatically generated by prefixing
+    the base name of the source file with "processed_".
+
+    Args:
+        storage_control (StorageControl): An instance with a method `copy_from_sftp_to_gcs`.
+        sftp_source_filenames (list): A list of file paths on the SFTP server.
+    """
+    num_files = len(sftp_source_filenames)
+    logger.info(f"Starting sftp_helper for {num_files} file(s).")
+    all_blobs = []
+    for sftp_source_filename in sftp_source_filenames:
+        sftp_source_filename = sftp_source_filename["path"]
+        if (
+            sftp_source_filename
+            == "./receive/AO1600pdp_AO1600_AR_DEIDENTIFIED_STUDYID_20250228030226.csv"
+        ):
+            logger.debug(f"Processing source file: {sftp_source_filename}")
+
+            # Extract the base filename.
+            base_filename = os.path.basename(sftp_source_filename)
+            dest_filename = f"{base_filename}"
+            logger.debug(f"Destination filename will be: {dest_filename}")
+
+            try:
+                storage_control.copy_from_sftp_to_gcs(
+                    sftp_vars["SFTP_HOST"],
+                    22,
+                    sftp_vars["SFTP_USER"],
+                    sftp_vars["SFTP_PASSWORD"],
+                    sftp_source_filename,
+                    get_sftp_bucket_name(env_vars["ENV"]),
+                    dest_filename,
+                )
+                all_blobs.append(dest_filename)
+                logger.info(
+                    f"Successfully processed '{sftp_source_filename}' as '{dest_filename}'."
+                )
+                return all_blobs
+            except Exception as e:
+                logger.error(
+                    f"Error processing '{sftp_source_filename}': {e}", exc_info=True
+                )
+    return all_blobs
 
 
 @app.post("/execute-pdp-pull", response_model=PdpPullResponse)
 def execute_pdp_pull(
     req: PdpPullRequest,
     current_username: Annotated[str, Depends(get_current_username)],
     storage_control: Annotated[StorageControl, Depends(StorageControl)],
+    api_key_enduser_tuple: str = Security(get_api_key),
 ) -> Any:
     """Performs the PDP pull of the file."""
     storage_control.create_bucket_if_not_exists(get_sftp_bucket_name(env_vars["ENV"]))
-    sftp_helper(storage_control, "sftp_file.csv", "write_out_file.csv")
+    files = storage_control.list_sftp_files(
+        sftp_vars["SFTP_HOST"], 22, sftp_vars["SFTP_USER"], sftp_vars["SFTP_PASSWORD"]
+    )
+    all_blobs = sftp_helper(storage_control, files)
+    valid_pdp_ids = []
+    invalid_ids = []
+
+    for blobs in all_blobs:
+        signed_urls = split_csv_and_generate_signed_urls(
+            bucket_name=get_sftp_bucket_name(env_vars["ENV"]), source_blob_name=blobs
+        )
+
+        temp_valid_pdp_ids, temp_invalid_ids = fetch_institution_ids(
+            pdp_ids=list(signed_urls.keys()),
+            backend_api_key=next(
+                key for key in api_key_enduser_tuple if key is not None
+            ),
+        )
+        valid_pdp_ids.append(temp_valid_pdp_ids)
+        invalid_ids.append(temp_invalid_ids)
+
     return {
-        "pdp_inst_generated": [],
-        "pdp_inst_not_found": [],
+        "sftp_files": files,
+        "pdp_inst_generated": valid_pdp_ids,
+        "pdp_inst_not_found": invalid_ids,
     }

src/worker/main_test.py

@@ -1,12 +1,11 @@
 """Test file for the main.py file and constituent API functions."""
 
 import pytest
-import json
+from typing import Any
 
 from fastapi.testclient import TestClient
 from .main import app
-import uuid
-from .authn import get_current_username
+from .authn import get_current_username, get_api_key
 from unittest import mock
 from .utilities import StorageControl
 
@@ -21,22 +20,26 @@ def get_current_username_override():
     def storage_control_override():
         return MOCK_STORAGE
 
+    def get_api_key_override():
+        return ("valid_api_key", "end_user")
+
     app.dependency_overrides[StorageControl] = storage_control_override
 
     app.dependency_overrides[get_current_username] = get_current_username_override
 
+    app.dependency_overrides[get_api_key] = get_api_key_override
     client = TestClient(app, root_path="/workers/api/v1")
     yield client
     app.dependency_overrides.clear()
 
 
-def test_get_root(client: TestClient):
+def test_get_root(client: TestClient) -> Any:
     """Test GET /."""
     response = client.get("/")
     assert response.status_code == 200
 
 
-def test_retrieve_token(client: TestClient):
+def test_retrieve_token(client: TestClient) -> Any:
     """Test POST /token."""
     response = client.post(
         "/token",
@@ -46,14 +49,23 @@ def test_retrieve_token(client: TestClient):
     assert response.status_code == 200
 
 
-def test_execute_pdp_pull(client: TestClient):
+def test_execute_pdp_pull(client: TestClient) -> Any:
     """Test POST /execute-pdp-pull."""
-    MOCK_STORAGE.copy_from_sftp_to_gcs.return_value = None
+    MOCK_STORAGE.copy_from_sftp_to_gcs.side_effect = (
+        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"},
+    ]
 
     response = client.post("/execute-pdp-pull", 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": [],
     }