Skip to content

Commit 942828d

Browse files
authored
feat(investigations): Add response serializers (#121576)
This splits out the response serializers from #121403 to keep the total pr size down, and restructures them into separate folders. <!-- Describe your PR here. -->
1 parent ec97449 commit 942828d

8 files changed

Lines changed: 993 additions & 0 deletions

File tree

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
__all__ = (
2+
"InvestigationBlockSerializer",
3+
"InvestigationBlockSerializerResponse",
4+
"InvestigationDetailsSerializer",
5+
"InvestigationDetailsSerializerResponse",
6+
"InvestigationParameterSerializer",
7+
"InvestigationParameterSerializerResponse",
8+
"InvestigationSerializer",
9+
"InvestigationSerializerResponse",
10+
)
11+
12+
13+
from .block import InvestigationBlockSerializer, InvestigationBlockSerializerResponse
14+
from .investigation import (
15+
InvestigationDetailsSerializer,
16+
InvestigationDetailsSerializerResponse,
17+
InvestigationSerializer,
18+
InvestigationSerializerResponse,
19+
)
20+
from .parameter import (
21+
InvestigationParameterSerializer,
22+
InvestigationParameterSerializerResponse,
23+
)
Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
from __future__ import annotations
2+
3+
from collections import defaultdict
4+
from collections.abc import Mapping, MutableMapping, Sequence
5+
from datetime import datetime
6+
from typing import Any, TypedDict, override
7+
8+
from django.contrib.auth.models import AnonymousUser
9+
from django.db.models import prefetch_related_objects
10+
11+
from sentry.api.serializers import Serializer
12+
from sentry.investigations.models import (
13+
InvestigationBlock,
14+
InvestigationBlockDependency,
15+
InvestigationBlockExecution,
16+
InvestigationBlockExecutionStatus,
17+
InvestigationBlockKind,
18+
InvestigationBlockParameter,
19+
)
20+
from sentry.users.models.user import User
21+
from sentry.users.services.user.model import RpcUser
22+
23+
24+
class InvestigationBlockExecutionSerializerResponse(TypedDict):
25+
id: str
26+
status: str
27+
executor: str
28+
schemaVersion: int
29+
startedAt: datetime | None
30+
completedAt: datetime | None
31+
error: Any | None
32+
33+
34+
class InvestigationBlockSerializerResponse(TypedDict):
35+
id: str
36+
position: int
37+
kind: str
38+
title: str
39+
content: str
40+
generationPrompt: str
41+
generatedContent: str
42+
output: Any | None
43+
outputStatus: str
44+
currentExecution: InvestigationBlockExecutionSerializerResponse | None
45+
config: dict[str, Any]
46+
display: dict[str, Any]
47+
dependencies: list[str]
48+
parameterKeys: list[str]
49+
version: int
50+
staleAt: datetime | None
51+
createdBy: str | None
52+
lastEditedBy: str | None
53+
54+
55+
class InvestigationBlockSerializer(Serializer[InvestigationBlockSerializerResponse]):
56+
"""
57+
Serializes a block, hiding output the viewer may not see.
58+
59+
``accessible_project_ids`` is the set of projects the viewer can read. A
60+
block's persisted output is withheld unless every project that contributed
61+
to it is in that set, so it must be supplied by the caller.
62+
"""
63+
64+
def __init__(self, accessible_project_ids: set[int]) -> None:
65+
self.accessible_project_ids = accessible_project_ids
66+
67+
@override
68+
def get_attrs(
69+
self,
70+
item_list: Sequence[InvestigationBlock],
71+
user: User | RpcUser | AnonymousUser,
72+
**kwargs: Any,
73+
) -> MutableMapping[InvestigationBlock, dict[str, Any]]:
74+
dependencies: MutableMapping[int, list[str]] = defaultdict(list)
75+
for link in (
76+
InvestigationBlockDependency.objects.filter(block__in=item_list)
77+
.values_list("block_id", "depends_on_id")
78+
.order_by("id")
79+
):
80+
dependencies[link[0]].append(str(link[1]))
81+
82+
parameter_keys: MutableMapping[int, list[str]] = defaultdict(list)
83+
for block_id, key in (
84+
InvestigationBlockParameter.objects.filter(block__in=item_list)
85+
.values_list("block_id", "parameter__key")
86+
.order_by("parameter__position")
87+
):
88+
parameter_keys[block_id].append(key)
89+
90+
prefetch_related_objects(
91+
item_list,
92+
"current_execution__data_projects",
93+
"content_execution__data_projects",
94+
"result_execution__data_projects",
95+
)
96+
97+
return {
98+
block: {
99+
"dependencies": dependencies[block.id],
100+
"parameter_keys": parameter_keys[block.id],
101+
}
102+
for block in item_list
103+
}
104+
105+
@override
106+
def serialize(
107+
self,
108+
obj: InvestigationBlock,
109+
attrs: Mapping[Any, Any],
110+
user: User | RpcUser | AnonymousUser,
111+
**kwargs: Any,
112+
) -> InvestigationBlockSerializerResponse:
113+
def is_accessible(execution: InvestigationBlockExecution | None) -> bool:
114+
if execution is None:
115+
return True
116+
return {project.id for project in execution.data_projects.all()}.issubset(
117+
self.accessible_project_ids
118+
)
119+
120+
execution = obj.current_execution
121+
result_execution = obj.result_execution
122+
content_execution = obj.content_execution
123+
content_restricted = bool(
124+
obj.kind == InvestigationBlockKind.TEXT
125+
and content_execution is not None
126+
and not is_accessible(content_execution)
127+
)
128+
if execution is None:
129+
output = None
130+
output_status = "notRun"
131+
else:
132+
visible_execution = (
133+
result_execution if obj.kind == InvestigationBlockKind.QUERY else execution
134+
)
135+
if not is_accessible(visible_execution):
136+
output = None
137+
output_status = "restricted"
138+
else:
139+
output = visible_execution.result if visible_execution is not None else None
140+
output_status = (
141+
"available"
142+
if execution.status == InvestigationBlockExecutionStatus.COMPLETED
143+
else execution.status
144+
)
145+
if content_restricted:
146+
output = None
147+
output_status = "restricted"
148+
149+
content = obj.content
150+
generated_content = obj.generated_content
151+
if content_restricted:
152+
content = ""
153+
generated_content = ""
154+
155+
return {
156+
"id": str(obj.id),
157+
"position": obj.position,
158+
"kind": obj.kind,
159+
"title": obj.title,
160+
"content": content,
161+
"generationPrompt": obj.prompt,
162+
"generatedContent": generated_content,
163+
"output": output,
164+
"outputStatus": output_status,
165+
"currentExecution": (
166+
{
167+
"id": str(execution.id),
168+
"status": execution.status,
169+
"executor": execution.executor,
170+
"schemaVersion": execution.result_schema_version,
171+
"startedAt": execution.started_at,
172+
"completedAt": execution.completed_at,
173+
"error": execution.error,
174+
}
175+
if execution is not None
176+
else None
177+
),
178+
"config": obj.config,
179+
"display": obj.display,
180+
"dependencies": attrs["dependencies"],
181+
"parameterKeys": attrs["parameter_keys"],
182+
"version": obj.version,
183+
"staleAt": obj.stale_at,
184+
"createdBy": str(obj.created_by_id) if obj.created_by_id is not None else None,
185+
"lastEditedBy": (
186+
str(obj.last_edited_by_id) if obj.last_edited_by_id is not None else None
187+
),
188+
}

0 commit comments

Comments
 (0)