feat(bzlmod): support bazel downloader when downloading wheels (#1827)

This introduces 3 attributes and the minimal code to be able to download wheels using the bazel downloader for the host platform. This is not yet adding support for targeting a different platform but just allows us to get the wheels for the host platform instead of using `pip`. All of this is achieved by calling the PyPI's SimpleAPI (Artifactory should work as well) and getting the all URLs for packages from there. Then we use the `sha256` information within the requirements files to match the entries found on SimpleAPI and then pass the `url`, `sha256` and the `filename` to `whl_library`, which uses `repository_ctx.download`. If we cannot find any suitable artifact to use, we fallback to legacy `pip` behaviour. Testing notes: * Most of the code has unit tests, but the `pypi_index.bzl` extension could have more. * You can see the lock file for what the output of all of this code would be on your platform. * Thanks to @dougthor42 for testing this using the credentials helper against a private registry that needs authentication to be accessed. Work towards #1357
bazelbuild · Apr 5, 2024 · 4a615be · 4a615be
1 parent b9f39bf
commit 4a615be
Show file tree

Hide file tree

Showing 18 changed files with 1,361 additions and 77 deletions.
diff --git a/.bazelrc b/.bazelrc
@@ -30,5 +30,4 @@ build:rtd --stamp
 # Some bzl files contain repos only available under bzlmod
 build:rtd --enable_bzlmod
 
-# Disabled due to https://github.com/bazelbuild/bazel/issues/20942
-build --lockfile_mode=off
+build --lockfile_mode=update
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,12 @@ A brief description of the categories of changes:
 
 ### Changed
 
+* (bzlmod): The `MODULE.bazel.lock` `whl_library` rule attributes are now
+  sorted in the attributes section. We are also removing values that are not
+  default in order to reduce the size of the lock file.
+* (deps): Bumped bazel_features to 1.9.1 to detect optional support
+  non-blocking downloads.
+
 ### Fixed
 
 * (whl_library): Fix the experimental_target_platforms overriding for platform
@@ -48,12 +54,15 @@ A brief description of the categories of changes:
 * (gazelle) Added a new `python_default_visibility` directive to control the
   _default_ visibility of generated targets. See the [docs][python_default_visibility]
   for details.
-
 * (wheel) Add support for `data_files` attributes in py_wheel rule
   ([#1777](https://github.com/bazelbuild/rules_python/issues/1777))
-
 * (py_wheel) `bzlmod` installations now provide a `twine` setup for the default
   Python toolchain in `rules_python` for version 3.11.
+* (bzlmod) New `experimental_index_url`, `experimental_extra_index_urls` and
+  `experimental_index_url_overrides` to `pip.parse` for using the bazel
+  downloader. If you see any issues, report in
+  [#1357](https://github.com/bazelbuild/rules_python/issues/1357). The URLs for
+  the whl and sdist files will be written to the lock file.
 
 [0.XX.0]: https://github.com/bazelbuild/rules_python/releases/tag/0.XX.0
 [python_default_visibility]: gazelle/README.md#directive-python_default_visibility

diff --git a/MODULE.bazel b/MODULE.bazel
@@ -4,7 +4,7 @@ module(
     compatibility_level = 1,
 )
 
-bazel_dep(name = "bazel_features", version = "1.1.1")
+bazel_dep(name = "bazel_features", version = "1.9.1")
 bazel_dep(name = "bazel_skylib", version = "1.3.0")
 bazel_dep(name = "platforms", version = "0.0.4")
 
@@ -58,6 +58,7 @@ register_toolchains("@pythons_hub//:all")
 
 pip = use_extension("//python/extensions:pip.bzl", "pip")
 pip.parse(
+    experimental_index_url = "https://pypi.org/simple",
     hub_name = "rules_python_publish_deps",
     python_version = "3.11",
     requirements_darwin = "//tools/publish:requirements_darwin.txt",
@@ -69,7 +70,7 @@ use_repo(pip, "rules_python_publish_deps")
 # ===== DEV ONLY DEPS AND SETUP BELOW HERE =====
 bazel_dep(name = "stardoc", version = "0.6.2", dev_dependency = True, repo_name = "io_bazel_stardoc")
 bazel_dep(name = "rules_bazel_integration_test", version = "0.20.0", dev_dependency = True)
-bazel_dep(name = "rules_testing", version = "0.5.0", dev_dependency = True)
+bazel_dep(name = "rules_testing", version = "0.6.0", dev_dependency = True)
 bazel_dep(name = "rules_cc", version = "0.0.9", dev_dependency = True)
 
 # Extra gazelle plugin deps so that WORKSPACE.bzlmod can continue including it for e2e tests.
@@ -83,6 +84,8 @@ dev_pip = use_extension(
     dev_dependency = True,
 )
 dev_pip.parse(
+    envsubst = ["PIP_INDEX_URL"],
+    experimental_index_url = "${PIP_INDEX_URL:-https://pypi.org/simple}",
     experimental_requirement_cycles = {
         "sphinx": [
             "sphinx",
@@ -98,6 +101,8 @@ dev_pip.parse(
     requirements_lock = "//docs/sphinx:requirements.txt",
 )
 dev_pip.parse(
+    envsubst = ["PIP_INDEX_URL"],
+    experimental_index_url = "${PIP_INDEX_URL:-https://pypi.org/simple}",
     hub_name = "pypiserver",
     python_version = "3.11",
     requirements_lock = "//examples/wheel:requirements_server.txt",

diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
@@ -94,6 +94,20 @@ use_repo(pip, "whl_mods_hub")
 # Alternatively, `python_interpreter_target` can be used to directly specify
 # the Python interpreter to run to resolve dependencies.
 pip.parse(
+    # We can use `envsubst in the above
+    envsubst = ["PIP_INDEX_URL"],
+    # Use the bazel downloader to query the simple API for downloading the sources
+    # Note, that we can use envsubst for this value.
+    experimental_index_url = "${PIP_INDEX_URL:-https://pypi.org/simple}",
+    # One can also select a particular index for a particular package.
+    # This ensures that the setup is resistant against confusion attacks.
+    # experimental_index_url_overrides = {
+    #    "my_package": "https://different-index-url.com",
+    # },
+    # Or you can specify extra indexes like with `pip`:
+    # experimental_extra_index_urls = [
+    #    "https://different-index-url.com",
+    # ],
     experimental_requirement_cycles = {
         "sphinx": [
             "sphinx",

diff --git a/internal_deps.bzl b/internal_deps.bzl
@@ -57,18 +57,9 @@ def rules_python_internal_deps():
 
     http_archive(
         name = "rules_testing",
-        sha256 = "b84ed8546f1969d700ead4546de9f7637e0f058d835e47e865dcbb13c4210aed",
-        strip_prefix = "rules_testing-0.5.0",
-        url = "https://github.com/bazelbuild/rules_testing/releases/download/v0.5.0/rules_testing-v0.5.0.tar.gz",
-    )
-
-    http_archive(
-        name = "rules_license",
-        urls = [
-            "https://mirror.bazel.build/github.com/bazelbuild/rules_license/releases/download/0.0.7/rules_license-0.0.7.tar.gz",
-            "https://github.com/bazelbuild/rules_license/releases/download/0.0.7/rules_license-0.0.7.tar.gz",
-        ],
-        sha256 = "4531deccb913639c30e5c7512a054d5d875698daeb75d8cf90f284375fe7c360",
+        sha256 = "02c62574631876a4e3b02a1820cb51167bb9cdcdea2381b2fa9d9b8b11c407c4",
+        strip_prefix = "rules_testing-0.6.0",
+        url = "https://github.com/bazelbuild/rules_testing/releases/download/v0.6.0/rules_testing-v0.6.0.tar.gz",
     )
 
     http_archive(
@@ -221,3 +212,10 @@ def rules_python_internal_deps():
         ],
         sha256 = "4531deccb913639c30e5c7512a054d5d875698daeb75d8cf90f284375fe7c360",
     )
+
+    http_archive(
+        name = "bazel_features",
+        sha256 = "d7787da289a7fb497352211ad200ec9f698822a9e0757a4976fd9f713ff372b3",
+        strip_prefix = "bazel_features-1.9.1",
+        url = "https://github.com/bazel-contrib/bazel_features/releases/download/v1.9.1/bazel_features-v1.9.1.tar.gz",
+    )
diff --git a/internal_setup.bzl b/internal_setup.bzl
@@ -14,6 +14,7 @@
 
 """Setup for rules_python tests and tools."""
 
+load("@bazel_features//:deps.bzl", "bazel_features_deps")
 load("@bazel_skylib//:workspace.bzl", "bazel_skylib_workspace")
 load("@cgrindel_bazel_starlib//:deps.bzl", "bazel_starlib_dependencies")
 load("@com_google_protobuf//:protobuf_deps.bzl", "protobuf_deps")
@@ -42,3 +43,4 @@ def rules_python_internal_setup():
     bazel_integration_test_rules_dependencies()
     bazel_starlib_dependencies()
     bazel_binaries(versions = SUPPORTED_BAZEL_VERSIONS)
+    bazel_features_deps()
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
@@ -22,6 +22,7 @@ load("//python/pip_install:requirements_parser.bzl", parse_requirements = "parse
 load("//python/pip_install/private:generate_group_library_build_bazel.bzl", "generate_group_library_build_bazel")
 load("//python/pip_install/private:generate_whl_library_build_bazel.bzl", "generate_whl_library_build_bazel")
 load("//python/pip_install/private:srcs.bzl", "PIP_INSTALL_PY_SRCS")
+load("//python/private:auth.bzl", "AUTH_ATTRS", "get_auth")
 load("//python/private:envsubst.bzl", "envsubst")
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:parse_whl_name.bzl", "parse_whl_name")
@@ -187,7 +188,7 @@ def use_isolated(ctx, attr):
 
     return use_isolated
 
-def _parse_optional_attrs(rctx, args):
+def _parse_optional_attrs(rctx, args, extra_pip_args = None):
     """Helper function to parse common attributes of pip_repository and whl_library repository rules.
 
     This function also serializes the structured arguments as JSON
@@ -196,6 +197,7 @@ def _parse_optional_attrs(rctx, args):
     Args:
         rctx: Handle to the rule repository context.
         args: A list of parsed args for the rule.
+        extra_pip_args: The pip args to pass.
     Returns: Augmented args list.
     """
 
@@ -212,7 +214,7 @@ def _parse_optional_attrs(rctx, args):
 
     # Check for None so we use empty default types from our attrs.
     # Some args want to be list, and some want to be dict.
-    if rctx.attr.extra_pip_args != None:
+    if extra_pip_args != None:
         args += [
             "--extra_pip_args",
             json.encode(struct(arg = [
@@ -759,24 +761,64 @@ def _whl_library_impl(rctx):
         "--requirement",
         rctx.attr.requirement,
     ]
-
-    args = _parse_optional_attrs(rctx, args)
+    extra_pip_args = []
+    extra_pip_args.extend(rctx.attr.extra_pip_args)
 
     # Manually construct the PYTHONPATH since we cannot use the toolchain here
     environment = _create_repository_execution_environment(rctx, python_interpreter)
 
-    repo_utils.execute_checked(
-        rctx,
-        op = "whl_library.ResolveRequirement({}, {})".format(rctx.attr.name, rctx.attr.requirement),
-        arguments = args,
-        environment = environment,
-        quiet = rctx.attr.quiet,
-        timeout = rctx.attr.timeout,
-    )
+    whl_path = None
+    if rctx.attr.whl_file:
+        whl_path = rctx.path(rctx.attr.whl_file)
+
+        # Simulate the behaviour where the whl is present in the current directory.
+        rctx.symlink(whl_path, whl_path.basename)
+        whl_path = rctx.path(whl_path.basename)
+    elif rctx.attr.urls:
+        filename = rctx.attr.filename
+        urls = rctx.attr.urls
+        if not filename:
+            _, _, filename = urls[0].rpartition("/")
+
+        if not (filename.endswith(".whl") or filename.endswith("tar.gz") or filename.endswith(".zip")):
+            if rctx.attr.filename:
+                msg = "got '{}'".format(filename)
+            else:
+                msg = "detected '{}' from url:\n{}".format(filename, urls[0])
+            fail("Only '.whl', '.tar.gz' or '.zip' files are supported, {}".format(msg))
+
+        result = rctx.download(
+            url = urls,
+            output = filename,
+            sha256 = rctx.attr.sha256,
+            auth = get_auth(rctx, urls),
+        )
+
+        if not result.success:
+            fail("could not download the '{}' from {}:\n{}".format(filename, urls, result))
+
+        if filename.endswith(".whl"):
+            whl_path = rctx.path(rctx.attr.filename)
+        else:
+            # It is an sdist and we need to tell PyPI to use a file in this directory
+            # and not use any indexes.
+            extra_pip_args.extend(["--no-index", "--find-links", "."])
+
+    args = _parse_optional_attrs(rctx, args, extra_pip_args)
 
-    whl_path = rctx.path(json.decode(rctx.read("whl_file.json"))["whl_file"])
-    if not rctx.delete("whl_file.json"):
-        fail("failed to delete the whl_file.json file")
+    if not whl_path:
+        repo_utils.execute_checked(
+            rctx,
+            op = "whl_library.ResolveRequirement({}, {})".format(rctx.attr.name, rctx.attr.requirement),
+            arguments = args,
+            environment = environment,
+            quiet = rctx.attr.quiet,
+            timeout = rctx.attr.timeout,
+        )
+
+        whl_path = rctx.path(json.decode(rctx.read("whl_file.json"))["whl_file"])
+        if not rctx.delete("whl_file.json"):
+            fail("failed to delete the whl_file.json file")
 
     if rctx.attr.whl_patches:
         patches = {}
@@ -890,14 +932,18 @@ if __name__ == "__main__":
     )
     return contents
 
-whl_library_attrs = {
+# NOTE @aignas 2024-03-21: The usage of dict({}, **common) ensures that all args to `dict` are unique
+whl_library_attrs = dict({
     "annotation": attr.label(
         doc = (
             "Optional json encoded file containing annotation to apply to the extracted wheel. " +
             "See `package_annotation`"
         ),
         allow_files = True,
     ),
+    "filename": attr.string(
+        doc = "Download the whl file to this filename. Only used when the `urls` is passed. If not specified, will be auto-detected from the `urls`.",
+    ),
     "group_deps": attr.string_list(
         doc = "List of dependencies to skip in order to break the cycles within a dependency group.",
         default = [],
@@ -911,7 +957,18 @@ whl_library_attrs = {
     ),
     "requirement": attr.string(
         mandatory = True,
-        doc = "Python requirement string describing the package to make available",
+        doc = "Python requirement string describing the package to make available, if 'urls' or 'whl_file' is given, then this only needs to include foo[any_extras] as a bare minimum.",
+    ),
+    "sha256": attr.string(
+        doc = "The sha256 of the downloaded whl. Only used when the `urls` is passed.",
+    ),
+    "urls": attr.string_list(
+        doc = """\
+The list of urls of the whl to be downloaded using bazel downloader. Using this
+attr makes `extra_pip_args` and `download_only` ignored.""",
+    ),
+    "whl_file": attr.label(
+        doc = "The whl file that should be used instead of downloading or building the whl.",
     ),
     "whl_patches": attr.label_keyed_string_dict(
         doc = """a label-keyed-string dict that has
@@ -933,9 +990,8 @@ whl_library_attrs = {
             for repo in all_requirements
         ],
     ),
-}
-
-whl_library_attrs.update(**common_attrs)
+}, **common_attrs)
+whl_library_attrs.update(AUTH_ATTRS)
 
 whl_library = repository_rule(
     attrs = whl_library_attrs,

diff --git a/python/private/BUILD.bazel b/python/private/BUILD.bazel
@@ -119,6 +119,18 @@ bzl_library(
     srcs = ["parse_whl_name.bzl"],
 )
 
+bzl_library(
+    name = "pypi_index_bzl",
+    srcs = ["pypi_index.bzl"],
+    deps = [
+        ":auth_bzl",
+        ":normalize_name_bzl",
+        ":text_util_bzl",
+        "//python/pip_install:requirements_parser_bzl",
+        "//python/private/bzlmod:bazel_features_bzl",
+    ],
+)
+
 bzl_library(
     name = "py_cc_toolchain_bzl",
     srcs = [
@@ -260,6 +272,9 @@ bzl_library(
     name = "whl_target_platforms_bzl",
     srcs = ["whl_target_platforms.bzl"],
     visibility = ["//:__subpackages__"],
+    deps = [
+        "parse_whl_name_bzl",
+    ],
 )
 
 bzl_library(