Skip to content

Commit

Permalink
Refactoring for changed file names (#132)
Browse files Browse the repository at this point in the history
* Refactoring changed file names

* Roll version

* Update src/unstract/sdk/utils/tool_utils.py

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: Gayathri <[email protected]>

* Update tests/test_fs_permanent.py

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: Gayathri <[email protected]>

* Update tests/test_fs_permanent.py

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: Gayathri <[email protected]>

* Update src/unstract/sdk/utils/tool_utils.py

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: Gayathri <[email protected]>

* Address review comments

---------

Signed-off-by: Gayathri <[email protected]>
Co-authored-by: Chandrasekharan M <[email protected]>
  • Loading branch information
1 parent a1399a6 commit f1402c0
Show file tree
Hide file tree
Showing 6 changed files with 59 additions and 18 deletions.
2 changes: 1 addition & 1 deletion src/unstract/sdk/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
__version__ = "0.54.0rc4"
__version__ = "0.54.0rc5"


def get_sdk_version():
Expand Down
19 changes: 18 additions & 1 deletion src/unstract/sdk/file_storage/env_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
45 changes: 35 additions & 10 deletions src/unstract/sdk/file_storage/helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 3 additions & 2 deletions src/unstract/sdk/tool/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions src/unstract/sdk/utils/tool_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down
3 changes: 1 addition & 2 deletions tests/test_fs_permanent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit f1402c0

Please sign in to comment.