konflux_update_rp_rpa_for_fbc.py

import asyncio
import sys
from pathlib import Path
from ruamel.yaml import YAML

from sqlalchemy import create_engine, text, Table, MetaData, select, Column, String, DateTime

yaml = YAML(typ="rt")
yaml.preserve_quotes = True
yaml.explicit_start = True
yaml.indent(mapping=2, sequence=4, offset=2)
yaml.width = 1000

DB_URI = "bigquery://openshift-art"
DB_CREDENTIALS_PATH = "/Users/yux/.kube/openshift-art-61fdabdf4001.json"
KONFLUX_RELEASE_DATA_PATH = "/Volumes/Workspaces/projects/openshift/konflux-release-data"


async def get_bundle_package_names(dgks: list[str], major: int, minor: int) -> list[str]:
    engine = create_engine(DB_URI, credentials_path=DB_CREDENTIALS_PATH)
    metadata = MetaData()
    bundles_table = Table(
        'events.bundles', metadata,
        Column('name', String, primary_key=True),
        Column('group', String),
        Column('engine', String),
        Column('outcome', String),
        Column('bundle_package_name', String),
        Column('start_time', DateTime),
    )
    with engine.connect() as conn:
        def get_bundle_package_name(dgk):
            bundle_name = f'{dgk}-bundle'
            stmt = select(bundles_table.c.bundle_package_name).where(
                text('TIMESTAMP_TRUNC(start_time, DAY) >= TIMESTAMP("2025-05-01")'),
                bundles_table.c.group == f"openshift-{major}.{minor}",
            #     bundles_table.c.engine == "konflux"
            #     bundles_table.c.outcome == "success"
                bundles_table.c.name == bundle_name,
                bundles_table.c.bundle_package_name.isnot(None),
                bundles_table.c.bundle_package_name != ''
            ).order_by(
                bundles_table.c.start_time.desc()
            ).limit(1)
            result = conn.execute(stmt).first()
            if not result:
                raise ValueError(
                    f"Bundle {bundle_name} not found in the database.")
            assert result.bundle_package_name, f"Bundle {bundle_name} has no package name."
            bundle_name = str(result.bundle_package_name)
            return bundle_name
        results = await asyncio.gather(
            *[asyncio.to_thread(get_bundle_package_name, dgk) for dgk in dgks],
            return_exceptions=True
        )
        errors = {dgk: r for dgk, r in zip(dgks, results) if isinstance(r, Exception)}
        if errors:
            for dgk, error in errors.items():
                print(f"Error for {dgk}: {error}", file=sys.stderr)
            raise IOError("Failed to retrieve some bundle package names.")
        return results


async def run_for(major: int, minor: int):
    build_data_path = "/Volumes/Workspaces/projects/openshift/ocp-build-data"
    # run `git checkout openshift-{major}.{minor}`
    proc = await asyncio.create_subprocess_exec("git", "checkout", f"openshift-{major}.{minor}", cwd=build_data_path)
    await proc.wait()
    if proc.returncode != 0:
        raise RuntimeError(f"Failed to checkout openshift-{major}.{minor} branch in {build_data_path}")
    # run `git reset --hard upstream/openshift-{major}.{minor}`
    proc = await asyncio.create_subprocess_exec("git", "reset", "--hard", f"upstream/openshift-{major}.{minor}", cwd=build_data_path)
    await proc.wait()
    if proc.returncode != 0:
        raise RuntimeError(f"Failed to reset to upstream/openshift-{major}.{minor} in {build_data_path}")
    olm_operators = {}
    for meta_file in Path(build_data_path, "images").glob("*.yml"):
        dgk = meta_file.name.removesuffix(".yml")
        meta = yaml.load(meta_file.open("r"))
        if meta.get("mode") in ["disabled", "wip"]:
            continue
        if meta.get("update-csv"):
            olm_operators[dgk] = meta

    dgks = sorted(olm_operators.keys())
    bundle_package_names = await get_bundle_package_names(dgks, major=major, minor=minor)

    # Update RPAs
    rpa_paths = [Path(KONFLUX_RELEASE_DATA_PATH,
                      f"config/kflux-ocp-p01.7ayg.p1/product/ReleasePlanAdmission/ocp-art/ocp-art-fbc-stage-{major}-{minor}.yaml"),
                 Path(KONFLUX_RELEASE_DATA_PATH, f"config/kflux-ocp-p01.7ayg.p1/product/ReleasePlanAdmission/ocp-art/ocp-art-fbc-prod-{major}-{minor}.yaml")]
    for rpa_path in rpa_paths:
        rpa = yaml.load(rpa_path.open("r"))
        rpa_name = rpa_path.name.removesuffix(".yaml")
        rpa["metadata"]["name"] = rpa_name
        rpa["spec"]["applications"] = [f"fbc-openshift-{major}-{minor}"]
        rpa["spec"]["data"]["fbc"]["allowedPackages"] = sorted(
            bundle_package_names)
        print(
            f"Updated RPA for OCP {major}.{minor} with {len(bundle_package_names)} packages.")
        with rpa_path.open("w") as f:
            yaml.dump(rpa, f)

    # Update RPs
    rp_paths = [Path(KONFLUX_RELEASE_DATA_PATH,
                   f"tenants-config/cluster/kflux-ocp-p01/tenants/ocp-art-tenant/stage/ocp-art-fbc-stage-{major}-{minor}-rp.yaml"),
                   Path(KONFLUX_RELEASE_DATA_PATH, f"tenants-config/cluster/kflux-ocp-p01/tenants/ocp-art-tenant/prod/ocp-art-fbc-prod-{major}-{minor}-rp.yaml")]
    for rp_path in rp_paths:
        rp_name = rp_path.name.removesuffix("-rp.yaml")
        rp = yaml.load(rp_path.open("r"))
        rp["metadata"]["name"] = rp_name
        rp["metadata"]["labels"]["release.appstudio.openshift.io/releasePlanAdmission"] = rp_name  # RPA name matches RP name
        rp["spec"]["application"] = f"fbc-openshift-{major}-{minor}"
        if "stage" in rp_name:
            rp["spec"]["data"]["mapping"]["components"] = [
                {"name": f"fbc-ose-{major}-{minor}-{dgk}", "repository": "quay.io/openshift-art/stage-fbc-fragments",
                    "tags": [f"ocp__{major}.{minor}__{dgk}"]}
                for dgk in dgks
            ]
        with rp_path.open("w") as f:
            yaml.dump(rp, f)


async def main():
    major = 4
    for minor in range(20, 11, -1):
        await run_for(major, minor)
        print(f"Processed OCP {major}.{minor}.")

if __name__ == "__main__":
    asyncio.run(main())