MINOR: Typo for datetime attribute (#19420)

This commit is contained in:
Ayush Shah 2025-01-17 14:29:36 +05:30 committed by GitHub
parent 69d413ac34
commit 7d21a8295b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -11,6 +11,7 @@
""" """
Module containing AWS Client Module containing AWS Client
""" """
import datetime
from enum import Enum from enum import Enum
from functools import partial from functools import partial
from typing import Any, Callable, Dict, Optional, Type, TypeVar from typing import Any, Callable, Dict, Optional, Type, TypeVar
@ -47,11 +48,22 @@ class AWSAssumeRoleException(Exception):
""" """
class AWSAssumeRoleCredentialResponse(BaseModel):
AccessKeyId: str = Field()
SecretAccessKey: str = Field()
SessionToken: Optional[str] = Field(
default=None,
)
Expiration: Optional[datetime.datetime] = None
class AWSAssumeRoleCredentialWrapper(BaseModel): class AWSAssumeRoleCredentialWrapper(BaseModel):
accessKeyId: str = Field(alias="access_key") accessKeyId: str = Field()
secretAccessKey: CustomSecretStr = Field(alias="secret_key") secretAccessKey: CustomSecretStr = Field()
sessionToken: Optional[str] = Field(default=None, alias="token") sessionToken: Optional[str] = Field(
expiryTime: Optional[str] = Field(alias="expiry_time") default=None,
)
expiryTime: Optional[str] = Field()
AWSAssumeRoleCredentialFormat = TypeVar( AWSAssumeRoleCredentialFormat = TypeVar(
@ -102,12 +114,14 @@ class AWSClient:
) )
if resp: if resp:
credentials = resp.get("Credentials", {}) credentials: AWSAssumeRoleCredentialResponse = (
AWSAssumeRoleCredentialResponse(**resp.get("Credentials", {}))
)
creds_wrapper = AWSAssumeRoleCredentialWrapper( creds_wrapper = AWSAssumeRoleCredentialWrapper(
accessKeyId=credentials.get("AccessKeyId"), accessKeyId=credentials.AccessKeyId,
secretAccessKey=credentials.get("SecretAccessKey"), secretAccessKey=credentials.SecretAccessKey,
sessionToken=credentials.get("SessionToken"), sessionToken=credentials.SessionToken,
expiryTime=credentials.get("Expiration").isformat(), expiryTime=credentials.Expiration.isoformat(),
) )
if return_type == Dict: if return_type == Dict:
return creds_wrapper.model_dump(by_alias=True) return creds_wrapper.model_dump(by_alias=True)
@ -143,9 +157,11 @@ class AWSClient:
return Session( return Session(
aws_access_key_id=aws_access_key_id, aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key.get_secret_value() aws_secret_access_key=(
aws_secret_access_key.get_secret_value()
if aws_secret_access_key if aws_secret_access_key
else None, else None
),
aws_session_token=aws_session_token, aws_session_token=aws_session_token,
region_name=aws_region, region_name=aws_region,
profile_name=profile, profile_name=profile,