[ENG-11736] pass downloaded bytes, IP and completed flag to callback

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):

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'),

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):

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,

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,12 +141,13 @@ 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,
                                  metadata=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)

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):
@@ -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 = await self.write_stream(result)
+        self._send_hook('download_file', path=self.path, completed=completed)
 
     async def post(self):
         """Create a folder"""

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 = await self.write_stream(result)
+        self._send_hook('download_zip', path=self.path, completed=completed)

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,14 @@ def on_finish(self):
             'DELETE': lambda: 'delete'
         }[method]()
 
+        if action in {'download_file', 'download_zip'}:
+            completed = getattr(self, '_download_completed', status in {200, 302})
+            self._send_hook(action, completed=completed)
+            return
+
         self._send_hook(action)
 
-    def _send_hook(self, action):
+    def _send_hook(self, action, completed=False):
         source = None
         destination = None
 
@@ -281,4 +287,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)

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)

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