Refactor code for consistency; update string formatting and improve readability

Author: oto-macenauer-absa

src/__init__.py

@@ -1,2 +1 @@
 """Package initializer for EventGate source modules."""
-

src/event_gate_lambda.py

@@ -1,18 +1,18 @@
-# 
+#
 # Copyright 2024 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.
-# 
+#
 import base64
 import json
 import logging
@@ -28,8 +28,8 @@
 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')
+_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__)))
 
@@ -45,28 +45,28 @@
 logger.addHandler(logging.StreamHandler())
 logger.debug("Initialized LOGGER")
 
-with open(os.path.join(_CONF_DIR, 'api.yaml'), 'r') as file:
+with open(os.path.join(_CONF_DIR, "api.yaml"), "r") as file:
     API = file.read()
 logger.debug("Loaded API definition")
 
 TOPICS = {}
-with open(os.path.join(_CONF_DIR, 'topic_runs.json'), 'r') as file:
+with open(os.path.join(_CONF_DIR, "topic_runs.json"), "r") 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") 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") 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") 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)
 logger.debug("Initialized AWS S3 Client")
 
 if CONFIG["access_config"].startswith("s3://"):
-    name_parts = CONFIG["access_config"].split('/')
+    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"))
@@ -84,55 +84,49 @@
 writer_kafka.init(logger, CONFIG)
 writer_postgres.init(logger)
 
+
 def _error_response(status, err_type, message):
     return {
         "statusCode": status,
         "headers": {"Content-Type": "application/json"},
-        "body": json.dumps({
-            "success": False,
-            "statusCode": status,
-            "errors": [{"type": err_type, "message": message}]
-        })
+        "body": json.dumps(
+            {"success": False, "statusCode": status, "errors": [{"type": err_type, "message": message}]}
+        ),
     }
 
+
 def get_api():
-    return {
-        "statusCode": 200,
-        "body": API
-    }
+    return {"statusCode": 200, "body": API}
+
 
 def get_token():
     logger.debug("Handling GET Token")
