async def download_data(missing_data: dict,
communicator: BaseCommunicator) -> bool:
"""Download missing data for one data object.
:returns: status of the transfer: True for success, False for failure.
:rtype: bool
"""
data_id = missing_data["data_id"]
logger.debug("Starting data transfer for data_id %d",
missing_data["data_id"])
objects = None
for retry in range(1, RETRIES + 1):
try:
to_connector = connectors["local"]
from_connector = connectors[missing_data["connector_name"]]
from_storage_location_id = missing_data["from_storage_location_id"]
to_storage_location_id = missing_data["to_storage_location_id"]
logger.debug("Locking storage location with id %d",
from_storage_location_id)
if objects is None:
response = await communicator.send_command(
Message.command(ExecutorProtocol.GET_FILES_TO_DOWNLOAD,
from_storage_location_id))
objects = response.message_data
t = Transfer(from_connector, to_connector)
t.transfer_objects(missing_data["url"], objects)
await communicator.send_command(
Message.command(
ExecutorProtocol.DOWNLOAD_FINISHED,
to_storage_location_id,
))
return True
except DataTransferError:
logger.exception(
"Data transfer error downloading data with id {}, retry {}/{}".
format(data_id, retry, RETRIES))
except Exception:
logger.exception(
"Unknown error downloading data with id {}, retry {}/{}".
format(data_id, retry, RETRIES))
# None od the retries has been successfull, abort the download.
await communicator.send_command(
Message.command(ExecutorProtocol.DOWNLOAD_ABORTED,
to_storage_location_id))
return False
def test_max_thread(self):
t = Transfer(self.local, self.local)
with patch.object(Transfer, "_transfer_chunk") as transfer_mock:
t.transfer_objects("base", [{}] * 10)
self.assertEqual(len(transfer_mock.call_args_list), 10)
with patch.object(Transfer, "_transfer_chunk") as transfer_mock:
t.transfer_objects("base", [{}] * 20)
self.assertEqual(len(transfer_mock.call_args_list), 10)
with patch.object(Transfer, "_transfer_chunk") as transfer_mock:
t.transfer_objects("base", [{}] * 20, max_threads=20)
self.assertEqual(len(transfer_mock.call_args_list), 20)
async def download_data(missing_data: dict,
communicator: BaseCommunicator) -> bool:
"""Download missing data for one data object.
:returns: status of the transfer: True for success, False for failure.
:rtype: bool
"""
data_id = missing_data["data_id"]
to_connector = connectors[missing_data["to_connector"]]
dest_dir = os.path.join(to_connector.path, missing_data["url"])
os.makedirs(dest_dir, exist_ok=True)
logger.debug("Starting data transfer for data_id %d",
missing_data["data_id"])
objects = None
for retry in range(1, RETRIES + 1):
try:
from_connector = connectors[missing_data["connector_name"]]
from_storage_location_id = missing_data["from_storage_location_id"]
to_storage_location_id = missing_data["to_storage_location_id"]
logger.debug("Locking storage location with id %d",
from_storage_location_id)
if objects is None:
response = await communicator.send_command(
Message.command(ExecutorProtocol.GET_FILES_TO_DOWNLOAD,
from_storage_location_id))
objects = response.message_data
# Execute long running task in a threadpool.
loop = asyncio.get_event_loop()
with concurrent.futures.ThreadPoolExecutor() as pool:
t = Transfer(from_connector, to_connector)
await loop.run_in_executor(pool, t.transfer_objects,
missing_data["url"], objects)
await communicator.send_command(
Message.command(
ExecutorProtocol.DOWNLOAD_FINISHED,
to_storage_location_id,
))
return True
except DataTransferError:
logger.exception(
"Data transfer error downloading data with id {}, retry {}/{}".
format(data_id, retry, RETRIES))
except Exception:
logger.exception(
"Unknown error downloading data with id {}, retry {}/{}".
format(data_id, retry, RETRIES))
# None od the retries has been successfull, abort the download.
await communicator.send_command(
Message.command(ExecutorProtocol.DOWNLOAD_ABORTED,
to_storage_location_id))
return False
def test_retry_transfer(self):
t = Transfer(self.local, self.local)
mock: MagicMock = MagicMock(side_effect=[True, True])
t.transfer = retry_on_transfer_error(mock)
# with self.assertRaises(DataTransferError):
t.transfer_objects("test_url", [{
"path": "1"
}, {
"path": "2"
}],
max_threads=1)
self.assertEqual(len(mock.call_args_list), 2)
mock.reset_mock()
mock.side_effect = [DataTransferError, True, True]
start = time()
t.transfer_objects("test_url", [{
"path": "1"
}, {
"path": "2"
}],
max_threads=1)
end = time()
self.assertEqual(len(mock.call_args_list), 3)
self.assertGreater(end - start, 0.1)
mock.reset_mock()
mock.side_effect = [DataTransferError, DataTransferError, True, True]
start = time()
t.transfer_objects("test_url", [{
"path": "1"
}, {
"path": "2"
}],
max_threads=1)
end = time()
self.assertEqual(len(mock.call_args_list), 4)
self.assertGreater(end - start, 2 * 0.1)
mock.reset_mock()
mock.side_effect = [
DataTransferError, DataTransferError, DataTransferError
]
with self.assertRaises(DataTransferError):
t.transfer_objects("test_url", [{
"path": "1"
}, {
"path": "2"
}],
max_threads=1)
def test_exception(self):
t = Transfer(self.local, self.local)
with patch.object(Transfer, "_transfer_chunk") as transfer_mock:
transfer_mock.side_effect = [None, DataTransferError]
with self.assertRaises(DataTransferError):
t.transfer_objects("test_url", [{}, {}])
def test_ok(self):
t = Transfer(self.local, self.local)
with patch.object(Transfer, "_transfer_chunk") as transfer_mock:
t.transfer_objects("base", [{}])
transfer_mock.assert_called_once_with(Path("base"), [{}])