feat(dynamic-sampling):Add raw_query and user relation to custom_rules model (second try) (#58447)

Author: RaduW

fixtures/backup/model_dependencies/detailed.json

@@ -1103,6 +1103,11 @@
   "sentry.customdynamicsamplingrule": {
     "dangling": false,
     "foreign_keys": {
+      "created_by_id": {
+        "kind": "HybridCloudForeignKey",
+        "model": "sentry.user",
+        "nullable": true
+      },
       "organization": {
         "kind": "FlexibleForeignKey",
         "model": "sentry.organization",

fixtures/backup/model_dependencies/flat.json

@@ -150,7 +150,8 @@
     "sentry.project"
   ],
   "sentry.customdynamicsamplingrule": [
-    "sentry.organization"
+    "sentry.organization",
+    "sentry.user"
   ],
   "sentry.customdynamicsamplingruleproject": [
     "sentry.customdynamicsamplingrule",

migrations_lockfile.txt

@@ -9,5 +9,5 @@ feedback: 0003_feedback_add_env
 hybridcloud: 0005_add_missing_org_integration_scope
 nodestore: 0002_nodestore_no_dictfield
 replays: 0003_add_size_to_recording_segment
-sentry: 0577_drop_latest_incident_index
+sentry: 0578_add_query_and_users_to_custom_dynamic_sampling_rules
 social_auth: 0002_default_auto_field

src/sentry/api/endpoints/custom_rules.py

@@ -158,6 +158,7 @@ def post(self, request: Request, organization: Organization) -> Response:
                 organization_id=organization.id,
                 num_samples=NUM_SAMPLES_PER_CUSTOM_RULE,
                 sample_rate=1.0,
+                query=query,
             )
 
             # schedule update for affected project configs

src/sentry/migrations/0578_add_query_and_users_to_custom_dynamic_sampling_rules.py

@@ -0,0 +1,39 @@
+# Generated by Django 3.2.20 on 2023-10-20 07:59
+
+from django.db import migrations, models
+
+import sentry.db.models.fields.hybrid_cloud_foreign_key
+from sentry.new_migrations.migrations import CheckedMigration
+
+
+class Migration(CheckedMigration):
+    # This flag is used to mark that a migration shouldn't be automatically run in production. For
+    # the most part, this should only be used for operations where it's safe to run the migration
+    # after your code has deployed. So this should not be used for most operations that alter the
+    # schema of a table.
+    # Here are some things that make sense to mark as dangerous:
+    # - Large data migrations. Typically we want these to be run manually by ops so that they can
+    #   be monitored and not block the deploy for a long period of time while they run.
+    # - Adding indexes to large tables. Since this can take a long time, we'd generally prefer to
+    #   have ops run this and not block the deploy. Note that while adding an index is a schema
+    #   change, it's completely safe to run the operation after the code has deployed.
+    is_dangerous = False
+
+    dependencies = [
+        ("sentry", "0577_drop_latest_incident_index"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="customdynamicsamplingrule",
+            name="created_by_id",
+            field=sentry.db.models.fields.hybrid_cloud_foreign_key.HybridCloudForeignKey(
+                "sentry.User", blank=True, db_index=True, null=True, on_delete="CASCADE"
+            ),
+        ),
+        migrations.AddField(
+            model_name="customdynamicsamplingrule",
+            name="query",
+            field=models.TextField(null=True),
+        ),
+    ]

src/sentry/models/dynamicsampling.py

@@ -8,12 +8,12 @@
 
 from sentry.backup.scopes import RelocationScope
 from sentry.db.models import FlexibleForeignKey, Model, region_silo_only_model
+from sentry.db.models.fields.hybrid_cloud_foreign_key import HybridCloudForeignKey
 from sentry.utils import json
 
 if TYPE_CHECKING:
     from sentry.models.project import Project
 
-
 # max number of custom rules that can be created per organization
 MAX_CUSTOM_RULES = 2000
 CUSTOM_RULE_START = 3000
@@ -104,6 +104,9 @@ class CustomDynamicSamplingRule(Model):
     end_date = models.DateTimeField()
     num_samples = models.IntegerField()
     condition_hash = models.CharField(max_length=40)
+    # the raw query field from the request
+    query = models.TextField(null=True)
+    created_by_id = HybridCloudForeignKey("sentry.User", on_delete="CASCADE", null=True, blank=True)
 
     @property
     def external_rule_id(self) -> int:
@@ -160,6 +163,7 @@ def update_or_create(
         organization_id: int,
         num_samples: int,
         sample_rate: float,
+        query: str,
     ) -> "CustomDynamicSamplingRule":
 
         from sentry.models.project import Project
@@ -195,6 +199,7 @@ def update_or_create(
                     condition_hash=rule_hash,
                     is_active=True,
                     is_org_level=is_org_level,
+                    query=query,
                 )
 
                 rule.save()

src/sentry/testutils/helpers/backups.py

@@ -408,6 +408,7 @@ def create_exhaustive_organization(
             organization_id=org.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod event.type:transaction",
         )
 
         # Environment*

tests/sentry/api/endpoints/test_custom_rules.py

@@ -55,6 +55,7 @@ def setUp(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=1.0,
+            query="event.type:transaction, environment:prod",
         )
 
         # create an org rule
@@ -76,24 +77,9 @@ def setUp(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=1.0,
+            query="event.type:transaction, environment:dev",
         )
 
-        # # create a condition with empty query
-        # now = timezone.now()
-        # self.empty_condition = {"op": "and", "inner": []}
-        # start = now - timedelta(hours=2)
-        # end = now + timedelta(hours=2)
-        #
-        # CustomDynamicSamplingRule.update_or_create(
-        #     condition=self.empty_condition,
-        #     start=start,
-        #     end=end,
-        #     project_ids=[self.known_projects[0].id],
-        #     organization_id=self.organization.id,
-        #     num_samples=100,
-        #     sample_rate=1.0,
-        # )
-
     def test_finds_project_rule(self):
         """
         Tests that the endpoint finds the rule when the query matches and

tests/sentry/backup/snapshots/ReleaseTests/test_at_head.pysnap

@@ -672,6 +672,8 @@ source: tests/sentry/backup/test_releases.py
     is_org_level: false
     num_samples: 100
     organization: 4552811202150400
+    created_by_id: null
+    query: environment:prod event.type:transaction
     rule_id: 1
     sample_rate: 0.5
     start_date: '2023-10-17T01:50:40.016Z'

tests/sentry/backup/test_models.py

@@ -247,6 +247,7 @@ def test_custom_dynamic_sampling(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod",
         )
         return self.import_export_then_validate()
 

tests/sentry/dynamic_sampling/test_generate_rules.py

@@ -738,6 +738,7 @@ def test_generate_rules_return_custom_rules(get_blended_sample_rate, default_old
         organization_id=default_old_project.organization.id,
         num_samples=100,
         sample_rate=0.5,
+        query="environment:prod1",
     )
     # and an organization rule
     condition = {"op": "eq", "name": "environment", "value": "prod2"}
@@ -749,6 +750,7 @@ def test_generate_rules_return_custom_rules(get_blended_sample_rate, default_old
         organization_id=default_old_project.organization.id,
         num_samples=100,
         sample_rate=0.5,
+        query="environment:prod2",
     )
 
     rules = generate_rules(default_old_project)

tests/sentry/models/test_dynamicsampling.py

@@ -31,6 +31,7 @@ def test_update_or_create(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod",
         )
 
         end2 = timezone.now() + timedelta(hours=1)
@@ -42,6 +43,7 @@ def test_update_or_create(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod",
         )
 
         assert rule.id == updated_rule.id
@@ -64,6 +66,7 @@ def create_rule_for_env(env_idx: int):
                 organization_id=self.organization.id,
                 num_samples=100,
                 sample_rate=0.5,
+                query=f"environment:prod{env_idx}",
             )
 
         rule_ids = set()
@@ -116,6 +119,7 @@ def create_rule(is_old: bool, idx: int):
                 organization_id=self.organization.id,
                 num_samples=100,
                 sample_rate=0.5,
+                query=f"environment:prod{idx}",
             )
 
         for i in range(10):
@@ -160,6 +164,7 @@ def test_get_rule_for_org(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod",
         )
 
         rule = CustomDynamicSamplingRule.get_rule_for_org(
@@ -204,6 +209,7 @@ def create_rule(
                 organization_id=org_id,
                 num_samples=100,
                 sample_rate=0.5,
+                query=f"environment:prod{idx[0]}",
             )
 
         valid_project_rule = create_rule([self.project.id, self.second_project.id])
@@ -243,6 +249,7 @@ def test_separate_projects_create_different_rules(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod",
         )
 
         end2 = timezone.now() + timedelta(hours=1)
@@ -254,6 +261,7 @@ def test_separate_projects_create_different_rules(self):
             organization_id=self.organization.id,
             num_samples=100,
             sample_rate=0.5,
+            query="environment:prod",
         )
 
         assert rule.id != second_rule.id