feat: StorageProvider auto-selection on StorageClient API #6

Merged
CoreRasurae merged 1 commits from feat/storage_auto_selection-#5 into master 2025-09-29 20:30:13 +00:00
7 changed files with 548 additions and 22 deletions
+2
View File
@@ -5,3 +5,5 @@
## v0.1.0
* Expose Pre-Signed URL handling methods via StorageClient API
* Expose close() method via StorageClient API
* Enable StorageProvider auto-selection on StorageClient API
+202 -22
View File
@@ -35,6 +35,32 @@ class StorageClient:
self.provider = provider
self.logger = logging.getLogger(__name__)
def _helper_find_provider_if_needed(self, path: str):
"""
Helper method that checks if the current provider supports the filesystem with a given path format and if not,
tries to temporarily instantiate a suitable provider, if one can be found.
"""
source_provider: StorageProviderInterface = None
# Check if the current provider supports the path
if not self.provider.supports_path(path):
# If not, try to get the appropriate provider for the path
from clevercloud_storage_framework.factory import StorageClientFactory
# Extract SAS token if present in the path
sas_token = None
if "?" in path:
sas_token = "?" + path.split("?", 1)[1]
# Get provider for the path, passing the SAS token if available
if sas_token:
source_provider = StorageClientFactory.get_provider_for_path(
path, sas_token
)
else:
source_provider = StorageClientFactory.get_provider_for_path(path)
return self.provider if source_provider is None else source_provider
CoreRasurae marked this conversation as resolved
Review

Idiomatic Python might prefer return source_provider or self.provider

Idiomatic Python might prefer `return source_provider or self.provider`
Review

I will take note of those suggestions for later, but i have some urgency on getting this merged, as CleverSwarm Endpoints S3 support depends on this.

  • Regarding the first suggestion yes, it makes the code more maintainable, however the code repetition is not that large and it may be easier to track in debug sessions. If a bug comes up in that particular section i will consider refactoring it your way.

  • As for the second, you are also right, i will probably open an issue with this suggestions.

