From 7d3a010bb63c68455b49e93e3e9ec20ef92a4dc3 Mon Sep 17 00:00:00 2001
From: Aldo Borrero <82811+aldoborrero@users.noreply.github.com>
Date: Wed, 14 Feb 2024 18:23:18 +0100
Subject: [PATCH] feat: incorporate update.nix and update.py logic with
 improved semantics

---
 maintainers/scripts/update.py         |  57 ------
 maintainers/scripts/update.sh         |  19 --
 maintainers/scripts/update/update.nix |  63 +++++++
 maintainers/scripts/update/update.py  | 244 ++++++++++++++++++++++++++
 maintainers/scripts/update/update.sh  |  12 ++
 5 files changed, 319 insertions(+), 76 deletions(-)
 delete mode 100755 maintainers/scripts/update.py
 delete mode 100755 maintainers/scripts/update.sh
 create mode 100644 maintainers/scripts/update/update.nix
 create mode 100644 maintainers/scripts/update/update.py
 create mode 100755 maintainers/scripts/update/update.sh

diff --git a/maintainers/scripts/update.py b/maintainers/scripts/update.py
deleted file mode 100755
index e31115f6..00000000
--- a/maintainers/scripts/update.py
+++ /dev/null
@@ -1,57 +0,0 @@
-#!/usr/bin/env nix-shell
-#! nix-shell -i python3 -p "python311.withPackages (ps: with ps; [ click ])"
-
-import json
-import os
-import subprocess
-
-import click
-
-
-def get_nixpkgs_path(path):
-    result = subprocess.run(
-        ["nix", "eval", "-f", path, "inputs.nixpkgs.outPath", "--json"],
-        capture_output=True,
-        text=True,
-    )
-    result.check_returncode()
-    return json.loads(result.stdout)
-
-
-def resolve_nixpkgs(nixpkgs, path):
-    return nixpkgs if nixpkgs else get_nixpkgs_path(path)
-
-
-def build_nix_args(update_nix, path, attr_path, commit):
-    args = [
-        update_nix,
-        "--arg",
-        "include-overlays",
-        f"[(import {path}).overlays.default]",
-        "--argstr",
-        "path",
-        attr_path,
-    ]
-    if commit:
-        args.append("--argstr", "commit", "true")
-    return args
-
-
-@click.command()
-@click.option("--commit", is_flag=True, help="Commit the changes")
-@click.option(
-    "--nixpkgs",
-    help="Override the nixpkgs flake input with this path, it will be used for finding update.nix",
-    default=None,
-)
-@click.argument("attr_path")
-def main(commit, nixpkgs, attr_path):
-    path = os.getcwd()
-    nixpkgs_path = resolve_nixpkgs(nixpkgs, path)
-    update_nix_bin = os.path.join(nixpkgs_path, "maintainers/scripts/update.nix")
-    nix_args = build_nix_args(update_nix_bin, path, attr_path, commit)
-    os.execvp("nix-shell", ["nix-shell"] + nix_args)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/maintainers/scripts/update.sh b/maintainers/scripts/update.sh
deleted file mode 100755
index 2e9d56c4..00000000
--- a/maintainers/scripts/update.sh
+++ /dev/null
@@ -1,19 +0,0 @@
-#!/usr/bin/env bash
-
-set -ex
-
-rootDir="$(git rev-parse --show-toplevel)"
-
-nixpkgs=$(nix eval --raw -f "$rootDir" "inputs.nixpkgs.outPath")
-flake=$(nix eval --raw -f "$rootDir")
-
-nix-shell --show-trace "${nixpkgs}/maintainers/scripts/update.nix" \
-  --arg include-overlays "[(import $rootDir).overlays.default]" \
-  --arg keep-going 'true' \
-  --arg predicate "(
-    let
-        prefix = \"$flake/pkgs/\";
-        prefixLen = builtins.stringLength prefix;
-    in
-        (path: pkg: (builtins.substring 0 prefixLen pkg.meta.position) == prefix)
-  )"
diff --git a/maintainers/scripts/update/update.nix b/maintainers/scripts/update/update.nix
new file mode 100644
index 00000000..e2fbff77
--- /dev/null
+++ b/maintainers/scripts/update/update.nix
@@ -0,0 +1,63 @@
+{
+  flakePath ? ./.,
+  system ? "x86_64-linux",
+  input ? "nixpkgs",
+  max-workers ? null,
+  keep-going ? null,
+  commit ? null,
+  no-confirm ? null,
+}: let
+  # Import the flake.
+  flake = builtins.getFlake (builtins.toString flakePath);
+
+  # Access nixpkgs from the flake inputs.
+  pkgs = flake.inputs.${input}.legacyPackages.${system};
+  inherit (pkgs) lib;
+
+  filterPkgsWithUpdateScript = pkgs:
+    lib.filterAttrs (
+      _name: pkg:
+        lib.isDerivation pkg && lib.hasAttrByPath ["passthru" "updateScript"] pkg
+    )
+    pkgs;
+
+  # Assuming the flake provides packages in the standard flake output schema.
+  flakePkgs = flake.packages.${system};
+
+  filteredPackages = filterPkgsWithUpdateScript flakePkgs;
+
+  packageData = name: package: {
+    inherit name;
+    pname = lib.getName package;
+    oldVersion = lib.getVersion package;
+    inherit (package.passthru) updateScript;
+    attrPath = name; # Assuming you have a way to determine the attribute path here.
+  };
+
+  packagesJson = pkgs.writeText "packages.json" (builtins.toJSON (lib.mapAttrsToList packageData filteredPackages));
+
+  optionalArgs =
+    lib.optional (max-workers != null) "--max-workers=${toString max-workers}"
+    ++ lib.optional (keep-going == "true") "--keep-going"
+    ++ lib.optional (no-confirm == "true") "--no-confirm"
+    ++ lib.optional (commit == "true") "--commit";
+
+  args = [packagesJson] ++ optionalArgs;
+in
+  pkgs.stdenv.mkDerivation {
+    name = "flake-packages-update-script";
+    buildCommand = ''
+      echo ""
+      echo "----------------------------------------------------------------"
+      echo ""
+      echo "Not possible to update packages using \`nix-build\`"
+      echo "Please use \`nix-shell\` with this derivation."
+      echo ""
+      echo "----------------------------------------------------------------"
+      exit 1
+    '';
+    shellHook = ''
+      unset shellHook # Prevent contamination in nested shells.
+      exec ${pkgs.python3}/bin/python ${./update.py} ${builtins.concatStringsSep " " args}
+    '';
+  }
diff --git a/maintainers/scripts/update/update.py b/maintainers/scripts/update/update.py
new file mode 100644
index 00000000..d5e000ef
--- /dev/null
+++ b/maintainers/scripts/update/update.py
@@ -0,0 +1,244 @@
+from typing import Dict, Generator, List, Optional, Tuple
+import argparse
+import asyncio
+import contextlib
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+
+class CalledProcessError(Exception):
+    process: asyncio.subprocess.Process
+
+class UpdateFailedException(Exception):
+    pass
+
+def eprint(*args, **kwargs):
+    print(*args, file=sys.stderr, **kwargs)
+
+async def check_subprocess(*args, **kwargs):
+    """
+    Emulate check argument of subprocess.run function.
+    """
+    process = await asyncio.create_subprocess_exec(*args, **kwargs)
+    returncode = await process.wait()
+
+    if returncode != 0:
+        error = CalledProcessError()
+        error.process = process
+
+        raise error
+
+    return process
+
+async def run_update_script(nixpkgs_root: str, merge_lock: asyncio.Lock, temp_dir: Optional[Tuple[str, str]], package: Dict, keep_going: bool):
+    worktree: Optional[str] = None
+
+    update_script_command = package['updateScript']
+
+    if temp_dir is not None:
+        worktree, _branch = temp_dir
+
+        # Ensure the worktree is clean before update.
+        await check_subprocess('git', 'reset', '--hard', '--quiet', 'HEAD', cwd=worktree)
+
+        # Update scripts can use $(dirname $0) to get their location but we want to run
+        # their clones in the git worktree, not in the main nixpkgs repo.
+        update_script_command = map(lambda arg: re.sub(r'^{0}'.format(re.escape(nixpkgs_root)), worktree, arg), update_script_command)
+
+    eprint(f" - {package['name']}: UPDATING ...")
+
+    try:
+        update_process = await check_subprocess(
+            'env',
+            f"UPDATE_NIX_NAME={package['name']}",
+            f"UPDATE_NIX_PNAME={package['pname']}",
+            f"UPDATE_NIX_OLD_VERSION={package['oldVersion']}",
+            f"UPDATE_NIX_ATTR_PATH={package['attrPath']}",
+            *update_script_command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=worktree,
+        )
+        update_info = await update_process.stdout.read()
+
+        await merge_changes(merge_lock, package, update_info, temp_dir)
+    except KeyboardInterrupt as e:
+        eprint('Cancelling…')
+        raise asyncio.exceptions.CancelledError()
+    except CalledProcessError as e:
+        eprint(f" - {package['name']}: ERROR")
+        eprint()
+        eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
+        eprint()
+        stderr = await e.process.stderr.read()
+        eprint(stderr.decode('utf-8'))
+        with open(f"{package['pname']}.log", 'wb') as logfile:
+            logfile.write(stderr)
+        eprint()
+        eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
+
+        if not keep_going:
+            raise UpdateFailedException(f"The update script for {package['name']} failed with exit code {e.process.returncode}")
+
+@contextlib.contextmanager
+def make_worktree() -> Generator[Tuple[str, str], None, None]:
+    with tempfile.TemporaryDirectory() as wt:
+        branch_name = f'update-{os.path.basename(wt)}'
+        target_directory = f'{wt}/nixpkgs'
+
+        subprocess.run(['git', 'worktree', 'add', '-b', branch_name, target_directory])
+        yield (target_directory, branch_name)
+        subprocess.run(['git', 'worktree', 'remove', '--force', target_directory])
+        subprocess.run(['git', 'branch', '-D', branch_name])
+
+async def commit_changes(name: str, merge_lock: asyncio.Lock, worktree: str, branch: str, changes: List[Dict]) -> None:
+    for change in changes:
+        # Git can only handle a single index operation at a time
+        async with merge_lock:
+            await check_subprocess('git', 'add', *change['files'], cwd=worktree)
+            commit_message = '{attrPath}: {oldVersion} -> {newVersion}'.format(**change)
+            if 'commitMessage' in change:
+                commit_message = change['commitMessage']
+            elif 'commitBody' in change:
+                commit_message = commit_message + '\n\n' + change['commitBody']
+            await check_subprocess('git', 'commit', '--quiet', '-m', commit_message, cwd=worktree)
+            await check_subprocess('git', 'cherry-pick', branch)
+
+async def check_changes(package: Dict, worktree: str, update_info: str):
+    if 'commit' in package['supportedFeatures']:
+        changes = json.loads(update_info)
+    else:
+        changes = [{}]
+
+    # Try to fill in missing attributes when there is just a single change.
+    if len(changes) == 1:
+        # Dynamic data from updater take precedence over static data from passthru.updateScript.
+        if 'attrPath' not in changes[0]:
+            # update.nix is always passing attrPath
+            changes[0]['attrPath'] = package['attrPath']
+
+        if 'oldVersion' not in changes[0]:
+            # update.nix is always passing oldVersion
+            changes[0]['oldVersion'] = package['oldVersion']
+
+        if 'newVersion' not in changes[0]:
+            attr_path = changes[0]['attrPath']
+            obtain_new_version_process = await check_subprocess('nix-instantiate', '--expr', f'with import ./. {{}}; lib.getVersion {attr_path}', '--eval', '--strict', '--json', stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=worktree)
+            changes[0]['newVersion'] = json.loads((await obtain_new_version_process.stdout.read()).decode('utf-8'))
+
+        if 'files' not in changes[0]:
+            changed_files_process = await check_subprocess('git', 'diff', '--name-only', 'HEAD', stdout=asyncio.subprocess.PIPE, cwd=worktree)
+            changed_files = (await changed_files_process.stdout.read()).splitlines()
+            changes[0]['files'] = changed_files
+
+            if len(changed_files) == 0:
+                return []
+
+    return changes
+
+async def merge_changes(merge_lock: asyncio.Lock, package: Dict, update_info: str, temp_dir: Optional[Tuple[str, str]]) -> None:
+    if temp_dir is not None:
+        worktree, branch = temp_dir
+        changes = await check_changes(package, worktree, update_info)
+
+        if len(changes) > 0:
+            await commit_changes(package['name'], merge_lock, worktree, branch, changes)
+        else:
+            eprint(f" - {package['name']}: DONE, no changes.")
+    else:
+        eprint(f" - {package['name']}: DONE.")
+
+async def updater(nixpkgs_root: str, temp_dir: Optional[Tuple[str, str]], merge_lock: asyncio.Lock, packages_to_update: asyncio.Queue[Optional[Dict]], keep_going: bool, commit: bool):
+    while True:
+        package = await packages_to_update.get()
+        if package is None:
+            # A sentinel received, we are done.
+            return
+
+        if not ('commit' in package['supportedFeatures'] or 'attrPath' in package):
+            temp_dir = None
+
+        await run_update_script(nixpkgs_root, merge_lock, temp_dir, package, keep_going)
+
+async def start_updates(max_workers: int, keep_going: bool, commit: bool, packages: List[Dict]):
+    merge_lock = asyncio.Lock()
+    packages_to_update: asyncio.Queue[Optional[Dict]] = asyncio.Queue()
+
+    with contextlib.ExitStack() as stack:
+        temp_dirs: List[Optional[Tuple[str, str]]] = []
+
+        # Do not create more workers than there are packages.
+        num_workers = min(max_workers, len(packages))
+
+        nixpkgs_root_process = await check_subprocess('git', 'rev-parse', '--show-toplevel', stdout=asyncio.subprocess.PIPE)
+        nixpkgs_root = (await nixpkgs_root_process.stdout.read()).decode('utf-8').strip()
+
+        # Set up temporary directories when using auto-commit.
+        for i in range(num_workers):
+            temp_dir = stack.enter_context(make_worktree()) if commit else None
+            temp_dirs.append(temp_dir)
+
+        # Fill up an update queue,
+        for package in packages:
+            await packages_to_update.put(package)
+
+        # Add sentinels, one for each worker.
+        # A workers will terminate when it gets sentinel from the queue.
+        for i in range(num_workers):
+            await packages_to_update.put(None)
+
+        # Prepare updater workers for each temp_dir directory.
+        # At most `num_workers` instances of `run_update_script` will be running at one time.
+        updaters = asyncio.gather(*[updater(nixpkgs_root, temp_dir, merge_lock, packages_to_update, keep_going, commit) for temp_dir in temp_dirs])
+
+        try:
+            # Start updater workers.
+            await updaters
+        except asyncio.exceptions.CancelledError:
+            # When one worker is cancelled, cancel the others too.
+            updaters.cancel()
+        except UpdateFailedException as e:
+            # When one worker fails, cancel the others, as this exception is only thrown when keep_going is false.
+            updaters.cancel()
+            eprint(e)
+            sys.exit(1)
+
+def main(max_workers: int, keep_going: bool, commit: bool, no_confirm: bool, packages_path: str) -> None:
+    with open(packages_path) as f:
+        packages = json.load(f)
+
+    eprint()
+    eprint('Going to be running update for following packages:')
+    for package in packages:
+        eprint(f" - {package['name']}")
+    eprint()
+
+    if not no_confirm:
+        confirm = input('Press Enter key to continue...')
+        if confirm != '':
+            eprint('Aborting!')
+            sys.exit(130)
+
+    eprint('Running update for:')
+    asyncio.run(start_updates(max_workers, keep_going, commit, packages))
+    eprint('Packages updated!')
+    sys.exit()
+
+parser = argparse.ArgumentParser(description='Update packages')
+parser.add_argument('--max-workers', '-j', dest='max_workers', type=int, help='Number of updates to run concurrently', nargs='?', default=4)
+parser.add_argument('--keep-going', '-k', dest='keep_going', action='store_true', help='Do not stop after first failure')
+parser.add_argument('--commit', '-c', dest='commit', action='store_true', help='Commit the changes')
+parser.add_argument('--no-confirm', '-n', action='store_true', help='Skip the confirmation prompt and proceed with updates automatically')
+parser.add_argument('packages', help='JSON file containing the list of package names and their update scripts')
+
+if __name__ == '__main__':
+    args = parser.parse_args()
+
+    try:
+        main(args.max_workers, args.keep_going, args.commit, args.no_confirm, args.packages)
+    except KeyboardInterrupt as e:
+        # Let’s cancel outside of the main loop too.
+        sys.exit(130)
diff --git a/maintainers/scripts/update/update.sh b/maintainers/scripts/update/update.sh
new file mode 100755
index 00000000..57d376fa
--- /dev/null
+++ b/maintainers/scripts/update/update.sh
@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -e
+
+rootDir="$(git rev-parse --show-toplevel)"
+updateScript="maintainers/scripts/update/update.nix"
+
+nix-shell "${rootDir}/${updateScript}" \
+  --argstr flakePath "${rootDir}" \
+  --arg keep-going 'true' \
+  --arg commit "true" \
+  --arg max-workers "4"