def test_project_extension_query_adds_rate_limits():
extension = ProjectExtension(processor=ProjectExtensionProcessor(
project_column="project_id"))
raw_data = {'project': [2, 3]}
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query(
{"conditions": []},
TableSource("my_table", ColumnSet([])),
)
request_settings = RequestSettings(turbo=False,
consistent=False,
debug=False)
num_rate_limits_before_processing = len(
request_settings.get_rate_limit_params())
extension.get_processor().process_query(query, valid_data,
request_settings)
rate_limits = request_settings.get_rate_limit_params()
# make sure a rate limit was added by the processing
assert len(rate_limits) == num_rate_limits_before_processing + 1
most_recent_rate_limit = rate_limits[-1]
assert most_recent_rate_limit.bucket == '2'
assert most_recent_rate_limit.per_second_limit == 1000
assert most_recent_rate_limit.concurrent_limit == 1000
def get_extensions(self) -> Mapping[str, QueryExtension]:
return {
"project": ProjectExtension(
processor=ProjectExtensionProcessor(project_column="project_id")
),
"timeseries": TimeSeriesExtension(
default_granularity=3600,
default_window=timedelta(days=5),
timestamp_column="start_ts",
),
}
def get_extensions(self) -> Mapping[str, QueryExtension]:
return {
'project':
ProjectExtension(processor=ProjectExtensionProcessor()),
'timeseries':
TimeSeriesExtension(
default_granularity=3600,
default_window=timedelta(days=5),
timestamp_column='start_ts',
),
}
def test_project_extension_query_processing(raw_data: dict, expected_conditions: Sequence[Condition]):
extension = ProjectExtension(
processor=ProjectExtensionProcessor()
)
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({
"conditions": []
})
request_settings = RequestSettings(turbo=False, consistent=False, debug=False)
extension.get_processor().process_query(query, valid_data, request_settings)
assert query.get_conditions() == expected_conditions
def test_project_extension_query_processing(
raw_data: dict,
expected_conditions: Sequence[Condition],
expected_ast_conditions: Expression,
):
extension = ProjectExtension(
processor=ProjectExtensionProcessor(project_column="project_id")
)
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({"conditions": []}, TableSource("my_table", ColumnSet([])),)
request_settings = HTTPRequestSettings()
extension.get_processor().process_query(query, valid_data, request_settings)
assert query.get_conditions() == expected_conditions
assert query.get_condition_from_ast() == expected_ast_conditions
def test_project_extension_project_rate_limits_are_overridden():
extension = ProjectExtension(
processor=ProjectExtensionProcessor(project_column="project_id")
)
raw_data = {"project": [2, 3]}
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({"conditions": []}, TableSource("my_table", ColumnSet([])),)
request_settings = HTTPRequestSettings()
state.set_config("project_per_second_limit_2", 5)
state.set_config("project_concurrent_limit_2", 10)
extension.get_processor().process_query(query, valid_data, request_settings)
rate_limits = request_settings.get_rate_limit_params()
most_recent_rate_limit = rate_limits[-1]
assert most_recent_rate_limit.bucket == "2"
assert most_recent_rate_limit.per_second_limit == 5
assert most_recent_rate_limit.concurrent_limit == 10
def test_project_extension_project_rate_limits_are_overridden():
extension = ProjectExtension(
processor=ProjectExtensionProcessor()
)
raw_data = {
'project': [2, 3]
}
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({
'conditions': []
})
request_settings = RequestSettings(turbo=False, consistent=False, debug=False)
state.set_config('project_per_second_limit_2', 5)
state.set_config('project_concurrent_limit_2', 10)
extension.get_processor().process_query(query, valid_data, request_settings)
rate_limits = request_settings.get_rate_limit_params()
most_recent_rate_limit = rate_limits[-1]
assert most_recent_rate_limit.bucket == '2'
assert most_recent_rate_limit.per_second_limit == 5
assert most_recent_rate_limit.concurrent_limit == 10
