feat: add stage 1 AMI build caching based on input hash

.github/workflows/testinfra-ami-build.yml

@@ -91,10 +91,6 @@ jobs:
         with:
           endpoint: builders
 
-      - name: Generate random string
-        id: random
-        run: echo "random_string=$(openssl rand -hex 8)" >> $GITHUB_OUTPUT
-
       - name: Set PostgreSQL version environment variable
         run: |
           echo "POSTGRES_MAJOR_VERSION=${{ matrix.postgres_version }}" >> $GITHUB_ENV
@@ -107,29 +103,48 @@ jobs:
           echo 'postgres-version = "'$PG_VERSION'"' > common-nix.vars.pkr.hcl
           # Ensure there's a newline at the end of the file
           echo "" >> common-nix.vars.pkr.hcl
+          git add -f common-nix.vars.pkr.hcl
 
       - name: Build AMI stage 1
         env:
           AWS_MAX_ATTEMPTS: 10
           AWS_RETRY_MODE: adaptive
         run: |
           GIT_SHA=${{github.sha}}
-          nix run github:supabase/postgres/${GIT_SHA}#packer -- init amazon-arm64-nix.pkr.hcl
-          nix run github:supabase/postgres/${GIT_SHA}#packer -- build -var "git-head-version=${GIT_SHA}" -var "packer-execution-id=${EXECUTION_ID}" -var-file="development-arm.vars.pkr.hcl" -var-file="common-nix.vars.pkr.hcl" -var "ansible_arguments=" -var "postgres-version=${{ steps.random.outputs.random_string }}" -var "region=ap-southeast-1" -var 'ami_regions=["ap-southeast-1"]' -var "force-deregister=true" -var "ansible_arguments=-e postgresql_major=${POSTGRES_MAJOR_VERSION}" amazon-arm64-nix.pkr.hcl
+          nix run .#build-ami -- stage1 \
+            -var "git-head-version=${GIT_SHA}" \
+            -var "packer-execution-id=${EXECUTION_ID}" \
+            -var "ansible_arguments=-e postgresql_major=${POSTGRES_MAJOR_VERSION}" \
+            -var "region=ap-southeast-1" \
+            -var 'ami_regions=["ap-southeast-1"]' \
+            -var-file="development-arm.vars.pkr.hcl" \
+            -var-file="common-nix.vars.pkr.hcl" \
+            amazon-arm64-nix.pkr.hcl
 
       - name: Build AMI stage 2
+        id: build-stage2
         env:
           AWS_MAX_ATTEMPTS: 10
           AWS_RETRY_MODE: adaptive
+          PACKER_EXECUTION_ID: ${{ env.EXECUTION_ID }}
         run: |
           GIT_SHA=${{github.sha}}
-          nix run github:supabase/postgres/${GIT_SHA}#packer -- init stage2-nix-psql.pkr.hcl
-          nix run github:supabase/postgres/${GIT_SHA}#packer -- build -var "git-head-version=${GIT_SHA}" -var "packer-execution-id=${EXECUTION_ID}" -var "postgres_major_version=${POSTGRES_MAJOR_VERSION}" -var-file="development-arm.vars.pkr.hcl" -var-file="common-nix.vars.pkr.hcl"  -var "postgres-version=${{ steps.random.outputs.random_string }}" -var "region=ap-southeast-1" -var 'ami_regions=["ap-southeast-1"]' -var "force-deregister=true" -var "git_sha=${GITHUB_SHA}"  stage2-nix-psql.pkr.hcl
+          nix run .#build-ami -- stage2 \
+            -var "git-head-version=${GIT_SHA}" \
+            -var "packer-execution-id=${EXECUTION_ID}" \
+            -var "postgres_major_version=${POSTGRES_MAJOR_VERSION}" \
+            -var-file="development-arm.vars.pkr.hcl" \
+            -var-file="common-nix.vars.pkr.hcl" \
+            -var "postgres-version=${EXECUTION_ID}" \
+            -var 'ami_regions=["ap-southeast-1"]' \
+            -var "force-deregister=true" \
+            -var "git_sha=${GITHUB_SHA}" \
+            stage2-nix-psql.pkr.hcl
 
       - name: Run tests
         timeout-minutes: 10
         env:
