From a63a7f88e815e566dcbd19b03e83d612d0a66dbf Mon Sep 17 00:00:00 2001 From: Ogi Date: Wed, 5 Mar 2025 12:56:23 +0100 Subject: [PATCH] feat(demo-mode): artifact bundle sync (#86302) --- src/sentry/conf/server.py | 7 ++ src/sentry/demo_mode/tasks.py | 145 ++++++++++++++++++++++++++ src/sentry/options/defaults.py | 17 ++++ src/sentry/utils/demo_mode.py | 8 ++ tests/demo_mode/test_tasks.py | 146 +++++++++++++++++++++++++++ tests/sentry/utils/test_demo_mode.py | 17 +++- 6 files changed, 339 insertions(+), 1 deletion(-) create mode 100644 src/sentry/demo_mode/tasks.py create mode 100644 tests/demo_mode/test_tasks.py diff --git a/src/sentry/conf/server.py b/src/sentry/conf/server.py index f3b7e0395807b6..7a5ca4ee702134 100644 --- a/src/sentry/conf/server.py +++ b/src/sentry/conf/server.py @@ -838,6 +838,7 @@ def SOCIAL_AUTH_DEFAULT_USERNAME() -> str: "sentry.integrations.vsts.tasks", "sentry.integrations.vsts.tasks.kickoff_subscription_check", "sentry.integrations.tasks", + "sentry.demo_mode.tasks", ) # Enable split queue routing @@ -997,6 +998,7 @@ def SOCIAL_AUTH_DEFAULT_USERNAME() -> str: Queue("on_demand_metrics", routing_key="on_demand_metrics"), Queue("check_new_issue_threshold_met", routing_key="check_new_issue_threshold_met"), Queue("integrations_slack_activity_notify", routing_key="integrations_slack_activity_notify"), + Queue("demo_mode", routing_key="demo_mode"), ] from celery.schedules import crontab @@ -1278,6 +1280,11 @@ def SOCIAL_AUTH_DEFAULT_USERNAME() -> str: # Run every 1 minute "schedule": crontab(minute="*/1"), }, + "demo_mode_sync_artifact_bundles": { + "task": "sentry.demo_mode.tasks.sync_artifact_bundles", + # Run every hour + "schedule": crontab(minute="0", hour="*/1"), + }, } # Assign the configuration keys celery uses based on our silo mode. diff --git a/src/sentry/demo_mode/tasks.py b/src/sentry/demo_mode/tasks.py new file mode 100644 index 00000000000000..1b86bbd3d3606e --- /dev/null +++ b/src/sentry/demo_mode/tasks.py @@ -0,0 +1,145 @@ +from datetime import timedelta + +from django.db import router +from django.db.models import Q +from django.utils import timezone + +from sentry import options +from sentry.models.artifactbundle import ( + ArtifactBundle, + ProjectArtifactBundle, + ReleaseArtifactBundle, +) +from sentry.models.files import FileBlobOwner +from sentry.models.organization import Organization +from sentry.models.project import Project +from sentry.tasks.base import instrumented_task +from sentry.utils.db import atomic_transaction +from sentry.utils.demo_mode import get_demo_org, is_demo_mode_enabled + + +@instrumented_task( + name="sentry.demo_mode.tasks.sync_artifact_bundles", + queue="demo_mode", +) +def sync_artifact_bundles(): + + if ( + not options.get("sentry.demo_mode.sync_artifact_bundles.enable") + or not is_demo_mode_enabled() + ): + return + + source_org_id = options.get("sentry.demo_mode.sync_artifact_bundles.source_org_id") + source_org = Organization.objects.get(id=source_org_id) + + target_org = get_demo_org() + + lookback_days = options.get("sentry.demo_mode.sync_artifact_bundles.lookback_days") + + _sync_artifact_bundles(source_org, target_org, lookback_days) + + +def _sync_artifact_bundles(source_org: Organization, target_org: Organization, lookback_days=1): + if not source_org or not target_org: + return + + cutoff_date = timezone.now() - timedelta(days=lookback_days) + + artifact_bundles = ArtifactBundle.objects.filter( + Q(organization_id=source_org.id) | Q(organization_id=target_org.id), + date_uploaded__gte=cutoff_date, + ) + + source_artifact_bundles = artifact_bundles.filter(organization_id=source_org.id) + target_artifact_bundles = artifact_bundles.filter(organization_id=target_org.id) + + different_artifact_bundles = source_artifact_bundles.exclude( + bundle_id__in=target_artifact_bundles.values_list("bundle_id", flat=True) + ) + + for source_artifact_bundle in different_artifact_bundles: + _sync_artifact_bundle(source_artifact_bundle, target_org) + + +def _sync_artifact_bundle(source_artifact_bundle: ArtifactBundle, target_org: Organization): + with atomic_transaction( + using=( + router.db_for_write(ArtifactBundle), + router.db_for_write(FileBlobOwner), + router.db_for_write(ProjectArtifactBundle), + router.db_for_write(ReleaseArtifactBundle), + ) + ): + blobs = source_artifact_bundle.file.blobs.all() + for blob in blobs: + FileBlobOwner.objects.create( + organization_id=target_org.id, + blob_id=blob.id, + ) + + target_artifact_bundle = ArtifactBundle.objects.create( + organization_id=target_org.id, + bundle_id=source_artifact_bundle.bundle_id, + artifact_count=source_artifact_bundle.artifact_count, + date_last_modified=source_artifact_bundle.date_last_modified, + date_uploaded=source_artifact_bundle.date_uploaded, + file=source_artifact_bundle.file, + indexing_state=source_artifact_bundle.indexing_state, + ) + + _sync_project_artifact_bundle(source_artifact_bundle, target_artifact_bundle) + _sync_release_artifact_bundle(source_artifact_bundle, target_artifact_bundle) + + +def _sync_project_artifact_bundle( + source_artifact_bundle: ArtifactBundle, + target_artifact_bundle: ArtifactBundle, +): + source_project_artifact_bundle = ProjectArtifactBundle.objects.get( + artifact_bundle_id=source_artifact_bundle.id, + organization_id=source_artifact_bundle.organization_id, + ) + + target_project = _find_matching_project( + source_project_artifact_bundle.project_id, + target_artifact_bundle.organization_id, + ) + + if not target_project: + return + + ProjectArtifactBundle.objects.create( + project_id=target_project.id, + artifact_bundle_id=target_artifact_bundle.id, + organization_id=target_artifact_bundle.organization_id, + ) + + +def _sync_release_artifact_bundle( + source_artifact_bundle: ArtifactBundle, + target_artifact_bundle: ArtifactBundle, +): + source_release_artifact_bundle = ReleaseArtifactBundle.objects.filter( + artifact_bundle_id=source_artifact_bundle.id, + organization_id=source_artifact_bundle.organization_id, + ).first() + + if not source_release_artifact_bundle: + return + + ReleaseArtifactBundle.objects.create( + artifact_bundle_id=target_artifact_bundle.id, + organization_id=target_artifact_bundle.organization_id, + dist_name=source_release_artifact_bundle.dist_name, + release_name=source_release_artifact_bundle.release_name, + ) + + +def _find_matching_project(project_id, organization_id): + source_project = Project.objects.get(id=project_id) + + return Project.objects.get( + organization_id=organization_id, + slug=source_project.slug, + ) diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py index b8f145d1336745..7623f61c86733f 100644 --- a/src/sentry/options/defaults.py +++ b/src/sentry/options/defaults.py @@ -3101,3 +3101,20 @@ default=False, flags=FLAG_AUTOMATOR_MODIFIABLE, ) + +register( + "sentry.demo_mode.sync_artifact_bundles.enable", + type=Bool, + default=False, + flags=FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE, +) +register( + "sentry.demo_mode.sync_artifact_bundles.source_org_id", + default=None, + flags=FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE, +) +register( + "sentry.demo_mode.sync_artifact_bundles.lookback_days", + default=1, + flags=FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE, +) diff --git a/src/sentry/utils/demo_mode.py b/src/sentry/utils/demo_mode.py index e0d81e24ad056f..5938c2642dc55a 100644 --- a/src/sentry/utils/demo_mode.py +++ b/src/sentry/utils/demo_mode.py @@ -37,6 +37,14 @@ def is_demo_org(organization: Organization | None): return organization.id in options.get("demo-mode.orgs") +def get_demo_org(): + if not is_demo_mode_enabled(): + return None + + org_id = options.get("demo-mode.orgs")[0] + return Organization.objects.get(id=org_id) + + def get_demo_user(): if not is_demo_mode_enabled(): return None diff --git a/tests/demo_mode/test_tasks.py b/tests/demo_mode/test_tasks.py new file mode 100644 index 00000000000000..ee56aba8a8055d --- /dev/null +++ b/tests/demo_mode/test_tasks.py @@ -0,0 +1,146 @@ +from datetime import datetime, timedelta + +import pytest +from django.utils import timezone + +from sentry.demo_mode.tasks import _sync_artifact_bundles +from sentry.models.artifactbundle import ( + ArtifactBundle, + ProjectArtifactBundle, + ReleaseArtifactBundle, +) +from sentry.models.organization import Organization +from sentry.models.project import Project +from sentry.testutils.cases import TestCase + + +class SyncArtifactBundlesTest(TestCase): + + def setUp(self): + self.source_org = self.create_organization(slug="source_org") + self.target_org = self.create_organization(slug="target_org") + self.unrelated_org = self.create_organization(slug="unrelated_org") + + self.source_proj_foo = self.create_project(organization=self.source_org, slug="foo") + self.target_proj_foo = self.create_project(organization=self.target_org, slug="foo") + self.unrelated_proj_foo = self.create_project(organization=self.unrelated_org, slug="foo") + + self.source_proj_bar = self.create_project(organization=self.source_org, slug="bar") + self.target_proj_baz = self.create_project(organization=self.target_org, slug="baz") + + def set_up_artifact_bundle( + self, + organization: Organization, + project: Project, + date_uploaded: datetime | None = None, + ): + date_uploaded = date_uploaded or timezone.now() + artifact_bundle = self.create_artifact_bundle(org=organization, date_uploaded=date_uploaded) + project_artifact_bundle = ProjectArtifactBundle.objects.create( + organization_id=organization.id, + project_id=project.id, + artifact_bundle_id=artifact_bundle.id, + ) + + release_artifact_bundle = ReleaseArtifactBundle.objects.create( + organization_id=organization.id, + artifact_bundle_id=artifact_bundle.id, + dist_name="dist", + release_name="release", + ) + + return artifact_bundle, project_artifact_bundle, release_artifact_bundle + + def test_sync_artifact_bundles_no_bundles(self): + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + assert not ArtifactBundle.objects.all().exists() + + def test_sync_artifact_bundles_with_differences(self): + (source_artifact_bundle, _, __) = self.set_up_artifact_bundle( + self.source_org, self.source_proj_foo + ) + + assert not ArtifactBundle.objects.filter(organization_id=self.target_org.id).exists() + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + target_artifact_bundles = ArtifactBundle.objects.get(organization_id=self.target_org.id) + + assert target_artifact_bundles.bundle_id == source_artifact_bundle.bundle_id + + def test_sync_artifact_bundles_does_not_touch_other_orgs(self): + self.set_up_artifact_bundle(self.source_org, self.source_proj_foo) + self.set_up_artifact_bundle(self.unrelated_org, self.unrelated_proj_foo) + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + unrelated_artifact_bundles = ArtifactBundle.objects.filter( + organization_id=self.unrelated_org.id + ) + + assert unrelated_artifact_bundles.count() == 1 + + def test_sync_artifact_bundles_with_old_uploads(self): + self.set_up_artifact_bundle( + self.source_org, self.source_proj_foo, date_uploaded=timezone.now() - timedelta(days=2) + ) + + assert not ArtifactBundle.objects.filter(organization_id=self.target_org.id).exists() + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + assert not ArtifactBundle.objects.filter(organization_id=self.target_org.id).exists() + + def test_sync_artifact_bundles_only_once(self): + (source_artifact_bundle, _, __) = self.set_up_artifact_bundle( + self.source_org, self.source_proj_foo + ) + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + target_artifact_bundles = ArtifactBundle.objects.filter(organization_id=self.target_org.id) + + assert target_artifact_bundles.count() == 1 + assert target_artifact_bundles[0].bundle_id == source_artifact_bundle.bundle_id + + def test_sync_project_artifact_bundles(self): + self.set_up_artifact_bundle(self.source_org, self.source_proj_foo) + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + target_project_artifact_bundle = ProjectArtifactBundle.objects.get( + organization_id=self.target_org.id, + project_id=self.target_proj_foo.id, + ) + + assert target_project_artifact_bundle.project_id == self.target_proj_foo.id + assert target_project_artifact_bundle.organization_id == self.target_org.id + + def test_sync_release_artifact_bundles(self): + (_, __, source_release_bundle) = self.set_up_artifact_bundle( + self.source_org, self.source_proj_foo + ) + + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + target_release_bundle = ReleaseArtifactBundle.objects.get( + organization_id=self.target_org.id, + ) + + assert target_release_bundle.dist_name == source_release_bundle.dist_name + assert target_release_bundle.release_name == source_release_bundle.release_name + assert target_release_bundle.organization_id == self.target_org.id + + def test_sync_artifact_bunles_rolls_back_on_error(self): + self.set_up_artifact_bundle(self.source_org, self.source_proj_bar) + + with pytest.raises(Project.DoesNotExist): + _sync_artifact_bundles(source_org=self.source_org, target_org=self.target_org) + + assert not ArtifactBundle.objects.filter(organization_id=self.target_org.id).exists() + assert not ProjectArtifactBundle.objects.filter(organization_id=self.target_org.id).exists() + assert not ReleaseArtifactBundle.objects.filter(organization_id=self.target_org.id).exists() diff --git a/tests/sentry/utils/test_demo_mode.py b/tests/sentry/utils/test_demo_mode.py index 367213a2612502..50c6c4ac537b1b 100644 --- a/tests/sentry/utils/test_demo_mode.py +++ b/tests/sentry/utils/test_demo_mode.py @@ -3,7 +3,7 @@ from sentry.testutils.factories import Factories from sentry.testutils.helpers.options import override_options from sentry.testutils.pytest.fixtures import django_db_all -from sentry.utils.demo_mode import get_demo_user, is_demo_org, is_demo_user +from sentry.utils.demo_mode import get_demo_org, get_demo_user, is_demo_org, is_demo_user @override_options({"demo-mode.enabled": True, "demo-mode.users": [1]}) @@ -86,3 +86,18 @@ def test_get_demo_user_demo_mode_enabled(): with patch("sentry.utils.demo_mode.User.objects.get", return_value=user) as mock_user_get: assert get_demo_user() == user mock_user_get.assert_called_once_with(id=1) + + +@override_options({"demo-mode.enabled": False}) +@django_db_all +def test_get_demo_org_demo_mode_disabled(): + assert get_demo_org() is None + + +@override_options({"demo-mode.enabled": True, "demo-mode.orgs": [1]}) +@django_db_all +def test_get_demo_org_demo_mode_enabled(): + org = Factories.create_organization(id=1) + with patch("sentry.utils.demo_mode.Organization.objects.get", return_value=org) as mock_org_get: + assert get_demo_org() == org + mock_org_get.assert_called_once_with(id=1)