ahi-aws 0.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,25 @@
1
+ # Python artifacts
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.so
5
+ *.egg-info/
6
+ .dist/
7
+ build/
8
+ dist/
9
+
10
+ # Virtual environments
11
+ .venv/
12
+
13
+ # Tool caches
14
+ .pytest_cache/
15
+ .mypy_cache/
16
+ .ruff_cache/
17
+ .coverage
18
+ htmlcov/
19
+
20
+ # IDEs/OS
21
+ .DS_Store
22
+ .vscode/*.code-snippets
23
+
24
+ # Docs build
25
+ site/
@@ -0,0 +1,29 @@
1
+ # Changelog
2
+
3
+ ## [0.4.0](https://github.com/adaptivehealthintelligence/ahi-python-packages/compare/ahi-aws-v0.3.0...ahi-aws-v0.4.0) (2026-11-10)
4
+
5
+
6
+ ### Features
7
+
8
+ * **LHI-1728:** add ahi-aws auth helpers for SSO and attached roles ([7e2c7fe](https://github.com/adaptivehealthintelligence/ahi-python-packages/commit/7e2c7fed2ce86bbd95f0b0036120c9c5bf1cf050))
9
+
10
+ ## [0.3.0](https://github.com/adaptivehealthintelligence/ahi-python-packages/compare/ahi-aws-v0.2.1...ahi-aws-v0.3.0) (2026-10-21)
11
+
12
+
13
+ ### Features
14
+
15
+ * **ahi-aws:** add S3 storage helpers ([0a71bf6](https://github.com/adaptivehealthintelligence/ahi-python-packages/commit/0a71bf6e17bc43f13844ff2a47c5baad23924c33))
16
+
17
+ ## [0.2.1](https://github.com/adaptivehealthintelligence/ahi-python-packages/compare/ahi-aws-v0.2.0...ahi-aws-v0.2.1) (2026-06-28)
18
+
19
+
20
+ ### Bug Fixes
21
+
22
+ * **ahi-aws:** allow ahi-core 0.2 ([505ea28](https://github.com/adaptivehealthintelligence/ahi-python-packages/commit/505ea28d5bb3e154487d32af1e742fdf79299c83))
23
+
24
+ ## [0.2.0](https://github.com/adaptivehealthintelligence/ahi-python-packages/compare/ahi-aws-v0.1.0...ahi-aws-v0.2.0) (2026-06-27)
25
+
26
+
27
+ ### Features
28
+
29
+ * scaffold uv workspace with core, sql, aws packages ([8c42f2f](https://github.com/adaptivehealthintelligence/ahi-python-packages/commit/8c42f2f371284180bf0569785837640e39deb80b))
ahi_aws-0.4.0/PKG-INFO ADDED
@@ -0,0 +1,80 @@
1
+ Metadata-Version: 2.5
2
+ Name: ahi-aws
3
+ Version: 0.4.0
4
+ Summary: Thin, typed boto3 helpers for AHI: AWS auth (SSO and attached-role sessions), Redshift Data API and Athena query execution, and S3 file storage.
5
+ Project-URL: Homepage, https://github.com/adaptivehealthintelligence/ahi-python-packages
6
+ Project-URL: Documentation, https://adaptivehealthintelligence.github.io/ahi-python-packages/
7
+ Project-URL: Repository, https://github.com/adaptivehealthintelligence/ahi-python-packages
8
+ Author: Adaptive Health Intelligence
9
+ License: MIT
10
+ Keywords: ahi,athena,auth,aws,boto3,redshift,s3,sso,sts
11
+ Requires-Python: >=3.11
12
+ Requires-Dist: ahi-core<0.3,>=0.2
13
+ Requires-Dist: boto3>=1.34
14
+ Provides-Extra: types
15
+ Requires-Dist: boto3-stubs[athena,redshift-data,s3,sts]>=1.34; extra == 'types'
16
+ Description-Content-Type: text/markdown
17
+
18
+ # ahi-aws
19
+
20
+ Thin, typed `boto3` helpers for AHI. Today: AWS **auth** (one session, checked
21
+ once, with clear errors), running SQL on **Redshift** (Data API) and **Athena**,
22
+ with async query handling, and moving files and folders to and from **S3**.
23
+
24
+ ```python
25
+ from ahi.aws.auth import aws_session, check_identity
26
+
27
+ session = aws_session() # AWS_PROFILE / default profile / attached role
28
+ identity = check_identity(session) # raises an AuthError saying what to run
29
+ print(identity.arn, identity.method) # e.g. "sso", "iam-role", "container-role"
30
+ ```
31
+
32
+ ```python
33
+ from ahi.aws.redshift import RedshiftDataClient
34
+ from ahi.aws.athena import AthenaClient
35
+
36
+ rs = RedshiftDataClient(workgroup_name="analytics", database="dev")
37
+ rows = rs.run_query(
38
+ "select count(*) from events where day = :day",
39
+ parameters={"day": "2026-06-26"},
40
+ )
41
+ print(rows.scalar())
42
+
43
+ ath = AthenaClient(workgroup="primary", database="analytics")
44
+ for row in ath.run_query("select name from users where id = ?", parameters=[42]):
45
+ print(row["name"])
46
+ ```
47
+
48
+ ```python
49
+ from ahi.aws.auth import aws_session
50
+ from ahi.aws.s3 import S3Storage
51
+
52
+ session = aws_session()
53
+ storage = S3Storage("s3://my-bucket/some/prefix", client=session.client("s3"))
54
+
55
+ uploaded = storage.upload_dir("runs/run-01", prefix="run-01") # list[S3Object]
56
+ storage.download_dir("run-01", "local/run-01")
57
+ ```
58
+
59
+ ## Design notes
60
+
61
+ - **Never logs in for you.** No browser, no AWS CLI. Run `aws sso login`
62
+ beforehand (local SSO), or run where an AWS role is attached (EC2 instance
63
+ role, ECS task role); missing or expired credentials raise `AuthError`
64
+ with the command to run. No default region: set `AWS_REGION` or the
65
+ profile's `region`.
66
+ - **Parameterized queries by default.** Pass values via `parameters` (named
67
+ `:name` for Redshift, positional `?` for Athena) — never format them in.
68
+ - **One synchronous call.** `run_query` submits, polls with backoff (see
69
+ `ahi.core.poll_until`), and pages the full result set, returning a
70
+ `QueryResult` you can iterate, `len()`, or call `.scalar()` on.
71
+ - **Typed results from Redshift; strings from Athena** (Athena returns all values
72
+ as strings — we preserve that rather than guessing).
73
+ - **Failures raise** `QueryFailedError`; timeouts raise `PollTimeoutError`.
74
+ - **S3 transfers return `S3Object` records** (bucket, key, local path, size);
75
+ failures raise `UploadFailedError` / `DownloadFailedError`.
76
+
77
+ Install the optional `types` extra (`ahi-aws[types]`) for `boto3-stubs` during
78
+ development.
79
+
80
+ Part of [ahi-python-packages](https://github.com/adaptivehealthintelligence/ahi-python-packages).
@@ -0,0 +1,63 @@
1
+ # ahi-aws
2
+
3
+ Thin, typed `boto3` helpers for AHI. Today: AWS **auth** (one session, checked
4
+ once, with clear errors), running SQL on **Redshift** (Data API) and **Athena**,
5
+ with async query handling, and moving files and folders to and from **S3**.
6
+
7
+ ```python
8
+ from ahi.aws.auth import aws_session, check_identity
9
+
10
+ session = aws_session() # AWS_PROFILE / default profile / attached role
11
+ identity = check_identity(session) # raises an AuthError saying what to run
12
+ print(identity.arn, identity.method) # e.g. "sso", "iam-role", "container-role"
13
+ ```
14
+
15
+ ```python
16
+ from ahi.aws.redshift import RedshiftDataClient
17
+ from ahi.aws.athena import AthenaClient
18
+
19
+ rs = RedshiftDataClient(workgroup_name="analytics", database="dev")
20
+ rows = rs.run_query(
21
+ "select count(*) from events where day = :day",
22
+ parameters={"day": "2026-06-26"},
23
+ )
24
+ print(rows.scalar())
25
+
26
+ ath = AthenaClient(workgroup="primary", database="analytics")
27
+ for row in ath.run_query("select name from users where id = ?", parameters=[42]):
28
+ print(row["name"])
29
+ ```
30
+
31
+ ```python
32
+ from ahi.aws.auth import aws_session
33
+ from ahi.aws.s3 import S3Storage
34
+
35
+ session = aws_session()
36
+ storage = S3Storage("s3://my-bucket/some/prefix", client=session.client("s3"))
37
+
38
+ uploaded = storage.upload_dir("runs/run-01", prefix="run-01") # list[S3Object]
39
+ storage.download_dir("run-01", "local/run-01")
40
+ ```
41
+
42
+ ## Design notes
43
+
44
+ - **Never logs in for you.** No browser, no AWS CLI. Run `aws sso login`
45
+ beforehand (local SSO), or run where an AWS role is attached (EC2 instance
46
+ role, ECS task role); missing or expired credentials raise `AuthError`
47
+ with the command to run. No default region: set `AWS_REGION` or the
48
+ profile's `region`.
49
+ - **Parameterized queries by default.** Pass values via `parameters` (named
50
+ `:name` for Redshift, positional `?` for Athena) — never format them in.
51
+ - **One synchronous call.** `run_query` submits, polls with backoff (see
52
+ `ahi.core.poll_until`), and pages the full result set, returning a
53
+ `QueryResult` you can iterate, `len()`, or call `.scalar()` on.
54
+ - **Typed results from Redshift; strings from Athena** (Athena returns all values
55
+ as strings — we preserve that rather than guessing).
56
+ - **Failures raise** `QueryFailedError`; timeouts raise `PollTimeoutError`.
57
+ - **S3 transfers return `S3Object` records** (bucket, key, local path, size);
58
+ failures raise `UploadFailedError` / `DownloadFailedError`.
59
+
60
+ Install the optional `types` extra (`ahi-aws[types]`) for `boto3-stubs` during
61
+ development.
62
+
63
+ Part of [ahi-python-packages](https://github.com/adaptivehealthintelligence/ahi-python-packages).
@@ -0,0 +1,32 @@
1
+ [project]
2
+ name = "ahi-aws"
3
+ version = "0.4.0"
4
+ description = "Thin, typed boto3 helpers for AHI: AWS auth (SSO and attached-role sessions), Redshift Data API and Athena query execution, and S3 file storage."
5
+ readme = "README.md"
6
+ requires-python = ">=3.11"
7
+ license = { text = "MIT" }
8
+ authors = [{ name = "Adaptive Health Intelligence" }]
9
+ keywords = ["ahi", "aws", "boto3", "auth", "sso", "sts", "redshift", "athena", "s3"]
10
+ dependencies = [
11
+ # Independent versioning: depend on a compatible RANGE of ahi-core, not a pin.
12
+ "ahi-core>=0.2,<0.3",
13
+ "boto3>=1.34",
14
+ ]
15
+
16
+ [project.optional-dependencies]
17
+ # Opt-in type stubs for editor/CI type-checking of the boto3 surface we use.
18
+ types = [
19
+ "boto3-stubs[athena,redshift-data,s3,sts]>=1.34",
20
+ ]
21
+
22
+ [project.urls]
23
+ Homepage = "https://github.com/adaptivehealthintelligence/ahi-python-packages"
24
+ Documentation = "https://adaptivehealthintelligence.github.io/ahi-python-packages/"
25
+ Repository = "https://github.com/adaptivehealthintelligence/ahi-python-packages"
26
+
27
+ [build-system]
28
+ requires = ["hatchling>=1.24"]
29
+ build-backend = "hatchling.build"
30
+
31
+ [tool.hatch.build.targets.wheel]
32
+ packages = ["src/ahi"]
@@ -0,0 +1,44 @@
1
+ """Thin, typed boto3 helpers for AHI."""
2
+
3
+ from ahi.aws.athena import AthenaClient
4
+ from ahi.aws.auth import AwsIdentity, aws_session, check_identity
5
+ from ahi.aws.errors import (
6
+ AuthError,
7
+ AwsUnreachableError,
8
+ CredentialsNotFoundError,
9
+ CredentialsRejectedError,
10
+ DownloadFailedError,
11
+ ProfileNotFoundError,
12
+ QueryError,
13
+ QueryFailedError,
14
+ RegionNotConfiguredError,
15
+ S3Error,
16
+ SsoLoginRequiredError,
17
+ UploadFailedError,
18
+ )
19
+ from ahi.aws.redshift import RedshiftDataClient
20
+ from ahi.aws.results import QueryResult
21
+ from ahi.aws.s3 import S3Object, S3Storage
22
+
23
+ __all__ = [
24
+ "AthenaClient",
25
+ "AuthError",
26
+ "AwsIdentity",
27
+ "AwsUnreachableError",
28
+ "CredentialsNotFoundError",
29
+ "CredentialsRejectedError",
30
+ "DownloadFailedError",
31
+ "ProfileNotFoundError",
32
+ "QueryError",
33
+ "QueryFailedError",
34
+ "QueryResult",
35
+ "RedshiftDataClient",
36
+ "RegionNotConfiguredError",
37
+ "S3Error",
38
+ "S3Object",
39
+ "S3Storage",
40
+ "SsoLoginRequiredError",
41
+ "UploadFailedError",
42
+ "aws_session",
43
+ "check_identity",
44
+ ]
@@ -0,0 +1,167 @@
1
+ """Thin wrapper over Athena for running SQL.
2
+
3
+ Like Redshift's Data API, Athena is asynchronous: start an execution, poll its
4
+ state, then page the results. ``AthenaClient.run_query`` collapses that into one
5
+ synchronous call returning a :class:`~ahi.aws.results.QueryResult`.
6
+
7
+ from ahi.aws.athena import AthenaClient
8
+
9
+ client = AthenaClient(workgroup="primary", database="analytics")
10
+ result = client.run_query(
11
+ "select name from users where id = ?",
12
+ parameters=[42],
13
+ )
14
+ for row in result:
15
+ print(row["name"])
16
+
17
+ Athena returns every value as a string (or NULL); this wrapper preserves that
18
+ rather than guessing types.
19
+ """
20
+
21
+ from __future__ import annotations
22
+
23
+ from typing import TYPE_CHECKING, Any
24
+
25
+ from ahi.aws.errors import QueryFailedError
26
+ from ahi.aws.results import QueryResult
27
+ from ahi.core import get_logger, poll_until
28
+
29
+ if TYPE_CHECKING: # pragma: no cover - typing only
30
+ from collections.abc import Sequence
31
+
32
+ _log = get_logger("aws.athena")
33
+
34
+ # Terminal states reported by get_query_execution.
35
+ _TERMINAL_STATES = frozenset({"SUCCEEDED", "FAILED", "CANCELLED"})
36
+
37
+
38
+ class AthenaClient:
39
+ """Run queries against Athena.
40
+
41
+ Args:
42
+ workgroup: Athena workgroup. Result location and encryption are usually
43
+ configured on the workgroup, so ``output_location`` is optional.
44
+ database: Default Glue/Athena database (overridable per call).
45
+ catalog: Default data catalog (for example ``"AwsDataCatalog"``).
46
+ output_location: ``s3://`` prefix for results. Only needed when the
47
+ workgroup does not enforce one.
48
+ region_name: AWS region; defaults to the ambient boto3 configuration.
49
+ client: A pre-built ``athena`` boto3 client (injection point for tests).
50
+ """
51
+
52
+ def __init__(
53
+ self,
54
+ *,
55
+ workgroup: str = "primary",
56
+ database: str | None = None,
57
+ catalog: str | None = None,
58
+ output_location: str | None = None,
59
+ region_name: str | None = None,
60
+ client: Any | None = None,
61
+ ) -> None:
62
+ if client is None:
63
+ import boto3
64
+
65
+ client = boto3.client("athena", region_name=region_name)
66
+ self._client = client
67
+ self._workgroup = workgroup
68
+ self._database = database
69
+ self._catalog = catalog
70
+ self._output_location = output_location
71
+
72
+ def run_query(
73
+ self,
74
+ sql: str,
75
+ *,
76
+ parameters: Sequence[Any] | None = None,
77
+ database: str | None = None,
78
+ timeout: float = 300.0,
79
+ poll_interval: float = 1.0,
80
+ ) -> QueryResult:
81
+ """Execute ``sql`` and return its result.
82
+
83
+ Args:
84
+ sql: The SQL to run. Use positional ``?`` placeholders with
85
+ ``parameters`` rather than formatting values into the string.
86
+ parameters: Positional bind parameters, in order of ``?`` appearance.
87
+ database: Override the default database for this query.
88
+ timeout: Maximum seconds to wait for completion.
89
+ poll_interval: Initial delay between status polls.
90
+
91
+ Raises:
92
+ QueryFailedError: If the execution ends in ``FAILED``/``CANCELLED``.
93
+ PollTimeoutError: If it does not finish within ``timeout``.
94
+ """
95
+ query_id = self._start(sql, parameters, database)
96
+ _log.debug("started athena execution %s", query_id)
97
+
98
+ execution = poll_until(
99
+ lambda: self._client.get_query_execution(QueryExecutionId=query_id),
100
+ lambda e: e["QueryExecution"]["Status"]["State"] in _TERMINAL_STATES,
101
+ timeout=timeout,
102
+ interval=poll_interval,
103
+ )
104
+
105
+ status = execution["QueryExecution"]["Status"]
106
+ state = status["State"]
107
+ if state != "SUCCEEDED":
108
+ raise QueryFailedError(query_id, state, status.get("StateChangeReason"))
109
+ return self._collect(query_id)
110
+
111
+ def _start(
112
+ self,
113
+ sql: str,
114
+ parameters: Sequence[Any] | None,
115
+ database: str | None,
116
+ ) -> str:
117
+ args: dict[str, Any] = {"QueryString": sql, "WorkGroup": self._workgroup}
118
+
119
+ context: dict[str, str] = {}
120
+ db = database or self._database
121
+ if db is not None:
122
+ context["Database"] = db
123
+ if self._catalog is not None:
124
+ context["Catalog"] = self._catalog
125
+ if context:
126
+ args["QueryExecutionContext"] = context
127
+
128
+ if self._output_location is not None:
129
+ args["ResultConfiguration"] = {"OutputLocation": self._output_location}
130
+ if parameters:
131
+ args["ExecutionParameters"] = [str(value) for value in parameters]
132
+
133
+ return self._client.start_query_execution(**args)["QueryExecutionId"]
134
+
135
+ def _collect(self, query_id: str) -> QueryResult:
136
+ columns: list[str] = []
137
+ rows: list[dict[str, Any]] = []
138
+ next_token: str | None = None
139
+ first_page = True
140
+
141
+ while True:
142
+ kwargs: dict[str, Any] = {"QueryExecutionId": query_id}
143
+ if next_token:
144
+ kwargs["NextToken"] = next_token
145
+ page = self._client.get_query_results(**kwargs)
146
+ result_set = page.get("ResultSet", {})
147
+
148
+ if not columns:
149
+ column_info = result_set.get("ResultSetMetadata", {}).get("ColumnInfo", [])
150
+ columns = [col["Label"] for col in column_info]
151
+
152
+ page_rows = result_set.get("Rows", [])
153
+ # Athena repeats the column labels as the first data row of the first
154
+ # page for SELECT statements; skip it so it isn't returned as data.
155
+ if first_page and page_rows:
156
+ page_rows = page_rows[1:]
157
+ first_page = False
158
+
159
+ for row in page_rows:
160
+ values = [cell.get("VarCharValue") for cell in row.get("Data", [])]
161
+ rows.append(dict(zip(columns, values, strict=False)))
162
+
163
+ next_token = page.get("NextToken")
164
+ if not next_token:
165
+ break
166
+
167
+ return QueryResult(columns=columns, rows=rows, query_id=query_id)
@@ -0,0 +1,227 @@
1
+ """AWS credentials for AHI apps
2
+
3
+ This module provides :func:`aws_session`, which builds a boto3 session using
4
+ boto3's standard credential resolution, :func:`check_identity`, which checks it
5
+ with a single ``sts:GetCallerIdentity`` call, and :class:`AwsIdentity`, which
6
+ records who the credentials belong to and how they were found, great for provenance.
7
+
8
+ Caters to common ways of running AHI apps:
9
+
10
+ - **Local SSO** (a person at their machine): their own SSO login, done
11
+ beforehand with ``aws sso login``, picked up from ``AWS_PROFILE`` or the
12
+ default profile.
13
+ - **Attached role** (EC2 instance, ECS/Fargate task, scheduled jobs): the role
14
+ attached to the resource, found by boto3 with no profile and no login.
15
+
16
+ Hand ``session.client(...)`` to the other helpers:
17
+
18
+ from ahi.aws.athena import AthenaClient
19
+ from ahi.aws.auth import aws_session, check_identity
20
+
21
+ session = aws_session()
22
+ identity = check_identity(session)
23
+ log.info("AWS: %s via %s", identity.arn, identity.method)
24
+
25
+ athena = AthenaClient(workgroup="primary", client=session.client("athena"))
26
+
27
+ This module does not log in for you: no browser, no AWS CLI. When credentials are
28
+ missing or expired it raises an :class:`~ahi.aws.errors.AuthError` subclass
29
+ saying what to run (for example ``aws sso login``).
30
+ """
31
+
32
+ from __future__ import annotations
33
+
34
+ import os
35
+ from dataclasses import dataclass
36
+ from typing import Any
37
+
38
+ import boto3
39
+ from botocore.exceptions import (
40
+ BotoCoreError,
41
+ ClientError,
42
+ HTTPClientError,
43
+ NoCredentialsError,
44
+ NoRegionError,
45
+ PartialCredentialsError,
46
+ ProfileNotFound,
47
+ SSOTokenLoadError,
48
+ TokenRetrievalError,
49
+ UnauthorizedSSOTokenError,
50
+ )
51
+ from botocore.exceptions import ConnectionError as BotoConnectionError
52
+
53
+ from ahi.aws.errors import (
54
+ AuthError,
55
+ AwsUnreachableError,
56
+ CredentialsNotFoundError,
57
+ CredentialsRejectedError,
58
+ ProfileNotFoundError,
59
+ RegionNotConfiguredError,
60
+ SsoLoginRequiredError,
61
+ )
62
+ from ahi.core import get_logger
63
+
64
+ _log = get_logger("aws.auth")
65
+
66
+
67
+ @dataclass(frozen=True)
68
+ class AwsIdentity:
69
+ """The identity behind a set of AWS credentials
70
+
71
+ Returned by :func:`check_identity`. Log it, or record it with a run's provenance.
72
+
73
+ Attributes:
74
+ account: The AWS account id.
75
+ arn: The caller's ARN (an assumed-role ARN for SSO and attached roles).
76
+ user_id: The caller's unique id.
77
+ region: The region the identity check ran in.
78
+ profile: The named profile in effect, or ``None`` for the default
79
+ profile or when no profile is in use.
80
+ method: How botocore found the credentials, from
81
+ ``session.get_credentials().method``: for example ``"sso"`` for an
82
+ SSO login, ``"iam-role"`` for an EC2 instance role, or
83
+ ``"container-role"`` for an ECS/Fargate task role. ``None`` if it
84
+ could not tell.
85
+ """
86
+
87
+ account: str
88
+ arn: str
89
+ user_id: str
90
+ region: str | None
91
+ profile: str | None
92
+ method: str | None
93
+
94
+
95
+ def aws_session(*, profile: str | None = None, region: str | None = None) -> boto3.Session:
96
+ """Return a boto3 session using the standard credential resolution.
97
+
98
+ Credentials lookup order is the same as boto3's, but;
99
+ later sources are not checked:
100
+
101
+ 1. Environment variables: ``AWS_ACCESS_KEY_ID``, ``AWS_SECRET_ACCESS_KEY``
102
+ 2. The profile in use: ``AWS_PROFILE``, else ``[default]`` in ``~/.aws``.
103
+ 3. The attached role: ECS/Fargate task role, else EC2 instance role.
104
+
105
+ Passing ``profile`` skips this list and uses only that profile and authenticate via SSO.
106
+
107
+ Region lookup order;
108
+ 1. ``region`` argument to this function.
109
+ 2. Environment variables: ``AWS_REGION``, ``AWS_DEFAULT_REGION``.
110
+ 3. The profile's ``region`` setting.
111
+
112
+ Makes no network calls; use :func:`check_identity` to verify the session.
113
+
114
+ Args:
115
+ profile: AWS SSO profile name to use. This uses SSO method over other credential
116
+ sources (environment keys, ``AWS_PROFILE``, attached roles).
117
+ If you pass one, it must exist on the machine running the code, or this raises
118
+ ``ProfileNotFoundError``. On machines that use an attached role
119
+ (EC2, ECS), pass ``None``.
120
+ region: Region override.
121
+
122
+ Returns:
123
+ The session, ready for ``session.client(...)``.
124
+
125
+ Raises:
126
+ ProfileNotFoundError: If the profile is not configured.
127
+ RegionNotConfiguredError: If no region resolves.
128
+ """
129
+ region_name = region or os.environ.get("AWS_REGION") or None
130
+ try:
131
+ session = boto3.Session(profile_name=profile, region_name=region_name)
132
+ resolved_region = session.region_name
133
+ except ProfileNotFound as exc:
134
+ raise _to_auth_error(exc, profile=profile, method=None, region=None) from exc
135
+
136
+ profile_in_use = _profile_in_use(session)
137
+ if not resolved_region:
138
+ raise RegionNotConfiguredError(profile=profile_in_use)
139
+ _log.debug("AWS session: profile=%s region=%s", profile_in_use or "<default>", resolved_region)
140
+ return session
141
+
142
+
143
+ def check_identity(session: boto3.Session, *, client: Any | None = None) -> AwsIdentity:
144
+ """Return the identity behind the session's credentials, checked via an STS call.
145
+
146
+ Call this once at the start of a run so missing or expired credentials
147
+ fail early with a message that says what to do.
148
+
149
+ Args:
150
+ session: The session to check, usually from :func:`aws_session`.
151
+ client: A pre-built ``sts`` boto3 client (injection point for tests).
152
+ Defaults to ``session.client("sts")``.
153
+
154
+ Returns:
155
+ The caller's identity and how its credentials were found.
156
+
157
+ Raises:
158
+ SsoLoginRequiredError: If the SSO login is missing or expired.
159
+ CredentialsNotFoundError: If no credentials are found.
160
+ CredentialsRejectedError: If AWS rejects the credentials.
161
+ AwsUnreachableError: If AWS STS cannot be reached.
162
+ AuthError: For any other credential failure (the base of all the above).
163
+ """
164
+ profile = _profile_in_use(session)
165
+ method: str | None = None
166
+ region: str | None = session.region_name
167
+ try:
168
+ credentials = session.get_credentials()
169
+ method = credentials.method if credentials is not None else None
170
+ sts = client if client is not None else session.client("sts")
171
+ region = sts.meta.region_name
172
+ response = sts.get_caller_identity()
173
+ except (BotoCoreError, ClientError) as exc:
174
+ raise _to_auth_error(exc, profile=profile, method=method, region=region) from exc
175
+
176
+ identity = AwsIdentity(
177
+ account=response["Account"],
178
+ arn=response["Arn"],
179
+ user_id=response["UserId"],
180
+ region=region,
181
+ profile=profile,
182
+ method=method,
183
+ )
184
+ _log.debug("AWS identity: %s via %s", identity.arn, identity.method)
185
+ return identity
186
+
187
+
188
+ def _profile_in_use(session: boto3.Session) -> str | None:
189
+ """Return the profile in use, or ``None`` when it's the default profile."""
190
+ name = session.profile_name
191
+ return None if name == "default" else name
192
+
193
+
194
+ def _to_auth_error(
195
+ exc: BotoCoreError | ClientError,
196
+ *,
197
+ profile: str | None,
198
+ method: str | None,
199
+ region: str | None,
200
+ ) -> AuthError:
201
+ """Return the AHI auth error that matches a botocore failure.
202
+
203
+ Maps each botocore credential, region and network error to its ``AuthError``
204
+ subclass, adding context botocore doesn't have (the profile in effect, how
205
+ the credentials were found, the region) so the message can say what to do
206
+ next. Anything not recognised becomes a plain ``AuthError``.
207
+ """
208
+ reason = str(exc)
209
+ if isinstance(exc, ProfileNotFound):
210
+ return ProfileNotFoundError(profile=profile or exc.kwargs.get("profile"), reason=reason)
211
+ if isinstance(exc, SSOTokenLoadError | UnauthorizedSSOTokenError | TokenRetrievalError):
212
+ return SsoLoginRequiredError(profile=profile, reason=reason)
213
+ if isinstance(exc, NoCredentialsError | PartialCredentialsError):
214
+ return CredentialsNotFoundError(profile=profile, reason=reason)
215
+ if isinstance(exc, NoRegionError):
216
+ return RegionNotConfiguredError(profile=profile, reason=reason)
217
+ if isinstance(exc, BotoConnectionError | HTTPClientError):
218
+ return AwsUnreachableError(reason=reason, region=region, profile=profile)
219
+ if isinstance(exc, ClientError):
220
+ error = exc.response.get("Error", {})
221
+ return CredentialsRejectedError(
222
+ error.get("Code", "Unknown"),
223
+ reason=reason, # botocore's full message, which names the operation too
224
+ method=method,
225
+ profile=profile,
226
+ )
227
+ return AuthError(profile=profile, reason=reason)