11from  contextlib  import  contextmanager 
2- from  typing  import  Sequence 
2+ from  datetime  import  datetime , timedelta 
3+ from  typing  import  Any , Callable , Dict , Generator , Optional , Sequence , Union , cast 
34
45import  sentry_sdk 
5- from  django .http  import  HttpRequest 
6+ from  django .utils  import  timezone 
67from  django .utils .http  import  urlquote 
7- from  rest_framework .exceptions  import  APIException , ParseError 
8+ from  rest_framework .exceptions  import  APIException , ParseError , ValidationError 
9+ from  rest_framework .request  import  Request 
810from  sentry_relay .consts  import  SPAN_STATUS_CODE_TO_NAME 
911
10- from  sentry  import  features 
12+ from  sentry  import  features ,  quotas 
1113from  sentry .api .base  import  LINK_HEADER 
1214from  sentry .api .bases  import  NoProjects , OrganizationEndpoint 
1315from  sentry .api .helpers .teams  import  get_teams 
14- from  sentry .api .serializers .snuba  import  SnubaTSResultSerializer 
16+ from  sentry .api .serializers .snuba  import  BaseSnubaSerializer ,  SnubaTSResultSerializer 
1517from  sentry .discover .arithmetic  import  ArithmeticError , is_equation , strip_equation 
1618from  sentry .exceptions  import  InvalidSearchQuery 
17- from  sentry .models  import  Organization , Team 
19+ from  sentry .models  import  Organization , Project ,  Team 
1820from  sentry .models .group  import  Group 
1921from  sentry .search .events .constants  import  TIMEOUT_ERROR_MESSAGE 
2022from  sentry .search .events .fields  import  get_function_alias 
2123from  sentry .search .events .filter  import  get_filter 
2224from  sentry .snuba  import  discover 
2325from  sentry .utils  import  snuba 
26+ from  sentry .utils .cursors  import  Cursor 
2427from  sentry .utils .dates  import  get_interval_from_range , get_rollup_from_request , parse_stats_period 
2528from  sentry .utils .http  import  absolute_uri 
26- from  sentry .utils .snuba  import  MAX_FIELDS 
29+ from  sentry .utils .snuba  import  MAX_FIELDS ,  SnubaTSResult 
2730
2831
29- def  resolve_axis_column (column : str , index = 0 ) ->  str :
30-     return  get_function_alias (column ) if  not  is_equation (column ) else  f"equation[{ index }  ]" 
32+ def  resolve_axis_column (column : str , index : int  =  0 ) ->  str :
33+     return  cast (
34+         str , get_function_alias (column ) if  not  is_equation (column ) else  f"equation[{ index }  ]" 
35+     )
3136
3237
33- class  OrganizationEventsEndpointBase (OrganizationEndpoint ):
34-     def  has_feature (self , organization , request ) :
38+ class  OrganizationEventsEndpointBase (OrganizationEndpoint ):   # type: ignore 
39+     def  has_feature (self , organization :  Organization , request :  Request )  ->   bool :
3540        return  features .has (
3641            "organizations:discover-basic" , organization , actor = request .user 
3742        ) or  features .has ("organizations:performance-view" , organization , actor = request .user )
3843
39-     def  get_equation_list (self , organization : Organization , request : HttpRequest ) ->  Sequence [str ]:
44+     def  get_equation_list (self , organization : Organization , request : Request ) ->  Sequence [str ]:
4045        """equations have a prefix so that they can be easily included alongside our existing fields""" 
4146        return  [
4247            strip_equation (field ) for  field  in  request .GET .getlist ("field" )[:] if  is_equation (field )
4348        ]
4449
45-     def  get_field_list (self , organization : Organization , request : HttpRequest ) ->  Sequence [str ]:
50+     def  get_field_list (self , organization : Organization , request : Request ) ->  Sequence [str ]:
4651        return  [field  for  field  in  request .GET .getlist ("field" )[:] if  not  is_equation (field )]
4752
48-     def  get_snuba_filter (self , request , organization , params = None ):
49-         if  params  is  None :
50-             params  =  self .get_snuba_params (request , organization )
51-         query  =  request .GET .get ("query" )
52-         try :
53-             return  get_filter (query , params )
54-         except  InvalidSearchQuery  as  e :
55-             raise  ParseError (detail = str (e ))
56- 
57-     def  get_team_ids (self , request , organization ):
53+     def  get_team_ids (self , request : Request , organization : Organization ) ->  Sequence [int ]:
5854        if  not  request .user :
5955            return  []
6056
@@ -64,7 +60,9 @@ def get_team_ids(self, request, organization):
6460
6561        return  [team .id  for  team  in  teams ]
6662
67-     def  get_snuba_params (self , request , organization , check_global_views = True ):
63+     def  get_snuba_params (
64+         self , request : Request , organization : Organization , check_global_views : bool  =  True 
65+     ) ->  Dict [str , Any ]:
6866        with  sentry_sdk .start_span (op = "discover.endpoint" , description = "filter_params" ):
6967            if  (
7068                len (self .get_field_list (organization , request ))
@@ -75,7 +73,7 @@ def get_snuba_params(self, request, organization, check_global_views=True):
7573                    detail = f"You can view up to { MAX_FIELDS }   fields at a time. Please delete some and try again." 
7674                )
7775
78-             params  =  self .get_filter_params (request , organization )
76+             params :  Dict [ str ,  Any ]  =  self .get_filter_params (request , organization )
7977            params  =  self .quantize_date_params (request , params )
8078            params ["user_id" ] =  request .user .id  if  request .user  else  None 
8179            params ["team_id" ] =  self .get_team_ids (request , organization )
@@ -89,17 +87,27 @@ def get_snuba_params(self, request, organization, check_global_views=True):
8987
9088            return  params 
9189
92-     def  get_orderby (self , request ) :
93-         sort  =  request .GET .getlist ("sort" )
90+     def  get_orderby (self , request :  Request )  ->   Optional [ Sequence [ str ]] :
91+         sort :  Sequence [ str ]  =  request .GET .getlist ("sort" )
9492        if  sort :
9593            return  sort 
9694        # Deprecated. `sort` should be used as it is supported by 
9795        # more endpoints. 
98-         orderby  =  request .GET .getlist ("orderby" )
96+         orderby :  Sequence [ str ]  =  request .GET .getlist ("orderby" )
9997        if  orderby :
10098            return  orderby 
101- 
102-     def  get_snuba_query_args_legacy (self , request , organization ):
99+         return  None 
100+ 
101+     def  get_snuba_query_args_legacy (
102+         self , request : Request , organization : Organization 
103+     ) ->  Dict [
104+         str ,
105+         Union [
106+             Optional [datetime ],
107+             Sequence [Sequence [Union [str , str , Any ]]],
108+             Optional [Dict [str , Sequence [int ]]],
109+         ],
110+     ]:
103111        params  =  self .get_filter_params (request , organization )
104112        query  =  request .GET .get ("query" )
105113        try :
@@ -116,7 +124,7 @@ def get_snuba_query_args_legacy(self, request, organization):
116124
117125        return  snuba_args 
118126
119-     def  quantize_date_params (self , request , params ) :
127+     def  quantize_date_params (self , request :  Request , params :  Dict [ str ,  Any ])  ->   Dict [ str ,  Any ] :
120128        # We only need to perform this rounding on relative date periods 
121129        if  "statsPeriod"  not  in   request .GET :
122130            return  params 
@@ -133,7 +141,7 @@ def quantize_date_params(self, request, params):
133141        return  results 
134142
135143    @contextmanager  
136-     def  handle_query_errors (self ):
144+     def  handle_query_errors (self )  ->   Generator [ None ,  None ,  None ] :
137145        try :
138146            yield 
139147        except  discover .InvalidSearchQuery  as  error :
@@ -184,7 +192,7 @@ def handle_query_errors(self):
184192
185193
186194class  OrganizationEventsV2EndpointBase (OrganizationEventsEndpointBase ):
187-     def  build_cursor_link (self , request , name , cursor ) :
195+     def  build_cursor_link (self , request :  Request , name :  str , cursor :  Optional [ Cursor ])  ->   str :
188196        # The base API function only uses the last query parameter, but this endpoint 
189197        # needs all the parameters, particularly for the "field" query param. 
190198        querystring  =  "&" .join (
@@ -200,21 +208,33 @@ def build_cursor_link(self, request, name, cursor):
200208        else :
201209            base_url  =  base_url  +  "?" 
202210
203-         return  LINK_HEADER .format (
211+         return  cast ( str ,  LINK_HEADER ) .format (
204212            uri = base_url ,
205213            cursor = str (cursor ),
206214            name = name ,
207215            has_results = "true"  if  bool (cursor ) else  "false" ,
208216        )
209217
210-     def  handle_results_with_meta (self , request , organization , project_ids , results ):
218+     def  handle_results_with_meta (
219+         self ,
220+         request : Request ,
221+         organization : Organization ,
222+         project_ids : Sequence [int ],
223+         results : Dict [str , Any ],
224+     ) ->  Dict [str , Any ]:
211225        with  sentry_sdk .start_span (op = "discover.endpoint" , description = "base.handle_results" ):
212226            data  =  self .handle_data (request , organization , project_ids , results .get ("data" ))
213227            if  not  data :
214228                return  {"data" : [], "meta" : {}}
215229            return  {"data" : data , "meta" : results .get ("meta" , {})}
216230
217-     def  handle_data (self , request , organization , project_ids , results ):
231+     def  handle_data (
232+         self ,
233+         request : Request ,
234+         organization : Organization ,
235+         project_ids : Sequence [int ],
236+         results : Optional [Sequence [Any ]],
237+     ) ->  Optional [Sequence [Any ]]:
218238        if  not  results :
219239            return  results 
220240
@@ -240,7 +260,9 @@ def handle_data(self, request, organization, project_ids, results):
240260
241261        return  results 
242262
243-     def  handle_issues (self , results , project_ids , organization ):
263+     def  handle_issues (
264+         self , results : Sequence [Any ], project_ids : Sequence [int ], organization : Organization 
265+     ) ->  None :
244266        issue_ids  =  {row .get ("issue.id" ) for  row  in  results }
245267        issues  =  Group .issues_mapping (issue_ids , project_ids , organization )
246268        for  result  in  results :
@@ -249,16 +271,19 @@ def handle_issues(self, results, project_ids, organization):
249271
250272    def  get_event_stats_data (
251273        self ,
252-         request ,
253-         organization ,
254-         get_event_stats ,
255-         top_events = 0 ,
256-         query_column = "count()" ,
257-         params = None ,
258-         query = None ,
259-         allow_partial_buckets = False ,
260-         zerofill_results = True ,
261-     ):
274+         request : Request ,
275+         organization : Organization ,
276+         get_event_stats : Callable [
277+             [Sequence [str ], str , Dict [str , str ], int , bool , Optional [timedelta ]], SnubaTSResult 
278+         ],
279+         top_events : int  =  0 ,
280+         query_column : str  =  "count()" ,
281+         params : Optional [Dict [str , Any ]] =  None ,
282+         query : Optional [str ] =  None ,
283+         allow_partial_buckets : bool  =  False ,
284+         zerofill_results : bool  =  True ,
285+         comparison_delta : Optional [timedelta ] =  None ,
286+     ) ->  Dict [str , Any ]:
262287        with  self .handle_query_errors ():
263288            with  sentry_sdk .start_span (
264289                op = "discover.endpoint" , description = "base.stats_query_creation" 
@@ -287,11 +312,15 @@ def get_event_stats_data(
287312                except  InvalidSearchQuery :
288313                    sentry_sdk .set_tag ("user.invalid_interval" , request .GET .get ("interval" ))
289314                    date_range  =  params ["end" ] -  params ["start" ]
290-                     rollup  =  int (
291-                         parse_stats_period (
292-                             get_interval_from_range (date_range , False )
293-                         ).total_seconds ()
294-                     )
315+                     stats_period  =  parse_stats_period (get_interval_from_range (date_range , False ))
316+                     rollup  =  int (stats_period .total_seconds ()) if  stats_period  is  not   None  else  3600 
317+ 
318+                 if  comparison_delta  is  not   None :
319+                     retention  =  quotas .get_event_retention (organization = organization )
320+                     comparison_start  =  params ["start" ] -  comparison_delta 
321+                     if  retention  and  comparison_start  <  timezone .now () -  timedelta (days = retention ):
322+                         raise  ValidationError ("Comparison period is outside your retention window" )
323+ 
295324                # Backwards compatibility for incidents which uses the old 
296325                # column aliases as it straddles both versions of events/discover. 
297326                # We will need these aliases until discover2 flags are enabled for all 
@@ -308,7 +337,9 @@ def get_event_stats_data(
308337
309338                query_columns  =  [column_map .get (column , column ) for  column  in  columns ]
310339            with  sentry_sdk .start_span (op = "discover.endpoint" , description = "base.stats_query" ):
311-                 result  =  get_event_stats (query_columns , query , params , rollup , zerofill_results )
340+                 result  =  get_event_stats (
341+                     query_columns , query , params , rollup , zerofill_results , comparison_delta 
342+                 )
312343
313344        serializer  =  SnubaTSResultSerializer (organization , None , request .user )
314345
@@ -359,13 +390,13 @@ def get_event_stats_data(
359390
360391    def  serialize_multiple_axis (
361392        self ,
362-         serializer ,
363-         event_result ,
364-         columns ,
365-         query_columns ,
366-         allow_partial_buckets ,
367-         zerofill_results = True ,
368-     ):
393+         serializer :  BaseSnubaSerializer ,
394+         event_result :  SnubaTSResult ,
395+         columns :  Sequence [ str ] ,
396+         query_columns :  Sequence [ str ] ,
397+         allow_partial_buckets :  bool ,
398+         zerofill_results :  bool   =   True ,
399+     )  ->   Dict [ str ,  Any ] :
369400        # Return with requested yAxis as the key 
370401        result  =  {}
371402        equations  =  0 
@@ -387,10 +418,10 @@ def serialize_multiple_axis(
387418
388419
389420class  KeyTransactionBase (OrganizationEventsV2EndpointBase ):
390-     def  has_feature (self , request ,  organization ) :
421+     def  has_feature (self , organization :  Organization ,  request :  Request )  ->   bool :
391422        return  features .has ("organizations:performance-view" , organization , actor = request .user )
392423
393-     def  get_project (self , request , organization ) :
424+     def  get_project (self , request :  Request , organization :  Organization )  ->   Project :
394425        projects  =  self .get_projects (request , organization )
395426
396427        if  len (projects ) !=  1 :
0 commit comments