I will take note of those suggestions for later, but i have some urgency on getting this merged, as CleverSwarm Endpoints S3 support depends on this. - Regarding the first suggestion yes, it makes the code more maintainable, however the code repetition is not that large and it may be easier to track in debug sessions. If a bug comes up in that particular section i will consider refactoring it your way. - As for the second, you are also right, i will probably open an issue with this suggestions.
def get_chunk_size(self) -> int:
return self.provider.get_chunk_size()
@@ -49,6 +75,20 @@ class StorageClient:
Returns:
List[Dict[str, Any]]: A list of file metadata dictionaries
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return list(source_provider.list_files(path, recursive))
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
return list(self.provider.list_files(path, recursive))
def file_exists(self, path: str) -> bool:
@@ -61,6 +101,19 @@ class StorageClient:
Returns:
bool: True if the file exists, False otherwise
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return source_provider.file_exists(path)
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
return self.provider.file_exists(path)
def get_file_size(self, path: str) -> int:
@@ -73,6 +126,19 @@ class StorageClient:
Returns:
int: The size of the file in bytes
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return source_provider.get_file_size(path)
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
return self.provider.get_file_size(path)
def get_file_metadata(self, path: str) -> Dict[str, Any]:
@@ -85,6 +151,19 @@ class StorageClient:
Returns:
Dict[str, Any]: A dictionary of file metadata
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return source_provider.get_file_metadata(path)
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
return self.provider.get_file_metadata(path)
def read_file(
@@ -98,6 +177,45 @@ class StorageClient:
local_path (str, optional): The local path to save the file to. Defaults to None.
as_iterator (bool, optional): Whether to return an iterator of bytes chunks. Defaults to False.
Returns:
Union[bytes, str, Iterator[bytes]]:
- If local_path is provided: the local path where the file was saved (str)
- If as_iterator is True: an iterator of bytes chunks (Iterator[bytes])
- Otherwise: the file contents as bytes
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return self._helper_read_file(
source_provider, path, local_path, as_iterator
)
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
return self._helper_read_file(self.provider, path, local_path, as_iterator)
def _helper_read_file(
self,
provider: StorageProviderInterface,
path: str,
local_path: Optional[str] = None,
as_iterator: bool = False,
) -> (Union)[bytes, str, Iterator[bytes]]:
"""
Read a file and return its contents, save to a local file, or return as an iterator.
Args:
provider (StorageProviderInterface): The storage provider to be used for reading the file
path (str): The path to the file
local_path (str, optional): The local path to save the file to. Defaults to None.
as_iterator (bool, optional): Whether to return an iterator of bytes chunks. Defaults to False.
Returns:
Union[bytes, str, Iterator[bytes]]:
- If local_path is provided: the local path where the file was saved (str)
@@ -110,17 +228,17 @@ class StorageClient:
# Save to local file
with open(local_path, "wb") as f:
for chunk in self.provider.read_file(path):
for chunk in provider.read_file(path):
f.write(chunk)
return local_path
elif as_iterator:
# Return the iterator directly from the provider
return self.provider.read_file(path)
return provider.read_file(path)
else:
# Return as bytes
chunks = []
for chunk in self.provider.read_file(path):
for chunk in provider.read_file(path):
chunks.append(chunk)
return b"".join(chunks)
@@ -136,6 +254,39 @@ class StorageClient:
content (Union[bytes, str, BinaryIO, Iterator[bytes]]): The content to write.
Can be bytes, a string (file path or content), a file-like object, or an iterator of bytes chunks.
Raises:
ValueError: If content is not bytes, str, a file-like object, or an iterator of bytes
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return self._helper_write_file(source_provider, path, content)
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
return self._helper_write_file(self.provider, path, content)
def _helper_write_file(
self,
provider: StorageProviderInterface,
path: str,
content: Union[bytes, str, BinaryIO, Iterator[bytes]],
) -> None:
"""
Write content to a file.
Args:
provider (StorageProviderInterface): The storage provider to be used for writing the file
path (str): The path to write to
content (Union[bytes, str, BinaryIO, Iterator[bytes]]): The content to write.
Can be bytes, a string (file path or content), a file-like object, or an iterator of bytes chunks.
Raises:
ValueError: If content is not bytes, str, a file-like object, or an iterator of bytes
"""
@@ -143,16 +294,16 @@ class StorageClient:
# If content is a string, assume it's a local file path
if os.path.exists(content):
with open(content, "rb") as f:
self.provider.write_file(path, f)
provider.write_file(path, f)
else:
# Encode the string as bytes
self.provider.write_file(path, content.encode("utf-8"))
provider.write_file(path, content.encode("utf-8"))
elif isinstance(content, bytes) or hasattr(content, "read"):
# Pass through bytes or file-like object
self.provider.write_file(path, content)
provider.write_file(path, content)
elif hasattr(content, "__iter__") and hasattr(content, "__next__"):
# Handle iterator of bytes chunks
self.provider.write_file(path, content)
provider.write_file(path, content)
else:
# Invalid content type
raise ValueError(f"Unsupported content type: {type(content)}")
@@ -169,27 +320,17 @@ class StorageClient:
"""
# Check if the current provider supports the path
if not self.provider.supports_path(path):
# If not, try to get the appropriate provider for the path
from clevercloud_storage_framework.factory import StorageClientFactory
try:
# Extract SAS token if present in the path
sas_token = None
if "?" in path:
sas_token = "?" + path.split("?", 1)[1]
# Get provider for the path, passing the SAS token if available
if sas_token:
source_provider = StorageClientFactory.get_provider_for_path(
path, sas_token
)
else:
source_provider = StorageClientFactory.get_provider_for_path(path)
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
source_provider.delete_file(path)
return
except StorageError:
# Fall back to the current provider if no suitable provider is found
pass
finally:
if "source_provider" in locals() and source_provider:
source_provider.close()
# Use the current provider if it supports the path or no other provider was found
self.provider.delete_file(path)
@@ -302,6 +443,9 @@ class StorageClient:
)
except StorageError:
source_provider = self.provider
finally:
if source_provider:
source_provider.close()
# Delete the source file using the source provider
source_provider.delete_file(source_path)
@@ -313,6 +457,20 @@ class StorageClient:
Args:
path (str): The path to the directory to create
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
source_provider.create_directory(path)
return
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
self.provider.create_directory(path)
def delete_directory(self, path: str, recursive: bool = False) -> None:
@@ -323,6 +481,19 @@ class StorageClient:
path (str): The path to the directory to delete
recursive (bool, optional): Whether to delete recursively. Defaults to False.
"""
# Obtain a suitable provider for the path if it exists
source_provider = self._helper_find_provider_if_needed(path)
if source_provider != self.provider:
try:
return source_provider.delete_directory(path, recursive)
except StorageError:
# Always fallback to the current provider
pass
finally:
if source_provider:
source_provider.close()
self.provider.delete_directory(path, recursive)
def get_provider_name(self) -> str:
@@ -356,3 +527,12 @@ class StorageClient:
UnsupportedError: If the provider does not handle pre-signed URLs
"""
return self.provider.generate_presigned_url(path, expiration)
def close(self):
"""
Close any open resources.
This method should be called when the client is no longer needed
to ensure proper cleanup of resources.
"""
self.provider.close()
@@ -315,3 +315,13 @@ class StorageProviderInterface(ABC):
bool: True if this provider can accept files from the source protocol, False otherwise
"""
pass
@abstractmethod
def close(self):
"""
Close any open resources.
This method should be called when the provider is no longer needed
to ensure proper cleanup of resources.
"""
pass
@@ -1001,3 +1001,6 @@ class EFSStorageProvider(StorageProviderInterface):
# If it doesn't start with efs://, it's not an EFS path
return False
def close(self):
pass
@@ -419,3 +419,6 @@ class GlacierStorageProvider(StorageProviderInterface):
# Glacier can only accept files from local and s3 sources
# EFS is not supported due to Glacier's specific upload requirements
return source_protocol in ["file", "s3"]
def close(self):
pass
@@ -710,3 +710,6 @@ class LocalStorageProvider(StorageProviderInterface):
# For other paths, assume they're local if they don't match other providers
return True
def close(self):
pass
+325
View File
@@ -0,0 +1,325 @@
"""
Test suite for StorageClient auto-selection functionality
=======================================================
This module contains tests for the auto-selection functionality added to StorageClient,
which automatically selects appropriate providers based on path format.
"""
import pytest
from unittest.mock import patch, MagicMock, call
from clevercloud_storage_framework.client import StorageClient
from clevercloud_storage_framework.exceptions import StorageError
class TestStorageClientAutoSelection:
"""Test suite for StorageClient auto-selection functionality."""
@pytest.fixture
def mock_provider(self):
"""Create a mock storage provider."""
provider = MagicMock()
provider.get_provider_name.return_value = "mock"
provider.supports_path.return_value = True
provider.close.return_value = None
return provider
@pytest.fixture
def mock_s3_provider(self):
"""Create a mock S3 provider."""
provider = MagicMock()
provider.get_provider_name.return_value = "s3"
provider.supports_path.return_value = True
provider.close.return_value = None
return provider
@pytest.fixture
def client(self, mock_provider):
"""Create a StorageClient with a mock provider."""
return StorageClient(mock_provider)
def test_helper_find_provider_if_needed_supported_path(self, client, mock_provider):
"""Test _helper_find_provider_if_needed when current provider supports the path."""
# Setup
path = "mock://bucket/file.txt"
mock_provider.supports_path.return_value = True
# Execute
result = client._helper_find_provider_if_needed(path)
# Verify
assert result == mock_provider
mock_provider.supports_path.assert_called_once_with(path)
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_helper_find_provider_if_needed_unsupported_path_no_sas(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test _helper_find_provider_if_needed when current provider doesn't support path and no SAS token."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
# Execute
result = client._helper_find_provider_if_needed(path)
# Verify
assert result == mock_s3_provider
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_helper_find_provider_if_needed_unsupported_path_with_sas(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test _helper_find_provider_if_needed when current provider doesn't support path and SAS token present."""
# Setup
path = "s3://bucket/file.txt?sp=rl&st=2023-01-01T00:00:00Z&se=2023-01-02T00:00:00Z&sv=2022-02-02&sr=b&sig=test"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
# Execute
result = client._helper_find_provider_if_needed(path)
# Verify
assert result == mock_s3_provider
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path, "?sp=rl&st=2023-01-01T00:00:00Z&se=2023-01-02T00:00:00Z&sv=2022-02-02&sr=b&sig=test")
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_list_files_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test list_files with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.list_files.return_value = [{"name": "file1.txt", "size": 100}]
# Execute
result = client.list_files(path, recursive=True)
# Verify
assert result == [{"name": "file1.txt", "size": 100}]
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.list_files.assert_called_once_with(path, True)
mock_s3_provider.close.assert_called_once()
mock_provider.list_files.assert_not_called()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_list_files_auto_selection_fallback(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test list_files with auto-selection when alternative provider fails and falls back."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.list_files.side_effect = StorageError("S3 error")
mock_provider.list_files.return_value = [{"name": "file2.txt", "size": 200}]
# Execute
result = client.list_files(path, recursive=False)
# Verify
assert result == [{"name": "file2.txt", "size": 200}]
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.list_files.assert_called_once_with(path, False)
mock_s3_provider.close.assert_called_once()
mock_provider.list_files.assert_called_once_with(path, False)
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_file_exists_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test file_exists with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.file_exists.return_value = True
# Execute
result = client.file_exists(path)
# Verify
assert result is True
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.file_exists.assert_called_once_with(path)
mock_s3_provider.close.assert_called_once()
mock_provider.file_exists.assert_not_called()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_file_exists_auto_selection_fallback(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test file_exists with auto-selection when alternative provider fails and falls back."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.file_exists.side_effect = StorageError("S3 error")
mock_provider.file_exists.return_value = False
# Execute
result = client.file_exists(path)
# Verify
assert result is False
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.file_exists.assert_called_once_with(path)
mock_s3_provider.close.assert_called_once()
mock_provider.file_exists.assert_called_once_with(path)
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_get_file_size_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test get_file_size with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.get_file_size.return_value = 1024
# Execute
result = client.get_file_size(path)
# Verify
assert result == 1024
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.get_file_size.assert_called_once_with(path)
mock_s3_provider.close.assert_called_once()
mock_provider.get_file_size.assert_not_called()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_get_file_metadata_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test get_file_metadata with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.get_file_metadata.return_value = {"size": 1024, "modified": "2023-01-01"}
# Execute
result = client.get_file_metadata(path)
# Verify
assert result == {"size": 1024, "modified": "2023-01-01"}
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.get_file_metadata.assert_called_once_with(path)
mock_s3_provider.close.assert_called_once()
mock_provider.get_file_metadata.assert_not_called()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_read_file_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test read_file with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.read_file.return_value = iter([b"test content"])
# Execute
result = client.read_file(path, as_iterator=True)
# Verify
assert list(result) == [b"test content"]
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.read_file.assert_called_once_with(path)
mock_s3_provider.close.assert_called_once()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_write_file_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test write_file with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/file.txt"
content = b"test content"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
# Execute
client.write_file(path, content)
# Verify
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.write_file.assert_called_once_with(path, content)
mock_s3_provider.close.assert_called_once()
mock_provider.write_file.assert_not_called()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_create_directory_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test create_directory with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/folder/"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
# Execute
client.create_directory(path)
# Verify
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.create_directory.assert_called_once_with(path)
mock_s3_provider.close.assert_called_once()
mock_provider.create_directory.assert_not_called()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_delete_directory_auto_selection_success(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test delete_directory with auto-selection when alternative provider succeeds."""
# Setup
path = "s3://bucket/folder/"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
# Execute
client.delete_directory(path, recursive=True)
# Verify
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_s3_provider.delete_directory.assert_called_once_with(path, True)
mock_s3_provider.close.assert_called_once()
mock_provider.delete_directory.assert_not_called()
def test_close_method(self, client, mock_provider):
"""Test the close method."""
# Execute
client.close()
# Verify
mock_provider.close.assert_called_once()
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_auto_selection_with_none_provider(self, mock_get_provider, client, mock_provider):
"""Test auto-selection when get_provider_for_path returns None."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = None
mock_provider.list_files.return_value = [{"name": "file.txt", "size": 100}]
# Execute
result = client.list_files(path)
# Verify
assert result == [{"name": "file.txt", "size": 100}]
mock_provider.supports_path.assert_called_once_with(path)
mock_get_provider.assert_called_once_with(path)
mock_provider.list_files.assert_called_once_with(path, False)
@patch('clevercloud_storage_framework.factory.StorageClientFactory.get_provider_for_path')
def test_auto_selection_provider_close_on_exception(self, mock_get_provider, client, mock_provider, mock_s3_provider):
"""Test that provider is closed even when an exception occurs."""
# Setup
path = "s3://bucket/file.txt"
mock_provider.supports_path.return_value = False
mock_get_provider.return_value = mock_s3_provider
mock_s3_provider.list_files.side_effect = StorageError("S3 error")
mock_provider.list_files.return_value = [{"name": "file.txt", "size": 100}]
# Execute
result = client.list_files(path)
# Verify
assert result == [{"name": "file.txt", "size": 100}]
mock_s3_provider.close.assert_called_once()
mock_provider.list_files.assert_called_once_with(path, False)