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.
- ahi_aws-0.4.0/.gitignore +25 -0
- ahi_aws-0.4.0/CHANGELOG.md +29 -0
- ahi_aws-0.4.0/PKG-INFO +80 -0
- ahi_aws-0.4.0/README.md +63 -0
- ahi_aws-0.4.0/pyproject.toml +32 -0
- ahi_aws-0.4.0/src/ahi/aws/__init__.py +44 -0
- ahi_aws-0.4.0/src/ahi/aws/athena.py +167 -0
- ahi_aws-0.4.0/src/ahi/aws/auth.py +227 -0
- ahi_aws-0.4.0/src/ahi/aws/errors.py +202 -0
- ahi_aws-0.4.0/src/ahi/aws/redshift.py +175 -0
- ahi_aws-0.4.0/src/ahi/aws/results.py +47 -0
- ahi_aws-0.4.0/src/ahi/aws/s3.py +277 -0
- ahi_aws-0.4.0/tests/conftest.py +41 -0
- ahi_aws-0.4.0/tests/test_athena.py +87 -0
- ahi_aws-0.4.0/tests/test_auth.py +389 -0
- ahi_aws-0.4.0/tests/test_redshift.py +138 -0
- ahi_aws-0.4.0/tests/test_results.py +21 -0
- ahi_aws-0.4.0/tests/test_s3.py +471 -0
ahi_aws-0.4.0/.gitignore
ADDED
|
@@ -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).
|
ahi_aws-0.4.0/README.md
ADDED
|
@@ -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)
|