-          AMI_NAME: "supabase-postgres-${{ steps.random.outputs.random_string }}"
+          AMI_ID: ${{ steps.build-stage2.outputs.stage2_ami_id }}
         run: |
           # TODO: use poetry for pkg mgmt
           pip3 install boto3 boto3-stubs[essential] docker ec2instanceconnectcli pytest pytest-testinfra[paramiko,docker] requests
@@ -138,30 +153,40 @@ jobs:
       - name: Cleanup resources on build cancellation
         if: ${{ cancelled() }}
         run: |
-          aws ec2 --region ap-southeast-1 describe-instances --filters "Name=tag:packerExecutionId,Values=${EXECUTION_ID}" --query "Reservations[].Instances[].InstanceId" --output text | xargs -r aws ec2 terminate-instances --region ap-southeast-1 --instance-ids
+          INSTANCE_IDS=$(aws ec2 --region ap-southeast-1 describe-instances --filters "Name=tag:packerExecutionId,Values=${EXECUTION_ID}" --query "Reservations[].Instances[].InstanceId" --output text)
+          if [ -n "$INSTANCE_IDS" ]; then
+            echo "Terminating packer build instances: $INSTANCE_IDS"
+            echo "$INSTANCE_IDS" | xargs -r aws ec2 terminate-instances --region ap-southeast-1 --instance-ids
+          else
+            echo "No packer build instances to clean up"
+          fi
 
       - name: Cleanup resources after build
         if: ${{ always() }}
         run: |
-          aws ec2 --region ap-southeast-1 describe-instances --filters "Name=tag:testinfra-run-id,Values=${EXECUTION_ID}" --query "Reservations[].Instances[].InstanceId" --output text | xargs -r aws ec2 terminate-instances --region ap-southeast-1 --instance-ids || true
+          INSTANCE_IDS=$(aws ec2 --region ap-southeast-1 describe-instances --filters "Name=tag:testinfra-run-id,Values=${EXECUTION_ID}" --query "Reservations[].Instances[].InstanceId" --output text)
+          if [ -n "$INSTANCE_IDS" ]; then
+            echo "Terminating testinfra instances: $INSTANCE_IDS"
+            echo "$INSTANCE_IDS" | xargs -r aws ec2 terminate-instances --region ap-southeast-1 --instance-ids || true
+          else
+            echo "No testinfra instances to clean up"
+          fi
 
-      - name: Cleanup AMIs
+      - name: Cleanup stage 2 AMI
         if: always()
         run: |
-          # Define AMI name patterns
-          STAGE1_AMI_NAME="supabase-postgres-ci-ami-test-stage-1"
-          STAGE2_AMI_NAME="${{ steps.random.outputs.random_string }}"
-
-          # Function to deregister AMIs by name pattern
-          deregister_ami_by_name() {
-            local ami_name_pattern=$1
-            local ami_ids=$(aws ec2 describe-images --region ap-southeast-1 --owners self --filters "Name=name,Values=${ami_name_pattern}" --query 'Images[*].ImageId' --output text)
-            for ami_id in $ami_ids; do
-              echo "Deregistering AMI: $ami_id"
-              aws ec2 deregister-image --region ap-southeast-1 --image-id $ami_id
+          STAGE2_AMI_IDS=$(aws ec2 describe-images \
+            --region ap-southeast-1 \
+            --owners self \
+            --filters "Name=tag:packerExecutionId,Values=${EXECUTION_ID}" \
+            --query 'Images[*].ImageId' \
+            --output text)
+
+          if [ -n "$STAGE2_AMI_IDS" ]; then
+            for ami_id in $STAGE2_AMI_IDS; do
+              echo "Deregistering stage 2 AMI: $ami_id"
+              aws ec2 deregister-image --region ap-southeast-1 --image-id $ami_id || true
             done
-          }
-
-          # Deregister AMIs
-          deregister_ami_by_name "$STAGE1_AMI_NAME"
-          deregister_ami_by_name "$STAGE2_AMI_NAME"
+          else
+            echo "No stage 2 AMI to clean up"
+          fi

amazon-arm64-nix.pkr.hcl

@@ -92,6 +92,12 @@ variable "force-deregister" {
   default = false
 }
 
