Send investigation rule notification email (#59684)

Co-authored-by: Ogi <86684834+obostjancic@users.noreply.github.com>

migrations_lockfile.txt

@@ -9,5 +9,5 @@ feedback: 0003_feedback_add_env
 hybridcloud: 0008_add_externalactorreplica
 nodestore: 0002_nodestore_no_dictfield
 replays: 0003_add_size_to_recording_segment
-sentry: 0592_delete_relocation_hybrid_cloud_foreign_keys
+sentry: 0593_add_notification_flag_to_dynamic_sampling_custom_rule
 social_auth: 0002_default_auto_field

src/sentry/api/endpoints/custom_rules.py

@@ -159,6 +159,7 @@ class CustomRulesEndpoint(OrganizationEndpoint):
                 num_samples=NUM_SAMPLES_PER_CUSTOM_RULE,
                 sample_rate=1.0,
                 query=query,
+                created_by_id=request.user.id,
             )
 
             # schedule update for affected project configs

src/sentry/conf/server.py

@@ -762,6 +762,7 @@ CELERY_IMPORTS = (
     "sentry.dynamic_sampling.tasks.recalibrate_orgs",
     "sentry.dynamic_sampling.tasks.sliding_window_org",
     "sentry.dynamic_sampling.tasks.utils",
+    "sentry.dynamic_sampling.tasks.custom_rule_notifications",
     "sentry.utils.suspect_resolutions.get_suspect_resolutions",
     "sentry.utils.suspect_resolutions_releases.get_suspect_resolutions_releases",
     "sentry.tasks.derive_code_mappings",
@@ -1124,6 +1125,11 @@ CELERYBEAT_SCHEDULE_REGION = {
         # Run every 10 minutes
         "schedule": crontab(minute="*/10"),
     },
+    "custom_rule_notifications": {
+        "task": "sentry.dynamic_sampling.tasks.custom_rule_notifications",
+        # Run every 10 minutes
+        "schedule": crontab(minute="*/10"),
+    },
     "weekly-escalating-forecast": {
         "task": "sentry.tasks.weekly_escalating_forecast.run_escalating_forecast",
         # TODO: Change this to run weekly once we verify the results

src/sentry/dynamic_sampling/tasks/common.py

@@ -120,6 +120,34 @@ class ContextIterator(Protocol):
         ...
 
 
+class _SimpleContextIterator(Iterator[Any]):
+    def __init__(self, inner: Iterator[Any]):
+        self.inner = inner
+        self.log_state = DynamicSamplingLogState()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return next(self.inner)
+
+    def get_current_state(self) -> DynamicSamplingLogState:
+        return self.log_state
+
+    def set_current_state(self, state: DynamicSamplingLogState) -> None:
+        self.log_state = state
+
+
+def to_context_iterator(inner: Iterator[Any]) -> ContextIterator:
+    """
+    Adds a LogState to a simple iterator turning it into a ContextIterator
+
+    Note: (RaduW) this was probably a mistake, ContextIterator, _SimpleContextIterator and TimedIterator
+    should have been rolled into one type TimedIterator.
+    """
+    return _SimpleContextIterator(inner)
+
+
 class TimedIterator(Iterator[Any]):
     """
     An iterator that wraps an existing ContextIterator.

src/sentry/dynamic_sampling/tasks/custom_rule_notifications.py

@@ -0,0 +1,163 @@
+"""
+Task for sending notifications when custom rules have gathered enough samples.
+"""
+from datetime import datetime, timezone
+from typing import Any, Dict, List
+
+from django.http import QueryDict
+
+from sentry.constants import ObjectStatus
+from sentry.dynamic_sampling.tasks.common import TimedIterator, to_context_iterator
+from sentry.dynamic_sampling.tasks.constants import MAX_TASK_SECONDS
+from sentry.dynamic_sampling.tasks.task_context import DynamicSamplingLogState, TaskContext
+from sentry.dynamic_sampling.tasks.utils import dynamic_sampling_task_with_context
+from sentry.models.dynamicsampling import CustomDynamicSamplingRule
+from sentry.models.user import User
+from sentry.silo import SiloMode
+from sentry.snuba import discover
+from sentry.tasks.base import instrumented_task
+from sentry.utils.email import MessageBuilder
+
+MIN_SAMPLES_FOR_NOTIFICATION = 10
+
+
+@instrumented_task(
+    name="sentry.dynamic_sampling.tasks.custom_rule_notifications",
+    queue="dynamicsampling",
+    default_retry_delay=5,
+    max_retries=5,
+    soft_time_limit=2 * 60 * 60,  # 2hours
+    time_limit=2 * 60 * 60 + 5,
+    silo_mode=SiloMode.REGION,
+)
+@dynamic_sampling_task_with_context(max_task_execution=MAX_TASK_SECONDS)
+def custom_rule_notifications(context: TaskContext) -> None:
+    """
+    Iterates through all active CustomRules and sends a notification to the rule creator
+    whenever enough samples have been collected.
+    """
+    log_state = DynamicSamplingLogState()
+    cur_org_id = None
+    now = datetime.now(tz=timezone.utc)
+    # just for protection filter out rules that are outside the time range (in case the
+    # task that deactivates rules is not working)
+    custom_active_rules = (
+        CustomDynamicSamplingRule.objects.filter(
+            is_active=True,
+            notification_sent=False,
+            start_date__lte=now,
+            end_date__gte=now,
+        )
+        .order_by("organization_id")
+        .iterator()
+    )
+
+    custom_active_rules = to_context_iterator(custom_active_rules)
+
+    for rule in TimedIterator(context, custom_active_rules, name="custom_rule_notifications"):
+        if rule.organization_id != cur_org_id:
+            cur_org_id = rule.organization_id
+            log_state.num_orgs += 1
+        log_state.num_rows_total += 1
+
+        num_samples = get_num_samples(rule)
+
+        if num_samples >= MIN_SAMPLES_FOR_NOTIFICATION:
+            send_notification(rule, num_samples)
+            rule.notification_sent = True
+            rule.save()
+
+
+def get_num_samples(rule: CustomDynamicSamplingRule) -> int:
+    """
+    Returns the number of samples accumulated for the given rule.
+    """
+    projects = rule.projects.all()
+
+    if not projects:
+        # org rule get all projects for org
+        projects = rule.organization.project_set.filter(status=ObjectStatus.ACTIVE)
+
+    params: Dict[str, Any] = {
+        "start": rule.start_date,
+        "end": rule.end_date,
+        "project_id": [p.id for p in projects],
+        "project_objects": projects,
+        "organization_id": rule.organization.id,
+    }
+
+    result = discover.query(
+        selected_columns=["count()"],
+        params=params,
+        query=rule.query,
+        referrer="dynamic_sampling.tasks.custom_rule_notifications",
+    )
+
+    return result["data"][0]["count"]
+
+
+def send_notification(rule: CustomDynamicSamplingRule, num_samples: int) -> None:
+    """
+    Notifies the rule creator that samples have been gathered.
+    """
+    subject_template = "We've collected {num_samples} samples for the query: {query} you made"
+
+    user_id = rule.created_by_id
+    if not user_id:
+        return
+
+    projects = rule.projects.all()
+    project_ids = [p.id for p in projects]
+
+    creator = User.objects.get_from_cache(id=user_id)
+    params = {
+        "query": rule.query,
+        "num_samples": num_samples,
+        "start_date": rule.start_date.strftime("%Y-%m-%d %H:%M:%S"),
+        "end_date": rule.end_date.strftime("%Y-%m-%d %H:%M:%S"),
+        "name": creator.name,
+        "email": creator.email,
+        "user_name": creator.username,
+        "display_name": creator.get_display_name(),
+        "discover_link": create_discover_link(rule, project_ids),
+    }
+
+    subject = subject_template.format(**params)
+
+    msg = MessageBuilder(
+        subject=subject,
+        template="sentry/emails/dyn-sampling-custom-rule-fulfilled.txt",
+        html_template="sentry/emails/dyn-sampling-custom-rule-fulfilled.html",
+        context=params,
+    )
+
+    msg.send_async([creator.email])
+
+
+def create_discover_link(rule: CustomDynamicSamplingRule, projects: List[int]) -> str:
+    """
+    Creates a discover link for the given rule.
+    It will point to a discover query using the same query as the rule
+    and the same time range as the rule.
+    """
+    if len(projects) == 0:
+        projects = [-1]
+
+    project_ids = [str(p) for p in projects]
+
+    q = QueryDict(mutable=True)
+    q["start"] = rule.start_date.strftime("%Y-%m-%dT%H:%M:%S")
+    q["end"] = rule.end_date.strftime("%Y-%m-%dT%H:%M:%S")
+    q.setlist("field", ["title", "event.type", "project", "user.display", "timestamp"])
+    q.setlist("project", project_ids)
+    q["name"] = "All Events"
+    q["query"] = rule.query if rule.query else ""
+    q["utc"] = "true"
+    q["yAxis"] = "count()"
+    q["sort"] = "-timestamp"
+
+    query_string = q.urlencode()
+    discover_url = rule.organization.absolute_url(
+        f"/organizations/{rule.organization.slug}/discover/results/", query=query_string
+    )
+    return discover_url

src/sentry/migrations/0593_add_notification_flag_to_dynamic_sampling_custom_rule.py

@@ -0,0 +1,31 @@
+# Generated by Django 3.2.23 on 2023-11-10 11:27
+
+from django.db import migrations, models
+
+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", "0592_delete_relocation_hybrid_cloud_foreign_keys"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="customdynamicsamplingrule",
+            name="notification_sent",
+            field=models.BooleanField(blank=True, null=True),
+        ),
+    ]

src/sentry/models/dynamicsampling.py

@@ -107,6 +107,7 @@ class CustomDynamicSamplingRule(Model):
     # 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)
+    notification_sent = models.BooleanField(null=True, blank=True)
 
     @property
     def external_rule_id(self) -> int:
@@ -164,6 +165,7 @@ class CustomDynamicSamplingRule(Model):
         num_samples: int,
         sample_rate: float,
         query: str,
+        created_by_id: Optional[int] = None,
     ) -> "CustomDynamicSamplingRule":
 
         from sentry.models.project import Project
@@ -200,6 +202,8 @@ class CustomDynamicSamplingRule(Model):
                     is_active=True,
                     is_org_level=is_org_level,
                     query=query,
+                    notification_sent=False,
+                    created_by_id=created_by_id,
                 )
 
                 rule.save()

src/sentry/snuba/referrer.py

@@ -475,6 +475,9 @@ class Referrer(Enum):
         "dynamic_sampling.counters.fetch_projects_with_transaction_totals"
     )
     DYNAMIC_SAMPLING_COUNTERS_FETCH_ACTIVE_ORGS = "dynamic_sampling.counters.fetch_active_orgs"
+    DYNAMIC_SAMPLING_TASKS_CUSTOM_RULE_NOTIFICATIONS = (
+        "dynamic_sampling.tasks.custom_rule_notifications"
+    )
     ESCALATING_GROUPS = "sentry.issues.escalating"
     EVENTSTORE_GET_EVENT_BY_ID_NODESTORE = "eventstore.backend.get_event_by_id_nodestore"
     EVENTSTORE_GET_EVENTS = "eventstore.get_events"

src/sentry/templates/sentry/emails/dyn-sampling-custom-rule-fulfilled.html

@@ -0,0 +1,10 @@
+<h3>Hello {{ display_name }},</h3>
+
+<p>We've collected {{ num_samples }} samples for the following search query:</p>
+<pre>
+{{ query }}
+</pre>
+
+<p>
+You can view the samples <a href="{{ discover_link|safe }}">here</a>.
+</p>

src/sentry/templates/sentry/emails/dyn-sampling-custom-rule-fulfilled.txt

@@ -0,0 +1,7 @@
+Hello {{ display_name }},
+
+We've collected {{ num_samples }} samples for the following search query:
+{{ query }}
+
+You can view the samples here:
+{{ discover_link|safe }}