ahi-aws 0.4.0__py3-none-any.whl

ahi/aws/__init__.py

@@ -0,0 +1,44 @@
+"""Thin, typed boto3 helpers for AHI."""
+
+from ahi.aws.athena import AthenaClient
+from ahi.aws.auth import AwsIdentity, aws_session, check_identity
+from ahi.aws.errors import (
+    AuthError,
+    AwsUnreachableError,
+    CredentialsNotFoundError,
+    CredentialsRejectedError,
+    DownloadFailedError,
+    ProfileNotFoundError,
+    QueryError,
+    QueryFailedError,
+    RegionNotConfiguredError,
+    S3Error,
+    SsoLoginRequiredError,
+    UploadFailedError,
+)
+from ahi.aws.redshift import RedshiftDataClient
+from ahi.aws.results import QueryResult
+from ahi.aws.s3 import S3Object, S3Storage
+
+__all__ = [
+    "AthenaClient",
+    "AuthError",
+    "AwsIdentity",
+    "AwsUnreachableError",
+    "CredentialsNotFoundError",
+    "CredentialsRejectedError",
+    "DownloadFailedError",
+    "ProfileNotFoundError",
+    "QueryError",
+    "QueryFailedError",
+    "QueryResult",
+    "RedshiftDataClient",
+    "RegionNotConfiguredError",
+    "S3Error",
+    "S3Object",
+    "S3Storage",
+    "SsoLoginRequiredError",
+    "UploadFailedError",
+    "aws_session",
+    "check_identity",
+]

ahi/aws/athena.py

