def test_load_explore_json_into_cache(self, mock_update_job):
async_query_manager.init_app(app)
table = self.get_table(name="birth_names")
user = security_manager.find_user("gamma")
form_data = {
"datasource": f"{table.id}__table",
"viz_type": "dist_bar",
"time_range_endpoints": ["inclusive", "exclusive"],
"granularity_sqla": "ds",
"time_range": "No filter",
"metrics": ["count"],
"adhoc_filters": [],
"groupby": ["gender"],
"row_limit": 100,
}
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": user.id,
"status": "pending",
"errors": [],
}
with mock.patch.object(
async_queries, "ensure_user_is_set"
) as ensure_user_is_set:
load_explore_json_into_cache(job_metadata, form_data)
ensure_user_is_set.assert_called_once_with(user.id)
mock_update_job.assert_called_once_with(
job_metadata, "done", result_url=mock.ANY
)
def test_load_explore_json_into_cache(self, mock_update_job):
async_query_manager.init_app(app)
table = get_table_by_name("birth_names")
form_data = {
"datasource": f"{table.id}__table",
"viz_type": "dist_bar",
"time_range_endpoints": ["inclusive", "exclusive"],
"granularity_sqla": "ds",
"time_range": "No filter",
"metrics": ["count"],
"adhoc_filters": [],
"groupby": ["gender"],
"row_limit": 100,
}
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": 1,
"status": "pending",
"errors": [],
}
load_explore_json_into_cache(job_metadata, form_data)
mock_update_job.assert_called_with(job_metadata,
"done",
result_url=mock.ANY)
def test_chart_data_async_cached_sync_response(self):
"""
Chart data API: Test chart data query returns results synchronously
when results are already cached.
"""
async_query_manager.init_app(app)
class QueryContext:
result_format = ChartDataResultFormat.JSON
result_type = ChartDataResultType.FULL
cmd_run_val = {
"query_context": QueryContext(),
"queries": [{
"query": "select * from foo"
}],
}
with mock.patch.object(ChartDataCommand,
"run",
return_value=cmd_run_val) as patched_run:
self.query_context_payload[
"result_type"] = ChartDataResultType.FULL
rv = self.post_assert_metric(CHART_DATA_URI,
self.query_context_payload, "data")
self.assertEqual(rv.status_code, 200)
data = json.loads(rv.data.decode("utf-8"))
patched_run.assert_called_once_with(force_cached=True)
self.assertEqual(data,
{"result": [{
"query": "select * from foo"
}]})
def test_explore_json_async(self):
tbl_id = self.table_ids.get("birth_names")
form_data = {
"datasource": f"{tbl_id}__table",
"viz_type": "dist_bar",
"time_range_endpoints": ["inclusive", "exclusive"],
"granularity_sqla": "ds",
"time_range": "No filter",
"metrics": ["count"],
"adhoc_filters": [],
"groupby": ["gender"],
"row_limit": 100,
}
async_query_manager.init_app(app)
self.login(username="******")
rv = self.client.post(
"/superset/explore_json/", data={"form_data": json.dumps(form_data)},
)
data = json.loads(rv.data.decode("utf-8"))
keys = list(data.keys())
self.assertEqual(rv.status_code, 202)
self.assertCountEqual(
keys, ["channel_id", "job_id", "user_id", "status", "errors", "result_url"]
)
def test_chart_data_async_results_type(self):
"""
Chart data API: Test chart data query non-JSON format (async)
"""
async_query_manager.init_app(app)
self.query_context_payload["result_type"] = "results"
rv = self.post_assert_metric(CHART_DATA_URI, self.query_context_payload, "data")
self.assertEqual(rv.status_code, 200)
def test_chart_data_cache_key_error(self):
"""
Chart data cache API: Test chart data async cache request with invalid cache key
"""
async_query_manager.init_app(app)
rv = self.get_assert_metric(f"{CHART_DATA_URI}/test-cache-key",
"data_from_cache")
self.assertEqual(rv.status_code, 404)
def test_chart_data_async_invalid_token(self):
"""
Chart data API: Test chart data query (async)
"""
async_query_manager.init_app(app)
test_client.set_cookie(
"localhost", app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME"],
"foo")
rv = test_client.post(CHART_DATA_URI, json=self.query_context_payload)
self.assertEqual(rv.status_code, 401)
def test_chart_data_async(self):
self.logout()
async_query_manager.init_app(app)
self.login("admin")
rv = self.post_assert_metric(CHART_DATA_URI, self.query_context_payload, "data")
self.assertEqual(rv.status_code, 202)
data = json.loads(rv.data.decode("utf-8"))
keys = list(data.keys())
self.assertCountEqual(
keys, ["channel_id", "job_id", "user_id", "status", "errors", "result_url"]
)
def test_events_last_id(self, mock_uuid4):
async_query_manager.init_app(app)
self.login(username="******")
with mock.patch.object(async_query_manager._redis, "xrange") as mock_xrange:
rv = self.fetch_events("1607471525180-0")
response = json.loads(rv.data.decode("utf-8"))
assert rv.status_code == 200
channel_id = app.config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] + self.UUID
mock_xrange.assert_called_with(channel_id, "1607471525180-1", "+", 100)
self.assertEqual(response, {"result": []})
def test_chart_data_cache_run_failed(self, cache_loader):
"""
Chart data cache API: Test chart data async cache request with run failure
"""
async_query_manager.init_app(app)
cache_loader.load.return_value = self.query_context_payload
rv = self.get_assert_metric(f"{CHART_DATA_URI}/test-cache-key",
"data_from_cache")
data = json.loads(rv.data.decode("utf-8"))
self.assertEqual(rv.status_code, 422)
self.assertEqual(data["message"], "Error loading data from cache")
def test_events_results(self, mock_uuid4):
async_query_manager.init_app(app)
self.login(username="******")
with mock.patch.object(async_query_manager._redis,
"xrange") as mock_xrange:
mock_xrange.return_value = [
(
"1607477697866-0",
{
"data":
'{"channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f", "job_id": "10a0bd9a-03c8-4737-9345-f4234ba86512", "user_id": "1", "status": "done", "errors": [], "result_url": "/api/v1/chart/data/qc-ecd766dd461f294e1bcdaa321e0e8463"}'
},
),
(
"1607477697993-0",
{
"data":
'{"channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f", "job_id": "027cbe49-26ce-4813-bb5a-0b95a626b84c", "user_id": "1", "status": "done", "errors": [], "result_url": "/api/v1/chart/data/qc-1bbc3a240e7039ba4791aefb3a7ee80d"}'
},
),
]
rv = self.fetch_events()
response = json.loads(rv.data.decode("utf-8"))
assert rv.status_code == 200
channel_id = app.config[
"GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] + self.UUID
mock_xrange.assert_called_with(channel_id, "-", "+", 100)
expected = {
"result": [
{
"channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f",
"errors": [],
"id": "1607477697866-0",
"job_id": "10a0bd9a-03c8-4737-9345-f4234ba86512",
"result_url":
"/api/v1/chart/data/qc-ecd766dd461f294e1bcdaa321e0e8463",
"status": "done",
"user_id": "1",
},
{
"channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f",
"errors": [],
"id": "1607477697993-0",
"job_id": "027cbe49-26ce-4813-bb5a-0b95a626b84c",
"result_url":
"/api/v1/chart/data/qc-1bbc3a240e7039ba4791aefb3a7ee80d",
"status": "done",
"user_id": "1",
},
]
}
self.assertEqual(response, expected)
def test_load_chart_data_into_cache(self, mock_update_job):
async_query_manager.init_app(app)
query_context = get_query_context("birth_names")
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": 1,
"status": "pending",
"errors": [],
}
load_chart_data_into_cache(job_metadata, query_context)
mock_update_job.assert_called_with(job_metadata,
"done",
result_url=mock.ANY)
def test_chart_data_cache_no_login(self, cache_loader):
"""
Chart data cache API: Test chart data async cache request (no login)
"""
self.logout()
async_query_manager.init_app(app)
cache_loader.load.return_value = self.query_context_payload
orig_run = ChartDataCommand.run
def mock_run(self, **kwargs):
assert kwargs["force_cached"] == True
# override force_cached to get result from DB
return orig_run(self, force_cached=False)
with mock.patch.object(ChartDataCommand, "run", new=mock_run):
rv = self.client.get(f"{CHART_DATA_URI}/test-cache-key", )
self.assertEqual(rv.status_code, 401)
def test_load_explore_json_into_cache_error(self, mock_update_job):
async_query_manager.init_app(app)
form_data = {}
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": 1,
"status": "pending",
"errors": [],
}
with pytest.raises(SupersetException):
load_explore_json_into_cache(job_metadata, form_data)
errors = ["The dataset associated with this chart no longer exists"]
mock_update_job.assert_called_with(job_metadata,
"error",
errors=errors)
def test_load_chart_data_into_cache_error(self, mock_update_job,
mock_run_command):
async_query_manager.init_app(app)
query_context = get_query_context("birth_names")
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": 1,
"status": "pending",
"errors": [],
}
with pytest.raises(ChartDataQueryFailedError):
load_chart_data_into_cache(job_metadata, query_context)
mock_run_command.assert_called_with(cache=True)
errors = [{"message": "Error: foo"}]
mock_update_job.assert_called_with(job_metadata,
"error",
errors=errors)
def test_explore_json_async_results_format(self):
tbl_id = self.table_ids.get("birth_names")
form_data = {
"datasource": f"{tbl_id}__table",
"viz_type": "dist_bar",
"time_range_endpoints": ["inclusive", "exclusive"],
"granularity_sqla": "ds",
"time_range": "No filter",
"metrics": ["count"],
"adhoc_filters": [],
"groupby": ["gender"],
"row_limit": 100,
}
async_query_manager.init_app(app)
self.login(username="******")
rv = self.client.post(
"/superset/explore_json/?results=true",
data={"form_data": json.dumps(form_data)},
)
self.assertEqual(rv.status_code, 200)
def test_load_chart_data_into_cache(self, mock_update_job):
async_query_manager.init_app(app)
query_context = get_query_context("birth_names")
user = security_manager.find_user("gamma")
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": user.id,
"status": "pending",
"errors": [],
}
with mock.patch.object(async_queries,
"ensure_user_is_set") as ensure_user_is_set:
load_chart_data_into_cache(job_metadata, query_context)
ensure_user_is_set.assert_called_once_with(user.id)
mock_update_job.assert_called_once_with(job_metadata,
"done",
result_url=mock.ANY)
def test_chart_data_cache(self, cache_loader):
"""
Chart data cache API: Test chart data async cache request
"""
async_query_manager.init_app(app)
cache_loader.load.return_value = self.query_context_payload
orig_run = ChartDataCommand.run
def mock_run(self, **kwargs):
assert kwargs["force_cached"] == True
# override force_cached to get result from DB
return orig_run(self, force_cached=False)
with mock.patch.object(ChartDataCommand, "run", new=mock_run):
rv = self.get_assert_metric(f"{CHART_DATA_URI}/test-cache-key",
"data_from_cache")
data = json.loads(rv.data.decode("utf-8"))
expected_row_count = self.get_expected_row_count("client_id_3")
self.assertEqual(rv.status_code, 200)
self.assertEqual(data["result"][0]["rowcount"], expected_row_count)
def test_load_chart_data_into_cache_error(self, mock_update_job, mock_run_command):
async_query_manager.init_app(app)
query_context = get_query_context("birth_names")
user = security_manager.find_user("gamma")
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": user.id,
"status": "pending",
"errors": [],
}
with pytest.raises(ChartDataQueryFailedError):
with mock.patch.object(
async_queries, "ensure_user_is_set"
) as ensure_user_is_set:
load_chart_data_into_cache(job_metadata, query_context)
ensure_user_is_set.assert_called_once_with(user.id)
mock_run_command.assert_called_once_with(cache=True)
errors = [{"message": "Error: foo"}]
mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors)
def test_load_explore_json_into_cache_error(self, mock_update_job):
async_query_manager.init_app(app)
user = security_manager.find_user("gamma")
form_data = {}
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": user.id,
"status": "pending",
"errors": [],
}
with pytest.raises(SupersetException):
with mock.patch.object(
async_queries, "ensure_user_is_set"
) as ensure_user_is_set:
load_explore_json_into_cache(job_metadata, form_data)
ensure_user_is_set.assert_called_once_with(user.id)
errors = ["The dataset associated with this chart no longer exists"]
mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors)
def test_soft_timeout_load_chart_data_into_cache(
self, mock_update_job, mock_run_command
):
async_query_manager.init_app(app)
user = security_manager.find_user("gamma")
form_data = {}
job_metadata = {
"channel_id": str(uuid4()),
"job_id": str(uuid4()),
"user_id": user.id,
"status": "pending",
"errors": [],
}
errors = ["A timeout occurred while loading chart data"]
with pytest.raises(SoftTimeLimitExceeded):
with mock.patch.object(
async_queries, "ensure_user_is_set",
) as ensure_user_is_set:
ensure_user_is_set.side_effect = SoftTimeLimitExceeded()
load_chart_data_into_cache(job_metadata, form_data)
ensure_user_is_set.assert_called_once_with(user.id, "error", errors=errors)
def configure_async_queries(self) -> None:
if feature_flag_manager.is_feature_enabled("GLOBAL_ASYNC_QUERIES"):
async_query_manager.init_app(self.superset_app)
def test_events_no_login(self):
async_query_manager.init_app(app)
rv = self.fetch_events()
assert rv.status_code == 401