feat(demo-mode): artifact bundle sync (#86302)

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.

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

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

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

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

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)