Skip to content

yggdrasil.aws

aws

yggdrasil AWS integration.

Mirror of the Databricks-side organization:

  • :class:AWSClient — top-level client. Owns both the static configuration (credentials, region, role, endpoint, SSO) and the lazy boto :class:Session. Singleton-cached per identity-bearing init kwarg.
  • :class:AWSService — abstract base for service objects (one per AWS service: S3, DynamoDB, ...).
  • :class:AWSResource — abstract base for individual entities (an S3 object, a DynamoDB row).
  • :class:AwsCredentials — wire-format credentials record (STS / Databricks Storage Credentials shape).

Filesystem

  • :class:S3Service (in :mod:yggdrasil.aws.fs.service) — thin S3 service object reachable as client.s3.
  • :class:S3Path (in :mod:yggdrasil.aws.fs.path) — :class:Path subclass over S3, registered for the s3:// / s3a:// / s3n:// URL schemes.

Quick start

>>> from yggdrasil.aws import AWSClient
>>> from yggdrasil.aws.fs.path import S3Path
>>>
>>> # Default chain (env / profile / instance metadata):
>>> p = S3Path("s3://my-bucket/data.parquet")
>>>
>>> # Explicit role:
>>> client = AWSClient(
...     role_arn="arn:aws:iam::1234:role/Reader",
...     region="us-east-1",
... )
>>> p = client.s3.path("s3://my-bucket/data.parquet")
>>>
>>> # IAM Identity Center (SSO) with external browser:
>>> client = AWSClient(
...     sso_start_url="https://example.awsapps.com/start",
...     sso_region="us-east-1",
...     sso_account_id="123456789012",
...     sso_role_name="DataReader",
... )

AWSAccount

AWSAccount(service: Optional[AWSService] = None, **kwargs: Any)

Bases: AWSResource

The AWS account a client authenticates as.

AWSClient.current().account AWSAccount(URL('https://us-east-1.console.aws.amazon.com/console/home?region=us-east-1')) AWSClient.current().account.account_id '123456789012'

explore_url property

explore_url: URL

AWS Console home for this account, pinned to its region.

caller_identity

caller_identity() -> dict

STS GetCallerIdentity — account, ARN, and user id of the caller.

AccountService

AccountService(client: Optional[AWSClient] = None)

Bases: AWSService

AWS account/identity service — STS-backed (GetCallerIdentity).

AWSBatch

AWSBatch(
    service: Optional[AWSService] = None,
    *,
    env: Optional[Mapping[str, str]] = None
)

Bases: AWSResource

The AWS Batch runtime context this process is running under.

>>> AWSClient.current().batch.is_batch
True
>>> AWSClient.current().batch
AWSBatch(URL('https://eu-west-1.console.aws.amazon.com/batch/home?region=eu-west-1#jobs/detail/abc'))

Off-Batch every field is None and :attr:is_batch is False; the captured env is snapshotted at construction so the resource is stable and picklable.

is_main_node property

is_main_node: bool

True for a single-node job, or the main node of a multi-node job.

explore_url property

explore_url: Optional[URL]

Console deep-link to this Batch job, or None off-Batch.

current classmethod

current(*, env: Optional[Mapping[str, str]] = None) -> 'AWSBatch'

The Batch context for the default client + live process env.

BatchService

BatchService(client: Optional[AWSClient] = None)

Bases: AWSService

AWS Batch service binding (the boto batch client lives on :attr:boto_client for future control-plane calls).

AWSClient

AWSClient(
    *,
    access_key_id: Any = ...,
    secret_access_key: Any = ...,
    session_token: Any = ...,
    region: Any = ...,
    profile: Any = ...,
    role_arn: Any = ...,
    role_session_name: Any = ...,
    external_id: Optional[str] = None,
    duration_seconds: int = 3600,
    endpoint_url: Any = ...,
    s3_addressing_style: Any = ...,
    sso_start_url: Any = ...,
    sso_region: Any = ...,
    sso_account_id: Any = ...,
    sso_role_name: Any = ...,
    refresher_key: Optional[str] = None,
    refresher: Optional[CredentialsRefresher] = None,
    singleton_ttl: Any = ...
)

Bases: Singleton

Merged AWS configuration + session + per-service client factory.

