feat(ci): split nix build workflow into separate extensions and checks jobs

Split the monolithic nix-build job into two workflows:
one for building PostgreSQL extensions and another for checks.

Building extensions can be resource-intensive and time-consuming, 
so isolating them allows for better resource allocation and parallelism.

Once they are built, the checks job can run tests and validations
on the already built extensions.

Author: jfroche

.github/workflows/nix-build.yml

@@ -15,7 +15,7 @@ permissions:
   packages: write
 
 jobs:
-  nix-matrix:
+  extensions-matrix:
     runs-on:
       group: self-hosted-runners-nix
       labels:
@@ -29,15 +29,74 @@ jobs:
         name: Generate Nix Matrix
         run: |
           set -Eeu
-          echo matrix="$(python scripts/github-matrix.py)" >> "$GITHUB_OUTPUT"
+          echo matrix="$(python scripts/github-matrix.py extensions)" >> "$GITHUB_OUTPUT"
 
-  build-run-image:
+  checks-matrix:
+    runs-on:
+      group: self-hosted-runners-nix
+      labels:
+        - aarch64-darwin
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+      - name: Checkout Repo
+        uses: actions/checkout@v4
+      - id: set-matrix
+        name: Generate Nix Matrix
+        run: |
+          set -Eeu
+          echo matrix="$(python scripts/github-matrix.py checks)" >> "$GITHUB_OUTPUT"
+
+  build-extensions:
+    name: ${{matrix.postgresql_version}}.${{ matrix.name }} (${{ matrix.system }})
+    needs: extensions-matrix
+    runs-on: ${{ matrix.runs_on.group && matrix.runs_on || matrix.runs_on.labels }}
+    concurrency:
+      group: runner-${{ matrix.runs_on.labels }}-${{ github.run_id }}-limit
+      cancel-in-progress: false
+    strategy:
+      fail-fast: false
+      matrix: ${{fromJSON(needs.extensions-matrix.outputs.matrix)}}
+    steps:
+      - name: Checkout Repo
+        uses: actions/checkout@v4
+      - name: aws-oidc
+        uses: aws-actions/[email protected]
+        with:
+          aws-region: us-east-2
+          role-to-assume: arn:aws:iam::279559813984:role/supabase-github-oidc-role # Shared Services
+          role-session-name: gha-oidc-${{ github.run_id }}
+      - name: aws-creds
+        uses: aws-actions/[email protected]
+        with:
+          disable-retry: true
+          aws-region: us-east-2
+          role-to-assume: arn:aws:iam::436098097459:role/nix-artifacts-deploy-role # supabase-dev
+          role-session-name: gha-oidc-${{ github.run_id }}
+          role-chaining: true
+          role-skip-session-tagging: true
+          role-duration-seconds: 18000
+      - name: Write creds files
+        run: |
+          umask 006
+          cat > /etc/nix/aws/nix-aws-credentials <<EOF
+          [ci-uploader]
+          aws_access_key_id = ${AWS_ACCESS_KEY_ID}
+          aws_secret_access_key = ${AWS_SECRET_ACCESS_KEY}
+          aws_session_token = ${AWS_SESSION_TOKEN}
+          EOF
+      - name: nix build
+        run: |
+           nix build -L .#${{ matrix.attr }}
+
+
+  build-checks:
     name: ${{ matrix.name }} (${{ matrix.system }})
-    needs: nix-matrix
+    needs: [checks-matrix, build-extensions]
     runs-on: ${{ matrix.runs_on.group && matrix.runs_on || matrix.runs_on.labels }}
     strategy:
       fail-fast: false
-      matrix: ${{fromJSON(needs.nix-matrix.outputs.matrix)}}
+      matrix: ${{fromJSON(needs.checks-matrix.outputs.matrix)}}
     steps:
       - name: Checkout Repo
         uses: actions/checkout@v4
@@ -56,7 +115,7 @@ jobs:
           role-session-name: gha-oidc-${{ github.run_id }}
           role-chaining: true
           role-skip-session-tagging: true
-          role-duration-seconds: 900 # TODO: switch to 18000 (5 hours)
+          role-duration-seconds: 18000
       - name: Write creds files
         run: |
           umask 006
@@ -68,13 +127,9 @@ jobs:
           EOF
       - name: nix build
         run: |
-           if ${{ matrix.already_cached }}; then
-             echo "${{ matrix.attr }} already cached, skipping build"
-             exit 0
-           fi
            nix build -L .#${{ matrix.attr }}
 
   run-tests:
-    needs: build-run-image
+    needs: build-checks
     if: ${{ success() }}
     uses: ./.github/workflows/test.yml

nix/packages/postgres.nix

