fix: removes broken filter in log method (#405)

Author: jtyoung84

src/aind_data_transfer_service/log_handler.py

@@ -1,8 +1,8 @@
 """Module to handle logging submit job requests"""
 
 import logging
-from typing import Any
 from enum import Enum
+from typing import Any
 
 
 class EventType(str, Enum):
@@ -13,20 +13,9 @@ class EventType(str, Enum):
     STAGE_FAILURE = "stage_failure"
 
 
-def log_stage_event(
-        message: str,
-        event_type: EventType,
-        **extra_fields: Any
+def log_submit_job_request(
+    content: Any, event_type: EventType | None = None
 ) -> None:
-    """Emit a structured log record for stage lifecycle events."""
-
-    logging.info(
-        message,
-        extra={"event_type": event_type.value, **extra_fields},
-    )
-
-
-def log_submit_job_request(content: Any) -> None:
     """
     Parses content object to log any lines with a subject_id and
     acquisition_name.
@@ -35,16 +24,22 @@ def log_submit_job_request(content: Any) -> None:
     ----------
     content : Any
       Pulled from request json, which may or may not return expected dict
+    event_type: EventType |  None
+      Type of event to log. Default is None.
     """
-    if isinstance(content, list) and all(
-        isinstance(row, dict) for row in content
+    upload_jobs = content.get("upload_jobs")
+    if (
+        upload_jobs is not None
+        and isinstance(upload_jobs, list)
+        and all(isinstance(row, dict) for row in upload_jobs)
     ):
-        for row in content:
+        for row in upload_jobs:
             subject_id = row.get("subject_id")
             acquisition_name = row.get("s3_prefix")
-            log_stage_event(
-                "Handling request",
-                event_type=EventType.STAGE_START,
-                subject_id=subject_id,
-                acquisition_name=acquisition_name,
-            )
+            extra_info = {
+                "subject_id": subject_id,
+                "acquisition_name": acquisition_name,
+            }
+            if event_type is not None:
+                extra_info["event_type"] = event_type
+            logging.info("Handling request", extra=extra_info)

src/aind_data_transfer_service/server.py

@@ -33,7 +33,6 @@
 )
 from aind_data_transfer_service.log_handler import (
     EventType,
-    log_stage_event,
     log_submit_job_request,
 )
 from aind_data_transfer_service.models.core import (
@@ -358,7 +357,9 @@ async def submit_jobs_v2(request: Request):
     logging.info("Received request to submit jobs v2")
     content = await request.json()
     try:
-        log_submit_job_request(content=content)
+        log_submit_job_request(
+            content=content, event_type=EventType.STAGE_START
+        )
         params = AirflowDagRunsRequestParameters(
             dag_ids=["transform_and_upload_v2", "run_list_of_jobs"],
             states=["running", "queued"],
@@ -395,24 +396,12 @@ async def submit_jobs_v2(request: Request):
                 url=os.getenv("AIND_AIRFLOW_SERVICE_URL"),
                 json={"conf": full_content},
             )
-            status_code = response.status_code
-            response_json = response.json()
-        if status_code > 300:
-            log_stage_event(
-                "Failed submit jobs v2 request",
-                event_type=EventType.STAGE_FAILURE,
-                dag_id=model.dag_id,
-                total_jobs=total_jobs,
-                status_code=status_code,
-            )
-        else:
-            log_stage_event(
-                "Completed submit jobs v2 request",
-                event_type=EventType.STAGE_COMPLETE,
-                dag_id=model.dag_id,
-                total_jobs=total_jobs,
-                status_code=status_code,
-            )
+        response.raise_for_status()
+        status_code = response.status_code
+        response_json = response.json()
+        log_submit_job_request(
+            content=content, event_type=EventType.STAGE_COMPLETE
+        )
         return JSONResponse(
             status_code=status_code,
             content={
@@ -422,6 +411,9 @@ async def submit_jobs_v2(request: Request):
         )
     except ValidationError as e:
         logging.warning(f"There were validation errors processing {content}")
+        log_submit_job_request(
+            content=content, event_type=EventType.STAGE_FAILURE
+        )
         return JSONResponse(
             status_code=406,
             content={
@@ -431,11 +423,8 @@ async def submit_jobs_v2(request: Request):
         )
     except Exception as e:
         logging.exception(e, exc_info=True)
-        log_stage_event(
-            "Failed submit jobs v2 request",
-            event_type=EventType.STAGE_FAILURE,
-            content=content,
-            error=str(e.args),
+        log_submit_job_request(
+            content=content, event_type=EventType.STAGE_FAILURE
         )
         return JSONResponse(
             status_code=500,

tests/test_log_handler.py

@@ -9,7 +9,6 @@
 from aind_data_transfer_service import CustomJsonFormatter
 from aind_data_transfer_service.log_handler import (
     EventType,
-    log_stage_event,
     log_submit_job_request,
 )
 
@@ -20,8 +19,12 @@ class TestLogHandler(unittest.TestCase):
     @patch("logging.info")
     def test_log_submit_job_request(self, mock_log: MagicMock):
         """Tests log_submit_job_request"""
-        content = [{"s3_prefix": "abc-123", "subject_id": "123456"}]
-        log_submit_job_request(content=content)
+        content = {
+            "upload_jobs": [{"s3_prefix": "abc-123", "subject_id": "123456"}]
+        }
+        log_submit_job_request(
+            content=content, event_type=EventType.STAGE_START
+        )
         mock_log.assert_called_once_with(
             "Handling request",
             extra={
@@ -31,26 +34,6 @@ def test_log_submit_job_request(self, mock_log: MagicMock):
             },
         )
 
-    @patch("logging.info")
-    def test_log_stage_event(self, mock_log: MagicMock):
-        """Tests log_stage_event emits structured extra fields."""
-
-        log_stage_event(
-            "Completed submit jobs v2 request",
-            event_type=EventType.STAGE_COMPLETE,
-            dag_id="run_list_of_jobs",
-            total_jobs=2,
-        )
-
-        mock_log.assert_called_once_with(
-            "Completed submit jobs v2 request",
-            extra={
-                "event_type": EventType.STAGE_COMPLETE,
-                "dag_id": "run_list_of_jobs",
-                "total_jobs": 2,
-            },
-        )
-
 
 class TestCustomJsonFormatter(unittest.TestCase):
     """Tests methods CustomJsonFormatter from init module"""

tests/test_server.py

@@ -1577,7 +1577,6 @@ def test_submit_v1_v2_jobs_406(
 
     @patch.dict(os.environ, EXAMPLE_ENV_VAR1, clear=True)
     @patch("httpx.AsyncClient.post")
-    @patch("aind_data_transfer_service.server.log_stage_event")
     @patch("aind_data_transfer_service.server.get_airflow_jobs")
     @patch("aind_data_transfer_service.server.get_project_names")
     @patch("aind_data_transfer_service.server.get_job_types")
@@ -1586,7 +1585,6 @@ def test_submit_v1_v2_jobs_200(
         mock_get_job_types: MagicMock,
         mock_get_project_names: MagicMock,
         mock_get_airflow_jobs: MagicMock,
-        mock_log_stage_event: MagicMock,
         mock_post: MagicMock,
     ):
         """Tests submit jobs success."""
@@ -1618,18 +1616,10 @@ def test_submit_v1_v2_jobs_200(
             params=expected_airflow_params, get_confs=True
         )
         self.assertEqual(1, mock_get_project_names.call_count)
-        self.assertEqual(4, len(captured.output))
-        mock_log_stage_event.assert_called_once_with(
-            "Completed submit jobs v2 request",
-            event_type="stage_complete",
-            dag_id="transform_and_upload_v2",
-            total_jobs=1,
-            status_code=200,
-        )
+        self.assertEqual(6, len(captured.output))
 
     @patch.dict(os.environ, EXAMPLE_ENV_VAR1, clear=True)
     @patch("httpx.AsyncClient.post")
-    @patch("aind_data_transfer_service.server.log_stage_event")
     @patch("aind_data_transfer_service.server.get_airflow_jobs")
     @patch("aind_data_transfer_service.server.get_project_names")
     @patch("aind_data_transfer_service.server.get_job_types")
@@ -1638,7 +1628,6 @@ def test_submit_v1_v2_jobs_500(
         mock_get_job_types: MagicMock,
         mock_get_project_names: MagicMock,
         mock_get_airflow_jobs: MagicMock,
-        mock_log_stage_event: MagicMock,
         mock_post: MagicMock,
     ):
         """Tests submit jobs failure."""
@@ -1670,14 +1659,7 @@ def test_submit_v1_v2_jobs_500(
             params=expected_airflow_params, get_confs=True
         )
         self.assertEqual(1, mock_get_project_names.call_count)
-        self.assertEqual(4, len(captured.output))
-        mock_log_stage_event.assert_called_once_with(
-            "Failed submit jobs v2 request",
-            event_type="stage_failure",
-            dag_id="transform_and_upload_v2",
-            total_jobs=1,
-            status_code=500,
-        )
+        self.assertEqual(7, len(captured.output))
 
     @patch.dict(os.environ, EXAMPLE_ENV_VAR1, clear=True)
     @patch("httpx.AsyncClient.post")
@@ -1773,7 +1755,7 @@ def test_submit_v2_jobs_200_basic_serialization(
                     url="/api/v2/submit_jobs", json=post_request_content_v2
                 )
         self.assertEqual(200, submit_job_response.status_code)
-        self.assertEqual(5, len(captured.output))
+        self.assertEqual(6, len(captured.output))
         mock_get_job_types.assert_called_once_with("v2")
         mock_get_airflow_jobs.assert_called_once()
         self.assertEqual(1, mock_get_project_names.call_count)
@@ -1822,7 +1804,7 @@ def test_validate_json(
         )
         mock_get_job_types.assert_called_once_with("v2")
         self.assertEqual(1, mock_get_project_names.call_count)
-        self.assertEqual(3, len(captured.output))
+        self.assertEqual(4, len(captured.output))
 
     @patch("aind_data_transfer_service.server.get_airflow_jobs")
     @patch("aind_data_transfer_service.server.get_project_names")