feat(osf): ENG-11735 record download events for files and zips

sh-andriy · sh-andriy · commit e0c8dd07e696 · 2026-07-22T17:10:52.000+03:00
diff --git a/osf/utils/download_telemetry.py b/osf/utils/download_telemetry.py
@@ -0,0 +1,150 @@
+import functools
+import logging
+
+from addons.osfstorage.settings import DEFAULT_REGION_NAME
+from framework.celery_tasks import app
+from framework.postcommit_tasks.handlers import enqueue_postcommit_task
+
+logger = logging.getLogger(__name__)
+
+# Set on every user until they pick something in their profile, so it says nothing
+# about where they actually are.
+UNSET_USER_TIMEZONE = 'Etc/UTC'
+
+
+def never_breaks_downloads(fn):
+    """Swallow and log anything this raises.
+
+    Wraps the whole capture, not just the write — gathering the values is as capable of
+    raising as storing them is, and neither is a reason for a download to fail.
+    """
+    @functools.wraps(fn)
+    def wrapped(*args, **kwargs):
+        try:
+            return fn(*args, **kwargs)
+        except Exception:
+            logger.exception('Failed to record a download event')
+    return wrapped
+
+
+@never_breaks_downloads
+def record_download(**kwargs):
+    """Enqueue a :class:`DownloadEvent` write."""
+    enqueue_postcommit_task(write_download_event, (), kwargs, celery=True)
+
+
+@app.task(max_retries=5, default_retry_delay=60)
+def write_download_event(
+    download_type,
+    resource_guid='',
+    path='',
+    file_id=None,
+    version_identifier=None,
+    size_bytes=None,
+    storage_region_id=None,
+    zip_completed=None,
+    user_guid=None,
+    ip=None,
+    source_area='',
+    tz='',
+):
+    """Resolve the expensive bits and write one row.
+
+    Callers hand over identifiers rather than loaded objects so that the download request
+    itself does no extra queries — everything that needs a lookup is resolved here.
+    """
+    from osf.models import BaseFileNode, DownloadEvent, OSFUser
+
+    user = OSFUser.load(user_guid) if user_guid else None
+    file_node = BaseFileNode.load(file_id) if file_id else None
+    file_version = _load_file_version(file_node, version_identifier)
+
+    if file_version is not None:
+        if size_bytes is None:
+            size_bytes = file_version.size
+        if storage_region_id is None:
+            storage_region_id = file_version.region_id
+
+    storage_region = _region_name(storage_region_id) or _resource_region_name(resource_guid)
+
+    if not path and file_node is not None:
+        path = getattr(file_node, 'materialized_path', '') or ''
+
+    DownloadEvent.objects.create(
+        download_type=download_type,
+        resource_guid=_truncate(resource_guid, 255),
+        path=path or '',
+        size_bytes=size_bytes if size_bytes is not None and size_bytes >= 0 else None,
+        zip_completed=zip_completed,
+        storage_region=_truncate(storage_region, 64),
+        user_region=_truncate(derive_user_region(tz, user, storage_region), 64),
+        ip=ip or None,
+        source_area=_truncate(source_area, 128),
+        user=user,
+    )
+
+
+def derive_user_region(tz, user, storage_region):
+    """Best available guess at where the user is, most to least trustworthy.
+
+    The live browser timezone is the only real signal; the rest are fallbacks so the
+    dashboard isn't mostly blank.  An empty string means we genuinely don't know, which
+    is more useful than a wrong guess.
+    """
+    if tz:
+        return tz
+
+    profile_timezone = getattr(user, 'timezone', '')
+    if profile_timezone and profile_timezone != UNSET_USER_TIMEZONE:
+        return profile_timezone
+
+    # Everything defaults to the US region, so it only tells us something when it's been
+    # deliberately changed.
+    if storage_region and storage_region != DEFAULT_REGION_NAME:
+        return storage_region
+
+    return ''
+
+
+def _load_file_version(file_node, version_identifier):
+    """The version that was served, for its size and region."""
+    if file_node is None:
+        return None
+
+    from osf.models import FileVersion
+
+    versions = FileVersion.objects.filter(basefilenode=file_node)
+    if version_identifier:
+        return versions.filter(identifier=version_identifier).first()
+    return versions.order_by('-created').first()
+
+
+def _region_name(region_id):
+    if not region_id:
+        return ''
+
+    from addons.osfstorage.models import Region
+
+    region = Region.objects.filter(id=region_id).first()
+    return region.name if region else ''
+
+
+def _resource_region_name(resource_guid):
+    """Where a zip was served from — zips have no single file version to read it off."""
+    if not resource_guid:
+        return ''
+
+    from osf.models import Guid
+
+    resource, _ = Guid.load_referent(resource_guid)
+    region = getattr(resource, 'osfstorage_region', None)
+    return getattr(region, 'name', '') or ''
+
+
+def _truncate(value, max_length):
+    """Keep user-controllable values inside their column.
+
+    ``source`` and ``tz`` arrive off the query string, so they're whatever the caller
+    put there.
+    """
+    return (value or '')[:max_length]
diff --git a/tests/test_download_telemetry.py b/tests/test_download_telemetry.py
@@ -0,0 +1,210 @@
+import time
+
+import pytest
+
+from addons.osfstorage.settings import DEFAULT_REGION_NAME
+from api_tests.utils import create_test_file
+from framework.auth import signing
+from osf.models import DownloadEvent
+from osf.utils.download_telemetry import derive_user_region, record_download
+from osf_tests.factories import AuthUserFactory, ProjectFactory
+from tests.base import OsfTestCase
+
+
+class TestZipDownloadTelemetry(OsfTestCase):
+    """Folder and project zips, recorded from the WaterButler callback.
+
+    Zips are requested straight from WaterButler, so this callback is the only place we
+    hear about them.
+    """
+
+    def setUp(self):
+        super().setUp()
+        self.user = AuthUserFactory()
+        self.node = ProjectFactory(creator=self.user)
+        self.url = self.node.api_url_for('create_waterbutler_log')
+
+    def build_payload(self, materialized='/', action='download_zip', **action_meta):
+        meta = dict(
+            bytes_downloaded=2048,
+            completed=True,
+            ip='198.51.100.7',
+            source='files',
+            tz='Europe/Kyiv',
+        )
+        meta.update(action_meta)
+        options = {
+            'auth': {'id': self.user._id},
+            'action': action,
+            'provider': 'osfstorage',
+            'time': time.time() + 1000,
+            'metadata': {
+                'nid': self.node._id,
+                'materialized': materialized,
+                'path': materialized,
+                'kind': 'folder',
+                'provider': 'osfstorage',
+            },
+            'action_meta': meta,
+        }
+        message, signature = signing.default_signer.sign_payload(options)
+        return {'payload': message, 'signature': signature}
+
+    def test_project_zip_is_recorded(self):
+        res = self.app.put(self.url, json=self.build_payload(materialized='/'))
+
+        assert res.status_code == 200
+        event = DownloadEvent.objects.get()
+        assert event.download_type == DownloadEvent.PROJECT
+        assert event.resource_guid == self.node._id
+        assert event.user == self.user
+        assert event.size_bytes == 2048
+        assert event.zip_completed is True
+        assert event.ip == '198.51.100.7'
+        assert event.source_area == 'files'
+        assert event.user_region == 'Europe/Kyiv'
+
+    def test_folder_zip_is_recorded_with_its_path(self):
+        self.app.put(self.url, json=self.build_payload(materialized='/data/raw/'))
+
+        event = DownloadEvent.objects.get()
+        assert event.download_type == DownloadEvent.FOLDER_ZIP
+        assert event.path == '/data/raw/'
+
+    def test_incomplete_zip_is_recorded_as_incomplete(self):
+        self.app.put(self.url, json=self.build_payload(completed=False))
+
+        assert DownloadEvent.objects.get().zip_completed is False
+
+    def test_mfr_render_is_not_recorded(self):
+        self.app.put(self.url, json=self.build_payload(is_mfr_render=True))
+
+        assert not DownloadEvent.objects.exists()
+
+    def test_single_file_action_is_not_recorded_here(self):
+        """Single files are caught at the redirect view — recording them here too would
+        double count every one of them."""
+        self.app.put(self.url, json=self.build_payload(action='download_file'))
+
+        assert not DownloadEvent.objects.exists()
+
+    def test_oversized_source_is_truncated_to_the_column(self):
+        self.app.put(self.url, json=self.build_payload(source='f' * 500))
+
+        assert len(DownloadEvent.objects.get().source_area) == 128
+
+    def test_callback_still_succeeds_when_recording_fails(self, ):
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(
+                'addons.base.views.record_download',
+                lambda **kwargs: (_ for _ in ()).throw(ValueError('boom')),
+            )
+            res = self.app.put(self.url, json=self.build_payload())
+
+        assert res.status_code == 200
+
+
+class TestSingleFileDownloadTelemetry(OsfTestCase):
+    """Single files, recorded at the redirect view before we 302 on to WaterButler."""
+
+    def setUp(self):
+        super().setUp()
+        self.user = AuthUserFactory()
+        self.node = ProjectFactory(creator=self.user)
+        self.file = create_test_file(self.node, self.user, size=4096)
+        self.guid = self.file.get_guid()._id
+
+    def test_download_is_recorded_with_link_tags(self):
+        res = self.app.get(
+            f'/download/{self.guid}/?source=file-detail&tz=Europe%2FKyiv',
+            auth=self.user.auth,
+        )
+
+        assert res.status_code == 302
+        event = DownloadEvent.objects.get()
+        assert event.download_type == DownloadEvent.FILE
+        assert event.resource_guid == self.node._id
+        assert event.user == self.user
+        assert event.source_area == 'file-detail'
+        assert event.user_region == 'Europe/Kyiv'
+
+    def test_size_and_region_come_from_the_file_version(self):
+        self.app.get(f'/download/{self.guid}/', auth=self.user.auth)
+
+        event = DownloadEvent.objects.get()
+        assert event.size_bytes == 4096
+        assert event.storage_region == self.node.osfstorage_region.name
+
+    def test_zip_completed_is_unset_for_single_files(self):
+        """Only zips stream through WaterButler, so nothing reports completion here."""
+        self.app.get(f'/download/{self.guid}/', auth=self.user.auth)
+
+        assert DownloadEvent.objects.get().zip_completed is None
+
+    def test_anonymous_download_is_recorded_without_a_user(self):
+        self.node.is_public = True
+        self.node.save()
+
+        self.app.get(f'/download/{self.guid}/')
+
+        event = DownloadEvent.objects.get()
+        assert event.user is None
+
+    def test_mfr_render_is_not_recorded(self):
+        self.app.get(f'/download/{self.guid}/?mode=render', auth=self.user.auth)
+
+        assert not DownloadEvent.objects.exists()
+
+    def test_download_still_succeeds_when_recording_fails(self):
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(
+                'addons.base.views.record_download',
+                lambda **kwargs: (_ for _ in ()).throw(ValueError('boom')),
+            )
+            res = self.app.get(f'/download/{self.guid}/', auth=self.user.auth)
+
+        assert res.status_code == 302
+
+
+class TestUserRegionDerivation:
+    """The fallback chain, most to least trustworthy."""
+
+    class FakeUser:
+        def __init__(self, timezone):
+            self.timezone = timezone
+
+    def test_live_browser_timezone_wins(self):
+        user = self.FakeUser('America/New_York')
+        assert derive_user_region('Europe/Kyiv', user, 'Germany') == 'Europe/Kyiv'
+
+    def test_falls_back_to_profile_timezone(self):
+        user = self.FakeUser('America/New_York')
+        assert derive_user_region('', user, 'Germany') == 'America/New_York'
+
+    def test_default_profile_timezone_is_not_a_signal(self):
+        """Every user has Etc/UTC until they change it, so it says nothing."""
+        user = self.FakeUser('Etc/UTC')
+        assert derive_user_region('', user, 'Germany') == 'Germany'
+
+    def test_falls_back_to_storage_region(self):
+        assert derive_user_region('', None, 'Germany') == 'Germany'
+
+    def test_default_storage_region_is_not_a_signal(self):
+        assert derive_user_region('', None, DEFAULT_REGION_NAME) == ''
+
+    def test_unknown_is_empty(self):
+        assert derive_user_region('', None, '') == ''
+
+
+@pytest.mark.django_db
+class TestRecordDownloadNeverRaises:
+
+    def test_enqueue_failure_is_swallowed(self, monkeypatch):
+        def explode(*args, **kwargs):
+            raise ValueError('boom')
+
+        monkeypatch.setattr('osf.utils.download_telemetry.enqueue_postcommit_task', explode)
+
+        record_download(download_type=DownloadEvent.FILE, resource_guid='abcde')
+
+        assert not DownloadEvent.objects.exists()
