1
import argparse
2
import json
3
import re
4
import subprocess
5
import sys
6
from copy import deepcopy
7
from dataclasses import dataclass
8
from functools import cache
9
from os import environ
10
from pathlib import Path
11
 
12
import requests
13
from tool_utils import b64digest
14
 
15
 
16
def cquery(rule):
17
    res = subprocess.run(
18
        [
19
            "bazel",
20
            "cquery",
21
            rule,
22
            "--output=files",
23
        ],
24
        capture_output=True,
25
        text=True,
26
        check=False,
27
    )
28
    if res.returncode:
29
        print(res.stdout)
30
        print(res.stderr)
31
        sys.exit(res.returncode)
32
    return res.stdout.strip()
33
 
34
 
35
def bazel_mod_tidy():
36
    res = subprocess.run(["bazel", "mod", "tidy"])
37
    if res.returncode:
38
        print(res.stdout)
39
        print(res.stderr)
40
        sys.exit(res.returncode)
41
 
42
 
43
@cache
44
def _bundle_version_info():
45
    with Path(cquery("@workerd//src/pyodide:bundle_version_info")).open() as f:
46
        return json.load(f)
47
 
48
 
49
def bundle_version_info():
50
    return deepcopy(_bundle_version_info())
51
 
52
 
53
def get_pyodide_bin_path(ver):
54
    return cquery(f"@workerd//src/pyodide:pyodide.capnp.bin@rule@{ver}")
55
 
56
 
57
def bundle_key(*, pyodide_version, pyodide_date, backport, **_kwds):
58
    return f"pyodide_{pyodide_version}_{pyodide_date}_{backport}.capnp.bin"
59
 
60
 
61
def bundle_url(**kwds):
62
    return "https://pyodide-capnp-bin.edgeworker.net/" + bundle_key(**kwds)
63
 
64
 
65
def get_backport(ver):
66
    info = bundle_version_info()[ver]
67
    backport = int(info["backport"])
68
    for b in range(backport + 1, backport + 20):
69
        info["backport"] = b
70
        url = bundle_url(**info)
71
        res = requests.head(url)
72
        if res.status_code == 404:
73
            return b
74
        res.raise_for_status()
75
 
76
 
77
def _get_replacer(backport, integrity):
78
    def replace_values(match):
79
        prefix = match.group(1)
80
        backport_key = match.group(2)
81
        middle = match.group(3)
82
        integrity_key = match.group(4)
83
        return (
84
            f'{prefix}{backport_key}"{backport}",{middle}{integrity_key}"{integrity}",'
85
        )
86
 
87
    return replace_values
88
 
89
 
90
@dataclass
91
class BundleInfo:
92
    version: str
93
    backport: int
94
    integrity: str
95
    path: Path
96
 
97
 
98
def update_python_metadata_bzl(bundles: list[BundleInfo]) -> None:
99
    """Update python_metadata.bzl file with new backport and integrity values."""
100
    metadata_path = (
101
        Path(__file__).parent.parent.parent / "build" / "python_metadata.bzl"
102
    )
103
    content = metadata_path.read_text()
104
 
105
    for info in bundles:
106
        # Find the version block and update backport and integrity
107
        version_pattern = rf'(\s+{{\s*\n\s*"name":\s*"{re.escape(info.version)}",.*?)("backport":\s*)"[^"]*",(.*?)("integrity":\s*)"[^"]*",'
108
 
109
        content = re.sub(
110
            version_pattern,
111
            _get_replacer(info.backport, info.integrity),
112
            content,
113
            flags=re.DOTALL,
114
        )
115
 
116
    metadata_path.write_text(content)
117
    bazel_mod_tidy()
118
 
119
 
120
def print_info(info: BundleInfo) -> None:
121
    print(f"Uploading version {info.version} backport {info.backport}")
122
    i = " " * 8
123
    print("Update python_metadata.bzl with:\n")
124
    print(i + f'"backport": "{info.backport}",')
125
    print(i + f'"integrity": "{info.integrity}",')
126
    print()
127
 
128
 
129
def make_bundles(update_released: bool) -> list[BundleInfo]:
130
    result = []
131
    for ver, info in bundle_version_info().items():
132
        if ver.startswith("dev"):
133
            continue
134
        if not update_released and info.get("released", False):
135
            continue
136
        path = Path(get_pyodide_bin_path(ver)).resolve()
137
        b = get_backport(ver)
138
        info["backport"] = b
139
        integrity = b64digest(path)
140
        result.append(BundleInfo(ver, b, integrity, path))
141
    return result
142
 
143
 
144
def upload_bundles(bundles: list[BundleInfo]):
145
    from boto3 import client
146
 
147
    s3 = client(
148
        "s3",
149
        endpoint_url=f"https://{environ['R2_ACCOUNT_ID']}.r2.cloudflarestorage.com",
150
        aws_access_key_id=environ["R2_ACCESS_KEY_ID"],
151
        aws_secret_access_key=environ["R2_SECRET_ACCESS_KEY"],
152
        region_name="auto",
153
    )
154
 
155
    for bundle in bundles:
156
        ver = bundle.version
157
        path = Path(get_pyodide_bin_path(ver)).resolve()
158
        b = get_backport(ver)
159
        info = bundle_version_info()[ver]
160
        info["backport"] = b
161
        key = bundle_key(**info)
162
        s3.upload_file(str(path), "pyodide-capnp-bin", key)
163
 
164
 
165
def main():
166
    parser = argparse.ArgumentParser(
167
        description="Upload Pyodide bundles and update metadata"
168
    )
169
    parser.add_argument(
170
        "--update-released",
171
        action="store_true",
172
        help="Update already released versions?",
173
    )
174
    args = parser.parse_args()
175
    bundles = make_bundles(args.update_released)
176
    for bundle in bundles:
177
        print_info(bundle)
178
    update_python_metadata_bzl(bundles)
179
    upload_bundles(bundles)
180
    return 0
181
 
182
 
183
if __name__ == "__main__":
184
    sys.exit(main())