diff --git a/updateindexfile.py b/updateindexfile.py index c06af49..45bf47d 100644 --- a/updateindexfile.py +++ b/updateindexfile.py @@ -1,108 +1,278 @@ import json +import os import subprocess +import sys import tempfile -import os +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from enum import StrEnum + +DRY_RUN = os.getenv("DRY_RUN", "").lower() in ("1", "true", "yes") -miniver = os.getenv("MINIVER") -splenv = os.getenv("SPLENV_REF") +MINIVER = os.getenv("MINIVER") +SPLENV_REF = os.getenv("SPLENV_REF") -folders = ["manifests/", "env/", "tables/"] BUCKET_NAME = "eups-prod" GCS_PREFIX = f"gs://{BUCKET_NAME}" INDEX_FILE = "index.json" +# Subfolders indexed within every conda platform folder. +PLATFORM_SUBFOLDERS = ["manifests/", "env/", "tables/"] +# The src tree has the same subfolders plus these. +SRC_SUBFOLDERS = PLATFORM_SUBFOLDERS + ["products/", "tags/"] -root_folders = [ +ROOT_FOLDERS = [ "stack/redhat/el7/conda-system", "stack/redhat/el8-arm/conda-system", "stack/osx/14-arm/conda-system", ] +MAX_WORKERS = 16 + + +class IndexStatus(StrEnum): + MISSING = "MISSING" # no index.json in prod yet + DIFFERS = "DIFFERS" # contents or order differs from prod + UNREADABLE = "UNREADABLE" # existing index could not be read + + +@dataclass +class IndexDiff: + dest: str + status: IndexStatus + count: int = 0 + reordered: bool = False + added: list[str] = field(default_factory=list) + removed: list[str] = field(default_factory=list) + -# If we have the MINIVER and SPLENV_REF defined, then we can return -# a string to filter by def filter_folders() -> str: - if miniver and splenv: - return f"miniconda3-{miniver}-{splenv}" + if MINIVER and SPLENV_REF: + return f"miniconda3-{MINIVER}-{SPLENV_REF}" return "" -def get_gcs_object_uris(target: str) -> list[str]: - indexdata = subprocess.run( - ["gcloud", "storage", "ls", target], capture_output=True, check=True, text=True - ) - return indexdata.stdout.split() +def list_recursive(prefix: str) -> list[str] | None: + """Return every object path under prefix (recursive), relative to the bucket. + Returns [] if the prefix matched nothing or None if the listing command itself failed. + Uses the ``**`` wildcard so gcloud returns a flat list of object URLs + (no directory placeholders, no per-directory headers), which is cheap to + group in Python. A prefix that matches nothing yields an empty list. + """ + target = f"{GCS_PREFIX}/{prefix.rstrip('/')}/**" + try: + result = subprocess.run( + ["gcloud", "storage", "ls", target], + capture_output=True, + check=True, + text=True, + ) + except subprocess.CalledProcessError as e: + stderr = e.stderr or "" + # error messages taken from gcloud storage ls docs. Could change in the future + if "matched no objects" in stderr: + return [] + print(f"ERROR: listing {target} failed: {stderr.strip()}") + return None + strip = f"{GCS_PREFIX}/" + return [u[len(strip) :] for u in result.stdout.splitlines() if u.startswith(strip)] -def copy_files(file: str, target: str): - copy = subprocess.run( - ["gcloud", "storage", "cp", file, target + INDEX_FILE], - capture_output=True, - check=True, - text=True, - ) - if copy.returncode == 0: - print(f"updated {INDEX_FILE}") + +def group_by_parent(objects: list[str]) -> dict[str, list[str]]: + """Map each folder prefix to the file names directly inside it.""" + by_parent: dict[str, list[str]] = defaultdict(list) + for obj in objects: + cut = obj.rfind("/") + parent, name = obj[: cut + 1], obj[cut + 1 :] + if name and name != INDEX_FILE: + by_parent[parent].append(name) + return by_parent + + +def existing_prefixes(objects: list[str]) -> set[str]: + """Every folder prefix that has at least one object beneath it.""" + prefixes: set[str] = set() + for obj in objects: + parts = obj.split("/") + for i in range(1, len(parts)): + prefixes.add("/".join(parts[:i]) + "/") + return prefixes + + +def platform_folders(root: str, objects: list[str]) -> set[str]: + """Discover the conda platform folders (one level below a root).""" + base = f"{root.rstrip('/')}/" + folders = set() + for obj in objects: + if not obj.startswith(base): + continue + seg = obj[len(base) :].split("/", 1)[0] + if seg: + folders.add(f"{base}{seg}/") + return folders + + +def upload_index(target: str, names: list[str]) -> None: + """Write the file list as index.json into the target folder.""" + dest = f"{GCS_PREFIX}/{target}{INDEX_FILE}" + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f: + tmp = f.name + json.dump(names, f) + try: + subprocess.run( + ["gcloud", "storage", "cp", tmp, dest], + capture_output=True, + check=True, + text=True, + ) + print(f"updated {dest}") + # Adding more detail to when an upload fails + except subprocess.CalledProcessError as e: + raise RuntimeError(f"upload to {dest} failed: {(e.stderr or '').strip()}") from e + finally: + os.remove(tmp) + + +def target_prefixes(root_objects: dict[str, list[str]], existing: set[str]) -> list[str]: + """Build the ordered, de-duplicated list of folders needing an index.""" + sub = filter_folders() + targets = [f"stack/src/{f}" for f in SRC_SUBFOLDERS] + for root, objects in root_objects.items(): + for folder in sorted(platform_folders(root, objects)): + if sub and sub not in folder: + continue + targets.append(folder) + targets.extend(f"{folder}{f}" for f in PLATFORM_SUBFOLDERS) + targets = [t for t in targets if t in existing] + return list(dict.fromkeys(targets)) -def update_helper(loc: str): - target = f"{GCS_PREFIX}/{loc}" - print(target) - # Using the gcloud cli tool was the most consistant way to get file names. - # SDK would give a mix of folders and files - indexdata = [] +def fetch_current_index(target: str) -> list[str] | None: + """Download the existing index.json for a target, or None if absent. Raise RuntimeError if the read itself failed""" + src = f"{GCS_PREFIX}/{target}{INDEX_FILE}" try: - indexdata = get_gcs_object_uris(target) - except subprocess.CalledProcessError: - print(f"{target} does not exist, skipping") - return - # makes a list of all the filenames in the target. It filters out folders - # because folders will be empty strings - # index = [file for i in indexdata for file in [i.split("/")[-1]] if file] - index = [i.split("/")[-1] for i in indexdata if i] - with tempfile.NamedTemporaryFile("w", delete_on_close=False) as f: - json.dump(index, f) - f.close() - print("Fetched files") - copy_files(f.name, target) - os.remove(f.name) - - -def get_list_of_folders() -> list[str]: - conda_folder = [] - for folder in root_folders: - target = f"{GCS_PREFIX}/{folder}" - indexdata = None - try: - indexdata = get_gcs_object_uris(target) - except subprocess.CalledProcessError: - print(f"{target} does not exist, skipping") + result = subprocess.run( + ["gcloud", "storage", "cat", src], + capture_output=True, + check=True, + text=True, + ) + except subprocess.CalledProcessError as e: + stderr = e.stderr or "" + if "does not exist" in stderr or "matched no objects" in stderr: + return None + raise RuntimeError(f"could not read {src}: {stderr.strip()}") from e + try: + return json.loads(result.stdout) + except json.JSONDecodeError: + print(f"WARNING: {src} is not valid JSON") + return None + + +def compare_index(target: str, names: list[str]) -> IndexDiff | None: + """Compare the generated index against the one currently in prod. Used for testing""" + dest = f"{GCS_PREFIX}/{target}{INDEX_FILE}" + try: + current = fetch_current_index(target) + except RuntimeError as e: + print(f"WARNING: {e}") + return IndexDiff(dest=dest, status=IndexStatus.UNREADABLE) + + if current is None: + return IndexDiff(dest=dest, status=IndexStatus.MISSING, count=len(names)) + + names = normalize_index(names) + current = normalize_index(current) + if current == names: + return None # up-to-date, nothing to report + + cur_set, new_set = set(current), set(names) + return IndexDiff( + dest=dest, + status=IndexStatus.DIFFERS, + reordered=cur_set == new_set, # same contents, different order + added=sorted(new_set - cur_set), + removed=sorted(cur_set - new_set), + ) + + +def normalize_index(names: list[str]) -> list[str]: + """Drop blank/whitespace-only entries and index.json for comparison.""" + return [n for n in names if n and n.strip() and n.strip() != INDEX_FILE] + + +def main() -> int: + # One recursive listing for src plus one per conda root, run concurrently. + listing_prefixes = ["stack/src"] + ROOT_FOLDERS + with ThreadPoolExecutor(max_workers=len(listing_prefixes)) as pool: + listings = list(pool.map(list_recursive, listing_prefixes)) + if any(objs is None for objs in listings): + print("Error: one or more listings failed; aborting and not uploading") + return 1 + + # Narrow the type: after the guard above, no entry is None. + listings = [objs for objs in listings if objs is not None] + + src_objects = listings[0] + root_objects = dict(zip(ROOT_FOLDERS, listings[1:])) + + all_objects = list(src_objects) + for objects in root_objects.values(): + all_objects.extend(objects) + + by_parent = group_by_parent(all_objects) + existing = existing_prefixes(all_objects) + targets = target_prefixes(root_objects, existing) + + action = compare_index if DRY_RUN else upload_index + if DRY_RUN: + print("[DRY-RUN] no objects will be uploaded") + + with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool: + futures = {pool.submit(action, t, sorted(by_parent.get(t, []))): t for t in targets} + results = [] + failures = 0 + for future in as_completed(futures): + target = futures[future] + try: + results.append(future.result()) + except Exception as e: + print(f"Error: {target} failed: {e}") + failures += 1 + + if not DRY_RUN: + if failures: + print(f"ERROR: {failures} of {len(targets)} uploads failed") + return 1 + return 0 + # Recap: only folders whose index would change. + diffs = [r for r in results if r] + print("\n" + "=" * 60) + print(f"[DRY-RUN] RECAP: {len(diffs)} of {len(targets)} indexes differ") + print("=" * 60) + + for r in diffs: + if r.status is IndexStatus.MISSING: + print(f"\n{r.dest}") + print(f" MISSING in prod, would create with {r.count} entries") continue - for j in indexdata: - conda_folder.append(j.removeprefix(f"{GCS_PREFIX}/")) - return conda_folder - - -def main(): - platforms = ["stack/src/"] - # Gets all of the folders - platforms.extend(get_list_of_folders()) - - # If miniver and splenv are set, filter out folders that are not included - platforms = [p for p in platforms if filter_folders() in p or "src/" in p] - for p in platforms: - if "src" in p: - # Src folder contains extra folders - srcfolders = folders + ["products/", "tags/"] - for f in srcfolders: - prefix = p + f - update_helper(prefix) - else: - update_helper(p) - for f in folders: - prefix = p + f - update_helper(prefix) + if r.status is IndexStatus.UNREADABLE: + print(f"\n{r.dest} UNREADABLE (could not compare)") + continue + + print(f"\n{r.dest} DIFFERS") + if r.reordered: + print(" (same contents, different order)") + for name in r.added: + print(f" + {name}") + for name in r.removed: + print(f" - {name}") + if failures: + print(f"WARNING: {failures} comparisons errored") + return 1 if failures else 0 if __name__ == "__main__": - main() + sys.exit(main())