Merge pull request #58717 from ClickHouse/improve-cherry-pick

Small fixes in different helpers

tests/ci/cache_utils.py

@@ -118,7 +118,12 @@ class Cache:
     def _download(self, url: str, ignore_error: bool = False) -> None:
         compressed_cache = self.temp_path / self.archive_name
         try:
-            download_build_with_progress(url, compressed_cache)
+            if url.startswith("file://"):
+                local_s3_cache = Path(url[7:])
+                if local_s3_cache.is_file():
+                    shutil.copy2(local_s3_cache, compressed_cache)
+            else:
+                download_build_with_progress(url, compressed_cache)
         except DownloadException as e:
             if not ignore_error:
                 raise CacheError(f"Failed to download {url}") from e

tests/ci/cherry_pick.py

@@ -32,6 +32,7 @@ from pathlib import Path
 from subprocess import CalledProcessError
 from typing import List, Optional
 
+import __main__
 from env_helper import TEMP_PATH
 from get_robot_token import get_best_robot_token
 from git_helper import git_runner, is_shallow
@@ -606,16 +607,18 @@ def parse_args():
 
 @contextmanager
 def clear_repo():
-    orig_ref = git_runner("git branch --show-current") or git_runner(
-        "git rev-parse HEAD"
-    )
+    def ref():
+        return git_runner("git branch --show-current") or git_runner(
+            "git rev-parse HEAD"
+        )
+
+    orig_ref = ref()
     try:
         yield
-    except (Exception, KeyboardInterrupt):
-        git_runner(f"git checkout -f {orig_ref}")
-        raise
-    else:
-        git_runner(f"git checkout -f {orig_ref}")
+    finally:
+        current_ref = ref()
+        if orig_ref != current_ref:
+            git_runner(f"git checkout -f {orig_ref}")
 
 
 @contextmanager
@@ -623,15 +626,14 @@ def stash():
     # diff.ignoreSubmodules=all don't show changed submodules
     need_stash = bool(git_runner("git -c diff.ignoreSubmodules=all diff HEAD"))
     if need_stash:
-        git_runner("git stash push --no-keep-index -m 'running cherry_pick.py'")
+        script = (
+            __main__.__file__ if hasattr(__main__, "__file__") else "unknown script"
+        )
+        git_runner(f"git stash push --no-keep-index -m 'running {script}'")
     try:
         with clear_repo():
             yield
-    except (Exception, KeyboardInterrupt):
-        if need_stash:
-            git_runner("git stash pop")
-        raise
-    else:
+    finally:
         if need_stash:
             git_runner("git stash pop")
 

tests/ci/git_helper.py

@@ -92,7 +92,7 @@ class Runner:
             return
         self._cwd = value
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any) -> str:
         return self.run(*args, **kwargs)