gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add support for AWS session token in S3 configuration (#1363)

Browse source 
<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->
feat: add support for AWS session token in S3 configuration

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
This commit is contained in:
Igor Ilic 2025-09-10 14:59:51 +02:00 • committed by GitHub
commit b137d680e2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 13 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cognee/infrastructure/files/storage/S3FileStorage.py
View file

@ -21,10 +21,15 @@ class S3FileStorage(Storage):
def __init__(self, storage_path: str): def __init__(self, storage_path: str):
self.storage_path = storage_path self.storage_path = storage_path
s3_config = get_s3_config() s3_config = get_s3_config()
if s3_config.aws_access_key_id is not None and s3_config.aws_secret_access_key is not None: if (
s3_config.aws_access_key_id is not None
and s3_config.aws_secret_access_key is not None
and s3_config.aws_session_token is not None
):
self.s3 = s3fs.S3FileSystem( self.s3 = s3fs.S3FileSystem(
key=s3_config.aws_access_key_id, key=s3_config.aws_access_key_id,
secret=s3_config.aws_secret_access_key, secret=s3_config.aws_secret_access_key,
token=s3_config.aws_session_token,
anon=False, anon=False,
endpoint_url=s3_config.aws_endpoint_url, endpoint_url=s3_config.aws_endpoint_url,
client_kwargs={"region_name": s3_config.aws_region}, client_kwargs={"region_name": s3_config.aws_region},

1
cognee/infrastructure/files/storage/s3_config.py
View file

@ -8,6 +8,7 @@ class S3Config(BaseSettings):
aws_endpoint_url: Optional[str] = None aws_endpoint_url: Optional[str] = None
aws_access_key_id: Optional[str] = None aws_access_key_id: Optional[str] = None
aws_secret_access_key: Optional[str] = None aws_secret_access_key: Optional[str] = None
aws_session_token: Optional[str] = None
model_config = SettingsConfigDict(env_file=".env", extra="allow") model_config = SettingsConfigDict(env_file=".env", extra="allow")

5
cognee/tasks/ingestion/resolve_data_directories.py
View file

@ -32,7 +32,10 @@ async def resolve_data_directories(
import s3fs import s3fs
fs = s3fs.S3FileSystem( fs = s3fs.S3FileSystem(
key=s3_config.aws_access_key_id, secret=s3_config.aws_secret_access_key, anon=False key=s3_config.aws_access_key_id,
secret=s3_config.aws_secret_access_key,
token=s3_config.aws_session_token,
anon=False,
) )
for item in data: for item in data:

2
distributed/pyproject.toml
View file

@ -40,7 +40,7 @@ dependencies = [
"aiohttp>=3.11.14,<4.0.0", "aiohttp>=3.11.14,<4.0.0",
"aiofiles>=23.2.1,<24.0.0", "aiofiles>=23.2.1,<24.0.0",
"rdflib>=7.1.4,<7.2.0", "rdflib>=7.1.4,<7.2.0",
"pypdf>=4.1.0,<6.0.0", "pypdf>=4.1.0,<7.0.0",
"jinja2>=3.1.3,<4", "jinja2>=3.1.3,<4",
"matplotlib>=3.8.3,<4", "matplotlib>=3.8.3,<4",
"networkx>=3.4.2,<4", "networkx>=3.4.2,<4",

2
pyproject.toml
View file

@ -41,7 +41,7 @@ dependencies = [
"aiohttp>=3.11.14,<4.0.0", "aiohttp>=3.11.14,<4.0.0",
"aiofiles>=23.2.1,<24.0.0", "aiofiles>=23.2.1,<24.0.0",
"rdflib>=7.1.4,<7.2.0", "rdflib>=7.1.4,<7.2.0",
"pypdf>=4.1.0,<6.0.0", "pypdf>=4.1.0,<7.0.0",
"jinja2>=3.1.3,<4", "jinja2>=3.1.3,<4",
"matplotlib>=3.8.3,<4", "matplotlib>=3.8.3,<4",
"networkx>=3.4.2,<4", "networkx>=3.4.2,<4",