Skip to content

Commit

Permalink
Set user-agent & retries for OIDC requests
Browse files Browse the repository at this point in the history
  • Loading branch information
takluyver committed Dec 10, 2024
1 parent 9c4506c commit ec859fb
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 33 deletions.
3 changes: 1 addition & 2 deletions twine/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from typing import TYPE_CHECKING, Callable, Optional, Type, cast
from urllib.parse import urlparse

import requests
from id import AmbientCredentialError # type: ignore
from id import detect_credential

Expand Down Expand Up @@ -77,7 +76,7 @@ def make_trusted_publishing_token(self) -> Optional[str]:
# Trusted publishing (OpenID Connect): get one token from the CI
# system, and exchange that for a PyPI token.
repository_domain = cast(str, urlparse(self.system).netloc)
session = requests.Session() # TODO: user agent & retries
session = utils.make_requests_session()

# Indices are expected to support `https://{domain}/_/oidc/audience`,
# which tells OIDC exchange clients which audience to use.
Expand Down
34 changes: 3 additions & 31 deletions twine/repository.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from typing import Any, Dict, List, Optional, Set, Tuple, cast
from typing import Any, Dict, List, Optional, Set, Tuple

import requests
import requests_toolbelt
import rich.progress
import urllib3
from requests import adapters
from requests_toolbelt.utils import user_agent
from rich import print

import twine
from twine import package as package_file
from twine.utils import make_requests_session

KEYWORDS_TO_NOT_FLATTEN = {"gpg_signature", "attestations", "content"}

Expand All @@ -47,7 +44,7 @@ def __init__(
) -> None:
self.url = repository_url

self.session = requests.session()
self.session = make_requests_session()
# requests.Session.auth should be Union[None, Tuple[str, str], ...]
# But username or password could be None
# See TODO for utils.RepositoryConfig
Expand All @@ -57,35 +54,10 @@ def __init__(
logger.info(f"username: {username if username else '<empty>'}")
logger.info(f"password: <{'hidden' if password else 'empty'}>")

self.session.headers["User-Agent"] = self._make_user_agent_string()
for scheme in ("http://", "https://"):
self.session.mount(scheme, self._make_adapter_with_retries())

# Working around https://github.com/python/typing/issues/182
self._releases_json_data: Dict[str, Dict[str, Any]] = {}
self.disable_progress_bar = disable_progress_bar

@staticmethod
def _make_adapter_with_retries() -> adapters.HTTPAdapter:
retry = urllib3.Retry(
allowed_methods=["GET"],
connect=5,
total=10,
status_forcelist=[500, 501, 502, 503],
)

return adapters.HTTPAdapter(max_retries=retry)

@staticmethod
def _make_user_agent_string() -> str:
user_agent_string = (
user_agent.UserAgentBuilder("twine", twine.__version__)
.include_implementation()
.build()
)

return cast(str, user_agent_string)

def close(self) -> None:
self.session.close()

Expand Down
26 changes: 26 additions & 0 deletions twine/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,11 @@

import requests
import rfc3986
import urllib3
from requests.adapters import HTTPAdapter
from requests_toolbelt.utils import user_agent

import twine
from twine import exceptions

# Shim for input to allow testing.
Expand Down Expand Up @@ -304,6 +308,28 @@ def get_userpass_value(
get_clientcert = functools.partial(get_userpass_value, key="client_cert")


def make_requests_session() -> requests.Session:
"""Prepare a requests Session with retries & twine's user-agent string."""
s = requests.Session()

retry = urllib3.Retry(
allowed_methods=["GET"],
connect=5,
total=10,
status_forcelist=[500, 501, 502, 503],
)

for scheme in ("http://", "https://"):
s.mount(scheme, HTTPAdapter(max_retries=retry))

s.headers["User-Agent"] = (
user_agent.UserAgentBuilder("twine", twine.__version__)
.include_implementation()
.build()
)
return s


class EnvironmentDefault(argparse.Action):
"""Get values from environment variable."""

Expand Down

0 comments on commit ec859fb

Please sign in to comment.