@@ -0,0 +1,167 @@
+"""Thin wrapper over Athena for running SQL.
+
+Like Redshift's Data API, Athena is asynchronous: start an execution, poll its
+state, then page the results. ``AthenaClient.run_query`` collapses that into one
+synchronous call returning a :class:`~ahi.aws.results.QueryResult`.
+
+    from ahi.aws.athena import AthenaClient
+
+    client = AthenaClient(workgroup="primary", database="analytics")
+    result = client.run_query(
+        "select name from users where id = ?",
+        parameters=[42],
+    )
+    for row in result:
+        print(row["name"])
+
+Athena returns every value as a string (or NULL); this wrapper preserves that
+rather than guessing types.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ahi.aws.errors import QueryFailedError
+from ahi.aws.results import QueryResult
+from ahi.core import get_logger, poll_until
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from collections.abc import Sequence
+
+_log = get_logger("aws.athena")
+
+# Terminal states reported by get_query_execution.
+_TERMINAL_STATES = frozenset({"SUCCEEDED", "FAILED", "CANCELLED"})
+
+
+class AthenaClient:
+    """Run queries against Athena.
+
+    Args:
+        workgroup: Athena workgroup. Result location and encryption are usually
+            configured on the workgroup, so ``output_location`` is optional.
+        database: Default Glue/Athena database (overridable per call).
+        catalog: Default data catalog (for example ``"AwsDataCatalog"``).
+        output_location: ``s3://`` prefix for results. Only needed when the
+            workgroup does not enforce one.
+        region_name: AWS region; defaults to the ambient boto3 configuration.
+        client: A pre-built ``athena`` boto3 client (injection point for tests).
+    """
+
+    def __init__(
+        self,
+        *,
+        workgroup: str = "primary",
+        database: str | None = None,
+        catalog: str | None = None,
+        output_location: str | None = None,
+        region_name: str | None = None,
+        client: Any | None = None,
+    ) -> None:
+        if client is None:
+            import boto3
+
+            client = boto3.client("athena", region_name=region_name)
+        self._client = client
+        self._workgroup = workgroup
+        self._database = database
+        self._catalog = catalog
+        self._output_location = output_location
+
+    def run_query(
+        self,
+        sql: str,
+        *,
+        parameters: Sequence[Any] | None = None,
+        database: str | None = None,
+        timeout: float = 300.0,
+        poll_interval: float = 1.0,
+    ) -> QueryResult:
+        """Execute ``sql`` and return its result.
+
+        Args:
+            sql: The SQL to run. Use positional ``?`` placeholders with
+                ``parameters`` rather than formatting values into the string.
+            parameters: Positional bind parameters, in order of ``?`` appearance.
+            database: Override the default database for this query.
+            timeout: Maximum seconds to wait for completion.
+            poll_interval: Initial delay between status polls.
+
+        Raises:
+            QueryFailedError: If the execution ends in ``FAILED``/``CANCELLED``.
+            PollTimeoutError: If it does not finish within ``timeout``.
+        """
+        query_id = self._start(sql, parameters, database)
+        _log.debug("started athena execution %s", query_id)
+
+        execution = poll_until(
+            lambda: self._client.get_query_execution(QueryExecutionId=query_id),
+            lambda e: e["QueryExecution"]["Status"]["State"] in _TERMINAL_STATES,
+            timeout=timeout,
+            interval=poll_interval,
+        )
+
+        status = execution["QueryExecution"]["Status"]
+        state = status["State"]
+        if state != "SUCCEEDED":
+            raise QueryFailedError(query_id, state, status.get("StateChangeReason"))
+        return self._collect(query_id)
+
+    def _start(
+        self,
+        sql: str,
+        parameters: Sequence[Any] | None,
+        database: str | None,
+    ) -> str:
+        args: dict[str, Any] = {"QueryString": sql, "WorkGroup": self._workgroup}
+
+        context: dict[str, str] = {}
+        db = database or self._database
+        if db is not None:
+            context["Database"] = db
+        if self._catalog is not None:
+            context["Catalog"] = self._catalog
+        if context:
+            args["QueryExecutionContext"] = context
+
+        if self._output_location is not None:
+            args["ResultConfiguration"] = {"OutputLocation": self._output_location}
+        if parameters:
+            args["ExecutionParameters"] = [str(value) for value in parameters]
+
+        return self._client.start_query_execution(**args)["QueryExecutionId"]
+
+    def _collect(self, query_id: str) -> QueryResult:
+        columns: list[str] = []
+        rows: list[dict[str, Any]] = []
+        next_token: str | None = None
+        first_page = True
+
+        while True:
+            kwargs: dict[str, Any] = {"QueryExecutionId": query_id}
+            if next_token:
+                kwargs["NextToken"] = next_token
+            page = self._client.get_query_results(**kwargs)
+            result_set = page.get("ResultSet", {})
+
+            if not columns:
+                column_info = result_set.get("ResultSetMetadata", {}).get("ColumnInfo", [])
+                columns = [col["Label"] for col in column_info]
+
+            page_rows = result_set.get("Rows", [])
+            # Athena repeats the column labels as the first data row of the first
+            # page for SELECT statements; skip it so it isn't returned as data.
+            if first_page and page_rows:
+                page_rows = page_rows[1:]
+            first_page = False
+
+            for row in page_rows:
+                values = [cell.get("VarCharValue") for cell in row.get("Data", [])]
+                rows.append(dict(zip(columns, values, strict=False)))
+
+            next_token = page.get("NextToken")
+            if not next_token:
+                break
+
+        return QueryResult(columns=columns, rows=rows, query_id=query_id)

ahi/aws/auth.py

@@ -0,0 +1,227 @@
+"""AWS credentials for AHI apps
+
+This module provides :func:`aws_session`, which builds a boto3 session using
+boto3's standard credential resolution, :func:`check_identity`, which checks it
+with a single ``sts:GetCallerIdentity`` call, and :class:`AwsIdentity`, which
+records who the credentials belong to and how they were found, great for provenance.
+
+Caters to common ways of running AHI apps:
+
+- **Local SSO** (a person at their machine): their own SSO login, done
+  beforehand with ``aws sso login``, picked up from ``AWS_PROFILE`` or the
+  default profile.
+- **Attached role** (EC2 instance, ECS/Fargate task, scheduled jobs): the role
+  attached to the resource, found by boto3 with no profile and no login.
+
+Hand ``session.client(...)`` to the other helpers:
+
+    from ahi.aws.athena import AthenaClient
+    from ahi.aws.auth import aws_session, check_identity
+
+    session = aws_session()
+    identity = check_identity(session)
+    log.info("AWS: %s via %s", identity.arn, identity.method)
+
+    athena = AthenaClient(workgroup="primary", client=session.client("athena"))
+
+This module does not log in for you: no browser, no AWS CLI. When credentials are
+missing or expired it raises an :class:`~ahi.aws.errors.AuthError` subclass
+saying what to run (for example ``aws sso login``).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import boto3
+from botocore.exceptions import (
+    BotoCoreError,
+    ClientError,
+    HTTPClientError,
+    NoCredentialsError,
+    NoRegionError,
+    PartialCredentialsError,
+    ProfileNotFound,
+    SSOTokenLoadError,
+    TokenRetrievalError,
+    UnauthorizedSSOTokenError,
+)
+from botocore.exceptions import ConnectionError as BotoConnectionError
+
+from ahi.aws.errors import (
+    AuthError,
+    AwsUnreachableError,
+    CredentialsNotFoundError,
+    CredentialsRejectedError,
+    ProfileNotFoundError,
+    RegionNotConfiguredError,
+    SsoLoginRequiredError,
+)
+from ahi.core import get_logger
+
+_log = get_logger("aws.auth")
+
+
+@dataclass(frozen=True)
+class AwsIdentity:
+    """The identity behind a set of AWS credentials
+
+    Returned by :func:`check_identity`. Log it, or record it with a run's provenance.
+
+    Attributes:
+        account: The AWS account id.
+        arn: The caller's ARN (an assumed-role ARN for SSO and attached roles).
+        user_id: The caller's unique id.
+        region: The region the identity check ran in.
+        profile: The named profile in effect, or ``None`` for the default
+            profile or when no profile is in use.
+        method: How botocore found the credentials, from
+            ``session.get_credentials().method``: for example ``"sso"`` for an
+            SSO login, ``"iam-role"`` for an EC2 instance role, or
+            ``"container-role"`` for an ECS/Fargate task role. ``None`` if it
+            could not tell.
+    """
+
+    account: str
+    arn: str
+    user_id: str
+    region: str | None
+    profile: str | None
+    method: str | None
+
+
+def aws_session(*, profile: str | None = None, region: str | None = None) -> boto3.Session:
+    """Return a boto3 session using the standard credential resolution.
+
+    Credentials lookup order is the same as boto3's, but;
+    later sources are not checked:
+
+    1. Environment variables: ``AWS_ACCESS_KEY_ID``, ``AWS_SECRET_ACCESS_KEY``
+    2. The profile in use: ``AWS_PROFILE``, else ``[default]`` in ``~/.aws``.
+    3. The attached role: ECS/Fargate task role, else EC2 instance role.
+
+    Passing ``profile`` skips this list and uses only that profile and authenticate via SSO.
+
+    Region lookup order;
+    1. ``region`` argument to this function.
+    2. Environment variables: ``AWS_REGION``, ``AWS_DEFAULT_REGION``.
+    3. The profile's ``region`` setting.
+
+    Makes no network calls; use :func:`check_identity` to verify the session.
+
+    Args:
+        profile: AWS SSO profile name to use. This uses SSO method over other credential
+            sources (environment keys, ``AWS_PROFILE``, attached roles).
+            If you pass one, it must exist on the machine running the code, or this raises
+            ``ProfileNotFoundError``. On machines that use an attached role
+            (EC2, ECS), pass ``None``.
+        region: Region override.
+
+    Returns:
+        The session, ready for ``session.client(...)``.
+
+    Raises:
+        ProfileNotFoundError: If the profile is not configured.
+        RegionNotConfiguredError: If no region resolves.
+    """
+    region_name = region or os.environ.get("AWS_REGION") or None
+    try:
+        session = boto3.Session(profile_name=profile, region_name=region_name)
+        resolved_region = session.region_name
+    except ProfileNotFound as exc:
+        raise _to_auth_error(exc, profile=profile, method=None, region=None) from exc
+
+    profile_in_use = _profile_in_use(session)
+    if not resolved_region:
+        raise RegionNotConfiguredError(profile=profile_in_use)
+    _log.debug("AWS session: profile=%s region=%s", profile_in_use or "<default>", resolved_region)
+    return session
+
+
+def check_identity(session: boto3.Session, *, client: Any | None = None) -> AwsIdentity:
+    """Return the identity behind the session's credentials, checked via an STS call.
+
+    Call this once at the start of a run so missing or expired credentials
+    fail early with a message that says what to do.
+
+    Args:
+        session: The session to check, usually from :func:`aws_session`.
+        client: A pre-built ``sts`` boto3 client (injection point for tests).
+            Defaults to ``session.client("sts")``.
+
+    Returns:
+        The caller's identity and how its credentials were found.
+
+    Raises:
+        SsoLoginRequiredError: If the SSO login is missing or expired.
+        CredentialsNotFoundError: If no credentials are found.
+        CredentialsRejectedError: If AWS rejects the credentials.
+        AwsUnreachableError: If AWS STS cannot be reached.
+        AuthError: For any other credential failure (the base of all the above).
+    """
+    profile = _profile_in_use(session)
+    method: str | None = None
+    region: str | None = session.region_name
+    try:
+        credentials = session.get_credentials()
+        method = credentials.method if credentials is not None else None
+        sts = client if client is not None else session.client("sts")
+        region = sts.meta.region_name
+        response = sts.get_caller_identity()
+    except (BotoCoreError, ClientError) as exc:
+        raise _to_auth_error(exc, profile=profile, method=method, region=region) from exc
+
+    identity = AwsIdentity(
+        account=response["Account"],
+        arn=response["Arn"],
+        user_id=response["UserId"],
+        region=region,
+        profile=profile,
+        method=method,
+    )
+    _log.debug("AWS identity: %s via %s", identity.arn, identity.method)
+    return identity
+
+
+def _profile_in_use(session: boto3.Session) -> str | None:
+    """Return the profile in use, or ``None`` when it's the default profile."""
+    name = session.profile_name
+    return None if name == "default" else name
+
+
+def _to_auth_error(
+    exc: BotoCoreError | ClientError,
+    *,
+    profile: str | None,
+    method: str | None,
+    region: str | None,
+) -> AuthError:
+    """Return the AHI auth error that matches a botocore failure.
+
+    Maps each botocore credential, region and network error to its ``AuthError``
+    subclass, adding context botocore doesn't have (the profile in effect, how
+    the credentials were found, the region) so the message can say what to do
+    next. Anything not recognised becomes a plain ``AuthError``.
+    """
+    reason = str(exc)
+    if isinstance(exc, ProfileNotFound):
+        return ProfileNotFoundError(profile=profile or exc.kwargs.get("profile"), reason=reason)
+    if isinstance(exc, SSOTokenLoadError | UnauthorizedSSOTokenError | TokenRetrievalError):
+        return SsoLoginRequiredError(profile=profile, reason=reason)
+    if isinstance(exc, NoCredentialsError | PartialCredentialsError):
+        return CredentialsNotFoundError(profile=profile, reason=reason)
+    if isinstance(exc, NoRegionError):
+        return RegionNotConfiguredError(profile=profile, reason=reason)
+    if isinstance(exc, BotoConnectionError | HTTPClientError):
+        return AwsUnreachableError(reason=reason, region=region, profile=profile)
+    if isinstance(exc, ClientError):
+        error = exc.response.get("Error", {})
+        return CredentialsRejectedError(
+            error.get("Code", "Unknown"),
+            reason=reason,  # botocore's full message, which names the operation too
+            method=method,
+            profile=profile,
+        )
+    return AuthError(profile=profile, reason=reason)

