From b4846097892433d5642fa45da2ea7b8352124895 Mon Sep 17 00:00:00 2001
From: Staci Mullins <63313398+stacimc@users.noreply.github.com>
Date: Thu, 31 Oct 2024 09:48:05 -0700
Subject: [PATCH] Add env in data refresh notifications (#5115)

---
 catalog/dags/data_refresh/dag_factory.py |  1 +
 catalog/dags/data_refresh/reporting.py   | 10 +++++++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/catalog/dags/data_refresh/dag_factory.py b/catalog/dags/data_refresh/dag_factory.py
index cb66c92c406..b4511808d02 100644
--- a/catalog/dags/data_refresh/dag_factory.py
+++ b/catalog/dags/data_refresh/dag_factory.py
@@ -295,6 +295,7 @@ def create_data_refresh_dag(
                 "after": after_record_count,
                 "media_type": data_refresh_config.media_type,
                 "dag_id": data_refresh_config.dag_id,
+                "env": target_environment,
             },
         )
 
diff --git a/catalog/dags/data_refresh/reporting.py b/catalog/dags/data_refresh/reporting.py
index a9e46c81e39..8128b55b4a1 100644
--- a/catalog/dags/data_refresh/reporting.py
+++ b/catalog/dags/data_refresh/reporting.py
@@ -18,7 +18,9 @@ def report_status(media_type: str, message: str, dag_id: str):
     return message
 
 
-def report_record_difference(before: dict, after: dict, media_type: str, dag_id: str):
+def report_record_difference(
+    before: dict, after: dict, media_type: str, dag_id: str, env: str
+):
     all_keys = before.keys() | after.keys()
     total_before = sum(before.values())
     total_after = sum(after.values())
@@ -37,7 +39,7 @@ def report_record_difference(before: dict, after: dict, media_type: str, dag_id:
 
     message = dedent(
         f"""
-        Data refresh for {media_type} complete! :tada:
+        {env.capitalize()} data refresh for {media_type} complete! :tada:
         _Note: All values are retrieved from elasticsearch_
         *Record count difference for `{media_type}`*: {total_before:,} → {total_after:,}
         *Change*: {count_diff:+,} ({percent_diff:+}% Δ)
@@ -46,6 +48,8 @@ def report_record_difference(before: dict, after: dict, media_type: str, dag_id:
     )
     message += breakdown_message
     slack.send_message(
-        text=message, dag_id=dag_id, username="Data refresh record difference"
+        text=message,
+        dag_id=dag_id,
+        username=f"{env.capitalize()} data refresh record difference",
     )
     return message