@@ -1,7 +1,7 @@
 { self, inputs, ... }:
 {
   perSystem =
-    { pkgs, ... }:
+    { pkgs, lib, ... }:
     let
       gitRev = "vcs=${self.shortRev or "dirty"}+${
         builtins.substring 0 8 (self.lastModifiedDate or self.lastModified or "19700101")
@@ -158,10 +158,9 @@
       #    install.
       #  - exts: an attrset containing all the extensions, mapped to their
       #    package names.
-      makePostgres = version: {
+      makePostgres = version: lib.recurseIntoAttrs {
         bin = makePostgresBin version;
         exts = makeOurPostgresPkgsSet version;
-        recurseForDerivations = true;
       };
       basePackages = {
         psql_15 = makePostgres "15";
@@ -171,5 +170,6 @@
     in
     {
       packages = inputs.flake-utils.lib.flattenTree basePackages;
+      legacyPackages = basePackages;
     };
 }

scripts/github-matrix.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import argparse
 import json
 import os
 import subprocess
@@ -47,6 +48,7 @@ class GitHubActionPackage(TypedDict):
     system: str
     already_cached: bool
     runs_on: RunsOnConfig
+    postgresql_version: NotRequired[str]
 
 
 BUILD_RUNNER_MAP: Dict[str, RunsOnConfig] = {
@@ -76,22 +78,23 @@ def get_worker_count() -> int:
         return 1
 
 
-def build_nix_eval_command(max_workers: int) -> List[str]:
+def build_nix_eval_command(max_workers: int, target: str) -> List[str]:
     """Build the nix-eval-jobs command with appropriate flags."""
-    return [
+    nix_eval_cmd = [
         "nix-eval-jobs",
         "--flake",
-        ".#checks",
+        f".#{target}",
         "--check-cache-status",
         "--force-recurse",
         "--quiet",
         "--workers",
         str(max_workers),
     ]
+    return nix_eval_cmd
 
 
 def parse_nix_eval_line(
-    line: str, drv_paths: Set[str]
+    line: str, drv_paths: Set[str], target: str
 ) -> Optional[GitHubActionPackage]:
     """Parse a single line of nix-eval-jobs output"""
     if not line.strip():
@@ -102,12 +105,11 @@ def parse_nix_eval_line(
         if data["drvPath"] in drv_paths:
             return None
         drv_paths.add(data["drvPath"])
-        print(f"Processing package: {data}", file=sys.stderr)
 
         runs_on_config = BUILD_RUNNER_MAP[data["system"]]
 
         return {
-            "attr": "checks." + data["attr"],
+            "attr": f"{target}.{data['attr']}",
             "name": data["name"],
             "system": data["system"],
             "already_cached": data.get("cacheStatus") != "notBuilt",
@@ -118,7 +120,9 @@ def parse_nix_eval_line(
         return None
 
 
-def run_nix_eval_jobs(cmd: List[str]) -> Generator[GitHubActionPackage, None, None]:
+def run_nix_eval_jobs(
+    cmd: List[str], target: str
+) -> Generator[GitHubActionPackage, None, None]:
     """Run nix-eval-jobs and yield parsed package data."""
     print(f"Running command: {' '.join(cmd)}", file=sys.stderr)
 
@@ -128,8 +132,8 @@ def run_nix_eval_jobs(cmd: List[str]) -> Generator[GitHubActionPackage, None, No
         drv_paths = set()
 
         for line in process.stdout:
-            package = parse_nix_eval_line(line, drv_paths)
-            if package:
+            package = parse_nix_eval_line(line, drv_paths, target)
+            if package and not package["already_cached"]:
                 yield package
 
         if process.returncode and process.returncode != 0:
@@ -138,11 +142,41 @@ def run_nix_eval_jobs(cmd: List[str]) -> Generator[GitHubActionPackage, None, No
             sys.exit(process.returncode)
 
 
+def is_extension_pkg(pkg: GitHubActionPackage) -> bool:
+    """Check if the package is a postgresql extension package."""
+    attrs = pkg["attr"].split(".")
+    return attrs[-2] == "exts"
+
+
 def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Generate GitHub Actions matrix for Nix builds"
+    )
+    parser.add_argument(
+        "target", choices=["checks", "extensions"], help="Type of matrix to generate"
+    )
+
+    args = parser.parse_args()
+
     max_workers = get_worker_count()
-    cmd = build_nix_eval_command(max_workers)
 
-    gh_action_packages = list(run_nix_eval_jobs(cmd))
+    if args.target == "checks":
+        flake_output = "checks"
+    else:
+        flake_output = "legacyPackages"
+
+    cmd = build_nix_eval_command(max_workers, flake_output)
+
+    gh_action_packages = list(run_nix_eval_jobs(cmd, flake_output))
+
+    if args.target == "extensions":
+        # filter to only include extension packages and add postgresql_version field
+        gh_action_packages = [
+            {**pkg, "postgresql_version": pkg["attr"].split(".")[-3]}
+            for pkg in gh_action_packages
+            if is_extension_pkg(pkg)
+        ]
+
     gh_output = {"include": gh_action_packages}
     print(json.dumps(gh_output))