ahi/aws/errors.py

@@ -0,0 +1,202 @@
+"""Errors raised by the AWS helpers."""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from ahi.core.errors import AhiError
+
+
+class QueryError(AhiError):
+    """Base class for query-execution failures."""
+
+
+class QueryFailedError(QueryError):
+    """Raised when a query reaches a non-successful terminal state.
+
+    Attributes:
+        query_id: The service-assigned statement/execution id.
+        state: The terminal state reported by the service (for example
+            ``"FAILED"`` or ``"ABORTED"``).
+        reason: The service-provided failure reason, if any.
+    """
+
+    def __init__(self, query_id: str, state: str, reason: str | None = None) -> None:
+        self.query_id = query_id
+        self.state = state
+        self.reason = reason
+        detail = f": {reason}" if reason else ""
+        super().__init__(f"Query {query_id} ended in state {state}{detail}")
+
+
+class S3Error(AhiError):
+    """Base class for S3 failures.
+
+    Attributes:
+        bucket: The bucket.
+        key: The object key, or the prefix being listed.
+        reason: The boto3/botocore message.
+        code: The S3 error code (for example ``"AccessDenied"``, or ``"404"`` for a
+            missing key), or ``None``.
+    """
+
+    _action: ClassVar[str] = "access"
+
+    def __init__(self, bucket: str, key: str, reason: str, *, code: str | None = None) -> None:
+        self.bucket = bucket
+        self.key = key
+        self.reason = reason
+        self.code = code
+        super().__init__(f"Failed to {self._action} s3://{bucket}/{key}: {reason}")
+
+
+class UploadFailedError(S3Error):
+    """Raised when an upload to S3 fails."""
+
+    _action = "upload"
+
+
+class DownloadFailedError(S3Error):
+    """Raised when listing or downloading from S3 fails."""
+
+    _action = "download"
+
+
+class AuthError(AhiError):
+    """Base class for AWS credential and region failures.
+
+    Raised as-is when botocore fails in a way no subclass covers. The message
+    says what went wrong, what botocore reported, and what to do about it, so
+    it can be shown to the person running the app as-is without hiding
+    botocore's own detail. The botocore error is also kept as ``__cause__``.
+
+    Attributes:
+        profile: The named AWS profile in effect, or ``None`` for the default
+            profile or when no profile is in use.
+        reason: The botocore message, or ``None`` when no botocore error is
+            behind the failure.
+        hint: What to do next (for example the ``aws sso login`` command to run).
+    """
+
+    def __init__(self, *, profile: str | None = None, reason: str | None = None) -> None:
+        self.profile = profile
+        self.reason = reason
+        self.hint = self._hint()
+        detail = f" (botocore: {reason.rstrip('.')})" if reason else ""
+        super().__init__(f"{self._problem()}{detail}. {self.hint}")
+
+    def _problem(self) -> str:
+        return "AWS credentials could not be resolved"
+
+    def _hint(self) -> str:
+        return "Check the AWS configuration with `aws configure list`."
+
+    def _login_command(self) -> str:
+        if self.profile:
+            return f"aws sso login --profile {self.profile}"
+        return "aws sso login"
+
+
+class ProfileNotFoundError(AuthError):
+    """Raised when ``AWS_PROFILE`` or ``profile=`` names a profile that is not configured."""
+
+    def _problem(self) -> str:
+        return f"AWS profile '{self.profile}' is not configured"
+
+    def _hint(self) -> str:
+        return (
+            "Check AWS_PROFILE (leave it unset to use a role attached to this machine), "
+            "or run `aws configure sso` to set the profile up."
+        )
+
+
+class SsoLoginRequiredError(AuthError):
+    """Raised when the AWS SSO login is missing or has expired."""
+
+    def _problem(self) -> str:
+        return "AWS SSO login is missing or expired"
+
+    def _hint(self) -> str:
+        return f"Run `{self._login_command()}`, then retry."
+
+
+class CredentialsNotFoundError(AuthError):
+    """Raised when no AWS credentials are found: no login, no keys, no attached role."""
+
+    def _problem(self) -> str:
+        return "No AWS credentials found"
+
+    def _hint(self) -> str:
+        return f"Log in with `{self._login_command()}`, or run where an AWS role is attached."
+
+
+class CredentialsRejectedError(AuthError):
+    """Raised when AWS rejects the credentials, for example because they expired.
+
+    Attributes:
+        code: The AWS error code (for example ``"ExpiredToken"`` or
+            ``"AccessDenied"``).
+        method: How botocore found the credentials (for example ``"sso"``), or
+            ``None``.
+        expired: Whether the code means the credentials have lapsed.
+    """
+
+    # AWS error codes that mean the credentials were valid once but have lapsed.
+    _expired_codes: ClassVar[frozenset[str]] = frozenset(
+        {"ExpiredToken", "ExpiredTokenException", "RequestExpired"}
+    )
+
+    def __init__(
+        self,
+        code: str,
+        *,
+        reason: str,
+        method: str | None = None,
+        profile: str | None = None,
+    ) -> None:
+        self.code = code
+        self.method = method
+        self.expired = code in self._expired_codes
+        super().__init__(profile=profile, reason=reason)
+
+    def _problem(self) -> str:
+        return f"AWS rejected the credentials ({self.code})"
+
+    def _hint(self) -> str:
+        source = self.method or "an unknown source"
+        if self.expired:
+            return (
+                f"The credentials from {source} have expired; "
+                f"run `{self._login_command()}`, then retry."
+            )
+        return f"Check the credentials from {source} are valid and allowed to call AWS."
+
+
+class RegionNotConfiguredError(AuthError):
+    """Raised when no AWS region resolves from the arguments, environment or profile."""
+
+    def _problem(self) -> str:
+        return "No AWS region configured"
+
+    def _hint(self) -> str:
+        return "Set AWS_REGION, add `region` to the AWS profile, or pass region=."
+
+
+class AwsUnreachableError(AuthError):
+    """Raised when AWS STS cannot be reached (network, proxy or SSL failure).
+
+    Attributes:
+        region: The region the call was made in, or ``None``.
+    """
+
+    def __init__(
+        self, *, reason: str, region: str | None = None, profile: str | None = None
+    ) -> None:
+        self.region = region
+        super().__init__(profile=profile, reason=reason)
+
+    def _problem(self) -> str:
+        return f"Cannot reach AWS STS in {self.region or 'the configured region'}"
+
+    def _hint(self) -> str:
+        return "Check network access to AWS (proxy, firewall, VPC endpoint)."

