Refactor slack notifications

loculus-project · Aug 29, 2024 · 4816478 · 4816478
1 parent 61826a5
commit 4816478
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 89 deletions.
diff --git a/ena-submission/scripts/create_assembly.py b/ena-submission/scripts/create_assembly.py
@@ -17,7 +17,7 @@
     ChromosomeType,
     MoleculeType,
 )
-from notifications import get_slack_config, notify
+from notifications import SlackConfig, send_slack_notification, slack_conn_init
 from submission_db_helper import (
     AssemblyTableEntry,
     Status,
@@ -61,28 +61,6 @@ class Config:
     slack_channel_id: str
 
 
-_last_notification_sent: datetime | None = None
-
-
-def send_slack_notification(config: Config, comment: str, time: datetime, time_threshold=12):
-    global _last_notification_sent  # noqa: PLW0603
-    slack_config = get_slack_config(
-        slack_hook_default=config.slack_hook,
-        slack_token_default=config.slack_token,
-        slack_channel_id_default=config.slack_channel_id,
-    )
-    if not slack_config.slack_hook:
-        logger.info("Could not find slack hook cannot send message")
-        return
-    if (
-        not _last_notification_sent
-        or time - timedelta(hours=time_threshold) > _last_notification_sent
-    ):
-        comment = f"{config.backend_url}: " + comment
-        notify(slack_config, comment)
-        _last_notification_sent = time
-
-
 def create_chromosome_list_object(
     unaligned_sequences: dict[str, str], seq_key: dict[str, str]
 ) -> str:
@@ -510,6 +488,7 @@ def assembly_table_update(db_config, config, retry_number=3, time_threshold=5):
 def assembly_table_handle_errors(
     db_config,
     config,
+    slack_config: SlackConfig,
     time_threshold=15,
     time_threshold_waiting=48,
     slack_time_threshold=12,
@@ -526,12 +505,15 @@ def assembly_table_handle_errors(
     )
     if len(entries_with_errors) > 0:
         error_msg = (
-            f"ENA Submission pipeline found {len(entries_with_errors)} entries in assembly_table in"
-            f" status HAS_ERRORS or SUBMITTING for over {time_threshold}m"
+            f"{config.backend_url}: ENA Submission pipeline found {len(entries_with_errors)} entries"
+            f" in assembly_table in status HAS_ERRORS or SUBMITTING for over {time_threshold}m"
         )
         logger.warning(error_msg)
         send_slack_notification(
-            config, error_msg, time=datetime.now(tz=pytz.utc), time_threshold=slack_time_threshold
+            error_msg,
+            slack_config,
+            time=datetime.now(tz=pytz.utc),
+            time_threshold=slack_time_threshold,
         )
         # TODO: Query ENA to check if assembly has in fact been created
         # If created update assembly_table
@@ -572,14 +554,19 @@ def create_assembly(log_level, config_file):
     logger.info(f"Config: {config}")
 
     db_config = db_init(config.db_password, config.db_username, config.db_host)
+    slack_config = slack_conn_init(
+        slack_hook_default=config.slack_hook,
+        slack_token_default=config.slack_token,
+        slack_channel_id_default=config.slack_channel_id,
+    )
 
     while True:
         submission_table_start(db_config)
         submission_table_update(db_config)
 
         assembly_table_create(db_config, config, retry_number=3)
         assembly_table_update(db_config, config)
-        assembly_table_handle_errors(db_config, config)
+        assembly_table_handle_errors(db_config, config, slack_config)
 
 
 if __name__ == "__main__":

diff --git a/ena-submission/scripts/create_project.py b/ena-submission/scripts/create_project.py
@@ -1,7 +1,7 @@
 import json
 import logging
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import datetime
 
 import click
 import pytz
@@ -18,7 +18,7 @@
     XmlAttribute,
     XrefType,
 )
-from notifications import get_slack_config, notify
+from notifications import SlackConfig, send_slack_notification, slack_conn_init
 from submission_db_helper import (
     ProjectTableEntry,
     Status,
@@ -108,29 +108,6 @@ def construct_project_set_object(
     return ProjectSet(project=[project_type])
 
 
-_last_notification_sent: datetime | None = None
-
-
-def send_slack_notification(config: Config, comment: str, time: datetime, time_threshold=12):
-    global _last_notification_sent  # noqa: PLW0603
-    slack_config = get_slack_config(
-        slack_hook_default=config.slack_hook,
-        slack_token_default=config.slack_token,
-        slack_channel_id_default=config.slack_channel_id,
-    )
-    if not slack_config.slack_hook:
-        logger.info("Could not find slack hook cannot send message")
-        return
-    if (
-        not _last_notification_sent
-        or time - timedelta(hours=time_threshold) > _last_notification_sent
-    ):
-        logger.warning(comment)
-        comment = f"{config.backend_url}: " + comment
-        notify(slack_config, comment)
-        _last_notification_sent = time
-
-
 def submission_table_start(db_config):
     """
     1. Find all entries in submission_table in state READY_TO_SUBMIT
@@ -334,7 +311,9 @@ def project_table_create(db_config, config, retry_number=3):
                 tries += 1
 
 
-def project_table_handle_errors(db_config, config, time_threshold=15, slack_time_threshold=12):
+def project_table_handle_errors(
+    db_config, config, slack_config: SlackConfig, time_threshold=15, slack_time_threshold=12
+):
     """
     - time_threshold: (minutes)
     - slack_time_threshold: (hours)
@@ -347,11 +326,14 @@ def project_table_handle_errors(db_config, config, time_threshold=15, slack_time
     )
     if len(entries_with_errors) > 0:
         error_msg = (
-            f"ENA Submission pipeline found {len(entries_with_errors)} entries in project_table in "
-            f"status HAS_ERRORS or SUBMITTING for over {time_threshold}m"
+            f"{config.backend_url}: ENA Submission pipeline found {len(entries_with_errors)} entries in project_table in "
+            f" in project_table in status HAS_ERRORS or SUBMITTING for over {time_threshold}m"
         )
         send_slack_notification(
-            config, error_msg, time=datetime.now(tz=pytz.utc), time_threshold=slack_time_threshold
+            error_msg,
+            slack_config,
+            time=datetime.now(tz=pytz.utc),
+            time_threshold=slack_time_threshold,
         )
         # TODO: Query ENA to check if project has in fact been created
         # If created update project_table
@@ -380,13 +362,18 @@ def create_project(log_level, config_file):
     logger.info(f"Config: {config}")
 
     db_config = db_init(config.db_password, config.db_username, config.db_host)
+    slack_config = slack_conn_init(
+        slack_hook_default=config.slack_hook,
+        slack_token_default=config.slack_token,
+        slack_channel_id_default=config.slack_channel_id,
+    )
 
     while True:
         submission_table_start(db_config)
         submission_table_update(db_config)
 
         project_table_create(db_config, config)
-        project_table_handle_errors(db_config, config)
+        project_table_handle_errors(db_config, config, slack_config)
 
 
 if __name__ == "__main__":

diff --git a/ena-submission/scripts/create_sample.py b/ena-submission/scripts/create_sample.py
@@ -2,7 +2,7 @@
 import logging
 import re
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import datetime
 
 import click
 import pytz
@@ -19,7 +19,7 @@
     XmlAttribute,
     XrefType,
 )
-from notifications import get_slack_config, notify
+from notifications import SlackConfig, send_slack_notification, slack_conn_init
 from submission_db_helper import (
     SampleTableEntry,
     Status,
@@ -151,29 +151,6 @@ def construct_sample_set_object(
     return SampleSetType(sample=[sample_type])
 
 
-_last_notification_sent: datetime | None = None
-
-
-def send_slack_notification(config: Config, comment: str, time: datetime, time_threshold=12):
-    global _last_notification_sent  # noqa: PLW0603
-    slack_config = get_slack_config(
-        slack_hook_default=config.slack_hook,
-        slack_token_default=config.slack_token,
-        slack_channel_id_default=config.slack_channel_id,
-    )
-    if not slack_config.slack_hook:
-        logger.info("Could not find slack hook cannot send message")
-        return
-    if (
-        not _last_notification_sent
-        or time - timedelta(hours=time_threshold) > _last_notification_sent
-    ):
-        logger.warning(comment)
-        comment = f"{config.backend_url}: " + comment
-        notify(slack_config, comment)
-        _last_notification_sent = time
-
-
 def submission_table_start(db_config):
     """
     1. Find all entries in submission_table in state SUBMITTED_PROJECT
@@ -363,7 +340,13 @@ def sample_table_create(db_config, config, retry_number=3):
                 tries += 1
 
 
-def sample_table_handle_errors(db_config, config, time_threshold=15, slack_time_threshold=12):
+def sample_table_handle_errors(
+    db_config,
+    config,
+    slack_config: SlackConfig,
+    time_threshold=15,
+    slack_time_threshold=12,
+):
     """
     - time_threshold: (minutes)
     - slack_time_threshold: (hours)
@@ -375,11 +358,14 @@ def sample_table_handle_errors(db_config, config, time_threshold=15, slack_time_
     )
     if len(entries_with_errors) > 0:
         error_msg = (
-            f"ENA Submission pipeline found {len(entries_with_errors)} entries in sample_table in "
-            f"status HAS_ERRORS or SUBMITTING for over {time_threshold}m"
+            f"{config.backend_url}: ENA Submission pipeline found {len(entries_with_errors)} entries"
+            f" in sample_table in status HAS_ERRORS or SUBMITTING for over {time_threshold}m"
         )
         send_slack_notification(
-            config, error_msg, time=datetime.now(tz=pytz.utc), time_threshold=slack_time_threshold
+            error_msg,
+            slack_config,
+            time=datetime.now(tz=pytz.utc),
+            time_threshold=slack_time_threshold,
         )
         # TODO: Query ENA to check if sample has in fact been created
         # If created update sample_table
@@ -408,13 +394,18 @@ def create_sample(log_level, config_file):
     logger.info(f"Config: {config}")
 
     db_config = db_init(config.db_password, config.db_username, config.db_host)
+    slack_config = slack_conn_init(
+        slack_hook_default=config.slack_hook,
+        slack_token_default=config.slack_token,
+        slack_channel_id_default=config.slack_channel_id,
+    )
 
     while True:
         submission_table_start(db_config)
         submission_table_update(db_config)
 
         sample_table_create(db_config, config)
-        sample_table_handle_errors(db_config, config)
+        sample_table_handle_errors(db_config, config, slack_config)
 
 
 if __name__ == "__main__":

diff --git a/ena-submission/scripts/notifications.py b/ena-submission/scripts/notifications.py
@@ -3,6 +3,7 @@
 import os
 import zipfile
 from dataclasses import dataclass
+from datetime import datetime, timedelta
 
 import requests
 from slack_sdk import WebClient
@@ -13,12 +14,13 @@ class SlackConfig:
     slack_hook: str
     slack_token: str
     slack_channel_id: str
+    last_notification_sent: datetime | None
 
 
 logger = logging.getLogger(__name__)
 
 
-def get_slack_config(
+def slack_conn_init(
     slack_hook_default: str, slack_token_default: str, slack_channel_id_default: str
 ):
     slack_hook = os.getenv("SLACK_HOOK")
@@ -33,10 +35,13 @@ def get_slack_config(
     if not slack_channel_id:
         slack_channel_id = slack_channel_id_default
 
+    last_notification_sent = None
+
     params = {
         "slack_hook": slack_hook,
         "slack_token": slack_token,
         "slack_channel_id": slack_channel_id,
+        "last_notification_sent": last_notification_sent,
     }
 
     return SlackConfig(**params)
@@ -61,3 +66,18 @@ def upload_file_with_comment(config: SlackConfig, file_path: str, comment: str):
         channel=config.slack_channel_id,
         initial_comment=comment,
     )
+
+
+def send_slack_notification(
+    comment: str, slack_config: SlackConfig, time: datetime, time_threshold=12
+):
+    if not slack_config.slack_hook:
+        logger.info("Could not find slack hook cannot send message")
+        return
+    if (
+        not slack_config.last_notification_sent
+        or time - timedelta(hours=time_threshold) > slack_config.last_notification_sent
+    ):
+        logger.warning(comment)
+        notify(slack_config, comment)
+        slack_config.last_notification_sent = time
diff --git a/ena-submission/scripts/test_ena_submission.py b/ena-submission/scripts/test_ena_submission.py
@@ -167,7 +167,7 @@ def test_create_files(self):
         # Check the file content
         with gzip.GzipFile(file_name_chromosome_list, "rb") as gz:
             content = gz.read()
-        print(content)
+
         self.assertEqual(
             content,
             b"test_accession.test_version_seg2\tseg2\tlinear-segmented\ntest_accession.test_version_seg3\tseg3\tlinear-segmented\n",