diff --git a/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_pipeline.py b/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_pipeline.py index b185d5f16..3e8fbca21 100644 --- a/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_pipeline.py +++ b/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_pipeline.py @@ -20,54 +20,34 @@ import pytest -from documentdb_tests.framework.assertions import assertProperties +from documentdb_tests.compatibility.tests.system.diagnostic.utils.explain_hash import ( + PLAN_CACHE_SHAPE_HASH, + QUERY_SHAPE_HASH, + assert_hash_relation, +) from documentdb_tests.framework.executor import execute_command from documentdb_tests.framework.parametrize import pytest_params -from documentdb_tests.framework.property_checks import Eq, Ne from documentdb_tests.framework.test_case import BaseTestCase pytestmark = pytest.mark.admin -def _get_plan_cache_shape_hash(collection, pipeline: list) -> str: - """Return queryPlanner.planCacheShapeHash from explain for an aggregate pipeline.""" - cmd: dict[str, Any] = { - "aggregate": collection.name, - "pipeline": pipeline, - "cursor": {}, - } - result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"}) - if isinstance(result, Exception): - raise RuntimeError(f"explain command failed: {result}") - # Standard format: queryPlanner at top level. - query_planner = result.get("queryPlanner", {}) - if "planCacheShapeHash" in query_planner: - return str(query_planner["planCacheShapeHash"]) - # Stages format ($out/$merge): nested under stages.0.$cursor.queryPlanner. - stages = result.get("stages", []) - if stages: - cursor_stage = stages[0].get("$cursor", {}) - qp = cursor_stage.get("queryPlanner", {}) - if "planCacheShapeHash" in qp: - return str(qp["planCacheShapeHash"]) - raise RuntimeError( - f"explain did not return planCacheShapeHash in expected location; got: {result}" - ) +# Explain reports planCacheShapeHash under queryPlanner for most pipelines, and +# under stages.0.$cursor.queryPlanner when the pipeline ends in $out or $merge. +PIPELINE_PLAN_CACHE_SHAPE_HASH = ( + PLAN_CACHE_SHAPE_HASH, + f"stages.0.$cursor.{PLAN_CACHE_SHAPE_HASH}", +) -def _get_query_shape_hash(collection, pipeline: list) -> str: - """Return the top-level queryShapeHash from explain for an aggregate pipeline.""" +def _explain_aggregate(collection, pipeline: list) -> Any: + """Run explain on an aggregate pipeline and return the raw result (or the exception).""" cmd: dict[str, Any] = { "aggregate": collection.name, "pipeline": pipeline, "cursor": {}, } - result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"}) - if isinstance(result, Exception): - raise RuntimeError(f"explain command failed: {result}") - if "queryShapeHash" not in result: - raise RuntimeError(f"explain did not return queryShapeHash; got: {result}") - return str(result["queryShapeHash"]) + return execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"}) @dataclass(frozen=True) @@ -361,14 +341,10 @@ def test_pipeline_plancacheshapehash_stage_sensitivity(collection, test): collection.insert_many( [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "sub": {"x": i}} for i in range(20)] ) - hash_a = _get_plan_cache_shape_hash(collection, list(test.pipeline_a)) - hash_b = _get_plan_cache_shape_hash(collection, list(test.pipeline_b)) - check = Eq(hash_b) if test.plan_cache_same else Ne(hash_b) - assertProperties( - {"planCacheShapeHash": hash_a}, - {"planCacheShapeHash": check}, - msg=test.msg, - raw_res=True, + result_a = _explain_aggregate(collection, list(test.pipeline_a)) + result_b = _explain_aggregate(collection, list(test.pipeline_b)) + assert_hash_relation( + result_a, result_b, PIPELINE_PLAN_CACHE_SHAPE_HASH, same=test.plan_cache_same, msg=test.msg ) @@ -378,12 +354,8 @@ def test_pipeline_queryshapehash_considers_all_stages(collection, test): collection.insert_many( [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "sub": {"x": i}} for i in range(20)] ) - hash_a = _get_query_shape_hash(collection, list(test.pipeline_a)) - hash_b = _get_query_shape_hash(collection, list(test.pipeline_b)) - check = Eq(hash_b) if test.query_shape_same else Ne(hash_b) - assertProperties( - {"queryShapeHash": hash_a}, - {"queryShapeHash": check}, - msg=test.msg, - raw_res=True, + result_a = _explain_aggregate(collection, list(test.pipeline_a)) + result_b = _explain_aggregate(collection, list(test.pipeline_b)) + assert_hash_relation( + result_a, result_b, QUERY_SHAPE_HASH, same=test.query_shape_same, msg=test.msg ) diff --git a/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_shape.py b/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_shape.py index 0de460033..05e7a7a59 100644 --- a/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_shape.py +++ b/documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_shape.py @@ -10,23 +10,26 @@ import pytest -from documentdb_tests.framework.assertions import assertProperties +from documentdb_tests.compatibility.tests.system.diagnostic.utils.explain_hash import ( + PLAN_CACHE_SHAPE_HASH, + QUERY_SHAPE_HASH, + assert_hash_relation, +) from documentdb_tests.framework.executor import execute_command from documentdb_tests.framework.parametrize import pytest_params -from documentdb_tests.framework.property_checks import Eq, Ne from documentdb_tests.framework.test_case import BaseTestCase pytestmark = pytest.mark.admin -def get_query_hash( +def explain_find( collection, find_filter: dict, sort: Optional[dict] = None, projection: Optional[dict] = None, collation: Optional[dict] = None, -) -> str: - """Return queryPlanner.planCacheShapeHash from explain for a find filter.""" +) -> Any: + """Run explain on a find and return the raw result (or the exception).""" cmd: dict[str, Any] = {"find": collection.name, "filter": find_filter} if sort is not None: cmd["sort"] = sort @@ -34,36 +37,7 @@ def get_query_hash( cmd["projection"] = projection if collation is not None: cmd["collation"] = collation - result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"}) - if isinstance(result, Exception): - raise RuntimeError(f"explain command failed: {result}") - query_planner = result.get("queryPlanner", {}) - if "planCacheShapeHash" not in query_planner: - raise RuntimeError(f"explain did not return queryPlanner.planCacheShapeHash; got: {result}") - return str(query_planner["planCacheShapeHash"]) - - -def get_query_shape_hash( - collection, - find_filter: dict, - sort: Optional[dict] = None, - projection: Optional[dict] = None, - collation: Optional[dict] = None, -) -> str: - """Return the top-level queryShapeHash from explain for a find filter.""" - cmd: dict[str, Any] = {"find": collection.name, "filter": find_filter} - if sort is not None: - cmd["sort"] = sort - if projection is not None: - cmd["projection"] = projection - if collation is not None: - cmd["collation"] = collation - result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"}) - if isinstance(result, Exception): - raise RuntimeError(f"explain command failed: {result}") - if "queryShapeHash" not in result: - raise RuntimeError(f"explain did not return queryShapeHash; got: {result}") - return str(result["queryShapeHash"]) + return execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"}) @dataclass(frozen=True) @@ -297,18 +271,13 @@ def test_explain_queryhash_same_shape(collection, test): collection.insert_many( [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)] ) - hash_a = get_query_hash( + result_a = explain_find( collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a ) - hash_b = get_query_hash( + result_b = explain_find( collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b ) - assertProperties( - {"planCacheShapeHash": hash_a}, - {"planCacheShapeHash": Eq(hash_b)}, - msg=test.msg, - raw_res=True, - ) + assert_hash_relation(result_a, result_b, PLAN_CACHE_SHAPE_HASH, same=True, msg=test.msg) DIFF_SHAPE_TESTS: list[QueryHashCase] = [ @@ -516,18 +485,13 @@ def test_explain_queryhash_different_shape(collection, test): collection.insert_many( [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)] ) - hash_a = get_query_hash( + result_a = explain_find( collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a ) - hash_b = get_query_hash( + result_b = explain_find( collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b ) - assertProperties( - {"planCacheShapeHash": hash_a}, - {"planCacheShapeHash": Ne(hash_b)}, - msg=test.msg, - raw_res=True, - ) + assert_hash_relation(result_a, result_b, PLAN_CACHE_SHAPE_HASH, same=False, msg=test.msg) @pytest.mark.parametrize("test", pytest_params(SAME_SHAPE_TESTS)) @@ -542,16 +506,13 @@ def test_explain_queryshapehash_same_shape(collection, test): collection.insert_many( [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)] ) - hash_a = get_query_shape_hash( + result_a = explain_find( collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a ) - hash_b = get_query_shape_hash( + result_b = explain_find( collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b ) - check = Eq(hash_b) if expected_same else Ne(hash_b) - assertProperties( - {"queryShapeHash": hash_a}, {"queryShapeHash": check}, msg=test.msg, raw_res=True - ) + assert_hash_relation(result_a, result_b, QUERY_SHAPE_HASH, same=expected_same, msg=test.msg) @pytest.mark.parametrize("test", pytest_params(DIFF_SHAPE_TESTS)) @@ -565,13 +526,10 @@ def test_explain_queryshapehash_different_shape(collection, test): collection.insert_many( [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)] ) - hash_a = get_query_shape_hash( + result_a = explain_find( collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a ) - hash_b = get_query_shape_hash( + result_b = explain_find( collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b ) - check = Eq(hash_b) if expected_same else Ne(hash_b) - assertProperties( - {"queryShapeHash": hash_a}, {"queryShapeHash": check}, msg=test.msg, raw_res=True - ) + assert_hash_relation(result_a, result_b, QUERY_SHAPE_HASH, same=expected_same, msg=test.msg) diff --git a/documentdb_tests/compatibility/tests/system/diagnostic/utils/explain_hash.py b/documentdb_tests/compatibility/tests/system/diagnostic/utils/explain_hash.py new file mode 100644 index 000000000..1b7f48d42 --- /dev/null +++ b/documentdb_tests/compatibility/tests/system/diagnostic/utils/explain_hash.py @@ -0,0 +1,78 @@ +"""Helpers for comparing the query hashes that explain reports. + +The hash tests run explain on two queries and assert that their hashes match +(same shape) or differ (different shape). The helpers here keep the explain +result raw and do every check, including that the hash is present at all, +through a single ``assertProperties`` call. A failed explain or a missing hash +field is therefore reported as an assertion failure instead of an error raised +before the assertion. +""" + +from __future__ import annotations + +from typing import Any, Union + +from documentdb_tests.framework.assertions import assertProperties +from documentdb_tests.framework.property_checks import Eq, Ne, NonEmptyStr + +PLAN_CACHE_SHAPE_HASH = "queryPlanner.planCacheShapeHash" +"""Dotted path of the plan cache shape hash in a standard explain result.""" + +QUERY_SHAPE_HASH = "queryShapeHash" +"""Dotted path of the top-level query shape hash in an explain result.""" + + +def _lookup(result: Any, path: str) -> Any: + """Return the value at a dotted path (dict keys and list indices), or None when absent.""" + value: Any = result + for part in path.split("."): + if isinstance(value, list) and part.isdigit() and int(part) < len(value): + value = value[int(part)] + elif isinstance(value, dict) and part in value: + value = value[part] + else: + return None + return value + + +def assert_hash_relation( + result_a: Any, + result_b: Any, + paths: Union[str, tuple[str, ...]], + same: bool, + msg: str, +) -> None: + """Assert both explain results report a hash, and that the two hashes match or differ. + + Args: + result_a: Raw explain result (or exception) for the first query. + result_b: Raw explain result (or exception) for the second query. + paths: Dotted path of the hash, or several candidate paths tried in + order (for explain output whose layout depends on the pipeline). + same: True to assert the hashes are equal, False to assert they differ. + msg: Failure message. + """ + candidates = (paths,) if isinstance(paths, str) else paths + for result in (result_a, result_b): + if isinstance(result, Exception): + # The failed explain is the result under test: assert on it and stop, + # so it is reported once rather than again as two missing hashes. + assertProperties(result, {candidates[0]: NonEmptyStr()}, msg=msg, raw_res=True) + return + hashes: dict[str, str] = {} + for key, result in (("hash_a", result_a), ("hash_b", result_b)): + for path in candidates: + value = _lookup(result, path) + if value is not None: + hashes[key] = str(value) + break + # Compare only when the second hash exists; its absence is reported on its own. + checks_a: list[Any] = [NonEmptyStr()] + if "hash_b" in hashes: + checks_a.append(Eq(hashes["hash_b"]) if same else Ne(hashes["hash_b"])) + assertProperties( + hashes, + {"hash_a": checks_a, "hash_b": NonEmptyStr()}, + msg=f"{msg} (hash at {' or '.join(candidates)})", + raw_res=True, + )