Fix sdc timestream logging multi-mission

HERMES-SOC · Jul 1, 2024 · f889122 · f889122
1 parent 38a73e2
commit f889122
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 51 deletions.
diff --git a/sdc_aws_utils/aws.py b/sdc_aws_utils/aws.py
@@ -262,9 +262,14 @@ def log_to_timestream(
         if not source_bucket and not destination_bucket:
             raise ValueError("A Source or Destination Buckets is required")
 
-        # Check environment
-        database_name = "sdc_aws_logs"
-        table_name = "sdc_aws_s3_bucket_log_table"
+        # Check environment variable for SWXSOC_MISSION
+        mission_name = os.getenv("SWXSOC_MISSION")
+        if not mission_name or mission_name == "hermes":
+            database_name = "sdc_aws_logs"
+            table_name = "sdc_aws_s3_bucket_log_table"
+        else:
+            database_name = f"{mission_name}_sdc_aws_logs"
+            table_name = f"{mission_name}_sdc_aws_s3_bucket_log_table"
         database_name = f"dev-{database_name}" if environment == "DEVELOPMENT" else database_name
         table_name = f"dev-{table_name}" if environment == "DEVELOPMENT" else table_name
 
@@ -321,13 +326,9 @@ def invoke_reprocessing_lambda(bucket: str, key: str, environment: str) -> None:
     """
     # Create the JSON structure
     data = {
-        "Records": [
-            {
-                "Sns": {
-                    "Message": json.dumps({"Records": [{"s3": {"bucket": {"name": bucket}, "object": {"key": key}}}]})
-                }
-            }
-        ]
+        "Records": [{
+            "Sns": {"Message": json.dumps({"Records": [{"s3": {"bucket": {"name": bucket}, "object": {"key": key}}}]})}
+        }]
     }
 
     # Initialize a boto3 client for Lambda

diff --git a/sdc_aws_utils/slack.py b/sdc_aws_utils/slack.py
@@ -25,12 +25,10 @@ def get_slack_client(slack_token: str) -> WebClient:
 
     # If the slack token is still not set, return None
     if not slack_token:
-        log.error(
-            {
-                "status": "ERROR",
-                "message": "Slack Token is not set",
-            }
-        )
+        log.error({
+            "status": "ERROR",
+            "message": "Slack Token is not set",
+        })
         return None
 
     # Initialize the slack client
@@ -123,40 +121,32 @@ def send_slack_notification(
     # Check if slack_message is a tuple
     if isinstance(slack_message, tuple):
         text = slack_message[0]
-        attachments = [
-            {
-                "color": color["purple"],
-                "blocks": [
-                    {
-                        "type": "section",
-                        "text": {
-                            "type": "mrkdwn",
-                            "text": f"{slack_message[1]}",
-                        },
-                    }
-                ],
-                "fallback": f"{slack_message[1]}",
-            }
-        ]
+        attachments = [{
+            "color": color["purple"],
+            "blocks": [{
+                "type": "section",
+                "text": {
+                    "type": "mrkdwn",
+                    "text": f"{slack_message[1]}",
+                },
+            }],
+            "fallback": f"{slack_message[1]}",
+        }]
         pretext = slack_message[0]
     else:
         text = slack_message
         pretext = slack_message
         if alert_type:
-            attachments = [
-                {
-                    "color": color[alert_type],
-                    "blocks": [
-                        {
-                            "type": "section",
-                            "text": {
-                                "type": "mrkdwn",
-                                "text": f"{slack_message}",
-                            },
-                        }
-                    ],
-                }
-            ]
+            attachments = [{
+                "color": color[alert_type],
+                "blocks": [{
+                    "type": "section",
+                    "text": {
+                        "type": "mrkdwn",
+                        "text": f"{slack_message}",
+                    },
+                }],
+            }]
             text = f"`{ts}` -"
 
     for i in range(slack_max_retries):
@@ -181,12 +171,10 @@ def send_slack_notification(
                 )
                 time.sleep(slack_retry_delay)
             else:  # If it's the last attempt, log the error and exit the loop
-                log.error(
-                    {
-                        "status": "ERROR",
-                        "message": f"Error sending Slack Notification (attempt {i + 1}): {e}",
-                    }
-                )
+                log.error({
+                    "status": "ERROR",
+                    "message": f"Error sending Slack Notification (attempt {i + 1}): {e}",
+                })
                 raise e