ci(unit): stop the full history clone on the db-and-spend shard, fetch the schema merge base by sha (#50120)

* ci(unit): fetch the merge base by sha for the db-and-spend schema check instead of cloning full history

* ci(unit): read the schema check arguments through one typed mapping

* ci(unit): give the schema check its own step with the ui-unit base fallback and register it with the merge-base step test

* ci(unit): keep the setup ceiling at 35 minutes, the install step never took over 33 seconds

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.github/workflows/test-unit.yml

@@ -977,10 +977,10 @@ jobs:
             dist: loadscope
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
-        timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
+        timeout-minutes: 3
         with:
           persist-credentials: false
-          fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
+          fetch-depth: 1
 
       - name: Detect relevant changes
         id: changes
@@ -1013,13 +1013,11 @@ jobs:
         with:
           artifact: ${{ needs.rust-bridge.outputs.artifact }}
 
-      - name: Install dependencies and check schema
+      - name: Install dependencies
         if: steps.changes.outputs.decision != 'skip'
-        timeout-minutes: 8
+        timeout-minutes: 5
         env:
           RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
-          UNIT_SHARD: ${{ matrix.shard }}
-          SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }}
         run: |
           diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
           if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
@@ -1031,9 +1029,21 @@ jobs:
             uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
           fi
           uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
-          if [ "$UNIT_SHARD" = "db-and-spend" ]; then
-            uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF"
-          fi
+
+      - name: Check schema.prisma against the merge base
+        if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request'
+        timeout-minutes: 3
+        env:
+          GH_TOKEN: ${{ github.token }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
+        run: |
+          base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref)
+          merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
+          test -n "$merge_base"
+          git fetch --no-tags --depth=1 origin "$merge_base"
+          uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base"
 
       - name: Cache Prisma binaries
         if: steps.changes.outputs.decision != 'skip'

scripts/check_schema_compatibility.py

@@ -52,14 +52,20 @@ def check_breaking_changes(old_schema: Schema, new_schema: Schema) -> tuple[str,
     return tuple(changes())
 
 
-def main() -> int:
-    parser: Final = argparse.ArgumentParser()
-    parser.add_argument("--base-ref", required=True)
-    options: Final = parser.parse_args()
-    base_ref: Final = TypeAdapter(str).validate_python(options.base_ref)
-    base: Final = subprocess.run(
+def _merge_base(base_ref: str) -> str:
+    return subprocess.run(
         ("git", "merge-base", "HEAD", base_ref), capture_output=True, text=True, check=True
     ).stdout.strip()
+
+
+def main() -> int:
+    parser: Final = argparse.ArgumentParser()
+    source: Final = parser.add_mutually_exclusive_group(required=True)
+    source.add_argument("--base-ref", help="compare against the merge base of HEAD and this ref (needs history)")
+    source.add_argument("--base-commit", help="compare against this commit as is (works in a shallow clone)")
+    options: Final = TypeAdapter(dict[str, str | None]).validate_python(vars(parser.parse_args()))
+    base_commit: Final = options["base_commit"]
+    base: Final = base_commit if base_commit else _merge_base(options["base_ref"] or "")
     previous: Final = subprocess.run(
         ("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True
     ).stdout

tests/unit/test_lint_workflow_diff_gates.py

@@ -383,10 +383,20 @@ def _guard_base_argument_and_checkout(run: StepRun) -> tuple[str, str]:
     return argv[argv.index("--base") + 1], _sha(run.runner.checkout, "HEAD")
 
 
+def _schema_base_commit_and_its_type_in_the_checkout(run: StepRun) -> tuple[str, str]:
+    argv: Final = _argv_recorded(run.runner, "uv")
+    base_commit: Final = argv[argv.index("--base-commit") + 1]
+    return base_commit, _git(run.runner.checkout, "cat-file", "-t", base_commit).strip()
+
+
 def _merge_base(upstream: StaleBaseUpstream, ref: str) -> str:
     return upstream.merge_base_by_ref[ref]
 
 
+def _merge_base_as_a_commit(upstream: StaleBaseUpstream, ref: str) -> tuple[str, str]:
+    return upstream.merge_base_by_ref[ref], "commit"
+
+
 def _ui_files_changed_since_merge_base(upstream: StaleBaseUpstream, ref: str) -> frozenset[str]:
     merge_base: Final = upstream.merge_base_by_ref[ref]
     changed: Final = _git(upstream.path, "diff", "--name-only", f"--relative={UI_DASHBOARD}", merge_base, upstream.head)
@@ -402,6 +412,7 @@ CONSUMERS: Final = MappingProxyType(
         "test-linting.yml:python": Consumer(_gate_base_exported, _merge_base),
         "test-linting.yml:ui": Consumer(_ui_files_queued_for_prettier, _ui_files_changed_since_merge_base),
         "test-unit.yml:ui-unit": Consumer(_ui_files_passed_to_vitest_related, _ui_files_changed_since_merge_base),
+        "test-unit.yml:unit": Consumer(_schema_base_commit_and_its_type_in_the_checkout, _merge_base_as_a_commit),
         "spend-logs-guard.yml:spend-logs-guard": Consumer(_guard_base_argument_and_checkout, _merge_base_and_tip),
         "cost-map-guard.yml:cost-map-guard": Consumer(_guard_base_argument_and_checkout, _merge_base_and_tip),
     }