ahi/aws/redshift.py

@@ -0,0 +1,175 @@
+"""Thin wrapper over the Redshift Data API for running SQL.
+
+The Redshift Data API is asynchronous: you submit a statement, poll its status,
+then page through the result set. ``RedshiftDataClient.run_query`` collapses that
+into one synchronous call returning a :class:`~ahi.aws.results.QueryResult`.
+
+It targets Redshift Serverless (``workgroup_name``) or a provisioned cluster
+(``cluster_identifier``). Authentication to the database is via the workgroup's
+default IAM auth, ``db_user`` (provisioned temp credentials), or ``secret_arn``
+(Secrets Manager) — whichever your setup uses.
+
+    from ahi.aws.redshift import RedshiftDataClient
+
+    client = RedshiftDataClient(workgroup_name="analytics", database="dev")
+    result = client.run_query(
+        "select count(*) from events where day = :day",
+        parameters={"day": "2026-06-26"},
+    )
+    print(result.scalar())
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ahi.aws.errors import QueryFailedError
+from ahi.aws.results import QueryResult
+from ahi.core import get_logger, poll_until
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from collections.abc import Mapping
+
+_log = get_logger("aws.redshift")
+
+# Terminal states reported by describe_statement.
+_TERMINAL_STATES = frozenset({"FINISHED", "FAILED", "ABORTED"})
+
+
+def _decode_field(field: Mapping[str, Any]) -> Any:
+    """Decode one Redshift Data API field into a native Python value."""
+    if field.get("isNull"):
+        return None
+    for key in ("stringValue", "longValue", "doubleValue", "booleanValue", "blobValue"):
+        if key in field:
+            return field[key]
+    return None
+
+
+class RedshiftDataClient:
+    """Run queries against Redshift via the Data API.
+
+    Args:
+        workgroup_name: Redshift Serverless workgroup (serverless target).
+        cluster_identifier: Provisioned cluster id (provisioned target).
+        database: Default database for queries (overridable per call).
+        db_user: Database user for provisioned temporary credentials.
+        secret_arn: Secrets Manager ARN holding database credentials.
+        region_name: AWS region; defaults to the ambient boto3 configuration.
+        client: A pre-built ``redshift-data`` boto3 client. When provided, the
+            connection arguments above are still used per request; mainly an
+            injection point for tests.
+    """
+
+    def __init__(
+        self,
+        *,
+        workgroup_name: str | None = None,
+        cluster_identifier: str | None = None,
+        database: str | None = None,
+        db_user: str | None = None,
+        secret_arn: str | None = None,
+        region_name: str | None = None,
+        client: Any | None = None,
+    ) -> None:
+        if workgroup_name is None and cluster_identifier is None:
+            raise ValueError("provide either workgroup_name or cluster_identifier")
+        if client is None:
+            import boto3
+
+            client = boto3.client("redshift-data", region_name=region_name)
+        self._client = client
+        self._workgroup_name = workgroup_name
+        self._cluster_identifier = cluster_identifier
+        self._database = database
+        self._db_user = db_user
+        self._secret_arn = secret_arn
+
+    def run_query(
+        self,
+        sql: str,
+        *,
+        parameters: Mapping[str, Any] | None = None,
+        database: str | None = None,
+        timeout: float = 300.0,
+        poll_interval: float = 1.0,
+    ) -> QueryResult:
+        """Execute ``sql`` and return its result.
+
+        Args:
+            sql: The SQL to run. Use named placeholders (``:name``) with
+                ``parameters`` rather than formatting values into the string.
+            parameters: Named bind parameters. Values are sent as strings, per
+                the Data API contract.
+            database: Override the default database for this query.
+            timeout: Maximum seconds to wait for completion.
+            poll_interval: Initial delay between status polls.
+
+        Raises:
+            QueryFailedError: If the statement ends in ``FAILED``/``ABORTED``.
+            PollTimeoutError: If it does not finish within ``timeout``.
+        """
+        statement_id = self._submit(sql, parameters, database)
+        _log.debug("submitted redshift statement %s", statement_id)
+
+        description = poll_until(
+            lambda: self._client.describe_statement(Id=statement_id),
+            lambda d: d["Status"] in _TERMINAL_STATES,
+            timeout=timeout,
+            interval=poll_interval,
+        )
+
+        status = description["Status"]
+        if status != "FINISHED":
+            raise QueryFailedError(statement_id, status, description.get("Error"))
+
+        if not description.get("HasResultSet"):
+            return QueryResult(query_id=statement_id)
+        return self._collect(statement_id)
+
+    def _submit(
+        self,
+        sql: str,
+        parameters: Mapping[str, Any] | None,
+        database: str | None,
+    ) -> str:
+        args: dict[str, Any] = {"Sql": sql}
+        if self._workgroup_name is not None:
+            args["WorkgroupName"] = self._workgroup_name
+        if self._cluster_identifier is not None:
+            args["ClusterIdentifier"] = self._cluster_identifier
+        db = database or self._database
+        if db is not None:
+            args["Database"] = db
+        if self._db_user is not None:
+            args["DbUser"] = self._db_user
+        if self._secret_arn is not None:
+            args["SecretArn"] = self._secret_arn
+        if parameters:
+            args["Parameters"] = [
+                {"name": name, "value": str(value)} for name, value in parameters.items()
+            ]
+        return self._client.execute_statement(**args)["Id"]
+
+    def _collect(self, statement_id: str) -> QueryResult:
+        columns: list[str] = []
+        rows: list[dict[str, Any]] = []
+        next_token: str | None = None
+
+        while True:
+            kwargs: dict[str, Any] = {"Id": statement_id}
+            if next_token:
+                kwargs["NextToken"] = next_token
+            page = self._client.get_statement_result(**kwargs)
+
+            if not columns:
+                columns = [col["name"] for col in page.get("ColumnMetadata", [])]
+            for record in page.get("Records", []):
+                values = [_decode_field(field) for field in record]
+                rows.append(dict(zip(columns, values, strict=False)))
+
+            next_token = page.get("NextToken")
+            if not next_token:
+                break
+
+        return QueryResult(columns=columns, rows=rows, query_id=statement_id)

ahi/aws/results.py

@@ -0,0 +1,47 @@
+"""The result type returned by the query helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class QueryResult:
+    """The outcome of a successful query.
+
+    Iterating a ``QueryResult`` yields row dicts, and ``len(result)`` is the row
+    count, so the common cases read naturally::
+
+        for row in client.run_query("select id, name from users"):
+            print(row["id"], row["name"])
+
+    Attributes:
+        columns: Column names in selection order.
+        rows: One dict per row, keyed by column name.
+        query_id: The service-assigned statement/execution id (useful for logs).
+    """
+
+    columns: list[str] = field(default_factory=list)
+    rows: list[dict[str, Any]] = field(default_factory=list)
+    query_id: str = ""
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        return iter(self.rows)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def scalar(self) -> Any:
+        """Return the single value of a one-row, one-column result.
+
+        Raises:
+            ValueError: If the result is not exactly one row and one column.
+        """
+        if len(self.rows) != 1 or len(self.columns) != 1:
+            raise ValueError(
+                f"scalar() requires exactly one row and one column, "
+                f"got {len(self.rows)} row(s) and {len(self.columns)} column(s)"
+            )
+        return self.rows[0][self.columns[0]]

