feat(investigations): add platform integration

arslnb · arslnb · commit e42180fcc0cc · 2026-08-06T00:07:36.000-07:00
diff --git a/src/sentry/backup/comparators.py b/src/sentry/backup/comparators.py
@@ -1005,6 +1005,30 @@ def get_default_comparators() -> dict[str, list[JSONScrubbingComparator]]:
             "insights.insightsstarredsegment": [
                 DateUpdatedComparator("date_updated", "date_added")
             ],
+            "investigations.investigation": [
+                DateUpdatedComparator("date_updated", "date_added"),
+                UUID4Comparator("uuid"),
+            ],
+            "investigations.investigationcell": [
+                DateUpdatedComparator("date_updated", "date_added"),
+                UUID4Comparator("uuid"),
+                IgnoredComparator("current_execution"),
+            ],
+            "investigations.investigationcellcomment": [
+                DateUpdatedComparator("date_updated", "date_added"),
+                UUID4Comparator("uuid"),
+            ],
+            "investigations.investigationfavoriteuser": [
+                DateUpdatedComparator("date_updated", "date_added")
+            ],
+            "investigations.investigationcellexecution": [
+                DateUpdatedComparator("date_updated", "date_added"),
+                UUID4Comparator("uuid", "request_id"),
+            ],
+            "investigations.investigationparameter": [
+                DateUpdatedComparator("date_updated", "date_added"),
+                UUID4Comparator("uuid"),
+            ],
             "monitors.monitor": [UUID4Comparator("guid")],
             "replays.organizationmemberreplayaccess": [
                 DateUpdatedComparator("date_updated", "date_added")
diff --git a/src/sentry/features/temporary.py b/src/sentry/features/temporary.py
@@ -112,6 +112,8 @@ def register_temporary_features(manager: FeatureManager) -> None:
     manager.add("organizations:expose-migrated-discover-queries", OrganizationFeature, FeatureHandlerStrategy.FLAGPOLE, api_expose=True)
     # Enable GenAI features such as Autofix and Issue Summary
     manager.add("organizations:gen-ai-features", OrganizationFeature, FeatureHandlerStrategy.FLAGPOLE, api_expose=True)
+    # Enable organization investigation notebooks.
+    manager.add("organizations:investigations", OrganizationFeature, FeatureHandlerStrategy.FLAGPOLE, api_expose=True)
     # Enable the 'translate' functionality for GenAI on the explore > traces page
     manager.add("organizations:gen-ai-search-agent-translate", OrganizationFeature, FeatureHandlerStrategy.FLAGPOLE, api_expose=True)
     # Enable AI search bar on the explore > metrics tab
diff --git a/src/sentry/testutils/factories.py b/src/sentry/testutils/factories.py
@@ -79,6 +79,21 @@
 from sentry.integrations.services.integration import RpcIntegration
 from sentry.integrations.types import ExternalProviders
 from sentry.integrations.utils.hostname import instance_hostname
+from sentry.investigations.models import (
+    Investigation,
+    InvestigationCell,
+    InvestigationCellComment,
+    InvestigationCellDependency,
+    InvestigationCellExecution,
+    InvestigationCellExecutionProject,
+    InvestigationCellParameter,
+    InvestigationCellReaction,
+    InvestigationCommentReaction,
+    InvestigationCommentTeamMention,
+    InvestigationCommentUserMention,
+    InvestigationParameter,
+    InvestigationPermissions,
+)
 from sentry.issue_detection.performance_problem import PerformanceProblem
 from sentry.issues.action_log.types import GroupActionType, GroupActorType
 from sentry.issues.grouptype import get_group_type_by_type_id
@@ -408,6 +423,83 @@ def _set_sample_rate_from_error_sampling(normalized_data: MutableMapping[str, An
 
 # TODO(dcramer): consider moving to something more scalable like factoryboy
 class Factories:
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation(organization, created_by=None, **kwargs):
+        return Investigation.objects.create(
+            organization=organization,
+            created_by_id=created_by.id if created_by else None,
+            **kwargs,
+        )
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_permissions(investigation, **kwargs):
+        return InvestigationPermissions.objects.create(investigation=investigation, **kwargs)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell(investigation, position=0, kind="text", **kwargs):
+        return InvestigationCell.objects.create(
+            investigation=investigation, position=position, kind=kind, **kwargs
+        )
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell_dependency(cell, depends_on):
+        return InvestigationCellDependency.objects.create(cell=cell, depends_on=depends_on)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_parameter(investigation, **kwargs):
+        return InvestigationParameter.objects.create(investigation=investigation, **kwargs)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell_parameter(cell, parameter, **kwargs):
+        return InvestigationCellParameter.objects.create(cell=cell, parameter=parameter, **kwargs)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell_execution(cell, **kwargs):
+        return InvestigationCellExecution.objects.create(cell=cell, **kwargs)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell_execution_project(execution, project):
+        return InvestigationCellExecutionProject.objects.create(
+            execution=execution, project=project
+        )
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell_comment(cell, author=None, **kwargs):
+        return InvestigationCellComment.objects.create(
+            cell=cell, author_id=author.id if author else None, **kwargs
+        )
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_cell_reaction(cell, user, **kwargs):
+        return InvestigationCellReaction.objects.create(cell=cell, user_id=user.id, **kwargs)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_comment_reaction(comment, user, **kwargs):
+        return InvestigationCommentReaction.objects.create(
+            comment=comment, user_id=user.id, **kwargs
+        )
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_comment_user_mention(comment, user):
+        return InvestigationCommentUserMention.objects.create(comment=comment, user_id=user.id)
+
+    @staticmethod
+    @assume_test_silo_mode(SiloMode.CELL)
+    def create_investigation_comment_team_mention(comment, team):
+        return InvestigationCommentTeamMention.objects.create(comment=comment, team=team)
+
     @staticmethod
     @assume_test_silo_mode(SiloMode.CELL)
     def create_organization(name=None, owner=None, cell: Cell | str | None = None, **kwargs):
diff --git a/src/sentry/testutils/fixtures.py b/src/sentry/testutils/fixtures.py
@@ -184,6 +184,45 @@ def organization_integration(self):
     def create_organization(self, *args, **kwargs):
         return Factories.create_organization(*args, **kwargs)
 
+    def create_investigation(self, *args, **kwargs):
+        return Factories.create_investigation(*args, **kwargs)
+
+    def create_investigation_permissions(self, *args, **kwargs):
+        return Factories.create_investigation_permissions(*args, **kwargs)
+
+    def create_investigation_cell(self, *args, **kwargs):
+        return Factories.create_investigation_cell(*args, **kwargs)
+
+    def create_investigation_cell_dependency(self, *args, **kwargs):
+        return Factories.create_investigation_cell_dependency(*args, **kwargs)
+
+    def create_investigation_parameter(self, *args, **kwargs):
+        return Factories.create_investigation_parameter(*args, **kwargs)
+
+    def create_investigation_cell_parameter(self, *args, **kwargs):
+        return Factories.create_investigation_cell_parameter(*args, **kwargs)
+
+    def create_investigation_cell_execution(self, *args, **kwargs):
+        return Factories.create_investigation_cell_execution(*args, **kwargs)
+
+    def create_investigation_cell_execution_project(self, *args, **kwargs):
+        return Factories.create_investigation_cell_execution_project(*args, **kwargs)
+
+    def create_investigation_cell_comment(self, *args, **kwargs):
+        return Factories.create_investigation_cell_comment(*args, **kwargs)
+
+    def create_investigation_cell_reaction(self, *args, **kwargs):
+        return Factories.create_investigation_cell_reaction(*args, **kwargs)
+
+    def create_investigation_comment_reaction(self, *args, **kwargs):
+        return Factories.create_investigation_comment_reaction(*args, **kwargs)
+
+    def create_investigation_comment_user_mention(self, *args, **kwargs):
+        return Factories.create_investigation_comment_user_mention(*args, **kwargs)
+
+    def create_investigation_comment_team_mention(self, *args, **kwargs):
+        return Factories.create_investigation_comment_team_mention(*args, **kwargs)
+
     def create_member(self, *args, **kwargs):
         return Factories.create_member(*args, **kwargs)
 
diff --git a/src/sentry/testutils/helpers/backups.py b/src/sentry/testutils/helpers/backups.py
@@ -58,6 +58,24 @@
 from sentry.integrations.models.data_forwarder_project import DataForwarderProject
 from sentry.integrations.models.integration import Integration
 from sentry.integrations.models.organization_integration import OrganizationIntegration
+from sentry.investigations.models import (
+    Investigation,
+    InvestigationCell,
+    InvestigationCellComment,
+    InvestigationCellDependency,
+    InvestigationCellExecution,
+    InvestigationCellExecutionProject,
+    InvestigationCellParameter,
+    InvestigationCellReaction,
+    InvestigationCommentReaction,
+    InvestigationCommentTeamMention,
+    InvestigationCommentUserMention,
+    InvestigationFavoriteUser,
+    InvestigationParameter,
+    InvestigationPermissions,
+    InvestigationPermissionsTeam,
+    InvestigationProject,
+)
 from sentry.models.activity import Activity
 from sentry.models.apiauthorization import ApiAuthorization
 from sentry.models.apidevicecode import ApiDeviceCode
@@ -841,6 +859,105 @@ def create_exhaustive_organization(
             overrides={"write_key": "test_override_write_key"},
         )
 
+        investigation = Investigation.objects.create(
+            organization=org,
+            created_by_id=owner_id,
+            title=f"Investigation in {slug}",
+            status="archived",
+            template_key="breached_metric",
+            template_version=2,
+            source_type="breached_metric",
+            source_ref={"groupId": str(group.id)},
+            filters={"environment": ["production"]},
+            version=2,
+        )
+        InvestigationProject.objects.create(investigation=investigation, project=project)
+        InvestigationFavoriteUser.objects.create(investigation=investigation, user_id=owner_id)
+        investigation_permissions = InvestigationPermissions.objects.create(
+            investigation=investigation,
+            is_editable_by_everyone=False,
+        )
+        InvestigationPermissionsTeam.objects.create(
+            permissions=investigation_permissions, team=team
+        )
+        parameter = InvestigationParameter.objects.create(
+            investigation=investigation,
+            key="environment",
+            label="Environment",
+            description="Environment filter",
+            type="string",
+            required=True,
+            constraints={"maxLength": 64},
+            default_value="development",
+            saved_value="production",
+            source="user",
+            position=1,
+            version=2,
+        )
+        upstream_cell = InvestigationCell.objects.create(
+            investigation=investigation,
+            created_by_id=owner_id,
+            last_edited_by_id=owner_id,
+            position=1,
+            kind="text",
+            title="Upstream",
+            content="Upstream evidence",
+            prompt="Create evidence",
+            generated_content="Generated evidence",
+            config={"language": "markdown"},
+            display={"type": "markdown"},
+            version=2,
+            stale_at=timezone.now(),
+            deleted_at=timezone.now(),
+        )
+        cell = InvestigationCell.objects.create(
+            investigation=investigation,
+            created_by_id=owner_id,
+            last_edited_by_id=owner_id,
+            position=2,
+            kind="query",
+            title="Query",
+            content="count()",
+            prompt="Build a query",
+            generated_content="count()",
+            config={"dataset": "spans"},
+            display={"type": "table"},
+            version=2,
+            stale_at=timezone.now(),
+            deleted_at=timezone.now(),
+        )
+        InvestigationCellDependency.objects.create(cell=cell, depends_on=upstream_cell)
+        InvestigationCellParameter.objects.create(cell=cell, parameter=parameter)
+        execution = InvestigationCellExecution.objects.create(
+            cell=cell,
+            triggered_by_id=owner_id,
+            executor="code_mode",
+            status="completed",
+            cell_version=2,
+            input_snapshot={"environment": "production"},
+            input_fingerprint="f" * 64,
+            result_schema_version=2,
+            result={"rows": [[1]]},
+            error={"previous": "transient"},
+            started_at=timezone.now(),
+            completed_at=timezone.now(),
+        )
+        InvestigationCellExecutionProject.objects.create(execution=execution, project=project)
+        cell.current_execution = execution
+        cell.save(update_fields=["current_execution"])
+        comment = InvestigationCellComment.objects.create(
+            cell=cell,
+            author_id=owner_id,
+            body="Please review this result",
+            deleted_at=timezone.now(),
+        )
+        InvestigationCellReaction.objects.create(cell=cell, user_id=owner_id, reaction="heart")
+        InvestigationCommentReaction.objects.create(
+            comment=comment, user_id=owner_id, reaction="eyes"
+        )
+        InvestigationCommentUserMention.objects.create(comment=comment, user_id=owner_id)
+        InvestigationCommentTeamMention.objects.create(comment=comment, team=team)
+
         return org
 
     @assume_test_silo_mode(SiloMode.CONTROL)
diff --git a/tests/sentry/backup/test_imports.py b/tests/sentry/backup/test_imports.py
@@ -33,6 +33,13 @@
 )
 from sentry.backup.scopes import ExportScope, ImportScope, RelocationScope
 from sentry.backup.services.import_export.model import RpcImportErrorKind
+from sentry.investigations.models import (
+    Investigation,
+    InvestigationCell,
+    InvestigationCellComment,
+    InvestigationCellExecution,
+    InvestigationParameter,
+)
 from sentry.models.apitoken import DEFAULT_EXPIRATION, ApiToken, generate_token
 from sentry.models.importchunk import (
     ControlImportChunk,
@@ -1463,7 +1470,15 @@ def test_colliding_member_invite_token(self, expected_models: list[type[Model]])
             with open(tmp_path, "rb") as tmp_file:
                 verify_models_in_output(expected_models, orjson.loads(tmp_file.read()))
 
-    @expect_models(COLLISION_TESTED, Monitor)
+    @expect_models(
+        COLLISION_TESTED,
+        Monitor,
+        Investigation,
+        InvestigationCell,
+        InvestigationCellComment,
+        InvestigationCellExecution,
+        InvestigationParameter,
+    )
     def test_colliding_monitor(self, expected_models: list[type[Model]]) -> None:
         owner = self.create_exhaustive_user("owner")
         invited = self.create_exhaustive_user("invited")
