Skip to content

Commit c3ad03b

Browse files
leeandherleeandher
committed
🏷️ correct typing issue
1 parent c5d13af commit c3ad03b

File tree

9 files changed

+57
-52
lines changed

9 files changed

+57
-52
lines changed

‎src/sentry/integrations/data_forwarding/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from sentry.integrations.data_forwarding.splunk.forwarder import SplunkForwarder
77
from sentry.integrations.types import DataForwarderProviderSlug
88

9-
FORWARDER_REGISTRY: Mapping[DataForwarderProviderSlug, type[BaseDataForwarder]] = {
9+
FORWARDER_REGISTRY: Mapping[str, type[BaseDataForwarder]] = {
1010
DataForwarderProviderSlug.SEGMENT.value: SegmentForwarder,
1111
DataForwarderProviderSlug.SQS.value: AmazonSQSForwarder,
1212
DataForwarderProviderSlug.SPLUNK.value: SplunkForwarder,

‎src/sentry/integrations/data_forwarding/amazon_sqs/forwarder.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from sentry.api.serializers import serialize
1313
from sentry.integrations.data_forwarding.base import BaseDataForwarder
1414
from sentry.integrations.types import DataForwarderProviderSlug
15-
from sentry.services.eventstore.models import Event
15+
from sentry.services.eventstore.models import Event, GroupEvent
1616

1717
logger = logging.getLogger(__name__)
1818

@@ -24,7 +24,9 @@ class AmazonSQSForwarder(BaseDataForwarder):
2424
provider = DataForwarderProviderSlug.SQS
2525
rate_limit = (0, 0)
2626

27-
def get_event_payload(self, event: Event, config: dict[str, Any]) -> dict[str, Any]:
27+
def get_event_payload(
28+
self, event: Event | GroupEvent, config: dict[str, Any]
29+
) -> dict[str, Any]:
2830
return serialize(event)
2931

3032
def is_unrecoverable_client_error(self, error: ClientError) -> bool:
@@ -61,9 +63,9 @@ def is_unrecoverable_client_error(self, error: ClientError) -> bool:
6163

6264
def forward_event(
6365
self,
66+
event: Event | GroupEvent,
6467
payload: dict[str, Any],
6568
config: dict[str, Any],
66-
event: Event,
6769
) -> bool:
6870
queue_url = config["queue_url"]
6971
region = config["region"]

‎src/sentry/integrations/data_forwarding/base.py‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from sentry import ratelimits
66
from sentry.integrations.models.data_forwarder_project import DataForwarderProject
77
from sentry.integrations.types import DataForwarderProviderSlug
8-
from sentry.services.eventstore.models import Event
8+
from sentry.services.eventstore.models import Event, GroupEvent
99

1010
logger = logging.getLogger(__name__)
1111

@@ -22,10 +22,10 @@ class BaseDataForwarder(ABC):
2222
Tuple of (Number of Requests, Window in Seconds)
2323
"""
2424

25-
def get_rl_key(self, event) -> str:
25+
def get_rl_key(self, event: Event | GroupEvent) -> str:
2626
return f"{self.provider.value}:{event.project.organization_id}"
2727

28-
def is_ratelimited(self, event: Event) -> bool:
28+
def is_ratelimited(self, event: Event | GroupEvent) -> bool:
2929
rl_key = self.get_rl_key(event)
3030
limit, window = self.rate_limit
3131
if limit and window and ratelimits.backend.is_limited(rl_key, limit=limit, window=window):
@@ -41,19 +41,25 @@ def is_ratelimited(self, event: Event) -> bool:
4141
return True
4242
return False
4343

44-
def initialize_variables(self, event: Event, config: dict[str, Any]) -> None:
44+
def initialize_variables(self, event: Event | GroupEvent, config: dict[str, Any]) -> None:
4545
"""This is only necessary for migrating Splunk plugin, needed for rate limiting"""
4646
return
4747

4848
@abstractmethod
49-
def forward_event(self, event: Event, payload: dict[str, Any], config: dict[str, Any]) -> bool:
49+
def forward_event(
50+
self, event: Event | GroupEvent, payload: dict[str, Any], config: dict[str, Any]
51+
) -> bool:
5052
raise NotImplementedError
5153

5254
@abstractmethod
53-
def get_event_payload(self, event: Event, config: dict[str, Any]) -> dict[str, Any]:
55+
def get_event_payload(
56+
self, event: Event | GroupEvent, config: dict[str, Any]
57+
) -> dict[str, Any]:
5458
raise NotImplementedError
5559

56-
def post_process(self, event: Event, data_forwarder_project: DataForwarderProject) -> None:
60+
def post_process(
61+
self, event: Event | GroupEvent, data_forwarder_project: DataForwarderProject
62+
) -> None:
5763
config = data_forwarder_project.get_config()
5864
if self.is_ratelimited(event):
5965
return

‎src/sentry/integrations/data_forwarding/segment/forwarder.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from sentry import VERSION, http
77
from sentry.integrations.data_forwarding.base import BaseDataForwarder
88
from sentry.integrations.types import DataForwarderProviderSlug
9-
from sentry.services.eventstore.models import Event
9+
from sentry.services.eventstore.models import Event, GroupEvent
1010

1111
logger = logging.getLogger(__name__)
1212

@@ -16,7 +16,9 @@ class SegmentForwarder(BaseDataForwarder):
1616
rate_limit = (200, 1)
1717
endpoint: ClassVar[str] = "https://api.segment.io/v1/track"
1818

19-
def get_event_payload(self, event: Event, config: dict[str, Any]) -> dict[str, Any]:
19+
def get_event_payload(
20+
self, event: Event | GroupEvent, config: dict[str, Any]
21+
) -> dict[str, Any]:
2022
context = {"library": {"name": "sentry", "version": VERSION}}
2123

2224
props = {
@@ -68,9 +70,9 @@ def get_event_payload(self, event: Event, config: dict[str, Any]) -> dict[str, A
6870

6971
def forward_event(
7072
self,
73+
event: Event | GroupEvent,
7174
payload: dict[str, Any],
7275
config: dict[str, Any],
73-
event: Event,
7476
) -> bool:
7577
# we currently only support errors
7678
if event.get_event_type() != "error":

‎src/sentry/integrations/data_forwarding/splunk/forwarder.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from sentry import tagstore
77
from sentry.integrations.data_forwarding.base import BaseDataForwarder
88
from sentry.integrations.types import DataForwarderProviderSlug
9-
from sentry.services.eventstore.models import Event
9+
from sentry.services.eventstore.models import Event, GroupEvent
1010
from sentry.shared_integrations.exceptions import ApiError, ApiHostError, ApiTimeoutError
1111
from sentry.utils.hashlib import md5_text
1212
from sentry_plugins.anonymizeip import anonymize_ip
@@ -24,7 +24,7 @@ class SplunkForwarder(BaseDataForwarder):
2424
host: str | None = None
2525
project_source: str | None = None
2626

27-
def get_host_for_splunk(self, event: Event) -> str | None:
27+
def get_host_for_splunk(self, event: Event | GroupEvent) -> str | None:
2828
host = event.get_tag("server_name")
2929
if host:
3030
return host
@@ -37,21 +37,21 @@ def get_host_for_splunk(self, event: Event) -> str | None:
3737

3838
return None
3939

40-
def initialize_variables(self, event):
41-
self.project_token = self.get_option("token", event.project)
42-
self.project_index = self.get_option("index", event.project)
43-
self.project_instance = self.get_option("instance", event.project)
40+
def initialize_variables(self, event: Event | GroupEvent, config: dict[str, Any]):
41+
self.project_token = config.get("token")
42+
self.project_index = config.get("index")
43+
self.project_instance = config.get("instance")
4444
self.host = self.get_host_for_splunk(event)
4545

4646
if self.project_instance and not self.project_instance.endswith("/services/collector"):
4747
self.project_instance = self.project_instance.rstrip("/") + "/services/collector"
4848

49-
self.project_source = self.get_option("source", event.project) or "sentry"
49+
self.project_source = config.get("source", "sentry")
5050

51-
def get_rl_key(self, event) -> str:
51+
def get_rl_key(self, event: Event | GroupEvent) -> str:
5252
return f"{self.provider.value}:{md5_text(self.project_token).hexdigest()}"
5353

54-
def get_event_payload_properties(self, event: Event) -> dict[str, Any]:
54+
def get_event_payload_properties(self, event: Event | GroupEvent) -> dict[str, Any]:
5555
props = {
5656
"event_id": event.event_id,
5757
"issue_id": event.group_id,
@@ -101,7 +101,9 @@ def get_event_payload_properties(self, event: Event) -> dict[str, Any]:
101101
props.update(user_payload)
102102
return props
103103

104-
def get_event_payload(self, event: Event, config: dict[str, Any]) -> dict[str, Any]:
104+
def get_event_payload(
105+
self, event: Event | GroupEvent, config: dict[str, Any]
106+
) -> dict[str, Any]:
105107
return {
106108
"time": int(event.datetime.strftime("%s")),
107109
"source": config.get("source", "sentry"),
@@ -111,9 +113,9 @@ def get_event_payload(self, event: Event, config: dict[str, Any]) -> dict[str, A
111113

112114
def forward_event(
113115
self,
116+
event: Event | GroupEvent,
114117
payload: dict[str, Any],
115118
config: dict[str, Any],
116-
event: Event,
117119
) -> bool:
118120
if not self.project_token or not self.project_index or not self.project_instance:
119121
return False

‎src/sentry/tasks/post_process.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
from sentry import features, options, projectoptions
1818
from sentry.exceptions import PluginError
19-
from sentry.integrations.data_forwarding.base import BaseDataForwarder
2019
from sentry.integrations.types import IntegrationProviderSlug
2120
from sentry.issues.grouptype import GroupCategory
2221
from sentry.issues.issue_occurrence import IssueOccurrence
@@ -1288,6 +1287,7 @@ def process_data_forwarding(job: PostProcessJob) -> None:
12881287
return
12891288

12901289
from sentry.integrations.data_forwarding import FORWARDER_REGISTRY
1290+
from sentry.integrations.data_forwarding.base import BaseDataForwarder
12911291
from sentry.integrations.models.data_forwarder_project import DataForwarderProject
12921292

12931293
data_forwarder_projects = DataForwarderProject.objects.filter(
@@ -1300,8 +1300,8 @@ def process_data_forwarding(job: PostProcessJob) -> None:
13001300
provider = data_forwarder_project.data_forwarder.provider
13011301
try:
13021302
# GroupEvent is compatible with Event for all operations forwarders need
1303-
forwarder: type[BaseDataForwarder] = FORWARDER_REGISTRY[provider]()
1304-
forwarder.post_process(event, data_forwarder_project)
1303+
forwarder: type[BaseDataForwarder] = FORWARDER_REGISTRY[provider]
1304+
forwarder().post_process(event, data_forwarder_project)
13051305
metrics.incr(
13061306
"data_forwarding.post_process",
13071307
tags={"provider": provider},

‎tests/sentry/integrations/data_forwarding/amazon_sqs/test_forwarder.py‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def setUp(self):
2929
project=self.project,
3030
is_enabled=True,
3131
)
32+
self.forwarder = AmazonSQSForwarder()
3233

3334
@patch("boto3.client")
3435
def test_simple_notification(self, mock_client):
@@ -42,9 +43,8 @@ def test_simple_notification(self, mock_client):
4243
project_id=self.project.id,
4344
)
4445

45-
result = AmazonSQSForwarder.forward_event(event, self.data_forwarder_project)
46+
self.forwarder.post_process(event, self.data_forwarder_project)
4647

47-
assert result is True
4848
mock_client.assert_called_once_with(
4949
service_name="sqs",
5050
region_name="us-east-1",
@@ -68,8 +68,7 @@ def test_token_error(self, mock_client):
6868
project_id=self.project.id,
6969
)
7070

71-
result = AmazonSQSForwarder.forward_event(event, self.data_forwarder_project)
72-
assert result is False
71+
self.forwarder.post_process(event, self.data_forwarder_project)
7372

7473
@patch("boto3.client")
7574
def test_message_group_error(self, mock_client):
@@ -88,8 +87,7 @@ def test_message_group_error(self, mock_client):
8887
project_id=self.project.id,
8988
)
9089

91-
result = AmazonSQSForwarder.forward_event(event, self.data_forwarder_project)
92-
assert result is False
90+
self.forwarder.post_process(event, self.data_forwarder_project)
9391

9492
@patch("boto3.client")
9593
def test_pass_message_group_id(self, mock_client):
@@ -104,9 +102,8 @@ def test_pass_message_group_id(self, mock_client):
104102
project_id=self.project.id,
105103
)
106104

107-
result = AmazonSQSForwarder.forward_event(event, self.data_forwarder_project)
105+
self.forwarder.post_process(event, self.data_forwarder_project)
108106

109-
assert result is True
110107
call_args = mock_client.return_value.send_message.call_args[1]
111108
assert call_args["MessageGroupId"] == "my_group"
112109
assert "MessageDeduplicationId" in call_args
@@ -124,9 +121,8 @@ def test_use_s3_bucket(self, mock_client):
124121
project_id=self.project.id,
125122
)
126123

127-
result = AmazonSQSForwarder.forward_event(event, self.data_forwarder_project)
124+
self.forwarder.post_process(event, self.data_forwarder_project)
128125

129-
assert result is True
130126
mock_client.return_value.put_object.assert_called_once()
131127
put_object_call = mock_client.return_value.put_object.call_args[1]
132128
assert put_object_call["Bucket"] == "my_bucket"

‎tests/sentry/integrations/data_forwarding/segment/test_forwarder.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ def setUp(self):
2323
project=self.project,
2424
is_enabled=True,
2525
)
26+
self.forwarder = SegmentForwarder()
2627

2728
@responses.activate
2829
def test_simple_notification(self):
@@ -39,9 +40,8 @@ def test_simple_notification(self):
3940
project_id=self.project.id,
4041
)
4142

42-
result = SegmentForwarder.forward_event(event, self.data_forwarder_project)
43+
self.forwarder.post_process(event, self.data_forwarder_project)
4344

44-
assert result is True
4545
assert len(responses.calls) == 1
4646

4747
request = responses.calls[0].request
@@ -75,5 +75,4 @@ def test_forward_event_http_error(self):
7575
project_id=self.project.id,
7676
)
7777

78-
result = SegmentForwarder.forward_event(event, self.data_forwarder_project)
79-
assert result is False
78+
self.forwarder.post_process(event, self.data_forwarder_project)

‎tests/sentry/integrations/data_forwarding/splunk/test_forwarder.py‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def setUp(self):
2929
project=self.project,
3030
is_enabled=True,
3131
)
32+
self.forwarder = SplunkForwarder()
3233

3334
@responses.activate
3435
def test_simple_notification(self):
@@ -38,9 +39,7 @@ def test_simple_notification(self):
3839
data={"message": "Hello world", "level": "warning"}, project_id=self.project.id
3940
)
4041

41-
result = SplunkForwarder.forward_event(event, self.data_forwarder_project)
42-
43-
assert result is True
42+
self.forwarder.post_process(event, self.data_forwarder_project)
4443
assert len(responses.calls) == 1
4544

4645
request = responses.calls[0].request
@@ -62,8 +61,7 @@ def test_dont_reraise_error(self):
6261
data={"message": "Hello world", "level": "warning"}, project_id=self.project.id
6362
)
6463

65-
result = SplunkForwarder.forward_event(event, self.data_forwarder_project)
66-
assert result is False
64+
self.forwarder.post_process(event, self.data_forwarder_project)
6765

6866
@responses.activate
6967
def test_reraise_error(self):
@@ -76,7 +74,7 @@ def test_reraise_error(self):
7674
)
7775

7876
with pytest.raises(ApiError):
79-
SplunkForwarder.forward_event(event, self.data_forwarder_project)
77+
self.forwarder.post_process(event, self.data_forwarder_project)
8078

8179
def test_http_payload(self):
8280
event = self.store_event(
@@ -91,7 +89,7 @@ def test_http_payload(self):
9189
)
9290

9391
config = self.data_forwarder_project.get_config()
94-
result = SplunkForwarder.get_event_payload(event, config)
92+
result = self.forwarder.get_event_payload(event, config)
9593
assert result["event"]["request_url"] == "http://example.com/"
9694
assert result["event"]["request_method"] == "POST"
9795
assert result["event"]["request_referer"] == "http://example.com/foo"
@@ -106,7 +104,7 @@ def test_error_payload(self):
106104
)
107105

108106
config = self.data_forwarder_project.get_config()
109-
result = SplunkForwarder.get_event_payload(event, config)
107+
result = self.forwarder.get_event_payload(event, config)
110108
assert result["event"]["type"] == "error"
111109
assert result["event"]["exception_type"] == "ValueError"
112110
assert result["event"]["exception_value"] == "foo bar"
@@ -126,7 +124,7 @@ def test_csp_payload(self):
126124
)
127125

128126
config = self.data_forwarder_project.get_config()
129-
result = SplunkForwarder.get_event_payload(event, config)
127+
result = self.forwarder.get_event_payload(event, config)
130128
assert result["event"]["type"] == "csp"
131129
assert result["event"]["csp_document_uri"] == "http://example.com/"
132130
assert result["event"]["csp_violated_directive"] == "style-src cdn.example.com"
@@ -140,7 +138,7 @@ def test_user_payload(self):
140138
)
141139

142140
config = self.data_forwarder_project.get_config()
143-
result = SplunkForwarder.get_event_payload(event, config)
141+
result = self.forwarder.get_event_payload(event, config)
144142
assert result["event"]["user_id"] == "1"
145143
assert result["event"]["user_email_hash"] == "b48def645758b95537d4424c84d1a9ff"
146144
assert result["event"]["user_ip_trunc"] == "127.0.0.0"

0 commit comments

Comments
 (0)