Coverage for cli/_image_uri.py: 88.46%

1""" 

2ECR image URI helpers backed by local AWS partition metadata. 

3 

4Lives in its own small module so both ``cli.images`` (which builds and 

5manages images) and ``cli.inference`` (which has to rewrite URIs to 

6target the local region of each deployed endpoint) can depend on it 

7without forming an import cycle. 

8 

9Static-analysis tools (CodeQL, pyright) flag deferred-import cycles 

10even when both imports happen inside method bodies, because the 

11resulting module-level dependency graph still has a cycle. Splitting 

12the helper out keeps the dependency graph a DAG: ``cli.images`` and 

13``cli.inference`` both depend on ``cli._image_uri``, and neither 

14depends on the other. Partition and DNS suffix resolution uses 

15botocore's bundled endpoint metadata and makes no AWS API calls. 

16""" 

17 

18from __future__ import annotations 

19 

20import re 

21from functools import cache 

22 

23import botocore.session 

24 

25# ECR registry host shape in any AWS partition: 

26# <account-id>.dkr.ecr.<region>.<partition-url-suffix> 

27_ECR_HOST_RE = re.compile( 

28 r"^(?P<account>\d+)\.dkr\.ecr\." 

29 r"(?P<region>[a-z]{2,4}(?:-[a-z0-9]+)+-[0-9]+)\." 

30 r"(?P<suffix>[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?)$" 

31) 

32 

33 

34@cache 

35def _partition_metadata(region: str) -> tuple[str, str]: 

36 """Return ``(partition, URL suffix)`` from botocore's local metadata.""" 

37 if not region: 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true

38 raise ValueError("AWS region must not be empty") 

39 

40 resolver = botocore.session.get_session().get_component("endpoint_resolver") 

41 partition = resolver.get_partition_for_region(region) 

42 if not partition: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true

43 raise ValueError(f"Could not resolve an AWS partition for region {region!r}") 

44 url_suffix = resolver.get_partition_dns_suffix(partition) 

45 if not url_suffix: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true

46 raise ValueError(f"Could not resolve the URL suffix for AWS partition {partition!r}") 

47 return str(partition), str(url_suffix) 

48 

49 

50def aws_partition(region: str) -> str: 

51 """Return the ARN partition for ``region`` using botocore metadata.""" 

52 return _partition_metadata(region)[0] 

53 

54 

55def aws_url_suffix(region: str) -> str: 

56 """Return the AWS DNS suffix for ``region`` using botocore metadata.""" 

57 return _partition_metadata(region)[1] 

58 

59 

60def ecr_registry_host(account_id: str, region: str) -> str: 

61 """Return the partition-correct private ECR registry hostname.""" 

62 return f"{account_id}.dkr.ecr.{region}.{aws_url_suffix(region)}" 

63 

64 

65def rewrite_image_uri_for_region(uri: str, region: str) -> str: 

66 """Rewrite an ECR image URI to target a specific region's replica. 

67 

68 Pure helper — no AWS calls. Detects ECR URIs by matching the 

69 ``<account>.dkr.ecr.<region>.<partition-url-suffix>`` host shape and 

70 validating that suffix against botocore's bundled endpoint metadata. 

71 It then replaces both the region and suffix for the target region, so 

72 same-partition and cross-partition rewrites cannot retain a stale DNS 

73 suffix. Non-ECR refs (Docker Hub, GHCR, etc.) are returned unchanged. 

74 

75 Args: 

76 uri: The image URI (with optional ``host/path:tag`` shape). 

77 region: Target AWS region for the rewrite. 

78 

79 Returns: 

80 The rewritten URI when the input is an ECR URI; otherwise the 

81 original input. 

82 """ 

83 if "://" in uri: 

84 # Not a bare image ref (looks like a URL with a scheme). 

85 return uri 

86 parts = uri.split("/", 1) 

87 host = parts[0] 

88 match = _ECR_HOST_RE.match(host) 

89 if match is None: 

90 return uri 

91 

92 source_region = match.group("region") 

93 if match.group("suffix") != aws_url_suffix(source_region): 

94 # The hostname has ECR-like labels but is not an AWS ECR endpoint. 

95 return uri 

96 

97 new_host = ecr_registry_host(match.group("account"), region) 

98 if len(parts) > 1: 

99 return f"{new_host}/{parts[1]}" 

100 return new_host