diff --git a/src/unstract/sdk/__init__.py b/src/unstract/sdk/__init__.py index b6a4ba8d..991333a3 100644 --- a/src/unstract/sdk/__init__.py +++ b/src/unstract/sdk/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.54.0rc4" +__version__ = "0.54.0rc5" def get_sdk_version(): diff --git a/src/unstract/sdk/file_storage/env_helper.py b/src/unstract/sdk/file_storage/env_helper.py index d8a1de2c..91e2850f 100644 --- a/src/unstract/sdk/file_storage/env_helper.py +++ b/src/unstract/sdk/file_storage/env_helper.py @@ -13,8 +13,24 @@ class EnvHelper: + ENV_CONFIG_FORMAT = ( + '{"provider": "gcs", ' '"credentials": {"token": "/path/to/google/creds.json"}}' + ) + @staticmethod def get_storage(storage_type: StorageType, env_name: str) -> FileStorage: + """Helper function for clients to pick up remote storage configuration + from env, initialise the file storage for the same and return the + instance. + + Args: + storage_type: Permanent / Temporary file storage + env_name: Name of the env which has the file storage config + + Returns: + FileStorage: FIleStorage instance initialised using the provider + and credentials configured in the env + """ try: file_storage_creds = json.loads(os.environ.get(env_name)) provider = FileStorageProvider( @@ -31,7 +47,8 @@ def get_storage(storage_type: StorageType, env_name: str) -> FileStorage: raise NotImplementedError() return file_storage except KeyError as e: - logger.error(f"Required credentials is missing in the env: {str(e)}") + logger.error(f"Required credentials are missing in the env: {str(e)}") + logger.error(f"The configuration format is {EnvHelper.ENV_CONFIG_FORMAT}") raise e except FileStorageError as e: raise e diff --git a/src/unstract/sdk/file_storage/helper.py b/src/unstract/sdk/file_storage/helper.py index 911edddc..ebe5eb32 100644 --- a/src/unstract/sdk/file_storage/helper.py +++ b/src/unstract/sdk/file_storage/helper.py @@ -70,21 +70,46 @@ def local_file_system_init() -> AbstractFileSystem: def skip_local_cache(func): + """Helper function/decorator for handling FileNotFound exception and making + sure that the error is not because of stale cache. + + Args: + func: The original function that is called in the context + + Returns: + NA + """ + def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except FileNotFoundError: - try: - # FileNotFound could have been caused by stale cache. - # Hence invalidate cache and retry again - args[0].fs.invalidate_cache() - return func(*args, **kwargs) - except Exception as e: - if isinstance(e, FileNotFoundError): - raise e - else: - raise FileOperationError(str(e)) from e + _handle_file_not_found(func, *args, **kwargs) except Exception as e: raise FileOperationError(str(e)) from e return wrapper + + +def _handle_file_not_found(func, *args, **kwargs): + """Helper function for handling FileNotFound exception and making sure that + the error is not because of stale cache. + + Args: + func: The original function that is called in the context + args: The context of the function call as an array + kwargs: args to the function being called in this context + + Returns: + NA + """ + try: + # FileNotFound could have been caused by stale cache. + # Hence invalidate cache and retry again + args[0].fs.invalidate_cache() + return func(*args, **kwargs) + except Exception as e: + if isinstance(e, FileNotFoundError): + raise e + else: + raise FileOperationError(str(e)) from e diff --git a/src/unstract/sdk/tool/base.py b/src/unstract/sdk/tool/base.py index 65a2b749..948e1231 100644 --- a/src/unstract/sdk/tool/base.py +++ b/src/unstract/sdk/tool/base.py @@ -14,7 +14,7 @@ ToolEnv, ToolExecKey, ) -from unstract.sdk.file_storage.fs_shared_temporary import SharedTemporaryFileStorage +from unstract.sdk.file_storage.shared_temporary import SharedTemporaryFileStorage from unstract.sdk.tool.mixin import ToolConfigHelper from unstract.sdk.tool.parser import ToolArgsParser from unstract.sdk.tool.stream import StreamMixin @@ -61,7 +61,8 @@ def __init__(self, log_level: LogLevel = LogLevel.INFO) -> None: "Please check your settings." ) self.workflow_filestorage = SharedTemporaryFileStorage( - provider=self.filestorage_provider, **self.filestorage_credentials + provider=self.filestorage_provider, + **self.filestorage_credentials, ) @classmethod diff --git a/src/unstract/sdk/utils/tool_utils.py b/src/unstract/sdk/utils/tool_utils.py index 486f04ab..0f1cbc74 100644 --- a/src/unstract/sdk/utils/tool_utils.py +++ b/src/unstract/sdk/utils/tool_utils.py @@ -9,8 +9,7 @@ import magic from unstract.sdk.exceptions import FileStorageError -from unstract.sdk.file_storage import FileStorage, FileStorageProvider -from unstract.sdk.file_storage.fs_shared_temporary import SharedTemporaryFileStorage +from unstract.sdk.file_storage import FileStorage, FileStorageProvider, SharedTemporaryFileStorage logger = logging.getLogger(__name__) diff --git a/tests/test_fs_permanent.py b/tests/test_fs_permanent.py index 39d4b851..37e3698d 100644 --- a/tests/test_fs_permanent.py +++ b/tests/test_fs_permanent.py @@ -5,8 +5,7 @@ import pytest from dotenv import load_dotenv -from unstract.sdk.file_storage import FileStorageProvider -from unstract.sdk.file_storage.fs_permanent import PermanentFileStorage +from unstract.sdk.file_storage import FileStorageProvider, PermanentFileStorage load_dotenv()