ahi/aws/s3.py

@@ -0,0 +1,277 @@
+"""Upload and download local files and folders to and from S3."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from boto3.exceptions import Boto3Error
+from boto3.s3.transfer import TransferConfig
+from botocore.exceptions import BotoCoreError, ClientError
+
+from ahi.aws.errors import DownloadFailedError, UploadFailedError
+from ahi.core import get_logger
+
+_log = get_logger("aws.s3")
+
+
+# S3 bucket naming rules (general purpose buckets).
+_BUCKET_RE = re.compile(r"[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]")
+
+_TRANSFER_ERRORS = (ClientError, BotoCoreError, Boto3Error)
+
+
+@dataclass(frozen=True)
+class S3Object:
+    """A file transferred between a local path and S3.
+
+    Attributes:
+        bucket: The bucket.
+        key: The full object key.
+        local_path: The local file.
+        size: Size in bytes.
+    """
+
+    bucket: str
+    key: str
+    local_path: Path
+    size: int
+
+    @property
+    def uri(self) -> str:
+        """The object's ``s3://bucket/key`` URI."""
+        return f"s3://{self.bucket}/{self.key}"
+
+
+class S3Storage:
+    """Upload and download files under an ``s3://bucket[/prefix]`` location.
+
+    Keys and prefixes passed to the methods are relative to that location.
+
+    Args:
+        location: The ``s3://bucket[/prefix]`` URI.
+        region_name: Region for the client built when ``client`` is not given.
+        client: A boto3 ``s3`` client, for example ``session.client("s3")``.
+        transfer_config: boto3 ``TransferConfig`` for multipart and concurrency (Manual Tune).
+
+    Raises:
+        ValueError: Invalid location, both ``client`` and ``region_name`` given,
+            or no region configured.
+    """
+
+    def __init__(
+        self,
+        location: str,
+        *,
+        region_name: str | None = None,
+        client: Any | None = None,
+        transfer_config: TransferConfig | None = None,
+    ) -> None:
+        if client is not None and region_name is not None:
+            raise ValueError("Pass region_name or client, not both; a client keeps its own region")
+        self._bucket, self._prefix = _parse_s3_uri(location)
+        self._client = client if client is not None else _build_client(region_name)
+        self._transfer_config = transfer_config
+
+    def __repr__(self) -> str:
+        return f"S3Storage({self.location!r})"
+
+    @property
+    def bucket(self) -> str:
+        """The bucket."""
+        return self._bucket
+
+    @property
+    def prefix(self) -> str:
+        """The key prefix (``""`` for the bucket root)."""
+        return self._prefix
+
+    @property
+    def location(self) -> str:
+        """The normalised ``s3://bucket[/prefix]`` URI."""
+        return self._uri(self._prefix)
+
+    def upload_file(self, path: str | os.PathLike[str], key: str | None = None) -> S3Object:
+        """Upload one file. ``key`` defaults to the file name.
+
+        Raises:
+            FileNotFoundError: ``path`` is not a file.
+            ValueError: ``key`` is empty or has ``.`` / ``..`` segments.
+            UploadFailedError: The upload failed.
+        """
+        source = Path(path)
+        if not source.is_file():
+            raise FileNotFoundError(f"Not a file: {source}")
+        return self._upload(source, _object_key(self._prefix, source.name if key is None else key))
+
+    def upload_dir(self, path: str | os.PathLike[str], prefix: str = "") -> list[S3Object]:
+        """Upload every file under a folder to ``prefix``, keeping its layout.
+
+        Raises:
+            NotADirectoryError: ``path`` is not a folder.
+            OSError: A folder under ``path`` can't be read (before any upload).
+            UploadFailedError: An upload failed.
+        """
+        root = Path(path)
+        if not root.is_dir():
+            raise NotADirectoryError(f"Not a directory: {root}")
+        base = _join_key(self._prefix, prefix)
+        # Walk first so local errors fail before any upload; as_posix() keeps
+        # Windows separators out of keys.
+        targets = [
+            (file, _object_key(base, file.relative_to(root).as_posix()))
+            for file in _walk_files(root)
+        ]
+        uploaded = [self._upload(file, key) for file, key in targets]
+        _log.info("uploaded %d file(s) from %s to %s", len(uploaded), root, self._uri(base))
+        return uploaded
+
+    def download_file(self, key: str, path: str | os.PathLike[str]) -> S3Object:
+        """Download one object to ``path``, creating folders and overwriting.
+
+        Raises:
+            IsADirectoryError: ``path`` is a folder.
+            ValueError: ``key`` is empty or has ``.`` / ``..`` segments.
+            DownloadFailedError: The download failed (``code == "404"`` if missing).
+        """
+        dest = Path(path)
+        if dest.is_dir():
+            raise IsADirectoryError(f"Is a directory: {dest}")
+        return self._download(_object_key(self._prefix, key), dest)
+
+    def download_dir(self, prefix: str, path: str | os.PathLike[str]) -> list[S3Object]:
+        """Download every object under ``prefix`` into a folder, keeping its layout.
+
+        Raises:
+            NotADirectoryError: ``path`` exists and is not a folder.
+            DownloadFailedError: Listing or a download failed, or a key can't be
+                mapped to a file inside ``path``.
+        """
+        root = Path(path)
+        if root.exists() and not root.is_dir():
+            raise NotADirectoryError(f"Not a directory: {root}")
+        base = _join_key(self._prefix, prefix)
+        list_prefix = f"{base}/" if base else ""
+        resolved_root = root.resolve()
+
+        targets: list[tuple[str, Path]] = []
+        for key in self._list_keys(list_prefix):
+            if key.endswith("/"):
+                continue
+            relative = key[len(list_prefix) :]
+            targets.append((key, self._local_target(root, resolved_root, key, relative)))
+
+        downloaded = [self._download(key, dest) for key, dest in targets]
+        _log.info("downloaded %d file(s) from %s to %s", len(downloaded), self._uri(base), root)
+        return downloaded
+
+    def _upload(self, source: Path, key: str) -> S3Object:
+        size = source.stat().st_size
+        try:
+            self._client.upload_file(str(source), self._bucket, key, Config=self._transfer_config)
+        except _TRANSFER_ERRORS as exc:
+            reason, code = _describe(exc)
+            raise UploadFailedError(self._bucket, key, reason, code=code) from exc
+        obj = S3Object(self._bucket, key, source, size)
+        _log.debug("uploaded %s to %s", source, obj.uri)
+        return obj
+
+    def _download(self, key: str, dest: Path) -> S3Object:
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            self._client.download_file(self._bucket, key, str(dest), Config=self._transfer_config)
+        except _TRANSFER_ERRORS as exc:
+            reason, code = _describe(exc)
+            raise DownloadFailedError(self._bucket, key, reason, code=code) from exc
+        obj = S3Object(self._bucket, key, dest, dest.stat().st_size)
+        _log.debug("downloaded %s to %s", obj.uri, dest)
+        return obj
+
+    def _list_keys(self, prefix: str) -> list[str]:
+        keys: list[str] = []
+        try:
+            paginator = self._client.get_paginator("list_objects_v2")
+            for page in paginator.paginate(Bucket=self._bucket, Prefix=prefix):
+                keys.extend(obj["Key"] for obj in page.get("Contents", []))
+        except _TRANSFER_ERRORS as exc:
+            reason, code = _describe(exc)
+            raise DownloadFailedError(self._bucket, prefix, reason, code=code) from exc
+        return sorted(keys)
+
+    def _local_target(self, root: Path, resolved_root: Path, key: str, relative: str) -> Path:
+        segments = relative.split("/")
+        dest = root.joinpath(*segments)
+        # resolve() also catches Windows drives and backslashes.
+        if any(s in ("", ".", "..") for s in segments) or not dest.resolve().is_relative_to(
+            resolved_root
+        ):
+            raise DownloadFailedError(
+                self._bucket, key, f"key can't be mapped to a file inside {resolved_root}"
+            )
+        return dest
+
+    def _uri(self, key: str) -> str:
+        return f"s3://{self._bucket}/{key}" if key else f"s3://{self._bucket}"
+
+
+def _build_client(region_name: str | None) -> Any:
+    import boto3
+
+    session = boto3.Session(region_name=region_name)
+    # Prevent silent fallback to us-east-1 (aws default)
+    if not session.region_name:
+        raise ValueError(
+            "No AWS region configured: pass region_name=, set AWS_DEFAULT_REGION, "
+            "add `region` to the AWS profile, or pass a client from a session with a region"
+        )
+    return session.client("s3")
+
+
+def _parse_s3_uri(uri: str) -> tuple[str, str]:
+    """Split ``s3://bucket[/prefix]`` into ``(bucket, prefix)``."""
+    scheme, sep, rest = uri.partition("://")
+    bucket, _, prefix = rest.partition("/")
+    if scheme != "s3" or not sep or not _BUCKET_RE.fullmatch(bucket):
+        raise ValueError(f"Expected an s3://bucket[/prefix] URI, got {uri!r}")
+    return bucket, _join_key(prefix)
+
+
+def _join_key(*parts: str) -> str:
+    """Join key parts with ``/``, dropping empty segments; reject ``.`` / ``..``."""
+    segments = [segment for part in parts for segment in part.split("/") if segment]
+    if any(segment in (".", "..") for segment in segments):
+        raise ValueError(f"Keys and prefixes can't contain '.' or '..' segments: {parts!r}")
+    return "/".join(segments)
+
+
+def _object_key(base: str, relative: str) -> str:
+    """Join ``relative`` onto ``base``; ``relative`` must not be empty."""
+    if not _join_key(relative):
+        raise ValueError(f"An object key can't be empty, got {relative!r}")
+    return _join_key(base, relative)
+
+
+def _walk_files(root: Path) -> Iterator[Path]:
+    """Yield files under ``root`` in sorted order, raising on unreadable folders
+    (``os.walk`` skips them by default)."""
+    for dirpath, dirnames, filenames in os.walk(root, onerror=_raise):
+        dirnames.sort()
+        for name in sorted(filenames):
+            yield Path(dirpath, name)
+
+
+def _raise(error: OSError) -> None:
+    raise error
+
+
+def _describe(exc: Exception) -> tuple[str, str | None]:
+    """Return ``(reason, code)``. ``S3UploadFailedError`` doesn't chain its
+    ``ClientError``, so also look on ``__context__``."""
+    for candidate in (exc, exc.__context__, getattr(exc, "last_exception", None)):
+        if isinstance(candidate, ClientError):
+            return str(candidate), candidate.response.get("Error", {}).get("Code")
+    return str(exc), None

