|
9 | 9 | logger = logging.getLogger("sentry.workflow_engine.process_data_source")
|
10 | 10 |
|
11 | 11 |
|
12 |
| -# TODO - @saponifi3d - change the text choices to an enum |
| 12 | +def bulk_fetch_enabled_detectors( |
| 13 | + source_ids: set[str], query_type: str |
| 14 | +) -> dict[str, list[Detector]]: |
| 15 | + """ |
| 16 | + Get all of the enabled detectors for a list of detector source ids and types. |
| 17 | + This will also prefetch all the subsequent data models for evaluating the detector. |
| 18 | + """ |
| 19 | + data_sources = ( |
| 20 | + DataSource.objects.filter( |
| 21 | + source_id__in=source_ids, |
| 22 | + type=query_type, |
| 23 | + detectors__enabled=True, |
| 24 | + ) |
| 25 | + .prefetch_related( |
| 26 | + Prefetch( |
| 27 | + "detectors", |
| 28 | + queryset=Detector.objects.filter(enabled=True) |
| 29 | + .select_related("workflow_condition_group") |
| 30 | + .prefetch_related("workflow_condition_group__conditions"), |
| 31 | + ) |
| 32 | + ) |
| 33 | + .distinct() |
| 34 | + ) |
| 35 | + |
| 36 | + result: dict[str, list[Detector]] = {} |
| 37 | + for data_source in data_sources: |
| 38 | + result[data_source.source_id] = list(data_source.detectors.all()) |
| 39 | + |
| 40 | + return result |
| 41 | + |
| 42 | + |
13 | 43 | # TODO - @saponifi3d - make query_type optional override, otherwise infer from the data packet.
|
14 | 44 | def process_data_sources[
|
15 | 45 | T
|
16 | 46 | ](data_packets: list[DataPacket[T]], query_type: str) -> list[tuple[DataPacket[T], list[Detector]]]:
|
17 | 47 | metrics.incr("workflow_engine.process_data_sources", tags={"query_type": query_type})
|
18 | 48 |
|
19 |
| - data_packet_ids = {packet.source_id for packet in data_packets} |
20 |
| - |
21 |
| - # Fetch all data sources and associated detectors for the given data packets |
22 |
| - with sentry_sdk.start_span(op="workflow_engine.process_data_sources.fetch_data_sources"): |
23 |
| - data_sources = DataSource.objects.filter( |
24 |
| - source_id__in=data_packet_ids, |
25 |
| - type=query_type, |
26 |
| - detectors__enabled=True, |
27 |
| - ).prefetch_related(Prefetch("detectors")) |
28 |
| - |
29 |
| - # Build a lookup dict for source_id to detectors |
30 |
| - source_id_to_detectors = {ds.source_id: list(ds.detectors.all()) for ds in data_sources} |
| 49 | + with sentry_sdk.start_span(op="workflow_engine.process_data_sources.get_enabled_detectors"): |
| 50 | + packet_source_ids = {packet.source_id for packet in data_packets} |
| 51 | + source_to_detector = bulk_fetch_enabled_detectors(packet_source_ids, query_type) |
31 | 52 |
|
32 | 53 | # Create the result tuples
|
33 | 54 | result = []
|
34 | 55 | for packet in data_packets:
|
35 |
| - detectors = source_id_to_detectors.get(packet.source_id) |
| 56 | + detectors: list[Detector] = source_to_detector.get(packet.source_id, []) |
36 | 57 |
|
37 | 58 | if detectors:
|
38 | 59 | data_packet_tuple = (packet, detectors)
|
|
0 commit comments