-    return {
-        "statusCode": 303,
-        "headers": {"Location": TOKEN_PROVIDER_URL}
-    }
-    
+    return {"statusCode": 303, "headers": {"Location": TOKEN_PROVIDER_URL}}
+
+
 def get_topics():
     logger.debug("Handling GET Topics")
     return {
         "statusCode": 200,
         "headers": {"Content-Type": "application/json"},
-        "body": json.dumps([topicName for topicName in TOPICS])
+        "body": json.dumps([topicName for topicName in 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")
-        
-    return {
-        "statusCode": 200,
-        "headers": {"Content-Type": "application/json"},
-        "body": json.dumps(TOPICS[topicName])
-    }
+
+    return {"statusCode": 200, "headers": {"Content-Type": "application/json"}, "body": json.dumps(TOPICS[topicName])}
+
 
 def post_topic_message(topicName, topicMessage, tokenEncoded):
     logger.debug(f"Handling POST {topicName}")
     try:
         token = jwt.decode(tokenEncoded, TOKEN_PUBLIC_KEY, algorithms=["RS256"])
     except Exception:
-         return _error_response(401, "auth", "Invalid or missing token")
+        return _error_response(401, "auth", "Invalid or missing token")
 
     if topicName not in TOPICS:
         return _error_response(404, "topic", f"Topic '{topicName}' not found")
@@ -144,8 +138,8 @@ def post_topic_message(topicName, topicMessage, tokenEncoded):
     try:
         validate(instance=topicMessage, schema=TOPICS[topicName])
     except ValidationError as e:
-         return _error_response(400, "validation", e.message)
-    
+        return _error_response(400, "validation", e.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)
@@ -163,31 +157,26 @@ def post_topic_message(topicName, topicMessage, tokenEncoded):
         return {
             "statusCode": 500,
             "headers": {"Content-Type": "application/json"},
-            "body": json.dumps({
-                "success": False,
-                "statusCode": 500,
-                "errors": errors
-            })
+            "body": json.dumps({"success": False, "statusCode": 500, "errors": errors}),
         }
 
     return {
         "statusCode": 202,
         "headers": {"Content-Type": "application/json"},
-        "body": json.dumps({
-            "success": True,
-            "statusCode": 202
-        })
+        "body": json.dumps({"success": True, "statusCode": 202}),
     }
 
+
 def extract_token(eventHeaders):
     # Initial implementation used bearer header directly
     if "bearer" in eventHeaders:
         return eventHeaders["bearer"]
-        
+
     if "Authorization" in eventHeaders and eventHeaders["Authorization"].startswith("Bearer "):
-        return eventHeaders["Authorization"][len("Bearer "):]
-        
-    return "" # Will result in 401
+        return eventHeaders["Authorization"][len("Bearer ") :]
+
+    return ""  # Will result in 401
+
 
 def lambda_handler(event, context):
     try:
@@ -201,7 +190,11 @@ def lambda_handler(event, context):
             if event["httpMethod"] == "GET":
                 return get_topic_schema(event["pathParameters"]["topic_name"].lower())
             if event["httpMethod"] == "POST":
-                return post_topic_message(event["pathParameters"]["topic_name"].lower(), json.loads(event["body"]), extract_token(event["headers"]))  
+                return post_topic_message(
+                    event["pathParameters"]["topic_name"].lower(),
+                    json.loads(event["body"]),
+                    extract_token(event["headers"]),
+                )
         if event["resource"].lower() == "/terminate":
             sys.exit("TERMINATING")
         return _error_response(404, "route", "Resource not found")

src/writer_eventbridge.py

@@ -2,17 +2,19 @@
 
 import boto3
 
+
 def init(logger, CONFIG):
     global _logger
     global EVENT_BUS_ARN
     global aws_eventbridge
-    
+
     _logger = logger
-    
-    aws_eventbridge = boto3.client('events')
+
+    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(topicName, message):
     if not EVENT_BUS_ARN:
         _logger.debug("No EventBus Arn - skipping")
@@ -24,17 +26,17 @@ def write(topicName, message):
             Entries=[
                 {
                     "Source": topicName,
-                    'DetailType': 'JSON',
-                    'Detail': json.dumps(message),
-                    'EventBusName': EVENT_BUS_ARN,
+                    "DetailType": "JSON",
+                    "Detail": json.dumps(message),
+                    "EventBusName": EVENT_BUS_ARN,
                 }
             ]
         )
         if response["FailedEntryCount"] > 0:
             _logger.error(str(response))
             return False
     except Exception as e:
-        _logger.error(f'The EventBridge writer failed with unknown error: {str(e)}')
+        _logger.error(f"The EventBridge writer failed with unknown error: {str(e)}")
         return False
-        
+
     return True

src/writer_kafka.py

@@ -3,43 +3,49 @@
 import boto3
 from confluent_kafka import Producer
 
+
 def init(logger, CONFIG):
     global _logger
     global kafka_producer
-    
+
     _logger = logger
-    
+
     producer_config = {"bootstrap.servers": CONFIG["kafka_bootstrap_server"]}
     if "kafka_sasl_kerberos_principal" in CONFIG and "kafka_ssl_key_path" in CONFIG:
-        producer_config.update({
-            "security.protocol": "SASL_SSL",
-            "sasl.mechanism": "GSSAPI",
-            "sasl.kerberos.service.name": "kafka",
-            "sasl.kerberos.keytab": CONFIG["kafka_sasl_kerberos_keytab_path"],
-            "sasl.kerberos.principal": CONFIG["kafka_sasl_kerberos_principal"],
-            "ssl.ca.location": CONFIG["kafka_ssl_ca_path"],
-            "ssl.certificate.location": CONFIG["kafka_ssl_cert_path"],
-            "ssl.key.location": CONFIG["kafka_ssl_key_path"],
-            "ssl.key.password": CONFIG["kafka_ssl_key_password"]
-        })
+        producer_config.update(
+            {
+                "security.protocol": "SASL_SSL",
+                "sasl.mechanism": "GSSAPI",
+                "sasl.kerberos.service.name": "kafka",
+                "sasl.kerberos.keytab": CONFIG["kafka_sasl_kerberos_keytab_path"],
+                "sasl.kerberos.principal": CONFIG["kafka_sasl_kerberos_principal"],
+                "ssl.ca.location": CONFIG["kafka_ssl_ca_path"],
+                "ssl.certificate.location": CONFIG["kafka_ssl_cert_path"],
+                "ssl.key.location": CONFIG["kafka_ssl_key_path"],
+                "ssl.key.password": CONFIG["kafka_ssl_key_password"],
+            }
+        )
         _logger.debug("producer will use SASL_SSL")
     kafka_producer = Producer(producer_config)
     _logger.debug("Initialized KAFKA writer")
 
+
 def write(topicName, message):
     try:
         _logger.debug(f"Sending to kafka {topicName}")
         error = []
-        kafka_producer.produce(topicName, 
-                               key="", 
-                               value=json.dumps(message).encode("utf-8"),
-                               callback = lambda err, msg: error.append(err) if err is not None else None)
+        kafka_producer.produce(
+            topicName,
+            key="",
+            value=json.dumps(message).encode("utf-8"),
+            callback=lambda err, msg: error.append(err) if err is not None else None,
+        )
         kafka_producer.flush()
         if error:
             _logger.error(str(error))
             return False
     except Exception as e:
-        _logger.error(f'The Kafka writer failed with unknown error: {str(e)}')
+        _logger.error(f"The Kafka writer failed with unknown error: {str(e)}")
         return False
-        
+
     return True