From 914fcacac0301e4acca9433661affb0b5c73bec1 Mon Sep 17 00:00:00 2001 From: Ihor Sokhan Date: Thu, 16 Jul 2026 19:39:33 +0300 Subject: [PATCH 1/2] pass downloaded bytes, IP and completed flag to callback --- tests/core/test_remote_logging.py | 71 +++++++++++++++++++ waterbutler/core/remote_logging.py | 11 ++- waterbutler/server/api/v0/core.py | 10 +-- waterbutler/server/api/v0/crud.py | 4 +- waterbutler/server/api/v0/zip.py | 2 +- .../server/api/v1/provider/__init__.py | 11 ++- 6 files changed, 96 insertions(+), 13 deletions(-) diff --git a/tests/core/test_remote_logging.py b/tests/core/test_remote_logging.py index 88ca8bd428..cdfb55fad4 100644 --- a/tests/core/test_remote_logging.py +++ b/tests/core/test_remote_logging.py @@ -3,6 +3,77 @@ from waterbutler.core import remote_logging +class TestLogToCallback: + + @pytest.mark.asyncio + @pytest.mark.parametrize('completed, expected', [(False, False), (True, True)]) + async def test_download_action_sets_completed_flag(self, monkeypatch, completed, expected): + captured = {} + + async def fake_send_signed_request(method, url, payload): + captured['payload'] = payload + return 200, b'success' + + monkeypatch.setattr(remote_logging.utils, 'send_signed_request', fake_send_signed_request) + + class DummySource: + auth = {'callback_url': 'https://example.com/callback'} + + def serialize(self): + return {'provider': 'osf'} + + source = DummySource() + request = { + 'request': { + 'method': 'GET', + 'url': 'https://example.com/file', + 'headers': {}, + }, + 'referrer': {'url': None}, + 'tech': {'ua': 'test-agent', 'ip': '127.0.0.1'}, + } + + await remote_logging.log_to_callback( + 'download_file', + source=source, + request=request, + completed=completed, + ) + + assert captured['payload']['action_meta']['completed'] is expected + + @pytest.mark.asyncio + async def test_non_download_action_omits_completed_flag(self, monkeypatch): + captured = {} + + async def fake_send_signed_request(method, url, payload): + captured['payload'] = payload + return 200, b'success' + + monkeypatch.setattr(remote_logging.utils, 'send_signed_request', fake_send_signed_request) + + class DummySource: + auth = {'callback_url': 'https://example.com/callback'} + + def serialize(self): + return {'provider': 'osf'} + + source = DummySource() + request = { + 'request': { + 'method': 'GET', + 'url': 'https://example.com/file', + 'headers': {}, + }, + 'referrer': {'url': None}, + 'tech': {'ua': 'test-agent', 'ip': '127.0.0.1'}, + } + + await remote_logging.log_to_callback('create', source=source, request=request) + + assert 'completed' not in captured['payload']['action_meta'] + + class TestScrubPayloadForKeen: def test_flat_dict(self): diff --git a/waterbutler/core/remote_logging.py b/waterbutler/core/remote_logging.py index bc2d256cfc..9debe58c47 100644 --- a/waterbutler/core/remote_logging.py +++ b/waterbutler/core/remote_logging.py @@ -18,7 +18,7 @@ @utils.async_retry(retries=5, backoff=5) async def log_to_callback(action, source=None, destination=None, start_time=None, errors=None, - request=None): + request=None, bytes_downloaded=0, completed=False): """PUT a logging payload back to the callback given by the auth provider.""" errors = errors or [] request = request or {} @@ -59,7 +59,10 @@ async def log_to_callback(action, source=None, destination=None, start_time=None is_mfr_render = (ref_url_domain == settings.MFR_DOMAIN or settings.MFR_IDENTIFYING_HEADER in request["request"]["headers"]) log_payload['action_meta']['is_mfr_render'] = is_mfr_render + log_payload['action_meta']['completed'] = completed + log_payload['action_meta']['bytes_downloaded'] = bytes_downloaded + log_payload['action_meta']['ip'] = request['tech']['ip'] resp_status, resp_data = await utils.send_signed_request('PUT', auth['callback_url'], log_payload) if resp_status // 100 != 2: @@ -217,12 +220,14 @@ async def _send_to_keen(payload, collection, project_id, write_key, action, doma def log_file_action(action, source, api_version, destination=None, request=None, - start_time=None, errors=None, bytes_downloaded=None, bytes_uploaded=None): + start_time=None, errors=None, bytes_downloaded=None, bytes_uploaded=None, + completed=False): """Kick off logging actions in the background. Returns array of asyncio.Tasks.""" request = request or {} return [ log_to_callback(action, source=source, destination=destination, - start_time=start_time, errors=errors, request=request,), + start_time=start_time, errors=errors, request=request, + bytes_downloaded=bytes_downloaded, completed=completed,), asyncio.ensure_future( log_to_keen(action, source=source, destination=destination, errors=errors, request=request, api_version=api_version, diff --git a/waterbutler/server/api/v0/core.py b/waterbutler/server/api/v0/core.py index 543d966a26..af702e1867 100644 --- a/waterbutler/server/api/v0/core.py +++ b/waterbutler/server/api/v0/core.py @@ -94,12 +94,13 @@ async def prepare(self): self.path = await self.provider.validate_path(**self.arguments) self.arguments['path'] = self.path # TODO Not this - def _send_hook(self, action, metadata=None, path=None): + def _send_hook(self, action, metadata=None, path=None, completed=False): source = LogPayload(self.arguments['nid'], self.provider, metadata=metadata, path=path) remote_logging.log_file_action(action, source=source, api_version='v0', request=remote_logging._serialize_request(self.request), bytes_downloaded=self.bytes_downloaded, - bytes_uploaded=self.bytes_uploaded) + bytes_uploaded=self.bytes_uploaded, + completed=completed) class BaseCrossProviderHandler(BaseHandler): @@ -140,7 +141,7 @@ def json(self): return self._json - def _send_hook(self, action, metadata): + def _send_hook(self, action, metadata, completed=False): source = LogPayload(self.json['source']['nid'], self.source_provider, path=self.json['source']['path']) destination = LogPayload(self.json['destination']['nid'], self.destination_provider, @@ -148,4 +149,5 @@ def _send_hook(self, action, metadata): remote_logging.log_file_action(action, source=source, destination=destination, api_version='v0', request=remote_logging._serialize_request(self.request), bytes_downloaded=self.bytes_downloaded, - bytes_uploaded=self.bytes_uploaded) + bytes_uploaded=self.bytes_uploaded, + completed=completed) diff --git a/waterbutler/server/api/v0/crud.py b/waterbutler/server/api/v0/crud.py index 07034155fa..e7271faecb 100644 --- a/waterbutler/server/api/v0/crud.py +++ b/waterbutler/server/api/v0/crud.py @@ -61,7 +61,7 @@ async def get(self): if isinstance(result, str): self.redirect(result) - self._send_hook('download_file', path=self.path) + self._send_hook('download_file', path=self.path, completed=True) return if getattr(result, 'partial', None): @@ -87,7 +87,7 @@ async def get(self): self.set_header('Content-Type', mime_types[ext]) await self.write_stream(result) - self._send_hook('download_file', path=self.path) + self._send_hook('download_file', path=self.path, completed=True) async def post(self): """Create a folder""" diff --git a/waterbutler/server/api/v0/zip.py b/waterbutler/server/api/v0/zip.py index 2ed0695c65..b52fe1f24f 100644 --- a/waterbutler/server/api/v0/zip.py +++ b/waterbutler/server/api/v0/zip.py @@ -21,4 +21,4 @@ async def get(self): result = await self.provider.zip(**self.arguments) await self.write_stream(result) - self._send_hook('download_zip', path=self.path) + self._send_hook('download_zip', path=self.path, completed=True) diff --git a/waterbutler/server/api/v1/provider/__init__.py b/waterbutler/server/api/v1/provider/__init__.py index 3cf8ed3184..10d85f2ace 100644 --- a/waterbutler/server/api/v1/provider/__init__.py +++ b/waterbutler/server/api/v1/provider/__init__.py @@ -246,6 +246,7 @@ def on_finish(self): 'zip' not in self.request.query_arguments))): return + completed = False # Done here just because method is defined action = { 'GET': lambda: 'download_file' if self.path.is_file else 'download_zip', @@ -254,9 +255,12 @@ def on_finish(self): 'DELETE': lambda: 'delete' }[method]() - self._send_hook(action) + if action in {'download_file', 'download_zip'}: + completed = status in {200, 302} - def _send_hook(self, action): + self._send_hook(action, completed=completed) + + def _send_hook(self, action, completed=False): source = None destination = None @@ -281,4 +285,5 @@ def _send_hook(self, action): remote_logging.log_file_action(action, source=source, destination=destination, api_version='v1', request=remote_logging._serialize_request(self.request), bytes_downloaded=self.bytes_downloaded, - bytes_uploaded=self.bytes_uploaded,) + bytes_uploaded=self.bytes_uploaded, + completed=completed) From f12ad8f813de4a60237ea160db4c508c49cb3b51 Mon Sep 17 00:00:00 2001 From: Ihor Sokhan Date: Fri, 17 Jul 2026 17:47:22 +0300 Subject: [PATCH 2/2] make completed depend on download state --- tests/server/api/v1/test_metadata_mixin.py | 14 ++++++++++++++ tests/server/api/v1/test_provider.py | 6 ++++-- waterbutler/server/api/v0/crud.py | 4 ++-- waterbutler/server/api/v0/zip.py | 4 ++-- waterbutler/server/api/v1/provider/__init__.py | 6 ++++-- waterbutler/server/api/v1/provider/metadata.py | 5 +++-- waterbutler/server/utils.py | 5 +++-- 7 files changed, 32 insertions(+), 12 deletions(-) diff --git a/tests/server/api/v1/test_metadata_mixin.py b/tests/server/api/v1/test_metadata_mixin.py index 48308e806e..09aafabb1d 100644 --- a/tests/server/api/v1/test_metadata_mixin.py +++ b/tests/server/api/v1/test_metadata_mixin.py @@ -120,6 +120,20 @@ async def test_download_file_headers_no_stream_name(self, http_request, mock_str handler.write_stream.assert_awaited_once() + @pytest.mark.asyncio + @pytest.mark.parametrize('write_stream_result, expected_completed', [(True, True), (False, False)]) + async def test_download_file_records_stream_completion(self, http_request, mock_stream, + write_stream_result, expected_completed): + + handler = mock_handler(http_request) + handler.provider.download = MockCoroutine(return_value=mock_stream) + handler.path = WaterButlerPath('/test_file') + handler.write_stream = MockCoroutine(return_value=write_stream_result) + + await handler.download_file() + + assert handler._download_completed is expected_completed + @pytest.mark.asyncio @pytest.mark.parametrize("given_arg,expected_name,filtered_name", [ (['résumé.doc'], 'r%C3%A9sum%C3%A9.doc', 'resume.doc'), diff --git a/tests/server/api/v1/test_provider.py b/tests/server/api/v1/test_provider.py index a39eb9911f..e9445635dc 100644 --- a/tests/server/api/v1/test_provider.py +++ b/tests/server/api/v1/test_provider.py @@ -165,15 +165,17 @@ async def test_data_received_stream(self, http_request): class TestProviderHandlerFinish: @pytest.mark.asyncio - async def test_on_finish_download_file(self, http_request): + @pytest.mark.parametrize('download_completed, expected_completed', [(True, True), (False, False)]) + async def test_on_finish_download_file(self, http_request, download_completed, expected_completed): handler = mock_handler(http_request) handler.request.method = 'GET' handler.path = WaterButlerPath('/file') + handler._download_completed = download_completed handler._send_hook = mock.Mock() assert handler.on_finish() is None - handler._send_hook.assert_called_once_with('download_file') + handler._send_hook.assert_called_once_with('download_file', completed=expected_completed) @pytest.mark.asyncio async def test_on_finish_download_zip(self, http_request): diff --git a/waterbutler/server/api/v0/crud.py b/waterbutler/server/api/v0/crud.py index e7271faecb..0957d00426 100644 --- a/waterbutler/server/api/v0/crud.py +++ b/waterbutler/server/api/v0/crud.py @@ -86,8 +86,8 @@ async def get(self): if ext in mime_types: self.set_header('Content-Type', mime_types[ext]) - await self.write_stream(result) - self._send_hook('download_file', path=self.path, completed=True) + completed = await self.write_stream(result) + self._send_hook('download_file', path=self.path, completed=completed) async def post(self): """Create a folder""" diff --git a/waterbutler/server/api/v0/zip.py b/waterbutler/server/api/v0/zip.py index b52fe1f24f..34c64a7447 100644 --- a/waterbutler/server/api/v0/zip.py +++ b/waterbutler/server/api/v0/zip.py @@ -20,5 +20,5 @@ async def get(self): result = await self.provider.zip(**self.arguments) - await self.write_stream(result) - self._send_hook('download_zip', path=self.path, completed=True) + completed = await self.write_stream(result) + self._send_hook('download_zip', path=self.path, completed=completed) diff --git a/waterbutler/server/api/v1/provider/__init__.py b/waterbutler/server/api/v1/provider/__init__.py index 10d85f2ace..a0defc6fee 100644 --- a/waterbutler/server/api/v1/provider/__init__.py +++ b/waterbutler/server/api/v1/provider/__init__.py @@ -256,9 +256,11 @@ def on_finish(self): }[method]() if action in {'download_file', 'download_zip'}: - completed = status in {200, 302} + completed = getattr(self, '_download_completed', status in {200, 302}) + self._send_hook(action, completed=completed) + return - self._send_hook(action, completed=completed) + self._send_hook(action) def _send_hook(self, action, completed=False): source = None diff --git a/waterbutler/server/api/v1/provider/metadata.py b/waterbutler/server/api/v1/provider/metadata.py index add78e150a..988a91f66a 100644 --- a/waterbutler/server/api/v1/provider/metadata.py +++ b/waterbutler/server/api/v1/provider/metadata.py @@ -73,6 +73,7 @@ async def download_file(self): ) if isinstance(stream, str): + self._download_completed = True return self.redirect(stream) if getattr(stream, 'partial', None): @@ -103,7 +104,7 @@ async def download_file(self): if ext in mime_types: self.set_header('Content-Type', mime_types[ext]) - await self.write_stream(stream) + self._download_completed = await self.write_stream(stream) if getattr(stream, 'partial', False) and isinstance(stream, ResponseStreamReader): await stream.response.release() @@ -133,4 +134,4 @@ async def download_folder_as_zip(self): result = await self.provider.zip(self.path) - await self.write_stream(result) + self._download_completed = await self.write_stream(result) diff --git a/waterbutler/server/utils.py b/waterbutler/server/utils.py index 83d2a01a29..e9afd243eb 100644 --- a/waterbutler/server/utils.py +++ b/waterbutler/server/utils.py @@ -123,7 +123,6 @@ def set_status(self, code, reason=None): async def write_stream(self, stream): try: - while True: chunk = await stream.read(settings.CHUNK_SIZE) if not chunk: @@ -138,4 +137,6 @@ async def write_stream(self, stream): except tornado.iostream.StreamClosedError: # Client has disconnected early. # No need for any exception to be raised - return + return False + + return True