diff --git a/snuba/split.py b/snuba/split.py index 20510430607..2b5196c2c7a 100644 --- a/snuba/split.py +++ b/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