From 91e6b5d7bac04ef0129a76711d00a88ef9bd2416 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A1nos=20Benjamin=20Antal?=
 <benjamin.antal@clickhouse.com>
Date: Thu, 7 Sep 2023 09:31:16 +0000
Subject: [PATCH] Make the logic robust over temporary errors

---
 tests/ci/docker_manifests_merge.py | 59 ++++++++++++++++++------------
 1 file changed, 35 insertions(+), 24 deletions(-)

diff --git a/tests/ci/docker_manifests_merge.py b/tests/ci/docker_manifests_merge.py
index dfc71a5b28f..3a9182c87e6 100644
--- a/tests/ci/docker_manifests_merge.py
+++ b/tests/ci/docker_manifests_merge.py
@@ -9,7 +9,11 @@ import subprocess
 from typing import List, Dict, Tuple
 from github import Github
 
-from clickhouse_helper import ClickHouseHelper, prepare_tests_results_for_clickhouse
+from clickhouse_helper import (
+    ClickHouseHelper,
+    prepare_tests_results_for_clickhouse,
+    CHException,
+)
 from commit_status_helper import format_description, get_commit, post_commit_status
 from docker_images_helper import IMAGES_FILE_PATH, get_image_names
 from env_helper import RUNNER_TEMP, GITHUB_WORKSPACE
@@ -183,7 +187,7 @@ def enrich_images(changed_images: Dict[str, str]) -> Dict[str, str]:
     )
 
     COMMIT_SHA_BATCH_SIZE = 100
-    MAX_COMMIT_BATCHES_TO_CHECK = 10
+    MAX_TRIES_COUNT = 10
     # Gets the sha of the last COMMIT_SHA_BATCH_SIZE commits after skipping some commits (see below)
     LAST_N_ANCESTOR_SHA_COMMAND = f"git log --format=format:'%H' --max-count={COMMIT_SHA_BATCH_SIZE} --skip={{}} --merges"
     git_runner = Runner()
@@ -204,38 +208,45 @@ def enrich_images(changed_images: Dict[str, str]) -> Dict[str, str]:
         """
 
     batch_count = 0
-    ch_helper = ClickHouseHelper()
+    try_count = 0
 
     while (
-        batch_count <= MAX_COMMIT_BATCHES_TO_CHECK and len(images_to_find_tags_for) != 0
+        try_count <= MAX_TRIES_COUNT and len(images_to_find_tags_for) != 0
     ):
         commit_shas = git_runner(
             LAST_N_ANCESTOR_SHA_COMMAND.format(batch_count * COMMIT_SHA_BATCH_SIZE)
         ).split("\n")
 
-        result = ch_helper.select_json_each_row(
-            "default",
-            GET_COMMIT_SHAS_QUERY,
-            {"commit_shas": commit_shas, "images": images_to_find_tags_for},
-        )
-        result.sort(key=lambda x: x["image_name"])
+        try:
+            result = ch_helper.select_json_each_row(
+                "default",
+                GET_COMMIT_SHAS_QUERY,
+                {"commit_shas": commit_shas, "images": images_to_find_tags_for},
+            )
+            result.sort(key=lambda x: x["image_name"])
 
-        logging.info(
-            "Found images for commits %s..%s: %s",
-            commit_shas[0],
-            commit_shas[-1],
-            "\n ".join(f"{im['image_name']}:{im['commit_sha']}" for im in result),
-        )
+            logging.info(
+                "Found images for commits %s..%s: %s",
+                commit_shas[0],
+                commit_shas[-1],
+                "\n ".join(f"{im['image_name']}:{im['commit_sha']}" for im in result),
+            )
 
-        for row in result:
-            image_name = row["image_name"]
-            commit_sha = row["commit_sha"]
-            # As we only get the SHAs of merge commits from master, the PR number will be always 0
-            tag = f"0-{commit_sha}"
-            new_changed_images[image_name] = tag
-            images_to_find_tags_for.remove(image_name)
+            for row in result:
+                image_name = row["image_name"]
+                commit_sha = row["commit_sha"]
+                # As we only get the SHAs of merge commits from master, the PR number will be always 0
+                tag = f"0-{commit_sha}"
+                new_changed_images[image_name] = tag
+                images_to_find_tags_for.remove(image_name)
+
+            batch_count += 1
+
+        except CHException as ex:
+            logging.warning("Request for ClickHouse failed: %s", ex)
+
+        try_count += 1
 
-        batch_count += 1
 
     return new_changed_images