rusackas commented on code in PR #44406:
URL: https://github.com/apache/superset/pull/44406#discussion_r4332207793
##########
tests/unit_tests/tasks/test_query_cancel.py:
##########
@@ -159,6 +159,26 @@ def
test_capture_registers_abort_handler_once_when_id_available() -> None:
assert ordered.index("set_cancellation") < ordered.index("on_abort")
+def test_capture_ignores_cursor_from_a_different_database() -> None:
Review Comment:
Added a test that goes through `notify_cursor` itself: the task's own
database registers the abort handler once across two cursors, and a different
database registers nothing.
##########
superset/common/query_context_processor.py:
##########
@@ -285,6 +279,18 @@ def get_df_payload_result(
# decision as the result cache, not the datasource-level default.
query_obj.force_query = force_query
query_obj.cache_timeout = timeout
+ # Resolved separately from the dataframe's force_query, against its own
+ # annotation_key. A forced GTF refresh's nonce marker is per (nonce,
+ # cache_key): the dataframe's marker is keyed on cache_key, so once
it's
+ # set, force_query goes False for every subsequent request carrying
that
+ # nonce, including a request for a *different* annotation access scope
+ # (its own annotation_key) that has never actually forced-refreshed its
+ # own entry. Reusing the dataframe's force_query there would read that
+ # scope's existing annotation cache entry instead of refreshing it.
+ annotation_force_query = (
+ self._resolve_forced_query(query_obj, annotation_key)
+ or timeout == CACHE_DISABLED_TIMEOUT
Review Comment:
Added a `get_df_payload_result` case with `get_cache_timeout=-1`,
`force=False` and a warm annotation entry: the cache is never read and the live
fetch's result is what lands in the payload.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -3167,11 +3915,211 @@ def payload(context: QueryContext, **kwargs: Any) ->
dict[str, Any]:
datasource="native-1", annotation_context=context
)
assert current_key != legacy_key
- assert context == {
- **legacy_context,
- "semantic_filter_protocol": "semantic-null-filters-v1",
- "source_versions": {"7": view.metadata_generation},
- }
+ assert context is not None
+ assert context["source_scope"][0]["semantic_filter_protocol"] == (
+ "semantic-null-filters-v1"
+ )
+ assert context["source_scope"][0]["metadata_version"] ==
view.metadata_generation
+
+
+def _annotated_payload_fixture(
+ *, status: str = "success"
+) -> tuple[QueryContextProcessor, Any, MagicMock]:
+ """A processor, annotated query object, and a dataframe cache hit."""
+ from superset.common.query_object import QueryObject
+
+ mock_query_context = MagicMock()
+ mock_query_context.force = False
+ mock_datasource = MagicMock()
+ mock_datasource.column_names = ["col1"]
+ processor = QueryContextProcessor(mock_query_context)
+ processor._qc_datasource = mock_datasource
+ query_obj = QueryObject(
+ datasource=mock_datasource,
+ columns=["col1"],
+ annotation_layers=[
+ {
+ "annotationType": "EVENT",
+ "sourceType": "NATIVE",
+ "name": "a",
+ "value": 1,
+ }
+ ],
+ )
+ cache = MagicMock()
+ cache.is_loaded = True
+ cache.applied_filter_columns = ["col1"]
+ cache.df = pd.DataFrame({"col1": [1, 2, 3]})
+ cache.query = ""
+ cache.status = status
+ cache.error_message = "dataframe failed" if status == QueryStatus.FAILED
else None
+ cache.is_cached = True
+ cache.applied_template_filters = []
+ cache.rejected_filter_columns = []
+ cache.bq_memory_limited = False
+ cache.result_persisted = False
+ return processor, query_obj, cache
+
+
+def _run_annotated_payload(
+ processor: QueryContextProcessor,
+ query_obj: Any,
+ cache: MagicMock,
+ *,
+ force_cached: bool = False,
+) -> Any:
+ with (
+ patch(
+ "superset.common.query_context_processor.QueryCacheManager"
+ ) as mock_cache_manager,
+ patch.object(query_obj, "validate", return_value=None),
+ patch.object(processor, "query_cache_key", return_value="df-key"),
+ patch.object(processor, "annotation_cache_key",
return_value="ann-key"),
+ patch.object(processor, "get_cache_timeout", return_value=3600),
+ ):
+ mock_cache_manager.get.return_value = cache
+ return processor.get_df_payload_result(query_obj,
force_cached=force_cached)
+
+
+def test_get_df_payload_result_skips_annotations_when_dataframe_failed() ->
None:
+ """A failed dataframe keeps its own error and never triggers the fetch."""
+ processor, query_obj, cache =
_annotated_payload_fixture(status=QueryStatus.FAILED)
+ with patch.object(processor, "_get_annotation_data_cached") as
mock_annotation:
+ result = _run_annotated_payload(processor, query_obj, cache)
+ mock_annotation.assert_not_called()
+ assert result.payload["annotation_data"] == {}
+ assert result.payload["error"] == "dataframe failed"
+
+
+def test_get_df_payload_result_force_cached_misses_without_annotation_entry()
-> None:
+ """Through the real helper, ``force_cached`` with no annotation entry
+ raises ``CacheLoadError`` instead of running the annotation SQL."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(processor, "get_annotation_data") as mock_live,
+ ):
+ cm.data_cache.get.return_value = None
+ with pytest.raises(CacheLoadError):
+ _run_annotated_payload(processor, query_obj, cache,
force_cached=True)
+ mock_live.assert_not_called()
+
+
+def test_get_df_payload_result_annotation_failure_marks_payload_failed() ->
None:
+ """An annotation error on a good dataframe fails the payload."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(
+ processor,
+ "get_annotation_data",
+ side_effect=QueryObjectValidationError("annotation boom"),
+ ),
+ ):
+ cm.data_cache.get.return_value = None
+ result = _run_annotated_payload(processor, query_obj, cache)
+ assert cache.status == QueryStatus.FAILED
+ assert result.payload["error"] == "annotation boom"
Review Comment:
Good catch. That test now asserts `isinstance(...,
QueryObjectValidationMessage)` and that `.cause` is the original exception.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -3167,11 +3915,211 @@ def payload(context: QueryContext, **kwargs: Any) ->
dict[str, Any]:
datasource="native-1", annotation_context=context
)
assert current_key != legacy_key
- assert context == {
- **legacy_context,
- "semantic_filter_protocol": "semantic-null-filters-v1",
- "source_versions": {"7": view.metadata_generation},
- }
+ assert context is not None
+ assert context["source_scope"][0]["semantic_filter_protocol"] == (
+ "semantic-null-filters-v1"
+ )
+ assert context["source_scope"][0]["metadata_version"] ==
view.metadata_generation
+
+
+def _annotated_payload_fixture(
+ *, status: str = "success"
+) -> tuple[QueryContextProcessor, Any, MagicMock]:
+ """A processor, annotated query object, and a dataframe cache hit."""
+ from superset.common.query_object import QueryObject
+
+ mock_query_context = MagicMock()
+ mock_query_context.force = False
+ mock_datasource = MagicMock()
+ mock_datasource.column_names = ["col1"]
+ processor = QueryContextProcessor(mock_query_context)
+ processor._qc_datasource = mock_datasource
+ query_obj = QueryObject(
+ datasource=mock_datasource,
+ columns=["col1"],
+ annotation_layers=[
+ {
+ "annotationType": "EVENT",
+ "sourceType": "NATIVE",
+ "name": "a",
+ "value": 1,
+ }
+ ],
+ )
+ cache = MagicMock()
+ cache.is_loaded = True
+ cache.applied_filter_columns = ["col1"]
+ cache.df = pd.DataFrame({"col1": [1, 2, 3]})
+ cache.query = ""
+ cache.status = status
+ cache.error_message = "dataframe failed" if status == QueryStatus.FAILED
else None
+ cache.is_cached = True
+ cache.applied_template_filters = []
+ cache.rejected_filter_columns = []
+ cache.bq_memory_limited = False
+ cache.result_persisted = False
+ return processor, query_obj, cache
+
+
+def _run_annotated_payload(
+ processor: QueryContextProcessor,
+ query_obj: Any,
+ cache: MagicMock,
+ *,
+ force_cached: bool = False,
+) -> Any:
+ with (
+ patch(
+ "superset.common.query_context_processor.QueryCacheManager"
+ ) as mock_cache_manager,
+ patch.object(query_obj, "validate", return_value=None),
+ patch.object(processor, "query_cache_key", return_value="df-key"),
+ patch.object(processor, "annotation_cache_key",
return_value="ann-key"),
+ patch.object(processor, "get_cache_timeout", return_value=3600),
+ ):
+ mock_cache_manager.get.return_value = cache
+ return processor.get_df_payload_result(query_obj,
force_cached=force_cached)
+
+
+def test_get_df_payload_result_skips_annotations_when_dataframe_failed() ->
None:
+ """A failed dataframe keeps its own error and never triggers the fetch."""
+ processor, query_obj, cache =
_annotated_payload_fixture(status=QueryStatus.FAILED)
+ with patch.object(processor, "_get_annotation_data_cached") as
mock_annotation:
+ result = _run_annotated_payload(processor, query_obj, cache)
+ mock_annotation.assert_not_called()
+ assert result.payload["annotation_data"] == {}
+ assert result.payload["error"] == "dataframe failed"
+
+
+def test_get_df_payload_result_force_cached_misses_without_annotation_entry()
-> None:
+ """Through the real helper, ``force_cached`` with no annotation entry
+ raises ``CacheLoadError`` instead of running the annotation SQL."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(processor, "get_annotation_data") as mock_live,
+ ):
+ cm.data_cache.get.return_value = None
+ with pytest.raises(CacheLoadError):
+ _run_annotated_payload(processor, query_obj, cache,
force_cached=True)
+ mock_live.assert_not_called()
+
+
+def test_get_df_payload_result_annotation_failure_marks_payload_failed() ->
None:
+ """An annotation error on a good dataframe fails the payload."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(
+ processor,
+ "get_annotation_data",
+ side_effect=QueryObjectValidationError("annotation boom"),
+ ),
+ ):
+ cm.data_cache.get.return_value = None
+ result = _run_annotated_payload(processor, query_obj, cache)
+ assert cache.status == QueryStatus.FAILED
+ assert result.payload["error"] == "annotation boom"
+
+
+def test_get_df_payload_result_completeness_error_escapes_annotation_fetch()
-> None:
+ """A semantic annotation's incomplete-result signal must propagate (as the
+ dataframe path does) so the async task is not reported successful."""
+ from superset.exceptions import SemanticResultCompletenessError
+
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(
+ processor,
+ "get_annotation_data",
+ side_effect=SemanticResultCompletenessError("unverified"),
+ ),
+ ):
+ cm.data_cache.get.return_value = None
+ with pytest.raises(SemanticResultCompletenessError):
+ _run_annotated_payload(processor, query_obj, cache)
+
+
+def test_get_df_payload_result_times_annotation_fetch_as_data_acquisition() ->
None:
+ """Annotation work lands in ``data_acquisition_ns`` even on a dataframe
+ cache hit (which otherwise reports no acquisition time)."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+
+ def slow_annotations(*_: Any, **__: Any) -> dict[str, Any]:
+ time.sleep(0.01)
+ return {}
+
+ with patch.object(
+ processor, "_get_annotation_data_cached", side_effect=slow_annotations
+ ):
+ result = _run_annotated_payload(processor, query_obj, cache)
+ assert result.timing.data_acquisition_ns is not None
+ assert result.timing.data_acquisition_ns >= 10_000_000
Review Comment:
Added a deterministic-clock case on a dataframe miss: the dataframe fetch
advances 30 ns, the annotation fetch 20 ns, and it asserts `data_acquisition_ns
== 50`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
