MUTE: fix issue after adding multi build to mute (#38910)

2 files changed, 35 insertions, 8 deletions

diff --git a/.github/scripts/telegram/parse_and_send_team_issues.py b/.github/scripts/telegram/parse_and_send_team_issues.py
index a5b44e01488..41a60838555 100644
--- a/.github/scripts/telegram/parse_and_send_team_issues.py
+++ b/.github/scripts/telegram/parse_and_send_team_issues.py
@@ -79,6 +79,13 @@ def _execute_ydb_query(query, description):
         return None
 
 
+def _sql_escape_literal(val) -> str:
+    """Escape a value for safe use inside YQL single-quoted string literals."""
+    if val is None:
+        return ""
+    return str(val).replace("'", "''")
+
+
 def _sql_build_type_clause(build_type) -> str:
     """Return YQL fragment ``AND build_type = '…'`` or empty string if ``all`` / unset."""
     if build_type is None:
@@ -88,7 +95,7 @@ def _sql_build_type_clause(build_type) -> str:
         return ""
     if not raw:
         raise ValueError(f"Invalid build_type: {build_type!r} (empty)")
-    escaped = raw.replace("'", "''")
+    escaped = _sql_escape_literal(raw)
     return f"\n    AND build_type = '{escaped}'"
 
 
@@ -108,6 +115,7 @@ def get_all_team_data(use_yesterday=False, build_type=DEFAULT_BUILD_TYPE, branch
     Args:
         use_yesterday: If True, use yesterday's data for development convenience.
         build_type: ``muted_tests_with_issue_and_area.build_type`` filter; ``"all"`` = no filter.
+        branch: Branch filter.
 
     Returns:
         dict: Keys are canonical team slugs from mart ``owner_team`` (effective owner,
@@ -139,6 +147,9 @@ def get_all_team_data(use_yesterday=False, build_type=DEFAULT_BUILD_TYPE, branch
         muted_tests_with_issue_and_area_table = ydb_wrapper.get_table_path("muted_tests_with_issue_and_area")
 
     bt_clause = _sql_build_type_clause(build_type)
+    eb = _sql_escape_literal(branch)
+    note_bt = "all build types" if str(build_type).strip().lower() == "all" else repr(build_type)
+    print(f"🔍 YDB stats slice: branch={branch!r}, build_type={note_bt}")
 
     # Single optimized query for all data from muted tests mart with issue/area enrichment.
     # Keep +today semantics event-like: tests whose mute_state_change_date is target date.
@@ -156,7 +167,7 @@ def get_all_team_data(use_yesterday=False, build_type=DEFAULT_BUILD_TYPE, branch
     FROM `{muted_tests_with_issue_and_area_table}`
     WHERE date_window >= Date('{start_date.strftime('%Y-%m-%d')}')
     AND date_window <= Date('{target_date.strftime('%Y-%m-%d')}')
-    AND branch = '{branch}'{bt_clause}
+    AND branch = '{eb}'{bt_clause}
     GROUP BY owner_team, date_window
     ORDER BY owner_team, date_window
     """
@@ -256,7 +267,12 @@ def get_muted_tests_stats(use_yesterday=False, build_type=DEFAULT_BUILD_TYPE, br
     return team_stats
 
 
-def get_monthly_trend_data(team_name=None, use_yesterday=False, build_type=DEFAULT_BUILD_TYPE, branch=DEFAULT_BRANCH):
+def get_monthly_trend_data(
+    team_name=None,
+    use_yesterday=False,
+    build_type=DEFAULT_BUILD_TYPE,
+    branch=DEFAULT_BRANCH,
+):
     """
     Get monthly trend data for a specific team.
     
@@ -730,6 +746,7 @@ def send_team_messages(teams, bot_token, delay=2, max_retries=5, retry_delay=10,
                         team_name=team_name,
                         use_yesterday=ydb_config.get('use_yesterday', False),
                         build_type=ydb_config.get('build_type', DEFAULT_BUILD_TYPE),
+                        branch=ydb_config.get('branch', DEFAULT_BRANCH),
                     )
                 else:
                     trend_data = None
@@ -1268,6 +1285,7 @@ def main():
         muted_stats = get_muted_tests_stats(
             use_yesterday=args.use_yesterday,
             build_type=args.build_type,
+            branch=args.branch,
         )
         if muted_stats:
             print(f"✅ Statistics loaded for {len(muted_stats)} teams")
diff --git a/.github/scripts/telegram/send_digest.py b/.github/scripts/telegram/send_digest.py
index 03e74f0ab93..2d3fe98a929 100644
--- a/.github/scripts/telegram/send_digest.py
+++ b/.github/scripts/telegram/send_digest.py
@@ -57,9 +57,10 @@ from ydb_wrapper import YDBWrapper
 
 sys.path.insert(0, os.path.dirname(__file__))
 from parse_and_send_team_issues import (
+    _sql_escape_literal,
+    get_all_team_data,
     load_team_channels,
     send_team_messages,
-    get_all_team_data,
 )
 
 
@@ -123,7 +124,7 @@ def _fetch_closed_unsent(w: YDBWrapper, profile_id: str) -> list:
         FROM `{queue_path}` AS q
         INNER JOIN `{issues_path}` AS i
             ON q.github_issue_number = i.issue_number
-        WHERE q.profile_id = '{profile_id.replace("'", "''")}'
+        WHERE q.profile_id = '{_sql_escape_literal(profile_id)}'
           AND q.sent_at IS NULL
           AND i.state = 'CLOSED'
         """,
@@ -158,7 +159,7 @@ def _fetch_unsent(w: YDBWrapper, profile_id: str) -> list:
         FROM `{queue_path}` AS q
         LEFT JOIN `{issues_path}` AS i
             ON q.github_issue_number = i.issue_number
-        WHERE q.profile_id = '{profile_id.replace("'", "''")}'
+        WHERE q.profile_id = '{_sql_escape_literal(profile_id)}'
           AND q.sent_at IS NULL
           AND (i.state IS NULL OR i.state != 'CLOSED')
         ORDER BY q.owner_team, q.github_issue_number
@@ -266,7 +267,10 @@ def run_digest(
         muted_stats   = None
         all_team_data = None
         try:
-            all_team_data = get_all_team_data(build_type=profile["build_type"], branch=profile["branch"])
+            all_team_data = get_all_team_data(
+                build_type=profile["build_type"],
+                branch=profile["branch"],
+            )
             if all_team_data:
                 muted_stats = {t: d["stats"] for t, d in all_team_data.items()}
         except Exception as exc:
@@ -279,11 +283,16 @@ def run_digest(
             muted_stats=muted_stats,
             include_plots=include_plots,
             ydb_config=(
-                {"use_yesterday": False, "build_type": profile["build_type"]}
+                {
+                    "use_yesterday": False,
+                    "build_type": profile["build_type"],
+                    "branch": profile["branch"],
+                }
                 if include_plots
                 else None
             ),
             all_team_data=all_team_data,
+            show_diff=True,
         )
 
         now = datetime.now(tz=timezone.utc)