feat(run-task): implement shallow git clones

Author: ahal

src/taskgraph/run-task/run-task

@@ -571,13 +571,17 @@ def git_fetch(
     ref: str,
     remote: str = "origin",
     tags: bool = False,
+    shallow: bool = False,
     env: Optional[Dict[str, str]] = None,
 ):
     args = ["git", "fetch"]
     if tags:
         # `--force` is needed to be able to update an existing outdated tag.
         args.extend(["--tags", "--force"])
 
+    if shallow:
+        args.append("--depth=1")
+
     args.extend([remote, ref])
     retry_required_command(b"vcs", args, cwd=destination_path, extra_env=env)
 
@@ -621,6 +625,19 @@ def _clean_git_checkout(destination_path):
     print_line(b"vcs", b"successfully cleaned git checkout!\n")
 
 
+def shortref(ref: str) -> str:
+    """Normalize a git ref to its short form.
+
+    Returns the ref unchanged if it's already in short form.
+    """
+    # Strip common ref prefixes
+    for prefix in ("refs/heads/", "refs/tags/"):
+        if ref.startswith(prefix):
+            return ref[len(prefix) :]
+
+    return ref
+
+
 def git_checkout(
     destination_path: str,
     head_repo: str,
@@ -630,6 +647,7 @@ def git_checkout(
     head_rev: Optional[str],
     ssh_key_file: Optional[Path],
     ssh_known_hosts_file: Optional[Path],
+    shallow: bool = False,
 ):
     env = {
         # abort if transfer speed is lower than 1kB/s for 1 minute
@@ -673,10 +691,18 @@ def git_checkout(
         args = [
             "git",
             "clone",
-            base_repo if base_repo else head_repo,
-            destination_path,
         ]
 
+        if shallow:
+            args.extend(["--depth=1", "--no-checkout"])
+
+        args.extend(
+            [
+                base_repo if base_repo else head_repo,
+                destination_path,
+            ]
+        )
+
         retry_required_command(b"vcs", args, extra_env=env)
 
     # For Github based repos, base_rev often doesn't refer to an ancestor of
@@ -685,7 +711,7 @@ def git_checkout(
     # that consumers can compute the merge-base or files modified between the
     # two as needed.
     if base_rev and base_rev != NULL_REVISION:
-        git_fetch(destination_path, base_rev, env=env)
+        git_fetch(destination_path, base_rev, shallow=shallow, env=env)
 
     # If a head_ref was provided, it might be tag, so we need to make sure we fetch
     # those. This is explicitly only done when base and head repo match,
@@ -698,16 +724,29 @@ def git_checkout(
 
     # If a head_ref isn't provided, we fetch all refs from head_repo, which may be slow.
     target = head_ref if head_ref else "+refs/heads/*:refs/remotes/work/*"
-    git_fetch(destination_path, target, remote=head_repo, tags=tags, env=env)
+    git_fetch(
+        destination_path,
+        target,
+        remote=head_repo,
+        tags=tags,
+        shallow=shallow,
+        env=env,
+    )
+
+    # If we have a shallow clone and specific commit, we need to fetch it too.
+    if shallow and head_rev and head_rev != head_ref:
+        git_fetch(
+            destination_path, head_rev, remote=head_repo, shallow=shallow, env=env
+        )
 
     args = [
         "git",
         "checkout",
         "-f",
     ]
 
-    if head_ref:
-        args.extend(["-B", head_ref])
+    if head_ref and head_ref != head_rev:
+        args.extend(["-B", shortref(head_ref)])
 
     # `git fetch` set `FETCH_HEAD` reference to the last commit of the desired branch
     args.append(head_rev if head_rev else "FETCH_HEAD")
@@ -884,11 +923,17 @@ def add_vcs_arguments(parser, project, name):
         f"--{project}-sparse-profile",
         help=f"Path to sparse profile for {name} checkout",
     )
+    parser.add_argument(
+        f"--{project}-shallow-clone",
+        action="store_true",
+        help=f"Use shallow clone for {name}",
+    )
 
 
 def collect_vcs_options(args, project, name):
     checkout = getattr(args, f"{project}_checkout")
     sparse_profile = getattr(args, f"{project}_sparse_profile")
+    shallow_clone = getattr(args, f"{project}_shallow_clone")
 
     env_prefix = project.upper()
 
@@ -933,6 +978,7 @@ def collect_vcs_options(args, project, name):
         "repo-type": repo_type,
         "ssh-secret-name": private_key_secret,
         "pip-requirements": pip_requirements,
+        "shallow-clone": shallow_clone,
     }
 
 
@@ -981,6 +1027,7 @@ def vcs_checkout_from_args(options):
                 head_rev,
                 ssh_key_file,
                 ssh_known_hosts_file,
+                shallow=options.get("shallow-clone", False),
             )
         elif options["repo-type"] == "hg":
             revision = hg_checkout(

test/test_scripts_run_task.py

@@ -151,9 +151,10 @@ def test_install_pip_requirements_with_uv(
 
 
 @pytest.mark.parametrize(
-    "env,extra_expected",
+    "args,env,extra_expected",
     [
         pytest.param(
+            {},
             {
                 "REPOSITORY_TYPE": "hg",
                 "BASE_REPOSITORY": "https://hg.mozilla.org/mozilla-central",
@@ -164,20 +165,38 @@ def test_install_pip_requirements_with_uv(
             {
                 "base-repo": "https://hg.mozilla.org/mozilla-unified",
             },
-        )
+            id="hg",
+        ),
+        pytest.param(
+            {"myrepo_shallow_clone": True},
+            {
+                "REPOSITORY_TYPE": "git",
+                "HEAD_REPOSITORY": "https://github.com/test/repo.git",
+                "HEAD_REV": "abc123",
+            },
+            {"shallow-clone": True},
+            id="git_with_shallow_clone",
+        ),
     ],
 )
-def test_collect_vcs_options(monkeypatch, run_task_mod, env, extra_expected):
+def test_collect_vcs_options(
+    monkeypatch,
+    run_task_mod,
+    args,
+    env,
+    extra_expected,
+):
     name = "myrepo"
     checkout = "checkout"
 
     monkeypatch.setattr(os, "environ", {})
     for k, v in env.items():
         monkeypatch.setenv(f"{name.upper()}_{k.upper()}", v)
 
-    args = Namespace()
-    setattr(args, f"{name}_checkout", checkout)
-    setattr(args, f"{name}_sparse_profile", False)
+    args.setdefault(f"{name}_checkout", checkout)
+    args.setdefault(f"{name}_shallow_clone", False)
+    args.setdefault(f"{name}_sparse_profile", False)
+    args = Namespace(**args)
 
     result = run_task_mod.collect_vcs_options(args, name, name)
 
@@ -193,6 +212,7 @@ def test_collect_vcs_options(monkeypatch, run_task_mod, env, extra_expected):
         "head-ref": env.get("HEAD_REF"),
         "head-rev": env.get("HEAD_REV"),
         "repo-type": env.get("REPOSITORY_TYPE"),
+        "shallow-clone": False,
         "ssh-secret-name": env.get("SSH_SECRET_NAME"),
         "sparse-profile": False,
         "store-path": env.get("HG_STORE_PATH"),
@@ -333,7 +353,9 @@ def mock_git_repo():
         )
 
         def _commit_file(message, filename):
-            with open(os.path.join(repo, filename), "w") as fout:
+            filepath = os.path.join(repo, filename)
+            os.makedirs(os.path.dirname(filepath), exist_ok=True)
+            with open(filepath, "w") as fout:
                 fout.write("test file content")
             subprocess.check_call(["git", "add", filename], cwd=repo_path)
             subprocess.check_call(["git", "commit", "-m", message], cwd=repo_path)
@@ -427,6 +449,104 @@ def test_git_checkout_with_commit(
     assert current_rev == mock_git_repo["branch"]
 
 
+def test_git_checkout_shallow(
+    mock_stdin,
+    run_task_mod,
+    mock_git_repo,
+    tmp_path,
+):
+    destination = tmp_path / "destination"
+
+    # Git ignores `--depth` when cloning from local directories, so use file://
+    # protocol to force shallow clone.
+    repo_url = f"file://{mock_git_repo['path']}"
+    base_rev = mock_git_repo["main"]
+    head_rev = mock_git_repo["branch"]
+
+    # Use shallow clone with head_ref != head_rev
+    run_task_mod.git_checkout(
+        destination_path=str(destination),
+        head_repo=repo_url,
+        base_repo=repo_url,
+        base_rev=base_rev,
+        head_ref="mybranch",
+        head_rev=head_rev,
+        ssh_key_file=None,
+        ssh_known_hosts_file=None,
+        shallow=True,
+    )
+    shallow_file = destination / ".git" / "shallow"
+    assert shallow_file.exists()
+
+    # Verify we're on the correct commit
+    final_rev = subprocess.check_output(
+        ["git", "rev-parse", "HEAD"],
+        cwd=str(destination),
+        universal_newlines=True,
+    ).strip()
+    assert final_rev == mock_git_repo["branch"]
+
+    # Verify both base_rev and head_rev are available.
+    for sha in [mock_git_repo["main"], mock_git_repo["branch"]]:
+        result = subprocess.run(
+            ["git", "cat-file", "-t", sha],
+            cwd=str(destination),
+            capture_output=True,
+            text=True,
+        )
+        assert result.returncode == 0, f"Commit {sha} should be available"
+        assert result.stdout.strip() == "commit"
+
+
+def test_git_fetch_shallow(
+    mock_stdin,
+    run_task_mod,
+    mock_git_repo,
+    tmp_path,
+):
+    destination = tmp_path / "destination"
+
+    # Git ignores `--depth` when cloning from local directories, so use file://
+    # protocol to force shallow clone.
+    repo_url = f"file://{mock_git_repo['path']}"
+
+    run_task_mod.run_command(
+        b"vcs",
+        [
+            "git",
+            "clone",
+            "--depth=1",
+            "--no-checkout",
+            repo_url,
+            str(destination),
+        ],
+    )
+    shallow_file = destination / ".git" / "shallow"
+    assert shallow_file.exists()
+
+    # Verify base_rev doesn't exist yet
+    base_rev = mock_git_repo["branch"]
+    result = subprocess.run(
+        ["git", "cat-file", "-t", base_rev],
+        cwd=str(destination),
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode != 0
+
+    run_task_mod.git_fetch(str(destination), base_rev, remote=repo_url, shallow=True)
+
+    # Verify base_rev is now available
+    result = subprocess.run(
+        ["git", "cat-file", "-t", base_rev],
+        cwd=str(destination),
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0
+    assert result.stdout.strip() == "commit"
+
+
 def test_display_python_version_should_output_python_versions_title(
     run_task_mod, capsys
 ):