-
Notifications
You must be signed in to change notification settings - Fork 83
feat(clp-py-utils): Add database-backed AWS credential management system. #1486
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: s3-credentials-manager
Are you sure you want to change the base?
Changes from 14 commits
f14178b
08682e0
5e0793e
a395256
535203c
9cac4c3
3906383
febfb75
5344d16
85ce7f9
517c995
b1f70d1
e7f8853
99c1dbf
d3808b2
58f05e0
417d73e
2576b2e
2391c19
4fc9a89
2b3258a
36c9a65
51c92fc
4be0a49
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,6 @@ | ||
| import os | ||
| import pathlib | ||
| from datetime import datetime | ||
| from enum import auto | ||
| from typing import Annotated, Any, ClassVar, Literal, Optional, Union | ||
|
|
||
|
|
@@ -11,6 +12,7 @@ | |
| model_validator, | ||
| PlainSerializer, | ||
| PrivateAttr, | ||
| SecretStr, | ||
| ) | ||
| from strenum import KebabCaseStrEnum, LowercaseStrEnum | ||
|
|
||
|
|
@@ -413,6 +415,96 @@ def validate_authentication(cls, data): | |
| return data | ||
|
|
||
|
|
||
| class AwsCredential(BaseModel): | ||
| """ | ||
| Represents a stored AWS credential retrieved from the database. | ||
|
|
||
| This model is used for credentials that are persisted in the aws_credentials table. | ||
| Credentials can be either static (access key + secret key) or configured for role | ||
| assumption (with role_arn set for Phase 2). | ||
| """ | ||
|
|
||
| id: int | ||
| name: Annotated[ | ||
| str, | ||
| Field( | ||
| min_length=1, | ||
| max_length=255, | ||
| pattern=r"^[a-zA-Z0-9_-]+$", | ||
| description="Credential name (alphanumeric, hyphens, underscores only; 1-255 characters)", | ||
| ), | ||
| ] | ||
|
|
||
| access_key_id: SecretStr | ||
| secret_access_key: SecretStr | ||
| role_arn: str | None = None | ||
|
|
||
| created_at: datetime = datetime.now() | ||
| updated_at: datetime = datetime.now() | ||
|
|
||
| def to_s3_credentials(self) -> S3Credentials: | ||
| """ | ||
| Converts to S3Credentials for use with boto3. | ||
|
|
||
| Note: This only works for static credentials. For temporary credentials | ||
| with session tokens, use the TemporaryCredential model instead. | ||
|
|
||
| :return: S3Credentials object with secrets revealed. | ||
| """ | ||
| return S3Credentials( | ||
| access_key_id=self.access_key_id.get_secret_value(), | ||
| secret_access_key=self.secret_access_key.get_secret_value(), | ||
| session_token=None, | ||
| ) | ||
|
|
||
|
|
||
| class TemporaryCredential(BaseModel): | ||
| """ | ||
| Represents cached temporary credentials (session tokens). | ||
|
|
||
| This model is used for credentials cached in the aws_temporary_credentials table. | ||
| These credentials can come from various sources: | ||
| - STS AssumeRole operations | ||
| - Resource-specific session tokens | ||
|
|
||
| The 'source' field tracks the origin of the session token, which can be: | ||
| - A role ARN: "arn:aws:iam::123456789012:role/MyRole" | ||
| - An S3 resource ARN: "arn:aws:s3:::bucket/path/*" | ||
| """ | ||
|
|
||
| id: int | ||
| long_term_key_id: int # Foreign key to aws_credentials table | ||
| access_key_id: SecretStr | ||
| secret_access_key: SecretStr | ||
| session_token: SecretStr | ||
| source: str # Role ARN or S3 resource ARN | ||
| expires_at: datetime | ||
| created_at: datetime | ||
|
|
||
| def to_s3_credentials(self) -> S3Credentials: | ||
| """ | ||
| Converts to S3Credentials for use with boto3. | ||
|
|
||
| :return: S3Credentials object with secrets revealed. | ||
| """ | ||
| return S3Credentials( | ||
| access_key_id=self.access_key_id.get_secret_value(), | ||
| secret_access_key=self.secret_access_key.get_secret_value(), | ||
| session_token=self.session_token.get_secret_value(), | ||
| ) | ||
|
|
||
| def is_expired(self, buffer_minutes: int = 5) -> bool: | ||
| """ | ||
| Checks if credential is expired or expiring soon. | ||
|
|
||
| :param buffer_minutes: Minutes of buffer before expiration to consider credential expired. | ||
| :return: True if expired or expiring within buffer_minutes. | ||
| """ | ||
| from datetime import timedelta | ||
|
|
||
| return datetime.now() >= self.expires_at - timedelta(minutes=buffer_minutes) | ||
|
||
|
|
||
|
|
||
| class S3Config(BaseModel): | ||
| region_code: NonEmptyStr | ||
| bucket: NonEmptyStr | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -13,6 +13,8 @@ | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ARCHIVE_TAGS_TABLE_SUFFIX = "archive_tags" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ARCHIVES_TABLE_SUFFIX = "archives" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| AWS_CREDENTIALS_TABLE_SUFFIX = "aws_credentials" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| AWS_TEMPORARY_CREDENTIALS_TABLE_SUFFIX = "aws_temporary_credentials" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| COLUMN_METADATA_TABLE_SUFFIX = "column_metadata" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| DATASETS_TABLE_SUFFIX = "datasets" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| FILES_TABLE_SUFFIX = "files" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -21,6 +23,8 @@ | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| TABLE_SUFFIX_MAX_LEN = max( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(ARCHIVE_TAGS_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(ARCHIVES_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(AWS_CREDENTIALS_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(AWS_TEMPORARY_CREDENTIALS_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(COLUMN_METADATA_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(DATASETS_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| len(FILES_TABLE_SUFFIX), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -110,6 +114,48 @@ def _create_column_metadata_table(db_cursor, table_prefix: str, dataset: str) -> | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def _create_aws_credentials_table(db_cursor, aws_credentials_table_name: str) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| db_cursor.execute( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| f""" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| CREATE TABLE IF NOT EXISTS `{aws_credentials_table_name}` ( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `id` INT NOT NULL AUTO_INCREMENT, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `name` VARCHAR(255) NOT NULL UNIQUE, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `access_key_id` VARCHAR(255) NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `secret_access_key` VARCHAR(255) NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `role_arn` VARCHAR(2048), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `created_at` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `updated_at` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| PRIMARY KEY (`id`), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| UNIQUE KEY `name_unique` (`name`) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) ROW_FORMAT=DYNAMIC | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
117
to
131
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion | 🟠 Major Remove redundant unique definition on You have both CREATE TABLE IF NOT EXISTS `{aws_credentials_table_name}` (
`id` INT NOT NULL AUTO_INCREMENT,
- `name` VARCHAR(255) NOT NULL UNIQUE,
+ `name` VARCHAR(255) NOT NULL UNIQUE,
`access_key_id` VARCHAR(255) NOT NULL,
`secret_access_key` VARCHAR(255) NOT NULL,
`role_arn` VARCHAR(2048),
`created_at` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
`updated_at` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
PRIMARY KEY (`id`),
- UNIQUE KEY `name_unique` (`name`)
) ROW_FORMAT=DYNAMIC📝 Committable suggestion
Suggested change
🤖 Prompt for AI Agents |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def _create_aws_temporary_credentials_table( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| db_cursor, aws_temporary_credentials_table_name: str, aws_credentials_table_name: str | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| db_cursor.execute( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| f""" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| CREATE TABLE IF NOT EXISTS `{aws_temporary_credentials_table_name}` ( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `id` INT NOT NULL AUTO_INCREMENT, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `long_term_key_id` INT NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `access_key_id` VARCHAR(255) NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `secret_access_key` VARCHAR(255) NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `session_token` VARCHAR(512) NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `source` VARCHAR(2048) NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `expires_at` DATETIME NOT NULL, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| `created_at` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| PRIMARY KEY (`id`), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| KEY `long_term_key_expires` (`long_term_key_id`, `expires_at`), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| KEY `source_expires` (`source`, `expires_at`), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| FOREIGN KEY (`long_term_key_id`) REFERENCES `{aws_credentials_table_name}` (`id`) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ON DELETE CASCADE | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) ROW_FORMAT=DYNAMIC | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
Eden-D-Zhang marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def _get_table_name(prefix: str, suffix: str, dataset: str | None) -> str: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| :param prefix: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -145,6 +191,34 @@ def create_datasets_table(db_cursor, table_prefix: str) -> None: | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def create_aws_credentials_table(db_cursor, table_prefix: str) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Creates the AWS credentials table for storing user-managed static credentials. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| :param db_cursor: The database cursor to execute the table creation. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| :param table_prefix: A string to prepend to the table name. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| aws_credentials_table_name = get_aws_credentials_table_name(table_prefix) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _create_aws_credentials_table(db_cursor, aws_credentials_table_name) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def create_aws_temporary_credentials_table(db_cursor, table_prefix: str) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Creates the AWS temporary credentials table for storing cached session tokens. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| This table caches session tokens from various sources (user-provided, role assumption, etc.) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| to enable efficient credential reuse. It references the aws_credentials table via foreign key. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| :param db_cursor: The database cursor to execute the table creation. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| :param table_prefix: A string to prepend to the table name. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| aws_credentials_table_name = get_aws_credentials_table_name(table_prefix) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| aws_temporary_credentials_table_name = get_aws_temporary_credentials_table_name(table_prefix) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _create_aws_temporary_credentials_table( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| db_cursor, aws_temporary_credentials_table_name, aws_credentials_table_name | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def add_dataset( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| db_conn, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| db_cursor, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -300,6 +374,14 @@ def get_archives_table_name(table_prefix: str, dataset: str | None) -> str: | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return _get_table_name(table_prefix, ARCHIVES_TABLE_SUFFIX, dataset) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def get_aws_credentials_table_name(table_prefix: str) -> str: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return _get_table_name(table_prefix, AWS_CREDENTIALS_TABLE_SUFFIX, None) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def get_aws_temporary_credentials_table_name(table_prefix: str) -> str: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return _get_table_name(table_prefix, AWS_TEMPORARY_CREDENTIALS_TABLE_SUFFIX, None) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def get_column_metadata_table_name(table_prefix: str, dataset: str | None) -> str: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return _get_table_name(table_prefix, COLUMN_METADATA_TABLE_SUFFIX, dataset) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fix tz-naive defaults and import; avoid import-time evaluation.
Use Field(default_factory=...) with UTC to prevent tz-naive values and import-time evaluation.
Also applies to: 442-443
🤖 Prompt for AI Agents