Refactor event_gate_lambda and writer modules for improved type hinting, logging, and error handling; enhance readability and maintainability

Author: oto-macenauer-absa

src/event_gate_lambda.py

@@ -18,8 +18,9 @@
 import logging
 import os
 import sys
+from typing import Any, Dict
+
 import urllib3
-from typing import Any
 
 import boto3
 import jwt
@@ -28,52 +29,57 @@
 from jsonschema import validate
 from jsonschema.exceptions import ValidationError
 
-# Resolve project root (parent directory of this file's directory)
-_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-_CONF_DIR = os.path.join(_PROJECT_ROOT, "conf")
-
 from . import writer_eventbridge, writer_kafka, writer_postgres
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 logger = logging.getLogger(__name__)
 log_level = os.environ.get("LOG_LEVEL", "INFO")
 logger.setLevel(log_level)
-logger.addHandler(logging.StreamHandler())
+if not logger.handlers:
+    logger.addHandler(logging.StreamHandler())
 logger.debug("Initialized LOGGER")
 
-with open(os.path.join(_CONF_DIR, "api.yaml"), "r") as file:
+# Resolve project root (parent directory of this file's directory)
+_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+_CONF_DIR = os.path.join(_PROJECT_ROOT, "conf")
+
+with open(os.path.join(_CONF_DIR, "api.yaml"), "r", encoding="utf-8") as file:
     API = file.read()
 logger.debug("Loaded API definition")
 
-TOPICS = {}
-with open(os.path.join(_CONF_DIR, "topic_runs.json"), "r") as file:
+TOPICS: Dict[str, Dict[str, Any]] = {}
+with open(os.path.join(_CONF_DIR, "topic_runs.json"), "r", encoding="utf-8") as file:
     TOPICS["public.cps.za.runs"] = json.load(file)
-with open(os.path.join(_CONF_DIR, "topic_dlchange.json"), "r") as file:
+with open(os.path.join(_CONF_DIR, "topic_dlchange.json"), "r", encoding="utf-8") as file:
     TOPICS["public.cps.za.dlchange"] = json.load(file)
-with open(os.path.join(_CONF_DIR, "topic_test.json"), "r") as file:
+with open(os.path.join(_CONF_DIR, "topic_test.json"), "r", encoding="utf-8") as file:
     TOPICS["public.cps.za.test"] = json.load(file)
 logger.debug("Loaded TOPICS")
 
-with open(os.path.join(_CONF_DIR, "config.json"), "r") as file:
+with open(os.path.join(_CONF_DIR, "config.json"), "r", encoding="utf-8") as file:
     CONFIG = json.load(file)
 logger.debug("Loaded main CONFIG")
 
-aws_s3 = boto3.Session().resource("s3", verify=False)
+aws_s3 = boto3.Session().resource("s3", verify=False)  # nosec Boto verify disabled intentionally
 logger.debug("Initialized AWS S3 Client")
 
 if CONFIG["access_config"].startswith("s3://"):
     name_parts = CONFIG["access_config"].split("/")
-    bucket_name = name_parts[2]
-    bucket_object = "/".join(name_parts[3:])
-    ACCESS = json.loads(aws_s3.Bucket(bucket_name).Object(bucket_object).get()["Body"].read().decode("utf-8"))
+    BUCKET_NAME = name_parts[2]
+    BUCKET_OBJECT_KEY = "/".join(name_parts[3:])
+    ACCESS = json.loads(
+        aws_s3.Bucket(BUCKET_NAME).Object(BUCKET_OBJECT_KEY).get()["Body"].read().decode("utf-8")
+    )
 else:
-    with open(CONFIG["access_config"], "r") as file:
+    with open(CONFIG["access_config"], "r", encoding="utf-8") as file:
         ACCESS = json.load(file)
 logger.debug("Loaded ACCESS definitions")
 
 TOKEN_PROVIDER_URL = CONFIG["token_provider_url"]
