Merge pull request #31 from neurostuff/fix/lambda_deploy

[FIX] the logging/result fetching endpoints

Author: jdkent

.gitignore

@@ -2,3 +2,4 @@
 .runner
 .pytest_cache
 _version.py
+__pycache__

compose_runner/aws_lambda/__pycache__/__init__.cpython-313.pyc

@@ -2,6 +2,8 @@
 
 import os
 import time
+import base64
+import json
 from typing import Any, Dict, List
 
 import boto3
@@ -11,9 +13,38 @@
 LOG_GROUP_ENV = "RUNNER_LOG_GROUP"
 DEFAULT_LOOKBACK_MS_ENV = "DEFAULT_LOOKBACK_MS"
 
+def _is_http_event(event: Any) -> bool:
+    return isinstance(event, dict) and "requestContext" in event
+
+
+def _extract_payload(event: Dict[str, Any]) -> Dict[str, Any]:
+    if not _is_http_event(event):
+        return event
+    body = event.get("body")
+    if not body:
+        return {}
+    if event.get("isBase64Encoded"):
+        body = base64.b64decode(body).decode("utf-8")
+    return json.loads(body)
+
+
+def _http_response(body: Dict[str, Any], status_code: int = 200) -> Dict[str, Any]:
+    return {
+        "statusCode": status_code,
+        "headers": {"Content-Type": "application/json"},
+        "body": json.dumps(body),
+    }
+
 
 def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
-    job_id = event["job_id"]
+    raw_event = event
+    event = _extract_payload(event)
+    job_id = event.get("job_id")
+    if not job_id:
+        message = "Request payload must include 'job_id'."
+        if _is_http_event(raw_event):
+            return _http_response({"status": "FAILED", "error": message}, status_code=400)
+        raise KeyError(message)
     next_token = event.get("next_token")
     start_time = event.get("start_time")
     end_time = event.get("end_time")
@@ -29,7 +60,7 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
 
     params: Dict[str, Any] = {
         "logGroupName": log_group,
-        "filterPattern": f'{{"job_id": "{job_id}"}}',
+        "filterPattern": f'"{job_id}"',
         "startTime": int(start_time),
     }
     if end_time is not None:
@@ -43,8 +74,11 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
         for item in response.get("events", [])
     ]
 
-    return {
+    body = {
         "job_id": job_id,
         "events": events,
         "next_token": response.get("nextToken"),
     }
+    if _is_http_event(raw_event):
+        return _http_response(body)
+    return body

compose_runner/aws_lambda/__pycache__/log_poll_handler.cpython-313.pyc

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import os
+import base64
+import json
 from datetime import datetime, timezone
 from typing import Any, Dict, List
 
@@ -19,11 +21,41 @@ def _serialize_dt(value: datetime) -> str:
     return value.astimezone(timezone.utc).isoformat()
 
 
+def _is_http_event(event: Any) -> bool:
+    return isinstance(event, dict) and "requestContext" in event
+
+
+def _extract_payload(event: Dict[str, Any]) -> Dict[str, Any]:
+    if not _is_http_event(event):
+        return event
+    body = event.get("body")
+    if not body:
+        return {}
+    if event.get("isBase64Encoded"):
+        body = base64.b64decode(body).decode("utf-8")
+    return json.loads(body)
+
+
+def _http_response(body: Dict[str, Any], status_code: int = 200) -> Dict[str, Any]:
+    return {
+        "statusCode": status_code,
+        "headers": {"Content-Type": "application/json"},
+        "body": json.dumps(body),
+    }
+
+
 def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
+    raw_event = event
+    event = _extract_payload(event)
     bucket = os.environ[RESULTS_BUCKET_ENV]
     prefix = os.environ.get(RESULTS_PREFIX_ENV)
 
-    job_id = event["job_id"]
+    job_id = event.get("job_id")
+    if not job_id:
+        message = "Request payload must include 'job_id'."
+        if _is_http_event(raw_event):
+            return _http_response({"status": "FAILED", "error": message}, status_code=400)
+        raise KeyError(message)
     expires_in = int(event.get("expires_in", DEFAULT_EXPIRES_IN))
 
     key_prefix = f"{prefix.rstrip('/')}/{job_id}" if prefix else job_id
@@ -51,9 +83,12 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
             }
         )
 
-    return {
+    body = {
         "job_id": job_id,
         "artifacts": artifacts,
         "bucket": bucket,
         "prefix": key_prefix,
     }
+    if _is_http_event(raw_event):
+        return _http_response(body)
+    return body

compose_runner/aws_lambda/__pycache__/results_handler.cpython-313.pyc

@@ -87,6 +87,16 @@ def filter_log_events(self, **kwargs):
     assert result["events"][0]["message"] == events_payload[0]["message"]
 
 
+def test_log_poll_handler_http_missing_job_id(monkeypatch):
+    monkeypatch.setenv("RUNNER_LOG_GROUP", "/aws/lambda/test")
+    http_event = _make_http_event({})
+    response = log_poll_handler.handler(http_event, DummyContext())
+    body = json.loads(response["body"])
+    assert response["statusCode"] == 400
+    assert body["status"] == "FAILED"
+    assert "job_id" in body["error"]
+
+
 def test_results_handler(monkeypatch):
     objects = [
         {"Key": "prefix/id/file1.nii.gz", "Size": 10, "LastModified": results_handler.datetime.now()}
@@ -109,8 +119,20 @@ def generate_presigned_url(self, client_method, Params, ExpiresIn):
     monkeypatch.setenv("RESULTS_PREFIX", "prefix")
     monkeypatch.setattr(results_handler, "_S3", FakeS3())
 
-    event = {"job_id": "id"}
-    result = results_handler.handler(event, DummyContext())
-    assert result["job_id"] == "id"
-    assert result["artifacts"][0]["url"] == "https://signed/url"
-    assert result["artifacts"][0]["filename"] == "file1.nii.gz"
+    event = _make_http_event({"job_id": "id"})
+    response = results_handler.handler(event, DummyContext())
+    body = json.loads(response["body"])
+    assert response["statusCode"] == 200
+    assert body["job_id"] == "id"
+    assert body["artifacts"][0]["url"] == "https://signed/url"
+    assert body["artifacts"][0]["filename"] == "file1.nii.gz"
+
+
+def test_results_handler_missing_job_id(monkeypatch):
+    monkeypatch.setenv("RESULTS_BUCKET", "bucket")
+    event = _make_http_event({})
+    response = results_handler.handler(event, DummyContext())
+    body = json.loads(response["body"])
+    assert response["statusCode"] == 400
+    assert body["status"] == "FAILED"
+    assert "job_id" in body["error"]