modify_distribution_matrix.py

# This script is used by CI to modify the deployment matrix for the extension distribution

import argparse
import json
import sys
import logging

# Define command-line arguments
parser = argparse.ArgumentParser(description="Filter a JSON file based on excluded duckdb_arch values and select an OS")
parser.add_argument("--input", required=True, help="Input JSON file path")
parser.add_argument("--exclude", required=True, help="Semicolon-separated list of excluded duckdb_arch values")
parser.add_argument("--opt_in", required=False, default="", help="Semicolon-separated list of opt-in duckdb_arch values")
parser.add_argument("--output", help="Output JSON file path")
parser.add_argument("--pretty", action="store_true", help="Pretty print the output JSON")
parser.add_argument("--reduced_ci_mode", required=True, help="Set to default/enabled/disabled, when enabled, filters out redundant archs for testing")
parser.add_argument("--select_os", help="Select an OS to include in the output JSON")
parser.add_argument("--deploy_matrix", action="store_true", help="Create a merged list used in deploy step")
parser.add_argument("--vcpkg_triplets_override", required=False, default="", help="Format: 'arch_name1:triplet1;arch_name2:triplet2'")
args = parser.parse_args()


# Parse the input file path, excluded arch values, and output file path
input_json_file_path = args.input
excluded_arch_values = args.exclude.split(";")
opt_in_arch_values = args.opt_in.split(";")
output_json_file_path = args.output
select_os = args.select_os
reduced_ci_mode = args.reduced_ci_mode

# Parse reduced CI mode
if reduced_ci_mode == "auto":
    # Note auto is off for now TODO: change?
    reduced_ci_mode = False
elif reduced_ci_mode == "enabled":
    reduced_ci_mode = True
elif reduced_ci_mode == "disabled":
    reduced_ci_mode = False
elif reduced_ci_mode is None:
    raise Exception("Unknown reduced_ci_mode value: None - must be auto/enabled/disabled.")
else:
    raise Exception("Unknown reduced_ci_mode value: " + reduced_ci_mode + " - must be auto/enabled/disabled.")

# Parse VCPKG triplets override
triplet_pairs = args.vcpkg_triplets_override.split(";")
triplet_dict = {}
for pair in triplet_pairs:
    stripped = pair.strip()
    if len(stripped) == 0:
        continue
    parts = stripped.split(":")
    if len(parts) != 2:
        raise Exception(f"Invalid 'vcpkg_triplets_override' entry: {stripped},\
 must be in format: 'arch_name1:triplet1'")
    triplet_dict[parts[0]] = parts[1]

# Read the input JSON file
with open(input_json_file_path, "r") as json_file:
    data = json.load(json_file)

def should_run(config, reduced_ci_mode, excluded_arch_values, opt_in_arch_values):
    arch = config["duckdb_arch"]
    if arch in excluded_arch_values:
        return False
    if reduced_ci_mode and not config["run_in_reduced_ci_mode"]:
        return False
    if config["opt_in"] and arch not in opt_in_arch_values:
        return False
    return True

# Function to filter entries based on duckdb_arch values
def filter_entries(data, excluded_arch_values, opt_in_arch_values):
    for os, config in data.items():
        if "include" in config:
            config["include"] = [entry for entry in config["include"] if should_run(entry, reduced_ci_mode, excluded_arch_values, opt_in_arch_values)]

            # Override VCPKG triplets
            for entry in config["include"]:
                override = triplet_dict.get(entry["duckdb_arch"])
                if override is not None:
                    entry["vcpkg_target_triplet"] = override
                    entry["vcpkg_host_triplet"] = override

        if not config["include"]:
            del config["include"]

    return data


# Filter the JSON data
filtered_data = filter_entries(data, excluded_arch_values, opt_in_arch_values)


# Select an OS if specified
if select_os:
    found = False
    for os in filtered_data.keys():
        if os == select_os:
            filtered_data = filtered_data[os]
            found = True
            break
    if found == False:
        logging.warning('A selection OS was provided but not found')
        filtered_data = []

# When deploy_matrix is specified, we only output a single merged include list with all the duckdb_archs
elif args.deploy_matrix:
    deploy_archs = []

    for os, config in filtered_data.items():
        if "include" in config:
            for item in config["include"]:
                deploy_archs.append({"duckdb_arch": item["duckdb_arch"]})

    filtered_data = {"include": deploy_archs}

# Determine the JSON formatting
indent = 2 if args.pretty else None

# If no output file is provided, print to stdout
if output_json_file_path:
    with open(output_json_file_path, "w") as output_json_file:
        if filtered_data:
            json.dump(filtered_data, output_json_file, indent=indent)
else:
    json.dump(filtered_data, sys.stdout, indent=indent)