Merge branch 'main' into fix/bigquery_null_equals_incremental

Author: sriramr98

dbt-athena/.changes/unreleased/Features-20260218-110731.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Add STS AssumeRole support for cross-account access
+time: 2026-02-18T11:07:31.136966+09:00
+custom:
+    Author: dtaniwaki
+    Issue: "1657"

dbt-athena/src/dbt/adapters/athena/connections.py

@@ -61,6 +61,10 @@ class AthenaCredentials(Credentials):
     aws_access_key_id: Optional[str] = None
     aws_secret_access_key: Optional[str] = None
     aws_session_token: Optional[str] = None
+    assume_role_arn: Optional[str] = None
+    assume_role_external_id: Optional[str] = None
+    assume_role_session_name: str = "dbt-athena"
+    assume_role_duration_seconds: int = 3600
     poll_interval: float = 1.0
     debug_query_state: bool = False
     _ALIASES = {"catalog": "database"}
@@ -99,6 +103,12 @@ def _connection_keys(self) -> Tuple[str, ...]:
             "poll_interval",
             "aws_profile_name",
             "aws_access_key_id",
+            "assume_role_arn",
+            # external_id is not a secret; it is a shared condition value to prevent confused deputy attacks.
+            # See: https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_create_for-user_externalid.html
+            "assume_role_external_id",
+            "assume_role_session_name",
+            "assume_role_duration_seconds",
             "endpoint_url",
             "s3_data_dir",
             "s3_data_naming",

dbt-athena/src/dbt/adapters/athena/session.py

@@ -1,13 +1,14 @@
 import json
 import threading
 import time
-from functools import cached_property
+from functools import cached_property, lru_cache
 from hashlib import md5
-from typing import Any, Dict
+from typing import Any, Dict, NamedTuple, Optional
 from uuid import UUID
 
 import boto3
 import boto3.session
+from botocore.exceptions import ClientError
 from dbt_common.exceptions import DbtRuntimeError
 from dbt_common.invocation import get_invocation_id
 
@@ -23,25 +24,99 @@
 spark_session_list: Dict[UUID, str] = {}
 spark_session_load: Dict[UUID, int] = {}
 
+# Refresh credentials this many seconds before actual expiration
+_EXPIRY_BUFFER_SECONDS = 300
+
+
+class _AssumeRoleParams(NamedTuple):
+    assume_role_arn: Optional[str]
+    assume_role_external_id: Optional[str]
+    assume_role_session_name: str
+    assume_role_duration_seconds: int
+    region_name: str
+    num_retries: int
+
+
+def _assume_role_session(
+    base_session: boto3.session.Session,
+    credentials: Any,
+) -> boto3.session.Session:
+    duration = credentials.assume_role_duration_seconds
+    # Valid range is 900–43200 seconds per AWS STS docs:
+    # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html#API_AssumeRole_RequestParameters
+    if not (900 <= duration <= 43200):
+        raise DbtRuntimeError(
+            f"assume_role_duration_seconds must be between 900 and 43200, got {duration}"
+        )
+    ttl = duration - _EXPIRY_BUFFER_SECONDS
+    key = _AssumeRoleParams(
+        assume_role_arn=credentials.assume_role_arn,
+        assume_role_external_id=credentials.assume_role_external_id,
+        assume_role_session_name=credentials.assume_role_session_name,
+        assume_role_duration_seconds=duration,
+        region_name=credentials.region_name,
+        num_retries=credentials.effective_num_retries,
+    )
+    # Increments every ttl seconds, causing lru_cache to treat each period as a distinct call
+    ttl_bucket = int(time.time() / ttl)
+    return _get_assume_role_session(base_session, key, ttl_bucket)
+
+
+@lru_cache(maxsize=1)
+def _get_assume_role_session(
+    base_session: boto3.session.Session,
+    key: _AssumeRoleParams,
+    _ttl_hash: int,  # artificial value that changes every ttl seconds to force cache invalidation
+) -> boto3.session.Session:
+    LOGGER.debug(f"Assuming role: {key.assume_role_arn}")
+    sts_client = base_session.client(
+        "sts",
+        config=get_boto3_config(key.num_retries),
+    )
+    kwargs: Dict[str, Any] = {
+        "RoleArn": key.assume_role_arn,
+        "RoleSessionName": key.assume_role_session_name,
+    }
+    if key.assume_role_external_id:
+        kwargs["ExternalId"] = key.assume_role_external_id
+    if key.assume_role_duration_seconds:
+        kwargs["DurationSeconds"] = key.assume_role_duration_seconds
+    try:
+        response = sts_client.assume_role(**kwargs)
+    except ClientError as e:
+        raise DbtRuntimeError(f"Failed to assume role {key.assume_role_arn}: {e}") from e
+    return boto3.session.Session(
+        aws_access_key_id=response["Credentials"]["AccessKeyId"],
+        aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
+        aws_session_token=response["Credentials"]["SessionToken"],
+        region_name=key.region_name,
+    )
+
 
 def get_boto3_session(connection: Connection) -> boto3.session.Session:
-    return boto3.session.Session(
+    base_session = boto3.session.Session(
         aws_access_key_id=connection.credentials.aws_access_key_id,
         aws_secret_access_key=connection.credentials.aws_secret_access_key,
         aws_session_token=connection.credentials.aws_session_token,
         region_name=connection.credentials.region_name,
         profile_name=connection.credentials.aws_profile_name,
     )
+    if connection.credentials.assume_role_arn:
+        return _assume_role_session(base_session, connection.credentials)
+    return base_session
 
 
 def get_boto3_session_from_credentials(credentials: Any) -> boto3.session.Session:
-    return boto3.session.Session(
+    base_session = boto3.session.Session(
         aws_access_key_id=credentials.aws_access_key_id,
         aws_secret_access_key=credentials.aws_secret_access_key,
         aws_session_token=credentials.aws_session_token,
         region_name=credentials.region_name,
         profile_name=credentials.aws_profile_name,
     )
+    if credentials.assume_role_arn:
+        return _assume_role_session(base_session, credentials)
+    return base_session
 
 
 class AthenaSparkSessionManager:

dbt-athena/tests/functional/conftest.py

@@ -22,5 +22,7 @@ def dbt_profile_target():
         "poll_interval": float(os.getenv("DBT_TEST_ATHENA_POLL_INTERVAL", "1.0")),
         "num_retries": int(os.getenv("DBT_TEST_ATHENA_NUM_RETRIES", "2")),
         "aws_profile_name": os.getenv("DBT_TEST_ATHENA_AWS_PROFILE_NAME") or None,
+        "assume_role_arn": os.getenv("DBT_TEST_ATHENA_ASSUME_ROLE_ARN") or None,
+        "assume_role_external_id": os.getenv("DBT_TEST_ATHENA_ASSUME_ROLE_EXTERNAL_ID") or None,
         "spark_work_group": os.getenv("DBT_TEST_ATHENA_SPARK_WORK_GROUP"),
     }