Refactor event writers to improve type hinting and logging; enhance error handling in write methods

Author: oto-macenauer-absa

pyproject.toml

@@ -9,3 +9,8 @@ omit = ["tests/*"]
 [tool.mypy]
 check_untyped_defs = true
 exclude = "tests"
+ignore_missing_imports = true
+python_version = "3.11"
+packages = ["src"]
+explicit_package_bases = true
+disable_error_code = ["import-untyped"]

src/event_gate_lambda.py

@@ -19,6 +19,7 @@
 import os
 import sys
 import urllib3
+from typing import Any
 
 import boto3
 import jwt
@@ -31,11 +32,7 @@
 _PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 _CONF_DIR = os.path.join(_PROJECT_ROOT, "conf")
 
-sys.path.append(os.path.join(os.path.dirname(__file__)))
-
-import writer_eventbridge
-import writer_kafka
-import writer_postgres
+from . import writer_eventbridge, writer_kafka, writer_postgres
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
@@ -77,7 +74,7 @@
 
 TOKEN_PROVIDER_URL = CONFIG["token_provider_url"]
 token_public_key_encoded = requests.get(CONFIG["token_public_key_url"], verify=False).json()["key"]
-TOKEN_PUBLIC_KEY = serialization.load_der_public_key(base64.b64decode(token_public_key_encoded))
+TOKEN_PUBLIC_KEY: Any = serialization.load_der_public_key(base64.b64decode(token_public_key_encoded))
 logger.debug("Loaded TOKEN_PUBLIC_KEY")
 
 writer_eventbridge.init(logger, CONFIG)

src/writer_eventbridge.py

@@ -14,9 +14,16 @@
 # limitations under the License.
 #
 import json
+import logging
+from typing import Optional, Tuple
 
 import boto3
 
+# Module globals for typing
+_logger: logging.Logger = logging.getLogger(__name__)
+EVENT_BUS_ARN: str = ""
+aws_eventbridge = None  # will hold boto3 client
+
 
 def init(logger, CONFIG):
     global _logger
@@ -30,10 +37,13 @@ def init(logger, CONFIG):
     _logger.debug("Initialized EVENTBRIDGE writer")
 
 
-def write(topicName, message):
+def write(topicName, message) -> Tuple[bool, Optional[str]]:
     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")
+        return True, None
 
     try:
         _logger.debug(f"Sending to eventBridge {topicName}")

src/writer_kafka.py

@@ -14,10 +14,16 @@
 # limitations under the License.
 #
 import json
+import logging
+from typing import Any, Dict, Optional, Tuple
 
 import boto3
 from confluent_kafka import Producer
 
+# Module globals for typing
+_logger: logging.Logger = logging.getLogger(__name__)
+kafka_producer: Optional[Producer] = None
+
 
 def init(logger, CONFIG):
     global _logger
@@ -45,10 +51,13 @@ def init(logger, CONFIG):
     _logger.debug("Initialized KAFKA writer")
 
 
-def write(topicName, message):
+def write(topicName, message) -> Tuple[bool, Optional[str]]:
     try:
+        if kafka_producer is None:
+            _logger.debug("Kafka producer not initialized - skipping")
+            return True, None
         _logger.debug(f"Sending to kafka {topicName}")
-        error = []
+        error: list[Any] = []
         kafka_producer.produce(
             topicName,
             key="",
@@ -57,10 +66,12 @@ def write(topicName, message):
         )
         kafka_producer.flush()
         if error:
-            _logger.error(str(error))
-            return False
+            msg = str(error)
+            _logger.error(msg)
+            return False, msg
     except Exception as e:
-        _logger.error(f"The Kafka writer failed with unknown error: {str(e)}")
-        return False
+        msg = f"The Kafka writer failed with unknown error: {str(e)}"
+        _logger.error(msg)
+        return False, msg
 
-    return True
+    return True, None

src/writer_postgres.py

@@ -15,14 +15,20 @@
 #
 import json
 import os
+import logging
+from typing import Any, Dict, Tuple, Optional
 
 import boto3
 from botocore.exceptions import ClientError
 
 try:
     import psycopg2  # noqa: F401
 except ImportError:  # pragma: no cover - environment without psycopg2
-    psycopg2 = None
+    psycopg2 = None  # type: ignore
+
+# Module level globals for typing
+_logger: logging.Logger = logging.getLogger(__name__)
+POSTGRES: Dict[str, Any] = {"database": ""}
 
 
 def init(logger):
@@ -201,23 +207,23 @@ def postgres_test_write(cursor, table, message):
     )
 
 
-def write(topicName, message):
+def write(topicName, message) -> Tuple[bool, Optional[str]]:
     try:
-        if not POSTGRES["database"]:
+        if not POSTGRES.get("database"):
             _logger.debug("No Postgres - skipping")
             return True, None
-        if psycopg2 is None:
+        if psycopg2 is None:  # type: ignore
             _logger.debug("psycopg2 not available - skipping actual Postgres write")
             return True, None
 
-        with psycopg2.connect(
+        with psycopg2.connect(  # type: ignore[attr-defined]
             database=POSTGRES["database"],
             host=POSTGRES["host"],
             user=POSTGRES["user"],
             password=POSTGRES["password"],
             port=POSTGRES["port"],
         ) as connection:
-            with connection.cursor() as cursor:
+            with connection.cursor() as cursor:  # type: ignore
                 if topicName == "public.cps.za.dlchange":
                     postgres_edla_write(cursor, "public_cps_za_dlchange", message)
                 elif topicName == "public.cps.za.runs":
@@ -229,7 +235,7 @@ def write(topicName, message):
                     _logger.error(msg)
                     return False, msg
 
-            connection.commit()
+            connection.commit()  # type: ignore
     except Exception as e:  # pragma: no cover - defensive (still tested though)
         err_msg = f"The Postgres writer with failed unknown error: {str(e)}"
         _logger.error(err_msg)