ref: Use explicit argument references in query splitter

snuba/split.py

@@ -12,8 +12,7 @@
 
 
 def split_query(query_func):
-    def wrapper(*args, **kwargs):
-        body = args[0]
+    def wrapper(body, *args, **kwargs):
         use_split = state.get_configs([
             ('use_split', 0),
         ])
@@ -34,13 +33,13 @@ def wrapper(*args, **kwargs):
                 and not body.get('aggregations')
                 and total_col_count > min_col_count
             ):
-                return col_split(*args, **kwargs)
+                return col_split(body, *args, **kwargs)
             elif orderby[:1] == ['-timestamp'] and remaining_offset < 1000:
-                return time_split(*args, **kwargs)
+                return time_split(body, *args, **kwargs)
 
-        return query_func(*args, **kwargs)
+        return query_func(body, *args, **kwargs)
 
-    def time_split(*args, **kwargs):
+    def time_split(body, *args, **kwargs):
         """
         If a query is:
             - ORDER BY timestamp DESC
@@ -52,8 +51,6 @@ def time_split(*args, **kwargs):
         into smaller increments, and start with the last one, so that we can potentially
         avoid querying the entire range.
         """
-        body = args[0]
-
         date_align, split_step = state.get_configs([
             ('date_align_seconds', 1),
             ('split_step', 3600),  # default 1 hour
@@ -77,7 +74,7 @@ def time_split(*args, **kwargs):
             # and set offset=0 so we can then trim them ourselves.
             body['offset'] = 0
             body['limit'] = limit - total_results + remaining_offset
-            result, status = query_func(*args, **kwargs)
+            result, status = query_func(body, *args, **kwargs)
 
             # If something failed, discard all progress and just return that
             if status != 200:
@@ -116,18 +113,16 @@ def time_split(*args, **kwargs):
 
         return overall_result, status
 
-    def col_split(*args, **kwargs):
+    def col_split(body, *args, **kwargs):
         """
         Split query in 2 steps if a large number of columns is being selected.
             - First query only selects event_id and project_id.
             - Second query selects all fields for only those events.
             - Shrink the date range.
         """
-        body = args[0]
-
         minimal_query = {**body, 'selected_columns': MIN_COLS}
 
-        result, status = query_func(minimal_query, *args[1:], **kwargs)
+        result, status = query_func(minimal_query, *args, **kwargs)
 
         # If something failed, just return
         if status != 200:
@@ -152,6 +147,6 @@ def col_split(*args, **kwargs):
             body['offset'] = 0
             body['limit'] = len(event_ids)
 
-        return query_func({**body, 'conditions': conditions}, *args[1:], **kwargs)
+        return query_func({**body, 'conditions': conditions}, *args, **kwargs)
 
     return wrapper