Coverage for cli / _image_uri.py: 100.00%

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: 

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

39 

40 # Public botocore surface only: ``Session.get_partition_for_region`` and 

41 # the bundled ``endpoints`` data file. The endpoint-resolver component 

42 # that used to answer both questions is internal, and fetching it through 

43 # ``get_component`` raises a DeprecationWarning (an error under the test 

44 # suite's warning filters). 

45 session = botocore.session.get_session() 

46 partition = session.get_partition_for_region(region) 

47 if not partition: 

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

49 partitions = session.get_data("endpoints").get("partitions", []) 

50 url_suffix = next( 

51 (entry.get("dnsSuffix") for entry in partitions if entry.get("partition") == partition), 

52 None, 

53 ) 

54 if not url_suffix: 

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

56 return str(partition), str(url_suffix) 

57 

58 

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

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

61 return _partition_metadata(region)[0] 

62 

63 

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

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

66 return _partition_metadata(region)[1] 

67 

68 

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

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

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

72 

73 

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

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

76 

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

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

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

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

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

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

83 

84 Args: 

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

86 region: Target AWS region for the rewrite. 

87 

88 Returns: 

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

90 original input. 

91 """ 

92 if "://" in uri: 

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

94 return uri 

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

96 host = parts[0] 

97 match = _ECR_HOST_RE.match(host) 

98 if match is None: 

99 return uri 

100 

101 source_region = match.group("region") 

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

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

104 return uri 

105 

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

107 if len(parts) > 1: 

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

109 return new_host