vitor-aignosi
981ac700d4
Implement MinIO Offload and Retention Features - Added configuration options for MinIO retention hours and offload threshold in README. - Introduced MinIO payload offloading for large DataFrame-derived payloads, storing them as parquet files. - Updated activities to utilize MinIO for data loading and cleanup, including new methods for offloading and retention management. - Refactored existing activities to integrate MinIO functionality, ensuring compatibility with previous workflows. - Removed the legacy MinioRepository class, consolidating MinIO operations under a new manager structure. - Updated requirements to use the latest version of the sientia-dataops-library.
309 lines
9.9 KiB
Python
309 lines
9.9 KiB
Python
import datetime
|
|
import os
|
|
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
from pytest import fixture, mark, raises
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.activities.postgres import Postgres
|
|
|
|
from laborious.activities.storage import Storage
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
}
|
|
}
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def storage(mock_minio_repository):
|
|
return Storage(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=AsyncMock(),
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___not_hasattr(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
storage = Storage(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
assert isinstance(storage, Postgres)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___none_minio_repository(mock_minio_repository, storage):
|
|
storage.minio_repository = None
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
storage.__init__(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___done_repository(mock_minio_repository, storage):
|
|
storage.__init__(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=AsyncMock(),
|
|
)
|
|
mock_minio_repository.assert_not_called()
|
|
assert storage.minio_repository is not None
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_minio_repository_not_initialized(storage):
|
|
storage.minio_repository = None
|
|
|
|
with raises(ValueError) as e:
|
|
await storage.query_to_minio({})
|
|
|
|
assert str(e.value) == 'Minio repository not initialized'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_not_data(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=None)
|
|
result = await storage.query_to_minio({})
|
|
|
|
storage.load_custom_query.assert_called_once_with({})
|
|
assert result['success'] is False
|
|
assert result['message'] == 'No data returned from query'
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.storage.pd.DataFrame')
|
|
@patch('laborious.activities.storage.now')
|
|
async def test_query_to_minio_success(now, dataframe, storage):
|
|
data = [{'a': 1}, {'a': 2}, {'a': 3}]
|
|
storage.load_custom_query = AsyncMock(return_value=data)
|
|
now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
|
|
storage.minio_repository.upload_file = AsyncMock(
|
|
return_value={
|
|
'minio_object_name': 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
|
|
}
|
|
)
|
|
storage.minio_repository.bucket = 'test'
|
|
|
|
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
|
|
|
|
dataframe.assert_called_once_with(data)
|
|
|
|
storage.minio_repository.upload_file.assert_called_once_with(
|
|
file_bytes=ANY,
|
|
relative_key='training_datasets/test_model/test_2024-01-01_00-00-00.parquet',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert result['success'] is True
|
|
assert (
|
|
result['object_key']
|
|
== 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
|
|
)
|
|
assert (
|
|
result['uri']
|
|
== 's3://test/sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_error(storage):
|
|
storage.send_notification = MagicMock()
|
|
storage.send_notification_async = AsyncMock()
|
|
|
|
storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
|
|
result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
|
|
assert result['success'] is False
|
|
assert result['message'] == 'test'
|
|
storage.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='ERROR_STORING_QUERY_TO_MINIO',
|
|
message='Error storing query to MinIO: test',
|
|
block='query_to_minio',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
def test_close(storage):
|
|
storage.minio_repository = MagicMock()
|
|
|
|
storage.close()
|
|
|
|
assert storage.minio_repository is None
|
|
|
|
|
|
def test___del__(storage):
|
|
storage.close = MagicMock()
|
|
|
|
storage.__del__()
|
|
|
|
storage.close.assert_called_once()
|
|
|
|
|
|
def test_estimate_payload_size_bytes(storage):
|
|
assert storage._estimate_payload_size_bytes({'x': 1}) > 0
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_query_with_minio_offload_no_rows(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=None)
|
|
result = await storage.load_query_with_minio_offload(
|
|
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
|
|
)
|
|
assert result['success'] is False
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_query_with_minio_offload_inline(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
|
|
result = await storage.load_query_with_minio_offload(
|
|
{**metadata, 'query': 'SELECT 1', 'model_name': 'my-model', 'key_prefix': 'predictions/s'}
|
|
)
|
|
assert result.get('success') is True
|
|
assert 'data' in result
|
|
assert result.get('object_key') is None
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.utils.models.minio_dataframe_payload.MinioDataFramePayload.estimate_size_bytes')
|
|
async def test_load_query_with_minio_offload_minio(mock_estimate, storage):
|
|
mock_estimate.return_value = 10**9
|
|
storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
|
|
storage.minio_repository.upload_file = AsyncMock(
|
|
return_value={
|
|
'minio_object_name': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet'
|
|
}
|
|
)
|
|
storage.minio_repository.bucket = 'test'
|
|
|
|
fixed = datetime.datetime(2024, 1, 15, 12, 30, 45)
|
|
with patch('laborious.utils.models.minio_dataframe_payload.now', return_value=fixed):
|
|
result = await storage.load_query_with_minio_offload(
|
|
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
|
|
)
|
|
|
|
assert result.get('success') is True
|
|
assert result.get('data') is None
|
|
assert (
|
|
result['object_key']
|
|
== 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet'
|
|
)
|
|
storage.minio_repository.upload_file.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
|
|
@patch('laborious.activities.storage.now')
|
|
async def test_cleanup_minio_objects_expired(mock_now, storage):
|
|
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
|
|
storage.minio_repository.list_objects = AsyncMock(
|
|
return_value=[
|
|
'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
|
|
'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet',
|
|
]
|
|
)
|
|
storage.minio_repository.delete_file = AsyncMock()
|
|
storage.send_notification_async = AsyncMock()
|
|
|
|
result = await storage.cleanup_minio_objects_expired(
|
|
{**metadata, 'prefixes': ['training_datasets/m']}
|
|
)
|
|
|
|
assert result['success'] is True
|
|
assert result['deleted_count'] == 1
|
|
storage.minio_repository.delete_file.assert_called_once_with(
|
|
object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
|
|
metadata=metadata['metadata'],
|
|
)
|