AWS CodeArtifact Support

3rdparty/python/requirements.txt

@@ -60,3 +60,7 @@ python-gnupg==0.4.9 # For validating signatures
 
 # Only used for release management
 PyGithub==2.4.0
+
+# For AWS CodeArtifact support.
+# TODO: Figure out why the backend-specific requirements.txt was not loading this into project venv.
+boto3>=1.35.62

3rdparty/python/user_reqs.lock

@@ -13,6 +13,7 @@
 //     "PyYAML<7.0,>=6.0",
 //     "ansicolors==1.1.8",
 //     "beautifulsoup4==4.11.1",
+//     "boto3>=1.35.62",
 //     "chevron==0.14.0",
 //     "debugpy==1.6.0",
 //     "fastapi==0.78.0",
@@ -57,7 +58,6 @@
   "allow_wheels": true,
   "build_isolation": true,
   "constraints": [],
-  "elide_unused_requires_dist": false,
   "excluded": [],
   "locked_resolves": [
     {
@@ -242,6 +242,53 @@
           "requires_python": ">=3.6.0",
           "version": "4.11.1"
         },
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "eb21380d73fec6645439c0d802210f72a0cdb3295b02953f246ff53f512faa8f",
+              "url": "https://files.pythonhosted.org/packages/2b/ed/464e1df3901fbfedd5a0786e551240216f0c867440fa6156595178227b3f/boto3-1.36.1-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "258ab77225a81d3cf3029c9afe9920cd9dec317689dfadec6f6f0a23130bb60a",
+              "url": "https://files.pythonhosted.org/packages/bf/04/0c6cea060653eee75f4348152dfc0aa0b241f7d1f99a530079ee44d61e4b/boto3-1.36.1.tar.gz"
+            }
+          ],
+          "project_name": "boto3",
+          "requires_dists": [
+            "botocore<1.37.0,>=1.36.1",
+            "botocore[crt]<2.0a0,>=1.21.0; extra == \"crt\"",
+            "jmespath<2.0.0,>=0.7.1",
+            "s3transfer<0.12.0,>=0.11.0"
+          ],
+          "requires_python": ">=3.8",
+          "version": "1.36.1"
+        },
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "dec513b4eb8a847d79bbefdcdd07040ed9d44c20b0001136f0890a03d595705a",
+              "url": "https://files.pythonhosted.org/packages/be/bb/5431f12e2dadd881fd023fb57e7e3ab82f7b697c38dc837fc8d70cca51bd/botocore-1.36.1-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "f789a6f272b5b3d8f8756495019785e33868e5e00dd9662a3ee7959ac939bb12",
+              "url": "https://files.pythonhosted.org/packages/39/aa/556720b3ee9629b7c4366b5a0d9797a84e83a97f78435904cbb9bdc41939/botocore-1.36.1.tar.gz"
+            }
+          ],
+          "project_name": "botocore",
+          "requires_dists": [
+            "awscrt==0.23.4; extra == \"crt\"",
+            "jmespath<2.0.0,>=0.7.1",
+            "python-dateutil<3.0.0,>=2.1",
+            "urllib3!=2.2.0,<3,>=1.25.4; python_version >= \"3.10\"",
+            "urllib3<1.27,>=1.25.4; python_version < \"3.10\""
+          ],
+          "requires_python": ">=3.8",
+          "version": "1.36.1"
+        },
         {
           "artifacts": [
             {
@@ -885,6 +932,24 @@
           "requires_python": ">=3.7",
           "version": "2.0.0"
         },
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980",
+              "url": "https://files.pythonhosted.org/packages/31/b4/b9b800c45527aadd64d5b442f9b932b00648617eb5d63d2c7a6587b7cafc/jmespath-1.0.1-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe",
+              "url": "https://files.pythonhosted.org/packages/00/2a/e867e8531cf3e36b41201936b7fa7ba7b5702dbef42922193f05c8976cd6/jmespath-1.0.1.tar.gz"
+            }
+          ],
+          "project_name": "jmespath",
+          "requires_dists": [],
+          "requires_python": ">=3.7",
+          "version": "1.0.1"
+        },
         {
           "artifacts": [
             {
@@ -1546,6 +1611,27 @@
           "requires_python": ">=3.8",
           "version": "2.32.3"
         },
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "8fa0aa48177be1f3425176dfe1ab85dcd3d962df603c3dbfc585e6bf857ef0ff",
+              "url": "https://files.pythonhosted.org/packages/5f/ce/22673f4a85ccc640735b4f8d12178a0f41b5d3c6eda7f33756d10ce56901/s3transfer-0.11.1-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "3f25c900a367c8b7f7d8f9c34edc87e300bde424f779dc9f0a8ae4f9df9264f6",
+              "url": "https://files.pythonhosted.org/packages/1a/aa/fdd958c626b00e3f046d4004363e7f1a2aba4354f78d65ceb3b217fa5eb8/s3transfer-0.11.1.tar.gz"
+            }
+          ],
+          "project_name": "s3transfer",
+          "requires_dists": [
+            "botocore<2.0a.0,>=1.36.0",
+            "botocore[crt]<2.0a.0,>=1.36.0; extra == \"crt\""
+          ],
+          "requires_python": ">=3.8",
+          "version": "0.11.1"
+        },
         {
           "artifacts": [
             {
@@ -2319,14 +2405,15 @@
   "only_wheels": [],
   "overridden": [],
   "path_mappings": {},
-  "pex_version": "2.29.0",
+  "pex_version": "2.24.1",
   "pip_version": "24.3.1",
   "prefer_older_binary": false,
   "requirements": [
     "PyGithub==2.4.0",
     "PyYAML<7.0,>=6.0",
     "ansicolors==1.1.8",
     "beautifulsoup4==4.11.1",
+    "boto3>=1.35.62",
     "chevron==0.14.0",
     "debugpy==1.6.0",
     "fastapi==0.78.0",

docs/notes/2.25.x.md

@@ -162,6 +162,8 @@ The version of Python used by Pants itself is now [3.11](https://docs.python.org
 
 The oldest [glibc version](https://www.sourceware.org/glibc/wiki/Glibc%20Timeline) supported by the published Pants wheels is now 2.28.  This should have no effect unless you are running on extremely old Linux distributions.  See <https://github.com/pypa/manylinux> for background context on Python wheels and C libraries.
 
+Added `PexKeyringConfigurationRequest`, a new API for plugins to supply credentials to Pex/Pip for when they access secured Python package indexes.
+
 
 ## Full Changelog
 

src/python/pants/backend/experimental/python/aws_codeartifact_support/BUILD

@@ -0,0 +1,4 @@
+# Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_sources()

src/python/pants/backend/experimental/python/aws_codeartifact_support/register.py

@@ -0,0 +1,10 @@
+# Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+from pants.backend.python.util_rules import aws_codeartifact
+
+
+def rules():
+    return aws_codeartifact.rules()

src/python/pants/backend/python/subsystems/aws_codeartifact.py

@@ -0,0 +1,44 @@
+# Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+from pants.option.option_types import BoolOption, StrOption
+from pants.option.subsystem import Subsystem
+from pants.util.strutil import help_text, softwrap
+
+
+class PythonAwsCodeartifact(Subsystem):
+    options_scope = "python-aws-codeartifact"
+
+    help = help_text(
+        """
+        AWS CodeArtifact configuration
+
+        These options are used to configure Pants to obtain and renew an AWS CodeArtifact token to access
+        a PyPi-compatible CodeArtifact repository.
+        """
+    )
+
+    enabled = BoolOption(
+        default=False,
+        help=softwrap(
+            """
+            If True, Pants will renew the AWS CodeArtifact token if it has expired. The token will be made
+            availbale to Pex/Pip automatically.
+            """
+        ),
+    )
+
+    domain = StrOption(
+        default="", help="AWS CodeArtifact domain containing the relevant repositories"
+    )
+
+    domain_owner = StrOption(
+        default=None,
+        help=softwrap(
+            """
+            If set, the value will be used as the `domainOwner` parameter passed to AWS CodeArtifact's GetAuthorizationToken API.
+            """
+        ),
+    )

src/python/pants/backend/python/util_rules/aws_codeartifact.py

@@ -0,0 +1,176 @@
+# Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+import datetime as dt
+import json
+import logging
+import urllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import boto3  # type: ignore[import-untyped]
+
+from pants.backend.python.subsystems.aws_codeartifact import PythonAwsCodeartifact
+from pants.backend.python.subsystems.repos import PythonRepos
+from pants.backend.python.util_rules.pex_cli import (
+    PexKeyringConfigurationRequest,
+    PexKeyringConfigurationResponse,
+)
+from pants.base.build_root import BuildRoot
+from pants.engine.internals.selectors import Get
+from pants.engine.rules import _uncacheable_rule, collect_rules, rule
+from pants.engine.unions import UnionRule
+from pants.util.frozendict import FrozenDict
+
+logger = logging.getLogger(__name__)
+
+
+# The duration before key expiration at which renewal will be triggered.
+_RENEWAL_WINDOW = dt.timedelta(minutes=30)
+
+
+@dataclass(frozen=True)
+class AuthToken:
+    token: str
+    expires: dt.datetime
+
+    def to_json_dict(self) -> dict[str, str]:
+        return {
+            "token": self.token,
+            "expires": self.expires.isoformat(),
+        }
+
+
+def _aws_codeartifact_dir() -> Path:
+    build_root: Path = BuildRoot().pathlib_path
+    return build_root / ".pants.d" / "aws" / "codeartifact"
+
+
+def _load_token() -> AuthToken | None:
+    aws_codeartifact_dir = _aws_codeartifact_dir()
+    aws_codeartifact_auth_cache = aws_codeartifact_dir / "auth-cache.json"
+    if not aws_codeartifact_auth_cache.exists():
+        return None
+
+    try:
+        raw_data = aws_codeartifact_auth_cache.read_bytes()
+        data = json.loads(raw_data)
+    except Exception as e:
+        logger.debug(f"CodeArtifact auth cache was not readable: {e}")
+        return None
+
+    if not isinstance(data, dict):
+        logger.debug("CodeArtifact auth cache was not a JSON object.")
+        return None
+
+    token: Any = data.get("token")
+    expires: Any = data.get("expires")
+    if token is None or expires is None:
+        logger.debug("CodeArtifact auth cache did not have all required fields.")
+        return None
+
+    if not isinstance(token, str) or not isinstance(expires, str):
+        logger.debug("CodeArtifact auth cache did not have all required fields with correct types.")
+        return None
+
+    return AuthToken(token=token, expires=dt.datetime.fromisoformat(expires))
+
+
+def _save_token(auth_token: AuthToken) -> None:
+    aws_codeartifact_dir = _aws_codeartifact_dir()
+    aws_codeartifact_dir.mkdir(parents=True, exist_ok=True)
+
+    aws_codeartifact_auth_cache = aws_codeartifact_dir / "auth-cache.json"
+    data = json.dumps(auth_token.to_json_dict()).encode()
+    aws_codeartifact_auth_cache.write_bytes(data)
+
+
+def _codeartifact_login(domain: str, domain_owner: str | None) -> AuthToken:
+    logger.debug("Logging in to AWS CodeArtifact.")
+    codeartifact = boto3.client("codeartifact")
+
+    kwargs = {}
+    if domain_owner:
+        kwargs["domainOwner"] = domain_owner
+    response = codeartifact.get_authorization_token(domain=domain, **kwargs)
+
+    logger.debug("Logged in to AWS CodeArtifact.")
+    return AuthToken(token=response["authorizationToken"], expires=response["expiration"])
+
+
+def _ensure_aws_codeartifact_login(options: PythonAwsCodeartifact) -> AuthToken:
+    auth_token = _load_token()
+    if auth_token is not None:
+        if dt.datetime.now(dt.timezone.utc) < auth_token.expires - _RENEWAL_WINDOW:
+            return auth_token
+
+    auth_token = _codeartifact_login(options.domain, options.domain_owner)
+    # TODO: Error handling and retry logic.
+    _save_token(auth_token)
+    return auth_token
+
+
+@dataclass(frozen=True)
+class _AwsCodeArtifactLogin:
+    auth_token: AuthToken | None
+
+
+# This rule is uncacheable so that it will every session to ensure the AWS CodeArtifact token is
+# renewed if need be.
+@_uncacheable_rule
+async def aws_codeartifact_login_rule(
+    codeartifact_subsystem: PythonAwsCodeartifact,
+) -> _AwsCodeArtifactLogin:
+    if codeartifact_subsystem.enabled:
+        auth_token = _ensure_aws_codeartifact_login(codeartifact_subsystem)
+        return _AwsCodeArtifactLogin(auth_token)
+    else:
+        return _AwsCodeArtifactLogin(None)
+
+
+class AwsCodeArtifactPexKeyringConfigurationRequest(PexKeyringConfigurationRequest):
+    pass
+
+
+@rule
+async def aws_code_artifact_pex_keyring_configuration_request(
+    _request: AwsCodeArtifactPexKeyringConfigurationRequest,
+    codeartifact_subsystem: PythonAwsCodeartifact,
+    python_repos: PythonRepos,
+) -> PexKeyringConfigurationResponse:
+    # Configure the AWS CodeArtifact token if any reposiory is in AWS CodeArtfact. We heurisitcally look
+    # for the string "codeartifact" in the package index URLs.
+    codeartifact_repo_urls = [
+        repo
+        for repo in [*python_repos.indexes, *python_repos.find_links]
+        if repo.find("codeartifact") >= 0
+    ]
+    if not codeartifact_repo_urls or not codeartifact_subsystem.enabled:
+        return PexKeyringConfigurationResponse(credentials=None)
+
+    auth_token_wrapper = await Get(_AwsCodeArtifactLogin)
+    auth_token = auth_token_wrapper.auth_token
+
+    credentials: dict[str, tuple[str, str]] = {}
+    if auth_token:
+        for codeartifact_repo_url in codeartifact_repo_urls:
+            parsed_url = urllib.parse.urlparse(codeartifact_repo_url)
+            username = parsed_url.username
+            hostname = parsed_url.hostname
+            if not hostname or not username:
+                # TODO: Warn or error if requried parts of the URL are missing?
+                continue
+            credentials[hostname] = (username, auth_token.token)
+
+    return PexKeyringConfigurationResponse(credentials=FrozenDict(credentials))
+
+
+def rules():
+    return [
+        *collect_rules(),
+        *PythonAwsCodeartifact.rules(),
+        UnionRule(PexKeyringConfigurationRequest, AwsCodeArtifactPexKeyringConfigurationRequest),
+    ]