+variable "input-hash" {
+  type    = string
+  default = ""
+  description = "Content hash of all input sources"
+}
+
 packer {
   required_plugins {
     amazon = {
@@ -106,7 +112,7 @@ source "amazon-ebssurrogate" "source" {
   profile = "${var.profile}"
   #access_key    = "${var.aws_access_key}"
   #ami_name = "${var.ami_name}-arm64-${formatdate("YYYY-MM-DD-hhmm", timestamp())}"
-  ami_name = "${var.ami_name}-${var.postgres-version}-stage-1"
+  ami_name = "${var.ami_name}-${var.postgres-version}-${var.input-hash}-stage-1"
   ami_virtualization_type = "hvm"
   ami_architecture = "arm64"
   ami_regions   = "${var.ami_regions}"
@@ -172,6 +178,7 @@ source "amazon-ebssurrogate" "source" {
     appType = "postgres"
     postgresVersion = "${var.postgres-version}-stage1"
     sourceSha = "${var.git-head-version}"
+    inputHash = "${var.input-hash}"
   }
 
   communicator = "ssh"

nix/packages/build-ami.nix

@@ -0,0 +1,158 @@
+{
+  lib,
+  stdenv,
+  writeShellApplication,
+  packer,
+  awscli2,
+  jq,
+  ...
+}:
+
+let
+  root = ../..;
+  packerSources = stdenv.mkDerivation {
+    name = "packer-sources";
+    src = lib.fileset.toSource {
+      inherit root;
+      fileset = lib.fileset.unions [
+        (root + "/ebssurrogate")
+        (root + "/ansible")
+        (root + "/migrations")
+        (root + "/scripts")
+        (root + "/amazon-arm64-nix.pkr.hcl")
+        (root + "/development-arm.vars.pkr.hcl")
+        (lib.fileset.maybeMissing (root + "/common-nix.vars.pkr.hcl"))
+      ];
+    };
+
+    phases = [
+      "unpackPhase"
+      "installPhase"
+    ];
+    installPhase = ''
+      mkdir -p $out
+      cp -r . $out/
+    '';
+  };
+in
+writeShellApplication {
+  name = "build-ami";
+
+  runtimeInputs = [
+    packer
+    awscli2
+    jq
+  ];
+
+  text = ''
+    set -euo pipefail
+
+    set -x
+
+    # Parse stage parameter
+    STAGE="''${1:-stage1}"
+    shift || true  # Remove first arg, ignore error if no args
+
+    REGION="''${AWS_REGION:-ap-southeast-1}"
+    POSTGRES_VERSION="''${POSTGRES_MAJOR_VERSION:-15}"
+    PACKER_SOURCES="${packerSources}"
+    INPUT_HASH=$(basename "$PACKER_SOURCES" | cut -d- -f1)
+
+    find_stage1_ami() {
+      set +e
+      local ami_output
+      ami_output=$(aws ec2 describe-images \
+        --region "$REGION" \
+        --owners self \
+        --filters \
+          "Name=tag:inputHash,Values=$INPUT_HASH" \
+          "Name=tag:postgresVersion,Values=$POSTGRES_VERSION-stage1" \
+          "Name=state,Values=available" \
+        --query 'Images[0].ImageId' \
+        --output text 2>&1)
+      local exit_code=$?
+      set -e
+
+      if [ $exit_code -ne 0 ] && [ $exit_code -ne 255 ]; then
+        echo "Error querying AWS: $ami_output"
+        exit 1
+      fi
+
+      if [ "$ami_output" = "None" ] || [ -z "$ami_output" ]; then
+        echo ""
+      else
+        echo "$ami_output"
+      fi
+    }
+
+    if [ "$STAGE" = "stage1" ]; then
+      echo "Building stage 1..."
+      echo "Checking for existing AMI..."
+
+      AMI_ID=$(find_stage1_ami)
+      if [ -n "$AMI_ID" ]; then
+        echo "Found existing AMI: $AMI_ID"
+        echo "STAGE1_AMI_ID=$AMI_ID"
+        exit 0
+      fi
+
+      echo "No cached AMI found"
+
+      cd "$PACKER_SOURCES"
+      packer init amazon-arm64-nix.pkr.hcl
+      packer build \
+        -var "input-hash=$INPUT_HASH" \
+        -var "postgres-version=$POSTGRES_VERSION" \
+        -var "region=$REGION" \
+        "$@"
+    elif [ "$STAGE" = "stage2" ]; then
+      echo "Building stage 2..."
+
+      STAGE1_AMI_ID=$(find_stage1_ami)
+      if [ -z "$STAGE1_AMI_ID" ]; then
+        echo "Error: Stage 1 AMI not found. Please build stage 1 first."
+        exit 1
+      fi
+
+      echo "Found stage 1 AMI: $STAGE1_AMI_ID"
+
+      packer init stage2-nix-psql.pkr.hcl
+      packer build \
+        -var "source_ami=$STAGE1_AMI_ID" \
+        -var "region=$REGION" \
+        "$@"
+
+      if [ -n "''${PACKER_EXECUTION_ID:-}" ]; then
+        STAGE2_AMI_ID=$(aws ec2 describe-images \
+          --region "$REGION" \
+          --owners self \
+          --filters \
+            "Name=tag:packerExecutionId,Values=''${PACKER_EXECUTION_ID}" \
+            "Name=state,Values=available" \
+          --query 'Images[0].ImageId' \
+          --output text)
+
+        if [ -n "$STAGE2_AMI_ID" ] && [ "$STAGE2_AMI_ID" != "None" ]; then
+          echo "STAGE2_AMI_ID=$STAGE2_AMI_ID"
+
+          if [ -n "''${GITHUB_OUTPUT:-}" ]; then
+            echo "stage2_ami_id=$STAGE2_AMI_ID" >> "$GITHUB_OUTPUT"
+          fi
+        fi
+      fi
+    else
+      echo "Error: Invalid stage '$STAGE'. Must be 'stage1' or 'stage2'"
+      exit 1
+    fi
+  '';
+
+  meta = {
+    description = "Build AMI if not cached based on input hash";
+    longDescription = ''
+      The input hash is computed from all source files that affect the build.
+      Before building, we verify the existence of an AMI with the same hash.
+      If found, the build is skipped. Otherwise, a new AMI is created and
+      tagged with the input hash for future cache hits.
+    '';
+  };
+}

nix/packages/default.nix

@@ -29,6 +29,7 @@
     {
       packages = (
         {
+          build-ami = pkgs.callPackage ./build-ami.nix { packer = self'.packages.packer; };
           build-test-ami = pkgs.callPackage ./build-test-ami.nix { };
           cleanup-ami = pkgs.callPackage ./cleanup-ami.nix { };
           dbmate-tool = pkgs.callPackage ./dbmate-tool.nix { inherit (self.supabase) defaults; };

stage2-nix-psql.pkr.hcl

@@ -51,6 +51,11 @@ variable "postgres_major_version" {
   default = ""
 }
 
+variable "source_ami" {
+  type    = string
+  description = "Source AMI ID from stage 1"
+}
+
 packer {
   required_plugins {
     amazon = {
@@ -64,15 +69,7 @@ source "amazon-ebs" "ubuntu" {
   ami_name      = "${var.ami_name}-${var.postgres-version}"
   instance_type = "c6g.4xlarge"
   region        = "${var.region}"
-  source_ami_filter {
-    filters = {
-      name   = "${var.ami_name}-${var.postgres-version}-stage-1"
-      root-device-type    = "ebs"
-      virtualization-type = "hvm"
-    }
-    most_recent = true
-    owners      = ["amazon", "self"]
-  }
+  source_ami    = "${var.source_ami}"
 
   communicator = "ssh"
   ssh_pty = true
@@ -107,6 +104,7 @@ source "amazon-ebs" "ubuntu" {
     appType = "postgres"
     postgresVersion = "${var.postgres-version}"
     sourceSha = "${var.git-head-version}"
+    packerExecutionId = "${var.packer-execution-id}"
   }
 }
 

testinfra/test_ami_nix.py

@@ -19,7 +19,7 @@
     + "@"
     + socket.gethostname(),
 )
-AMI_NAME = os.environ.get("AMI_NAME")
+AMI_ID = os.environ.get("AMI_ID")
 postgresql_schema_sql_content = """
 ALTER DATABASE postgres SET "app.settings.jwt_secret" TO  'my_jwt_secret_which_is_not_so_secret';
 ALTER DATABASE postgres SET "app.settings.jwt_exp" TO 3600;
@@ -224,13 +224,7 @@ def run_ssh_command(ssh, command, timeout=None):
 @pytest.fixture(scope="session")
 def host():
     ec2 = boto3.resource("ec2", region_name="ap-southeast-1")
-    images = list(
-        ec2.images.filter(
-            Filters=[{"Name": "name", "Values": [AMI_NAME]}],
-        )
-    )
-    assert len(images) == 1
-    image = images[0]
+    image = ec2.Image(AMI_ID)
 
     def gzip_then_base64_encode(s: str) -> str:
         return base64.b64encode(gzip.compress(s.encode())).decode()