diff --git a/great_ai/src/great_ai/large_file/__init__.py b/great_ai/src/great_ai/large_file/__init__.py index 437e068..576173b 100644 --- a/great_ai/src/great_ai/large_file/__init__.py +++ b/great_ai/src/great_ai/large_file/__init__.py @@ -1 +1 @@ -from .large_file import LargeFile, LargeFileS3 +from .large_file import LargeFile, LargeFileLocal, LargeFileMongo, LargeFileS3 diff --git a/great_ai/src/great_ai/large_file/large_file/large_file.py b/great_ai/src/great_ai/large_file/large_file/large_file.py index 57883dc..ce96ca8 100644 --- a/great_ai/src/great_ai/large_file/large_file/large_file.py +++ b/great_ai/src/great_ai/large_file/large_file/large_file.py @@ -2,7 +2,7 @@ import configparser import os import shutil import tempfile -from abc import ABC, abstractclassmethod, abstractmethod +from abc import ABC, abstractmethod from pathlib import Path from types import TracebackType from typing import IO, Any, List, Optional, Type, Union, cast @@ -48,6 +48,7 @@ class LargeFile(ABC): ``` """ + initialized = False cache_path = Path(".cache") max_cache_size: Optional[str] = "30GB" @@ -96,11 +97,11 @@ class LargeFile(ABC): credentials.default_section cls.configure_credentials(**credentials[credentials.default_section]) - @abstractclassmethod + @classmethod def configure_credentials( cls, ) -> None: - pass + cls.initialized = True def __enter__(self) -> IO: self._file: IO[Any] = ( @@ -236,7 +237,6 @@ class LargeFile(ABC): ), version=int(f.name.split(CACHE_NAME_VERSION_SEPARATOR)[-1]), remote_path=f, - origin="filesystem", ) for f in self.cache_path.glob( f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}*" diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_local.py b/great_ai/src/great_ai/large_file/large_file/large_file_local.py index 0b76749..ed39560 100644 --- a/great_ai/src/great_ai/large_file/large_file/large_file_local.py +++ b/great_ai/src/great_ai/large_file/large_file/large_file_local.py @@ -33,6 +33,7 @@ class LargeFileLocal(LargeFile): keep_last_n=keep_last_n, cache_only_mode=True, ) + super().configure_credentials() def _find_remote_instances(self) -> List[DataInstance]: return [] diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py b/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py index 836b828..f80c4ef 100644 --- a/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py +++ b/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py @@ -32,6 +32,7 @@ class LargeFileMongo(LargeFile): ) -> None: cls.connection_string = connection_string cls.database = database + super().configure_credentials() @cached_property def _client(self) -> GridFSBucket: diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_s3.py b/great_ai/src/great_ai/large_file/large_file/large_file_s3.py index ad15df7..35a02b7 100644 --- a/great_ai/src/great_ai/large_file/large_file/large_file_s3.py +++ b/great_ai/src/great_ai/large_file/large_file/large_file_s3.py @@ -66,6 +66,7 @@ class LargeFileS3(LargeFile): cls.secret_access_key = aws_secret_access_key cls.bucket_name = large_files_bucket_name cls.endpoint_url = aws_endpoint_url + super().configure_credentials() @cached_property def _client(self) -> boto3.client: @@ -99,7 +100,6 @@ class LargeFileS3(LargeFile): name=o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0], version=int(o["Key"].split(S3_NAME_VERSION_SEPARATOR)[-1]), remote_path=o["Key"], - origin="s3", ) for o in found_objects["Contents"] if o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0] == self._name diff --git a/great_ai/src/great_ai/large_file/models/data_instance.py b/great_ai/src/great_ai/large_file/models/data_instance.py index 7de533c..bfef26c 100644 --- a/great_ai/src/great_ai/large_file/models/data_instance.py +++ b/great_ai/src/great_ai/large_file/models/data_instance.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Any from pydantic import BaseModel @@ -7,4 +7,3 @@ class DataInstance(BaseModel): name: str version: int remote_path: Any - origin: Literal["filesystem", "mongodb", "s3"]