Holds every knob needed to mint a boto3 :class:Session directly on the instance — there is no separate AWSConfig class. Equality and hashing follow :meth:_singleton_key, which excludes :attr:refresher (callables aren't comparable) and lazy / transient session state. Use :attr:refresher_key as the discriminator when distinct refreshers must mint distinct clients.

Construction shapes:

  • Static credentials: pass access_key_id / secret_access_key / optional session_token.
  • Profile: pass profile (matches AWS_PROFILE); the session resolves through ~/.aws/credentials.
  • Assume-role: pass role_arn, optionally with role_session_name / external_id / duration_seconds. A refreshable credential provider drives STS AssumeRole on demand.
  • SSO (boto3-native): pass sso_start_url / sso_region / sso_account_id / sso_role_name for IAM Identity Center with external-browser device-code auth — boto3's SSOTokenProvider handles the device-code dance and token cache (typically primed by aws sso login).
  • Default chain: pass nothing. Runtime env vars (AWS_ACCESS_KEY_ID / AWS_REGION / AWS_PROFILE / AWS_ROLE_ARN / AWS_ROLE_SESSION_NAME / AWS_ENDPOINT_URL / AWS_S3_ADDRESSING_STYLE) are auto-detected so Lambda / Batch / ECS / EKS land with reasonable defaults; anything still empty falls through to boto3's own chain at session-build time.

session property

session: 'boto3.Session'

Lazily-built boto3 :class:Session. Cached.

  • :meth:has_refresher → :class:RefreshableCredentials driven by the user-supplied callback.
  • :meth:has_assume_role → :class:RefreshableCredentials driven by STS AssumeRole.
  • :meth:has_sso → boto3-native SSO token provider via a transient profile.
  • Otherwise → static / profile / default-chain creds.

s3 property

s3: 'S3Service'

The :class:S3Service bound to this client. Lazy + cached.

account_id property

account_id: str

Resolve the account ID via STS. Cached on the instance.

effective_region property

effective_region: Optional[str]

Configured region first, then boto session default.

explore_url property

explore_url: URL

AWS Console home for this client's region — clickable from code.

account property

account: 'AWSAccount'

The :class:AWSAccount resource for this client (STS-backed).

batch property

batch: 'AWSBatch'

The :class:AWSBatch runtime resource for this client.

Pure os.environ read — no network, no credentials. .is_batch gates whether the job-id / queue / array-index fields are meaningful.

to_singleton

to_singleton(ttl: Any = ...) -> 'Singleton'

Promote this instance into the per-class _INSTANCES cache.

Hot listing paths (iterdir / _ls / glob) build children with singleton_ttl=False so the bounded cache doesn't fill up with thousands of short-lived entries. When a caller decides one of those children is worth keeping around (handing it to a long-running worker, returning it from an API), :meth:to_singleton registers self into the cache so the next constructor call with the same key collapses to the same instance.

ttl defaults to the subclass's _SINGLETON_TTL (... = no caching, None = process lifetime, or a seconds count). When a different instance is already cached under this key, that pre-existing one wins and is returned unchanged — the cache is the source of truth.

invalidate_singleton

invalidate_singleton(remove_global: bool = True) -> None

Pop self from the per-class _INSTANCES cache.

Mutating ops on a Singleton-cached object (writes, deletes, schema invalidations on a Databricks table, put_object on an :class:S3Path) want to make sure the next caller asking for the same key gets a fresh build rather than collapsing onto this stale handle — that's what remove_global=True (the default) does. The pop is :meth:identity-guarded: only an entry that still points at self is removed, so a concurrent re-construction that already raced past this thread is left alone.

remove_global=False is a no-op. The keyword exists so subclass invalidators (invalidate_singleton, _invalidate_entity_tag_cache, …) can offer the same switch without branching at the call site.

has_refresher

has_refresher() -> bool

True iff a :attr:refresher callback is wired up.

Drives :meth:_build_session to mint a :class:RefreshableCredentials-backed session instead of a static one.

has_sso

has_sso() -> bool

True iff this client is configured for IAM Identity Center.

Either :attr:sso_start_url alone (token-cache flow primed by aws sso login) or the full role triple (sso_account_id + sso_role_name + sso_region) is enough — boto3 picks up whichever set is present once the profile is materialised at session-build time.

refresh_metadata

refresh_metadata() -> dict[str, Any]

Invoke :attr:refresher and return botocore-shaped metadata.

to_credentials

to_credentials() -> AwsCredentials

Snapshot the static credentials into an :class:AwsCredentials.

Returns the configured static fields; does NOT materialize assumed-role tokens — exporting a live STS token would defeat the auto-refresh that's the whole point of using a role.

to_url

to_url(scheme: Optional[str] = None) -> URL

Render this client as a URL.

Format: aws://[creds@]region/?profile=...&role_arn=...

  • Region goes in the host slot (a region is the closest AWS analog to a "host" — it parameterizes every endpoint).
  • Static creds go in user:password (when both set).
  • Everything else identity-bearing goes in the query string; the secret credential fields are emitted only via userinfo.

from_ classmethod

from_(obj: Any) -> TC

Coerce obj (str / URL / dict / AWSClient) to a client.

from_parsed_url classmethod

from_parsed_url(url: URL) -> TC

Parse an aws:// URL back into an :class:AWSClient.

from_credentials classmethod

from_credentials(
    creds: AwsCredentials,
    *,
    region: Optional[str] = None,
    endpoint_url: Optional[str] = None,
    refresher: Optional[CredentialsRefresher] = None,
    **kwargs: Any
) -> TC

Construct from a static :class:AwsCredentials.

Pass refresher for self-renewing temporary credentials.

from_refresher classmethod

from_refresher(
    refresher: CredentialsRefresher,
    *,
    region: Optional[str] = None,
    endpoint_url: Optional[str] = None,
    refresher_key: Optional[str] = None,
    **kwargs: Any
) -> TC

Build a self-refreshing :class:AWSClient from a credentials callback.

from_databricks_sql classmethod

from_databricks_sql(
    query: str,
    *,
    client: Optional["DatabricksClient"] = None,
    region: Optional[str] = None,
    endpoint_url: Optional[str] = None,
    columns: Optional[dict[str, str]] = None,
    **kwargs: Any
) -> TC

Build an :class:AWSClient that vends credentials from a Databricks SQL row.

current classmethod

current(*, reset: bool = False, **overrides: Any) -> TC

Process-global default :class:AWSClient.

reset=True rebuilds; overrides are passed to a fresh constructor. Mirrors :meth:DatabricksClient.current.

set_current classmethod

set_current(client: Optional['AWSClient']) -> None

Replace the process-global current client. Pass None to clear.

connect

connect(*, reset: bool = False) -> 'AWSClient'

Eagerly build the boto session. Idempotent.

close

close() -> None

Drop the cached session, all per-service clients, all service objects.

client

client(service: str, **overrides: Any) -> 'BaseClient'

Get a boto3 client for service. Cached per (service, overrides).

caller_identity

caller_identity() -> dict[str, Any]

Wrap STS GetCallerIdentity. Network call; not cached.

AWSResource

AWSResource(service: Optional[AWSService] = None, *args, **kwargs)

Bases: ExploreUrlRepr, ABC

Abstract base for AWS-backed entities.

Concrete resources (:class:~yggdrasil.aws.account.AWSAccount, :class:~yggdrasil.aws.fs.path.S3Bucket, …) override :attr:explore_url to return a Console deep-link; the inherited :class:ExploreUrlRepr then gives a clickable repr / _repr_html_ for free.

explore_url property

explore_url: 'Optional[URL]'

Web-UI deep-link for this resource, or None. Override in subclasses; the repr / HTML below key off it.

AWSService

AWSService(client: Optional[AWSClient] = None)

Bases: ABC

Abstract base for AWS service objects.

A service object binds an :class:AWSClient to a particular AWS service (S3, DynamoDB, …). Mirrors :class:DatabricksService: holds a client, delegates shared concerns (region, account_id, session) upstream, supports a current() singleton, and round-trips via URL.

Identity & singleton caching

Instances are cached per (class, client) in :attr:_INSTANCES. Pickling routes through :meth:__getnewargs__ so a service unpickled in the same process collapses to the live singleton. Subclasses add non-picklable handles to :attr:_TRANSIENT_STATE_ATTRS.

AwsCredentials dataclass

AwsCredentials(
    access_key_id: Optional[str] = None,
    access_point: Optional[str] = None,
    secret_access_key: Optional[str] = None,
    session_token: Optional[str] = None,
    expiration: Optional[str] = None,
)

AWS temporary credentials for API authentication.

Mirrors AWS STS's Credentials shape and the equivalent Databricks-issued temporary credentials. See https://docs.aws.amazon.com/STS/latest/APIReference/API_Credentials.html.

access_key_id class-attribute instance-attribute

access_key_id: Optional[str] = None

The access key ID that identifies the temporary credentials.

access_point class-attribute instance-attribute

access_point: Optional[str] = None

The Amazon Resource Name (ARN) of the S3 access point for temporary credentials related to the external location.

secret_access_key class-attribute instance-attribute

secret_access_key: Optional[str] = None

The secret access key that can be used to sign AWS API requests.

session_token class-attribute instance-attribute

session_token: Optional[str] = None

The token that users must pass to AWS API to use the temporary credentials.

expiration class-attribute instance-attribute

expiration: Optional[str] = None

ISO-8601 timestamp at which the credentials expire. Optional on construction; STS-issued creds always have one, long-lived keys don't.

is_complete

is_complete() -> bool

True iff at least access_key_id + secret_access_key are set.

to_botocore_metadata

to_botocore_metadata() -> Mapping[str, Optional[str]]

Render as the metadata dict botocore's RefreshableCredentials expects.

Used by :class:AWSClient when seeding refreshable creds from an initial static :class:AwsCredentials snapshot.

DatabricksSQLCredentialsRefresher dataclass

DatabricksSQLCredentialsRefresher(
    query: str,
    client: Optional["DatabricksClient"] = None,
    columns: Optional[dict[str, str]] = None,
    column_aliases: dict[str, tuple[str, ...]] = dict(),
)

Refresher that re-runs a Databricks SQL query for fresh AWS credentials.

Picklable by design — closure-free so botocore can hand the refresher to a Spark worker, a multiprocessing pool, or a cross-process job runner without cloudpickle. Serialization follows the project rule: the dataclass body holds plain picklable fields, the :class:DatabricksClient (when explicitly set) goes through its own URL-based round-trip, and the lazy DatabricksClient.current() fallback is resolved inside :meth:__call__ so the refresher stays constructible at module-import time without a live workspace.

Mapped to :class:AwsCredentials via column_aliases — each canonical field tries the user-provided columns override first, then walks the alias tuple in order. Per-call columns overrides win over the default :data:DATABRICKS_SQL_CREDENTIAL_COLUMNS.

query instance-attribute

query: str

SQL query that returns at least one row of credentials.

client class-attribute instance-attribute

client: Optional['DatabricksClient'] = None

Workspace to query. None resolves to :meth:DatabricksClient.current lazily on first call.

columns class-attribute instance-attribute

columns: Optional[dict[str, str]] = None

Per-call column-name overrides keyed by canonical field (access_key_id / secret_access_key / session_token / expiration). Wins over :attr:column_aliases.

column_aliases class-attribute instance-attribute

column_aliases: dict[str, tuple[str, ...]] = field(default_factory=dict)

Fallback alias tuples per canonical field. Populated from :data:DATABRICKS_SQL_CREDENTIAL_COLUMNS by :meth:AWSClient.from_databricks_sql; pre-snapshotted so the refresher survives a cross-process pickle without depending on the receiver's class-var defaults.

AwsCredentialsProvider

AwsCredentialsProvider(key: str)

Bases: ABC

Abstract refreshable AWS credentials provider.

Construct with a string key that uniquely identifies the credential scope. Two providers built with the same (cls, key) collapse to the same instance.

get_credentials abstractmethod

get_credentials(mode: Mode | None = None) -> AwsCredentials

Return a fresh :class:AwsCredentials. Called by botocore ~5 min before token expiry.

Subclasses that vend different credentials per read/write scope (e.g. the Databricks UC providers) use mode to pick the backing UC operation; providers that don't care about scope ignore it. None means "the subclass's default mode".

aws_client

aws_client(*, region: Optional[str] = None) -> 'AWSClient'

Return the cached :class:AWSClient for this provider / region.

First call seeds a botocore :class:RefreshableCredentials backed session by invoking self() once; subsequent calls with the same region return the same live client (and therefore share the connection pool, boto-client cache, and in-flight refresh state). Different region values mint different clients — boto region is a per-client concern.

in_aws_environment

in_aws_environment(env: Optional[Mapping[str, str]] = None) -> bool

True when running inside AWS-managed compute (Batch / ECS / Fargate / Lambda). Pure env-var probe — no IMDS round trip.