-token_public_key_encoded = requests.get(CONFIG["token_public_key_url"], verify=False).json()["key"]
+# Add timeout to avoid hanging requests
+response_json = requests.get(CONFIG["token_public_key_url"], verify=False, timeout=5).json()  # nosec external
+token_public_key_encoded = response_json["key"]
 TOKEN_PUBLIC_KEY: Any = serialization.load_der_public_key(base64.b64decode(token_public_key_encoded))
 logger.debug("Loaded TOKEN_PUBLIC_KEY")
 
@@ -82,7 +88,16 @@
 writer_postgres.init(logger)
 
 
-def _error_response(status, err_type, message):
+def _error_response(status: int, err_type: str, message: str) -> Dict[str, Any]:
+    """Build a standardized JSON error response body.
+
+    Args:
+        status: HTTP status code.
+        err_type: A short error classifier (e.g. 'auth', 'validation').
+        message: Human readable error description.
+    Returns:
+        A dictionary compatible with API Gateway Lambda Proxy integration.
+    """
     return {
         "statusCode": status,
         "headers": {"Content-Type": "application/json"},
@@ -92,55 +107,69 @@ def _error_response(status, err_type, message):
     }
 
 
-def get_api():
+def get_api() -> Dict[str, Any]:
+    """Return the OpenAPI specification text."""
     return {"statusCode": 200, "body": API}
 
 
-def get_token():
+def get_token() -> Dict[str, Any]:
+    """Return 303 redirect to token provider endpoint."""
     logger.debug("Handling GET Token")
     return {"statusCode": 303, "headers": {"Location": TOKEN_PROVIDER_URL}}
 
 
-def get_topics():
+def get_topics() -> Dict[str, Any]:
+    """Return list of available topic names."""
     logger.debug("Handling GET Topics")
     return {
         "statusCode": 200,
         "headers": {"Content-Type": "application/json"},
-        "body": json.dumps([topicName for topicName in TOPICS]),
+        "body": json.dumps(list(TOPICS)),
     }
 
 
-def get_topic_schema(topicName):
-    logger.debug(f"Handling GET TopicSchema({topicName})")
-    if topicName not in TOPICS:
-        return _error_response(404, "topic", f"Topic '{topicName}' not found")
+def get_topic_schema(topic_name: str) -> Dict[str, Any]:
+    """Return the JSON schema for a specific topic.
+
+    Args:
+        topic_name: The topic whose schema is requested.
+    """
+    logger.debug("Handling GET TopicSchema(%s)", topic_name)
+    if topic_name not in TOPICS:
+        return _error_response(404, "topic", f"Topic '{topic_name}' not found")
+
+    return {"statusCode": 200, "headers": {"Content-Type": "application/json"}, "body": json.dumps(TOPICS[topic_name])}
 
-    return {"statusCode": 200, "headers": {"Content-Type": "application/json"}, "body": json.dumps(TOPICS[topicName])}
 
+def post_topic_message(topic_name: str, topic_message: Dict[str, Any], token_encoded: str) -> Dict[str, Any]:
+    """Validate auth and schema; dispatch message to all writers.
 
-def post_topic_message(topicName, topicMessage, tokenEncoded):
-    logger.debug(f"Handling POST {topicName}")
+    Args:
+        topic_name: Target topic name.
+        topic_message: JSON message payload.
+        token_encoded: Encoded bearer JWT token string.
+    """
+    logger.debug("Handling POST %s", topic_name)
     try:
-        token = jwt.decode(tokenEncoded, TOKEN_PUBLIC_KEY, algorithms=["RS256"])
-    except Exception:
+        token = jwt.decode(token_encoded, TOKEN_PUBLIC_KEY, algorithms=["RS256"])  # type: ignore[arg-type]
+    except jwt.PyJWTError:  # type: ignore[attr-defined]
         return _error_response(401, "auth", "Invalid or missing token")
 
-    if topicName not in TOPICS:
-        return _error_response(404, "topic", f"Topic '{topicName}' not found")
+    if topic_name not in TOPICS:
+        return _error_response(404, "topic", f"Topic '{topic_name}' not found")
 
