From 5c03d7a54f8710976ecffbc40fc36026f1895037 Mon Sep 17 00:00:00 2001 From: Mark Story Date: Thu, 3 Aug 2023 13:33:58 -0400 Subject: [PATCH 1/3] feat(hybridcloud) Use control file for user/sentryapp avatars Going forward always write to ControlFile for User & sentry app avatars. This helps make application logic more consistent going forwards as new avatars are always stored in ControlFile. I've retained read paths for `File` backed avatars so that we preserve behavior in both single-tenant and self-hosted. We also retain the ability to write to a control silo specific storage bucket should an environment be configured in that way. --- src/sentry/models/avatars/base.py | 42 +++-- src/sentry/models/avatars/control_base.py | 34 +++-- src/sentry/models/avatars/user_avatar.py | 23 ++- src/sentry/models/files/control_fileblob.py | 52 ++++--- src/sentry/tasks/files.py | 14 +- .../api/endpoints/test_sentry_app_avatar.py | 48 +++++- .../sentry/api/endpoints/test_user_avatar.py | 143 +++++++++++++++--- tests/sentry/web/frontend/test_user_avatar.py | 14 ++ 8 files changed, 286 insertions(+), 84 deletions(-) diff --git a/src/sentry/models/avatars/base.py b/src/sentry/models/avatars/base.py index 18d87a458b20a6..0dd2c5bfa990a8 100644 --- a/src/sentry/models/avatars/base.py +++ b/src/sentry/models/avatars/base.py @@ -43,10 +43,11 @@ def save(self, *args, **kwargs): return super().save(*args, **kwargs) def get_file(self): - # Favor control_file_id if it exists and is set. - # Otherwise fallback to file_id. If still None, return. + # If we're getting a file, and the preferred write file + # type isn't present, move data over to new storage async. + file_id = getattr(self, self.file_write_fk(), None) file_class = self.file_class() - file_id = getattr(self, self.file_fk()) + if file_id is None: file_id = self.file_id file_class = File @@ -70,9 +71,6 @@ def get_file(self): self.update(**update) return None - def get_file_id(self): - return self.file_id - def delete(self, *args, **kwargs): file = self.get_file() if file: @@ -103,21 +101,32 @@ def get_cached_photo(self, size): cache.set(cache_key, photo) return photo - @classmethod - def file_class(cls): - from sentry.models import File - + def file_class(self): return File - @classmethod - def file_fk(cls) -> str: + def file_fk(self) -> str: + """ + Get the foreign key currently used by this record for blob storage. + Varies in ControlAvatarBase + """ + return "file_id" + + def file_write_fk(self) -> str: + """ + Get the foreign key that should be used for writes. + Varies in ControlAvatarBase + """ return "file_id" @classmethod def save_avatar(cls, relation, type, avatar=None, filename=None, color=None): if avatar: - with atomic_transaction(using=router.db_for_write(cls.file_class())): - photo = cls.file_class().objects.create(name=filename, type=cls.FILE_TYPE) + # Create an instance of the current class so we can + # access where new files should be stored. + dummy = cls() + file_class = dummy.file_class() + with atomic_transaction(using=router.db_for_write(file_class)): + photo = file_class.objects.create(name=filename, type=cls.FILE_TYPE) # XXX: Avatar may come in as a string instance in python2 # if it's not wrapped in BytesIO. if isinstance(avatar, str): @@ -141,11 +150,12 @@ def save_avatar(cls, relation, type, avatar=None, filename=None, color=None): file.delete() if photo: - setattr(instance, cls.file_fk(), photo.id) + if instance.file_fk() != instance.file_write_fk(): + setattr(instance, instance.file_fk(), None) + setattr(instance, instance.file_write_fk(), photo.id) instance.ident = uuid4().hex instance.avatar_type = [i for i, n in cls.AVATAR_TYPES if n == type][0] - instance.save() if photo and not created: diff --git a/src/sentry/models/avatars/control_base.py b/src/sentry/models/avatars/control_base.py index 36e7ebd6bc67ed..72e33f49fc31bc 100644 --- a/src/sentry/models/avatars/control_base.py +++ b/src/sentry/models/avatars/control_base.py @@ -1,6 +1,8 @@ +from typing import Type, Union + from sentry.db.models import BoundedBigIntegerField from sentry.models.avatars.base import AvatarBase -from sentry.models.files import ControlFileBlob, File +from sentry.models.files import ControlFile, File class ControlAvatarBase(AvatarBase): @@ -9,19 +11,25 @@ class ControlAvatarBase(AvatarBase): class Meta: abstract = True - @classmethod - def file_class(cls): - from sentry.models import ControlFile + def file_class(self) -> Union[Type[ControlFile], Type[File]]: + """ + Select the file class this avatar has used. + File classes can vary by the avatar as we have migrated + storage for saas, but self-hosted and single-tenant instances + did not have relations and storage migrated. + """ + if self.file_id: + return File + return ControlFile - if ControlFileBlob._storage_config(): - return ControlFile - return File + def file_fk(self) -> str: + if self.file_id: + return "file_id" + return "control_file_id" - @classmethod - def file_fk(cls) -> str: - if ControlFileBlob._storage_config(): - return "control_file_id" - return "file_id" + def file_write_fk(self) -> str: + """Prefer controlfile as user/sentryapp avatars are stored in control silo""" + return "control_file_id" - def get_file_id(self): + def get_file_id(self) -> int: return self.control_file_id or self.file_id diff --git a/src/sentry/models/avatars/user_avatar.py b/src/sentry/models/avatars/user_avatar.py index 5e7cec07dce756..52df2c8e58d3f5 100644 --- a/src/sentry/models/avatars/user_avatar.py +++ b/src/sentry/models/avatars/user_avatar.py @@ -1,5 +1,7 @@ from __future__ import annotations +from enum import IntEnum + from django.db import models from sentry.db.models import BaseManager, FlexibleForeignKey, control_silo_only_model @@ -7,6 +9,23 @@ from . import ControlAvatarBase +class UserAvatarType(IntEnum): + LETTER_AVATAR = 0 + UPLOAD = 1 + GRAVATAR = 2 + + def api_name(self): + return self.name.lower() + + @classmethod + def as_choices(cls) -> tuple[tuple[int, str], ...]: + return ( + (cls.LETTER_AVATAR, cls.LETTER_AVATAR.api_name()), + (cls.UPLOAD, cls.UPLOAD.api_name()), + (cls.GRAVATAR, cls.GRAVATAR.api_name()), + ) + + @control_silo_only_model class UserAvatar(ControlAvatarBase): """ @@ -14,12 +33,12 @@ class UserAvatar(ControlAvatarBase): and contains their preferences for avatar type. """ - AVATAR_TYPES = ((0, "letter_avatar"), (1, "upload"), (2, "gravatar")) + AVATAR_TYPES = UserAvatarType.as_choices() FILE_TYPE = "avatar.file" user = FlexibleForeignKey("sentry.User", unique=True, related_name="avatar") - avatar_type = models.PositiveSmallIntegerField(default=0, choices=AVATAR_TYPES) + avatar_type = models.PositiveSmallIntegerField(default=0, choices=UserAvatarType.as_choices()) objects = BaseManager(cache_fields=["user"]) diff --git a/src/sentry/models/files/control_fileblob.py b/src/sentry/models/files/control_fileblob.py index c9a83183487828..ab7088c8b302b7 100644 --- a/src/sentry/models/files/control_fileblob.py +++ b/src/sentry/models/files/control_fileblob.py @@ -1,3 +1,8 @@ +from __future__ import annotations + +from typing import Any, Dict + +from sentry import options from sentry.db.models import control_silo_only_model from sentry.models.files.abstractfileblob import AbstractFileBlob from sentry.models.files.control_fileblobowner import ControlFileBlobOwner @@ -5,6 +10,32 @@ from sentry.tasks.files import delete_file_control +def control_file_storage_config() -> Dict[str, Any] | None: + """ + When sentry is deployed in a siloed mode file relations + used by control silo models are stored separately from + region silo resources. + + While we consistently write to the ControlFile django + model for control silo resources, we can't ensure + that each deployment has separate control + region storage + backends. We coalesce those options here. None means use the + global default storage options. + """ + try: + # If these options exist, use them. Otherwise fallback to default behavior + storage_backend = options.get("filestore.control.backend") + storage_options = options.get("filestore.control.options") + if storage_backend: + return { + "backend": storage_backend, + "options": storage_options, + } + except UnknownOption: + pass + return None + + @control_silo_only_model class ControlFileBlob(AbstractFileBlob): class Meta: @@ -14,23 +45,6 @@ class Meta: FILE_BLOB_OWNER_MODEL = ControlFileBlobOwner DELETE_FILE_TASK = delete_file_control - # TODO(hybrid-cloud): This is a temporary measure to allow concurrent production deployments of filestore - # We override the behavior of only the control silo file storage implementation to use - # the new control instance while production runs in monolith mode @classmethod - def _storage_config(cls): - from sentry import options as options_store - - config = None - try: - # If these options exist, use them. Otherwise fallback to default behavior - backend = options_store.get("filestore.control.backend") - options = options_store.get("filestore.control.options") - if backend: - config = { - "backend": backend, - "options": options, - } - except UnknownOption: - pass - return config + def _storage_config(cls) -> Dict[str, Any] | None: + return control_file_storage_config() diff --git a/src/sentry/tasks/files.py b/src/sentry/tasks/files.py index b1de19a63d366e..17364ddeadf6bd 100644 --- a/src/sentry/tasks/files.py +++ b/src/sentry/tasks/files.py @@ -107,17 +107,15 @@ def copy_file_to_control_and_update_model( file_id: int, **kwargs, ): - from sentry.models.files import ControlFile, ControlFileBlob, File - - if ControlFileBlob._storage_config() is None: - return + from sentry.models.files import ControlFile, File lock = f"copy-file-lock-{model_name}:{model_id}" with locks.get(lock, duration=60, name="copy-file-lock").acquire(): # Short circuit duplicate copy calls - model = apps.get_app_config(app_name).get_model(model_name).objects.get(id=model_id) - if model.control_file_id: + model_class = apps.get_model(app_name, model_name) + instance = model_class.objects.get(id=model_id) + if instance.control_file_id: return file_model = File.objects.get(id=file_id) @@ -133,5 +131,5 @@ def copy_file_to_control_and_update_model( ) control_file.putfile(file_handle) - model.control_file_id = control_file.id - model.save() + instance.control_file_id = control_file.id + instance.save() diff --git a/tests/sentry/api/endpoints/test_sentry_app_avatar.py b/tests/sentry/api/endpoints/test_sentry_app_avatar.py index fd02772559a827..ea093ea134e533 100644 --- a/tests/sentry/api/endpoints/test_sentry_app_avatar.py +++ b/tests/sentry/api/endpoints/test_sentry_app_avatar.py @@ -2,6 +2,7 @@ from sentry import options as options_store from sentry.models import SentryAppAvatar +from sentry.models.files.control_file import ControlFile from sentry.testutils.cases import APITestCase from sentry.testutils.silo import control_silo_test @@ -37,7 +38,7 @@ def create_avatar(self, is_color): @control_silo_test(stable=True) -class SentryAppAvatarTest(SentryAppAvatarTestBase): +class SentryAppAvatarGetTest(SentryAppAvatarTestBase): def test_get(self): response = self.get_success_response(self.unpublished_app.slug) @@ -53,6 +54,30 @@ def test_get(self): assert simple_avatar["color"] is False assert response.data["uuid"] == str(self.unpublished_app.uuid) + def test_get_upload_control_file(self): + self.method = "put" + self.create_avatar(is_color=True) + + self.method = "get" + self.create_avatar(is_color=True) + response = self.get_success_response(self.unpublished_app.slug) + + assert response.status_code == 200, response.content + assert response.data["uuid"] == str(self.unpublished_app.uuid) + uploads = [ + avatar for avatar in response.data["avatars"] if avatar["avatarType"] == "upload" + ] + upload = uploads[0] + assert upload["avatarType"] == "upload" + assert upload["avatarUuid"] + + avatar = SentryAppAvatar.objects.filter(sentry_app_id=self.unpublished_app.id).first() + assert avatar + assert avatar.get_file_id() + assert avatar.control_file_id + file = avatar.get_file() + assert isinstance(file, ControlFile) + @control_silo_test(stable=True) class SentryAppAvatarPutTest(SentryAppAvatarTestBase): @@ -65,11 +90,25 @@ def test_upload(self): assert avatar.get_file_id() assert avatar.get_avatar_type_display() == "upload" color_avatar = self.get_avatar(resp) + assert color_avatar + assert color_avatar["avatarType"] == "upload" + assert color_avatar["avatarUuid"] is not None + assert color_avatar["color"] is True + + def test_upload_control_file(self): + resp = self.create_avatar(is_color=True) + + avatar = SentryAppAvatar.objects.get(sentry_app=self.unpublished_app, color=True) + assert avatar.get_file_id() + assert avatar.control_file_id + assert avatar.get_avatar_type_display() == "upload" + color_avatar = self.get_avatar(resp) + assert color_avatar assert color_avatar["avatarType"] == "upload" assert color_avatar["avatarUuid"] is not None assert color_avatar["color"] is True - def test_upload_control(self): + def test_upload_control_with_storage_options(self): with self.options( { "filestore.control.backend": options_store.get("filestore.backend"), @@ -82,6 +121,7 @@ def test_upload_control(self): assert avatar.get_file_id() assert avatar.get_avatar_type_display() == "upload" color_avatar = self.get_avatar(resp) + assert color_avatar assert color_avatar["avatarType"] == "upload" assert color_avatar["avatarUuid"] is not None assert color_avatar["color"] is True @@ -96,6 +136,7 @@ def test_upload_both(self): assert avatars[0].get_file_id() assert avatars[0].get_avatar_type_display() == "upload" color_avatar = self.get_avatar(resp) + assert color_avatar assert color_avatar["color"] is True assert color_avatar["avatarType"] == "upload" assert color_avatar["avatarUuid"] is not None @@ -103,6 +144,7 @@ def test_upload_both(self): assert avatars[1].get_file_id() assert avatars[1].get_avatar_type_display() == "upload" simple_avatar = self.get_avatar(resp, False) + assert simple_avatar assert simple_avatar["color"] is False assert simple_avatar["avatarType"] == "upload" assert simple_avatar["avatarUuid"] is not None @@ -129,10 +171,12 @@ def test_revert_to_default(self): color_avatar = self.get_avatar(response) simple_avatar = self.get_avatar(response, False) + assert color_avatar assert color_avatar["avatarType"] == "default" assert color_avatar["avatarUuid"] is not None assert color_avatar["color"] is True + assert simple_avatar assert simple_avatar["avatarType"] == "default" assert simple_avatar["avatarUuid"] is not None assert simple_avatar["color"] is False diff --git a/tests/sentry/api/endpoints/test_user_avatar.py b/tests/sentry/api/endpoints/test_user_avatar.py index 2c9aa57fd523cd..948d9f77afc7e3 100644 --- a/tests/sentry/api/endpoints/test_user_avatar.py +++ b/tests/sentry/api/endpoints/test_user_avatar.py @@ -1,17 +1,20 @@ from base64 import b64encode +from io import BytesIO +from unittest import mock from django.urls import reverse from sentry import options as options_store -from sentry.models import UserAvatar -from sentry.models.files import ControlFile +from sentry.models.avatars.user_avatar import UserAvatar, UserAvatarType +from sentry.models.files import ControlFile, File +from sentry.silo.base import SiloMode from sentry.testutils.cases import APITestCase -from sentry.testutils.silo import control_silo_test +from sentry.testutils.silo import assume_test_silo_mode, control_silo_test @control_silo_test(stable=True) class UserAvatarTest(APITestCase): - def test_get(self): + def test_get_letter_avatar(self): user = self.create_user(email="a@example.com") self.login_as(user=user) @@ -24,7 +27,58 @@ def test_get(self): assert response.data["avatar"]["avatarType"] == "letter_avatar" assert response.data["avatar"]["avatarUuid"] is None - def test_gravatar(self): + def test_get_gravatar(self): + user = self.create_user(email="a@example.com") + UserAvatar.objects.create(user=user, avatar_type=UserAvatarType.GRAVATAR) + + self.login_as(user=user) + + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + response = self.client.get(url, format="json") + + assert response.status_code == 200, response.content + assert response.data["id"] == str(user.id) + assert response.data["avatar"]["avatarType"] == "gravatar" + assert response.data["avatar"]["avatarUuid"] is None + + def test_get_upload_control_file(self): + user = self.create_user(email="a@example.com") + + photo = ControlFile.objects.create(name="test.png", type="avatar.file") + photo.putfile(BytesIO(b"test")) + UserAvatar.objects.create( + user=user, control_file_id=photo.id, avatar_type=UserAvatarType.UPLOAD + ) + + self.login_as(user=user) + + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + response = self.client.get(url, format="json") + + assert response.status_code == 200, response.content + assert response.data["id"] == str(user.id) + assert response.data["avatar"]["avatarType"] == "upload" + assert response.data["avatar"]["avatarUuid"] + + def test_get_upload_file(self): + user = self.create_user(email="a@example.com") + + with assume_test_silo_mode(SiloMode.REGION): + photo = File.objects.create(name="test.png", type="avatar.file") + photo.putfile(BytesIO(b"test")) + UserAvatar.objects.create(user=user, file_id=photo.id, avatar_type=UserAvatarType.UPLOAD) + + self.login_as(user=user) + + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + response = self.client.get(url, format="json") + + assert response.status_code == 200, response.content + assert response.data["id"] == str(user.id) + assert response.data["avatar"]["avatarType"] == "upload" + assert response.data["avatar"]["avatarUuid"] + + def test_put_gravatar(self): user = self.create_user(email="a@example.com") self.login_as(user=user) @@ -36,7 +90,7 @@ def test_gravatar(self): assert response.status_code == 200, response.content assert avatar.get_avatar_type_display() == "gravatar" - def test_upload(self): + def test_put_upload(self): user = self.create_user(email="a@example.com") self.login_as(user=user) @@ -55,14 +109,47 @@ def test_upload(self): assert response.status_code == 200, response.content assert avatar.get_avatar_type_display() == "upload" assert avatar.get_file_id() + assert avatar.control_file_id, "new files are control files" + assert ControlFile.objects.filter(id=avatar.control_file_id).exists() + + def test_put_upload_saves_to_control_file(self): + user = self.create_user(email="a@example.com") + + self.login_as(user=user) + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) - def test_transition_to_control_before_options_set(self): with self.tasks(): - user = self.create_user(email="a@example.com") + response = self.client.put( + url, + data={ + "avatar_type": "upload", + "avatar_photo": b64encode(self.load_fixture("avatar.jpg")), + }, + format="json", + ) - self.login_as(user=user) + avatar = UserAvatar.objects.get(user=user) + assert response.status_code == 200, response.content + assert avatar.get_file_id() + assert avatar.control_file_id + assert isinstance(avatar.get_file(), ControlFile) + assert ControlFile.objects.filter(id=avatar.control_file_id).exists() - url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + @mock.patch("sentry.tasks.files.copy_file_to_control_and_update_model.apply_async") + def test_get_upload_use_control_during_update(self, mock_task): + user = self.create_user(email="a@example.com") + + with assume_test_silo_mode(SiloMode.REGION): + photo = File.objects.create(name="test.png", type="avatar.file") + photo.putfile(BytesIO(b"test")) + UserAvatar.objects.create(user=user, file_id=photo.id, avatar_type=UserAvatarType.UPLOAD) + + self.login_as(user=user) + + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + + # Run in monolith as this transitional operation can only occur there. + with assume_test_silo_mode(SiloMode.MONOLITH): response = self.client.put( url, data={ @@ -72,24 +159,32 @@ def test_transition_to_control_before_options_set(self): format="json", ) - avatar = UserAvatar.objects.get(user=user) - assert response.status_code == 200, response.content - assert avatar.get_file_id() - assert isinstance(avatar.get_file(), ControlFile) + assert response.status_code == 200, response.content + assert response.data["id"] == str(user.id) + assert response.data["avatar"]["avatarType"] == "upload" + assert response.data["avatar"]["avatarUuid"] + assert mock_task.call_count == 1 - def test_transition_to_control_after_options_set(self): + avatar = UserAvatar.objects.get(user=user) + assert avatar + assert avatar.get_file_id() + assert avatar.file_id is None, "non-control file relation be removed." + assert avatar.control_file_id + assert isinstance(avatar.get_file(), ControlFile) + + def test_put_upload_saves_to_control_file_with_separate_storage(self): with self.options( { "filestore.control.backend": options_store.get("filestore.backend"), "filestore.control.options": options_store.get("filestore.options"), } ): - with self.tasks(): - user = self.create_user(email="a@example.com") + user = self.create_user(email="a@example.com") - self.login_as(user=user) + self.login_as(user=user) + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) - url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + with self.tasks(): response = self.client.put( url, data={ @@ -98,11 +193,11 @@ def test_transition_to_control_after_options_set(self): }, format="json", ) - - avatar = UserAvatar.objects.get(user=user) - assert response.status_code == 200, response.content - assert avatar.get_file_id() - assert isinstance(avatar.get_file(), ControlFile) + avatar = UserAvatar.objects.get(user=user) + assert response.status_code == 200, response.content + assert avatar.get_file_id() + assert isinstance(avatar.get_file(), ControlFile) + assert ControlFile.objects.filter(id=avatar.control_file_id).exists() def test_put_bad(self): user = self.create_user(email="a@example.com") diff --git a/tests/sentry/web/frontend/test_user_avatar.py b/tests/sentry/web/frontend/test_user_avatar.py index 90be0528d2a653..a6598dd7cc1dbe 100644 --- a/tests/sentry/web/frontend/test_user_avatar.py +++ b/tests/sentry/web/frontend/test_user_avatar.py @@ -3,6 +3,7 @@ from django.urls import reverse from sentry.models import File, UserAvatar +from sentry.models.files.control_file import ControlFile from sentry.silo import SiloMode from sentry.testutils.cases import TestCase from sentry.testutils.silo import assume_test_silo_mode, control_silo_test @@ -23,3 +24,16 @@ def test_headers(self): assert response["Cache-Control"] == FOREVER_CACHE assert response.get("Vary") is None assert response.get("Set-Cookie") is None + + def test_headers_control_file(self): + user = self.create_user(email="a@example.com") + photo = ControlFile.objects.create(name="test.png", type="avatar.file") + photo.putfile(BytesIO(b"test")) + avatar = UserAvatar.objects.create(user=user, control_file_id=photo.id) + + url = reverse("sentry-user-avatar-url", kwargs={"avatar_id": avatar.ident}) + response = self.client.get(url) + assert response.status_code == 200 + assert response["Cache-Control"] == FOREVER_CACHE + assert response.get("Vary") is None + assert response.get("Set-Cookie") is None From 4348b7230c1826e832438336007d81be4929b264 Mon Sep 17 00:00:00 2001 From: "getsantry[bot]" <66042841+getsantry[bot]@users.noreply.github.com> Date: Thu, 3 Aug 2023 19:31:19 +0000 Subject: [PATCH 2/3] :knife: regenerate mypy module blocklist --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b8037dff0aff2d..3ec616e7d5fa82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -760,7 +760,6 @@ module = [ "sentry.tasks.deliver_from_outbox", "sentry.tasks.derive_code_mappings", "sentry.tasks.digests", - "sentry.tasks.files", "sentry.tasks.groupowner", "sentry.tasks.integrations", "sentry.tasks.integrations.create_comment", From c346bf8e64b56e1965efa4394615f8be6c6172e3 Mon Sep 17 00:00:00 2001 From: Mark Story Date: Thu, 3 Aug 2023 16:14:11 -0400 Subject: [PATCH 3/3] Fix file_id not being wiped and fix precedence on reads. --- src/sentry/models/avatars/control_base.py | 4 +++ src/sentry/tasks/files.py | 1 + .../sentry/api/endpoints/test_user_avatar.py | 26 +++++++++++++++++++ tests/sentry/models/test_avatar.py | 7 ++--- 4 files changed, 35 insertions(+), 3 deletions(-) diff --git a/src/sentry/models/avatars/control_base.py b/src/sentry/models/avatars/control_base.py index 72e33f49fc31bc..71bdf402d0a778 100644 --- a/src/sentry/models/avatars/control_base.py +++ b/src/sentry/models/avatars/control_base.py @@ -18,11 +18,15 @@ def file_class(self) -> Union[Type[ControlFile], Type[File]]: storage for saas, but self-hosted and single-tenant instances did not have relations and storage migrated. """ + if self.control_file_id: + return ControlFile if self.file_id: return File return ControlFile def file_fk(self) -> str: + if self.control_file_id: + return "control_file_id" if self.file_id: return "file_id" return "control_file_id" diff --git a/src/sentry/tasks/files.py b/src/sentry/tasks/files.py index 17364ddeadf6bd..a327336630501d 100644 --- a/src/sentry/tasks/files.py +++ b/src/sentry/tasks/files.py @@ -132,4 +132,5 @@ def copy_file_to_control_and_update_model( control_file.putfile(file_handle) instance.control_file_id = control_file.id + instance.file_id = None instance.save() diff --git a/tests/sentry/api/endpoints/test_user_avatar.py b/tests/sentry/api/endpoints/test_user_avatar.py index 948d9f77afc7e3..15b7f1b734b84c 100644 --- a/tests/sentry/api/endpoints/test_user_avatar.py +++ b/tests/sentry/api/endpoints/test_user_avatar.py @@ -78,6 +78,32 @@ def test_get_upload_file(self): assert response.data["avatar"]["avatarType"] == "upload" assert response.data["avatar"]["avatarUuid"] + def test_get_prefers_control_file(self): + user = self.create_user(email="a@example.com") + with assume_test_silo_mode(SiloMode.REGION): + photo = File.objects.create(name="test.png", type="avatar.file") + photo.putfile(BytesIO(b"test")) + controlphoto = ControlFile.objects.create(name="control_test.png", type="avatar.file") + controlphoto.putfile(BytesIO(b"control test")) + + avatar = UserAvatar.objects.create( + user=user, + file_id=photo.id, + control_file_id=controlphoto.id, + avatar_type=UserAvatarType.UPLOAD, + ) + + self.login_as(user=user) + + url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"}) + response = self.client.get(url, format="json") + + assert response.status_code == 200, response.content + assert response.data["id"] == str(user.id) + assert response.data["avatar"]["avatarType"] == "upload" + assert response.data["avatar"]["avatarUuid"] + assert isinstance(avatar.get_file(), ControlFile) + def test_put_gravatar(self): user = self.create_user(email="a@example.com") diff --git a/tests/sentry/models/test_avatar.py b/tests/sentry/models/test_avatar.py index 7c0328fbef3679..d1485efd42ae01 100644 --- a/tests/sentry/models/test_avatar.py +++ b/tests/sentry/models/test_avatar.py @@ -43,9 +43,10 @@ def test_transition_to_control(self): ): with self.tasks(): assert isinstance(avatar.get_file(), File) - avatar = UserAvatar.objects.get(id=avatar.id) - assert avatar.control_file_id is not None - assert isinstance(avatar.get_file(), ControlFile) + avatar = UserAvatar.objects.get(id=avatar.id) + assert avatar.control_file_id is not None + assert avatar.file_id is None + assert isinstance(avatar.get_file(), ControlFile) @region_silo_test(stable=True)