ahi_aws-0.4.0.dist-info/METADATA

@@ -0,0 +1,80 @@
+Metadata-Version: 2.5
+Name: ahi-aws
+Version: 0.4.0
+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.
+Project-URL: Homepage, https://github.com/adaptivehealthintelligence/ahi-python-packages
+Project-URL: Documentation, https://adaptivehealthintelligence.github.io/ahi-python-packages/
+Project-URL: Repository, https://github.com/adaptivehealthintelligence/ahi-python-packages
+Author: Adaptive Health Intelligence
+License: MIT
+Keywords: ahi,athena,auth,aws,boto3,redshift,s3,sso,sts
+Requires-Python: >=3.11
+Requires-Dist: ahi-core<0.3,>=0.2
+Requires-Dist: boto3>=1.34
+Provides-Extra: types
+Requires-Dist: boto3-stubs[athena,redshift-data,s3,sts]>=1.34; extra == 'types'
+Description-Content-Type: text/markdown
+
+# ahi-aws
+
+Thin, typed `boto3` helpers for AHI. Today: AWS **auth** (one session, checked
+once, with clear errors), running SQL on **Redshift** (Data API) and **Athena**,
+with async query handling, and moving files and folders to and from **S3**.
+
+```python
+from ahi.aws.auth import aws_session, check_identity
+
+session = aws_session()  # AWS_PROFILE / default profile / attached role
+identity = check_identity(session)  # raises an AuthError saying what to run
+print(identity.arn, identity.method)  # e.g. "sso", "iam-role", "container-role"
+```
+
+```python
+from ahi.aws.redshift import RedshiftDataClient
+from ahi.aws.athena import AthenaClient
+
+rs = RedshiftDataClient(workgroup_name="analytics", database="dev")
+rows = rs.run_query(
+    "select count(*) from events where day = :day",
+    parameters={"day": "2026-06-26"},
+)
+print(rows.scalar())
+
+ath = AthenaClient(workgroup="primary", database="analytics")
+for row in ath.run_query("select name from users where id = ?", parameters=[42]):
+    print(row["name"])
+```
+
+```python
+from ahi.aws.auth import aws_session
+from ahi.aws.s3 import S3Storage
+
+session = aws_session()
+storage = S3Storage("s3://my-bucket/some/prefix", client=session.client("s3"))
+
+uploaded = storage.upload_dir("runs/run-01", prefix="run-01")  # list[S3Object]
+storage.download_dir("run-01", "local/run-01")
+```
+
+## Design notes
+
+- **Never logs in for you.** No browser, no AWS CLI. Run `aws sso login`
+  beforehand (local SSO), or run where an AWS role is attached (EC2 instance
+  role, ECS task role); missing or expired credentials raise `AuthError`
+  with the command to run. No default region: set `AWS_REGION` or the
+  profile's `region`.
+- **Parameterized queries by default.** Pass values via `parameters` (named
+  `:name` for Redshift, positional `?` for Athena) — never format them in.
+- **One synchronous call.** `run_query` submits, polls with backoff (see
+  `ahi.core.poll_until`), and pages the full result set, returning a
+  `QueryResult` you can iterate, `len()`, or call `.scalar()` on.
+- **Typed results from Redshift; strings from Athena** (Athena returns all values
+  as strings — we preserve that rather than guessing).
+- **Failures raise** `QueryFailedError`; timeouts raise `PollTimeoutError`.
+- **S3 transfers return `S3Object` records** (bucket, key, local path, size);
+  failures raise `UploadFailedError` / `DownloadFailedError`.
+
+Install the optional `types` extra (`ahi-aws[types]`) for `boto3-stubs` during
+development.
+
+Part of [ahi-python-packages](https://github.com/adaptivehealthintelligence/ahi-python-packages).

ahi_aws-0.4.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+ahi/aws/__init__.py,sha256=2fjWDJsepK7NKWggRPxsM_6Hf5YwImy5FH6ibVpYeZg,1071
+ahi/aws/athena.py,sha256=VuuBV1bkrya7LTKQZfg0C5mPX4AMtpMQvE9q1cZGth4,6072
+ahi/aws/auth.py,sha256=QquNlFN9yMTBQm-Af3sHd6D3Wipmo9u8dntyFSu6Xio,8685
+ahi/aws/errors.py,sha256=s9C2xby3NA-gSiJfSQgzHGwzLgfGY6LaTWvFdyUR38I,6619
+ahi/aws/redshift.py,sha256=z5M5V_55uNJu0nOcR0DPPRX2bPvLAZH7WX_iuRuSu9c,6625
+ahi/aws/results.py,sha256=9VuQKYJR6fgYOfE34BWfd6L2GkefE8t-u1b7PiZLk5w,1513
+ahi/aws/s3.py,sha256=kwNlxkRtfCXhm5lsXSJfnK6332mHbwuFhdkPfkH3qCs,10690
+ahi_aws-0.4.0.dist-info/METADATA,sha256=6gIOndW9j-utTglMYYqgCehnpDfC29si25dyYy1Qdf8,3392
+ahi_aws-0.4.0.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
+ahi_aws-0.4.0.dist-info/RECORD,,

ahi_aws-0.4.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.33.0
+Root-Is-Purelib: true
+Tag: py3-none-any