-    user = token["sub"]
-    if topicName not in ACCESS or user not in ACCESS[topicName]:
+    user = token.get("sub")
+    if topic_name not in ACCESS or user not in ACCESS[topic_name]:  # type: ignore[index]
         return _error_response(403, "auth", "User not authorized for topic")
 
     try:
-        validate(instance=topicMessage, schema=TOPICS[topicName])
-    except ValidationError as e:
-        return _error_response(400, "validation", e.message)
+        validate(instance=topic_message, schema=TOPICS[topic_name])
+    except ValidationError as exc:
+        return _error_response(400, "validation", exc.message)
 
-    # Run all writers independently (avoid short-circuit so failures in one don't skip others)
-    kafka_ok, kafka_err = writer_kafka.write(topicName, topicMessage)
-    eventbridge_ok, eventbridge_err = writer_eventbridge.write(topicName, topicMessage)
-    postgres_ok, postgres_err = writer_postgres.write(topicName, topicMessage)
+    kafka_ok, kafka_err = writer_kafka.write(topic_name, topic_message)
+    eventbridge_ok, eventbridge_err = writer_eventbridge.write(topic_name, topic_message)
+    postgres_ok, postgres_err = writer_postgres.write(topic_name, topic_message)
 
     errors = []
     if not kafka_ok:
@@ -164,37 +193,46 @@ def post_topic_message(topicName, topicMessage, tokenEncoded):
     }
 
 
-def extract_token(eventHeaders):
-    # Initial implementation used bearer header directly
-    if "bearer" in eventHeaders:
-        return eventHeaders["bearer"]
+def extract_token(event_headers: Dict[str, str]) -> str:
+    """Extract bearer token from headers.
 
-    if "Authorization" in eventHeaders and eventHeaders["Authorization"].startswith("Bearer "):
-        return eventHeaders["Authorization"][len("Bearer ") :]
+    Supports lowercase custom 'bearer' header or standard 'Authorization: Bearer <token>'.
+    Returns empty string if not present (caller handles auth error response).
+    """
+    if "bearer" in event_headers:
+        return event_headers["bearer"]
+    auth_header = event_headers.get("Authorization", "")
+    if auth_header.startswith("Bearer "):
+        return auth_header[len("Bearer ") :]
+    return ""
 
-    return ""  # Will result in 401
 
+def lambda_handler(event: Dict[str, Any], context: Any):  # pylint: disable=unused-argument,too-many-return-statements
+    """AWS Lambda entry point.
 
-def lambda_handler(event, context):
+    Dispatches based on API Gateway proxy 'resource' and 'httpMethod'.
+    """
     try:
-        if event["resource"].lower() == "/api":
+        resource = event.get("resource", "").lower()
+        if resource == "/api":
             return get_api()
-        if event["resource"].lower() == "/token":
+        if resource == "/token":
             return get_token()
-        if event["resource"].lower() == "/topics":
+        if resource == "/topics":
             return get_topics()
-        if event["resource"].lower() == "/topics/{topic_name}":
-            if event["httpMethod"] == "GET":
+        if resource == "/topics/{topic_name}":
+            method = event.get("httpMethod")
+            if method == "GET":
                 return get_topic_schema(event["pathParameters"]["topic_name"].lower())
-            if event["httpMethod"] == "POST":
+            if method == "POST":
                 return post_topic_message(
                     event["pathParameters"]["topic_name"].lower(),
                     json.loads(event["body"]),
-                    extract_token(event["headers"]),
+                    extract_token(event.get("headers", {})),
                 )
-        if event["resource"].lower() == "/terminate":
-            sys.exit("TERMINATING")
+        if resource == "/terminate":
+            sys.exit("TERMINATING")  # pragma: no cover - deliberate termination path
         return _error_response(404, "route", "Resource not found")
