fix: Improve command_static_url repack archive logic

This patch improves the logic of static URL fetch tasks
to better be able to determine if a file simply needs to
be renamed, or if it is an archive that needs to be repackaged.

src/taskgraph/run-task/fetch-content

@@ -435,11 +435,47 @@ def extract_archive(path, dest_dir):
     log("%s extracted in %.3fs" % (path, time.time() - t0))
 
 
+def should_try_to_repack_archive(orig: pathlib.Path, dest: pathlib.Path) -> bool:
+    """
+    Determines whether we should attempt to repack an archive based on the naming conventions used
+    in the original file path and the destination file path.
+    """
+    if orig == dest:
+        # There is nothing to do if the paths are exactly the same.
+        return False
+
+    if not orig.suffixes and dest.suffixes:
+        # If original path has no suffix, but the destination path does have a suffix, this could
+        # be a simple rename to add an extension, or it could be an attempt to package an archive.
+        #
+        # In this case, we will use the destination suffix to determine the course of action.
+        # If the final destination is a ".tar.zst" file then we will attempt to repack an archive.
+        return dest.suffixes[-2:] == [".tar", ".zst"]
+
+    if orig.suffix == dest.suffix:
+        # If the final suffix is the same, this is likely just a rename.
+        #
+        # This may be a case where multiple suffixes were detected unnecessarily due to the file name
+        # containing a semantic version. For example the file "python-3.8.10-amd64.exe" would be detected
+        # to have three suffixes [".8", ".10-amd64", ".exe"]. Changing this to "python.exe" should be a
+        # simple rename, rather than a repack.
+        #
+        # In this case, we will try to determine if the original path has a supported archive suffix.
+        # If the original path is detected to be an archive, we will try to repack, otherwise rename.
+        return archive_type(orig) is not None
+
+    # Otherwise, if the paths aren't the same, assume it's an archive and try to repack.
+
+    # It would be best to fail early if the repack fails than to fail during a test because a renamed
+    # file was incorrect type after all.
+    return True
+
+
 def repack_archive(
     orig: pathlib.Path, dest: pathlib.Path, strip_components=0, prefix=""
 ):
     assert orig != dest
-    log("Repacking as %s" % dest)
+    log(f"Repacking {orig} as {dest}")
     orig_typ, ifh = open_stream(orig)
     typ = archive_type(dest)
     if not typ:
@@ -534,6 +570,9 @@ def repack_archive(
                 # We only change compression here. The tar stream is unchanged.
                 ctx.copy_stream(ifh, fh)
 
+        else:
+            raise Exception(f"Attempt to repack an archive of unknown type {orig_typ}")
+
 
 def fetch_and_extract(url, dest_dir, extract=True, sha256=None, size=None):
     """Fetch a URL and extract it to a destination path.
@@ -773,8 +812,13 @@ def command_static_url(args):
         if gpg_sig_url:
             gpg_verify_path(dl_dest, gpg_key, gpg_signature)
 
-        if dl_dest != dest or args.strip_components or args.add_prefix:
-            repack_archive(dl_dest, dest, args.strip_components, args.add_prefix)
+        if dl_dest != dest:
+            if args.strip_components or args.add_prefix or should_try_to_repack_archive(dl_dest, dest):
+                repack_archive(dl_dest, dest, args.strip_components, args.add_prefix)
+            else:
+                log(f"Renaming {dl_dest} to {dest}")
+                dl_dest.rename(dest)
+
     except Exception:
         try:
             dl_dest.unlink()
@@ -783,7 +827,7 @@ def command_static_url(args):
 
         raise
 
-    if dl_dest != dest:
+    if dl_dest != dest and dl_dest.exists():
         log("Removing %s" % dl_dest)
         dl_dest.unlink()
 

test/test_scripts_fetch_content.py

@@ -1,4 +1,5 @@
 import os
+import pathlib
 import urllib.request
 from importlib.machinery import SourceFileLoader
 from importlib.util import module_from_spec, spec_from_loader
@@ -94,3 +95,34 @@ def getheader(field):
     except fetch_content_mod.IntegrityError:
         if not raises:
             raise
+
+
+@pytest.mark.parametrize(
+    "expected,orig,dest",
+    [
+        # Archives to repack
+        (True, pathlib.Path("archive"), pathlib.Path("archive.tar.zst")),
+        (True, pathlib.Path("archive.tar"), pathlib.Path("archive.tar.zst")),
+        (True, pathlib.Path("archive.tgz"), pathlib.Path("archive.tar.zst")),
+        (True, pathlib.Path("archive.zip"), pathlib.Path("archive.tar.zst")),
+        (True, pathlib.Path("archive.tar.xz"), pathlib.Path("archive.tar.zst")),
+        # Path is exactly the same
+        (False, pathlib.Path("archive"), pathlib.Path("archive")),
+        (False, pathlib.Path("archive.tar"), pathlib.Path("archive.tar")),
+        (False, pathlib.Path("archive.tgz"), pathlib.Path("archive.tgz")),
+        (False, pathlib.Path("archive.zip"), pathlib.Path("archive.zip")),
+        (False, pathlib.Path("file.txt"), pathlib.Path("file.txt")),
+        (False, pathlib.Path("archive.tar.zst"), pathlib.Path("archive.tar.zst")),
+        # Real edge cases
+        (False, pathlib.Path("python-3.8.10-amd64.exe"), pathlib.Path("python.exe")),
+        (
+            False,
+            pathlib.Path("9ee26e91-9b52-44ba-8d30-c0230dd587b2.bin"),
+            pathlib.Path("model.esen.intgemm.alphas.bin"),
+        ),
+    ],
+)
+def test_should_try_to_repack_archive(fetch_content_mod, orig, dest, expected):
+    assert (
+        fetch_content_mod.should_try_to_repack_archive(orig, dest) == expected
+    ), f"Failed for orig: {orig}, dest: {dest}, expected {expected} but received {not expected}"