-    except Exception as e:
-        logger.error(f"Unexpected exception: {e}")
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        logger.error("Unexpected exception: %s", exc)
         return _error_response(500, "internal", "Unexpected server error")

src/writer_eventbridge.py

@@ -1,69 +1,75 @@
-#
-# Copyright 2025 ABSA Group Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
+"""EventBridge writer module.
+
+Provides initialization and write functionality for publishing events to AWS EventBridge.
+"""
+
 import json
 import logging
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import boto3
+from botocore.exceptions import BotoCoreError, ClientError
 
-# Module globals for typing
-_logger: logging.Logger = logging.getLogger(__name__)
-EVENT_BUS_ARN: str = ""
-aws_eventbridge = None  # will hold boto3 client
+STATE: Dict[str, Any] = {"logger": logging.getLogger(__name__), "event_bus_arn": "", "client": None}
 
 
-def init(logger, CONFIG):
-    global _logger
-    global EVENT_BUS_ARN
-    global aws_eventbridge
+def init(logger: logging.Logger, config: Dict[str, Any]) -> None:
+    """Initialize the EventBridge writer.
 
-    _logger = logger
+    Args:
+        logger: Shared application logger.
+        config: Configuration dictionary (expects optional 'event_bus_arn').
+    """
+    STATE["logger"] = logger
+    STATE["client"] = boto3.client("events")
+    STATE["event_bus_arn"] = config.get("event_bus_arn", "")
+    STATE["logger"].debug("Initialized EVENTBRIDGE writer")
 
-    aws_eventbridge = boto3.client("events")
-    EVENT_BUS_ARN = CONFIG["event_bus_arn"] if "event_bus_arn" in CONFIG else ""
-    _logger.debug("Initialized EVENTBRIDGE writer")
 
+def write(topic_name: str, message: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
+    """Publish a message to EventBridge.
 
-def write(topicName, message) -> Tuple[bool, Optional[str]]:
-    if not EVENT_BUS_ARN:
-        _logger.debug("No EventBus Arn - skipping")
+    Args:
+        topic_name: Source topic name used as event Source.
+        message: JSON-serializable payload.
+    Returns:
+        Tuple of success flag and optional error message.
+    """
+    logger = STATE["logger"]
+    event_bus_arn = STATE["event_bus_arn"]
+    client = STATE["client"]
+
+    if not event_bus_arn:
+        logger.debug("No EventBus Arn - skipping")
         return True, None
-    if aws_eventbridge is None:  # defensive
-        _logger.debug("EventBridge client not initialized - skipping")
+    if client is None:  # defensive
+        logger.debug("EventBridge client not initialized - skipping")
         return True, None
 
     try:
-        _logger.debug(f"Sending to eventBridge {topicName}")
-        response = aws_eventbridge.put_events(
+        logger.debug("Sending to eventBridge %s", topic_name)
+        response = client.put_events(
             Entries=[
                 {
-                    "Source": topicName,
+                    "Source": topic_name,
                     "DetailType": "JSON",
                     "Detail": json.dumps(message),
-                    "EventBusName": EVENT_BUS_ARN,
+                    "EventBusName": event_bus_arn,
                 }
             ]
         )
-        if response["FailedEntryCount"] > 0:
+        if response.get("FailedEntryCount", 0) > 0:
             msg = str(response)
-            _logger.error(msg)
+            logger.error(msg)
             return False, msg
-    except Exception as e:
-        err_msg = f"The EventBridge writer failed with unknown error: {str(e)}"
-        _logger.error(err_msg)
+    except (BotoCoreError, ClientError) as err:
+        err_msg = f"The EventBridge writer failed: {err}"  # specific AWS error
+        logger.error(err_msg)
+        return False, err_msg
+    except Exception as err:  # pragma: no cover - unexpected failure path
+        err_msg = f"The EventBridge writer failed with unknown error: {err}" \
+            if not isinstance(err, (BotoCoreError, ClientError)) else str(err)
+        logger.error(err_msg)
         return False, err_msg
 
     return True, None