From 257fda72e5bdc28599a155a17b720abf4356b313 Mon Sep 17 00:00:00 2001
From: Paul Yang <github@ardorful.com>
Date: Fri, 26 Apr 2024 02:57:39 -0700
Subject: [PATCH] Add saved query dependency resolver.

---
 .../query/query_dependency_resolver.py        | 54 +++++++++++++++++++
 .../query/saved_query_dependency_resolver.py  |  0
 .../query/test_query_dependency_resolver.py   | 25 +++++++++
 3 files changed, 79 insertions(+)
 create mode 100644 metricflow-semantics/metricflow_semantics/query/query_dependency_resolver.py
 delete mode 100644 metricflow-semantics/metricflow_semantics/query/saved_query_dependency_resolver.py
 create mode 100644 metricflow-semantics/tests_metricflow_semantics/query/test_query_dependency_resolver.py

diff --git a/metricflow-semantics/metricflow_semantics/query/query_dependency_resolver.py b/metricflow-semantics/metricflow_semantics/query/query_dependency_resolver.py
new file mode 100644
index 0000000000..137e84d399
--- /dev/null
+++ b/metricflow-semantics/metricflow_semantics/query/query_dependency_resolver.py
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+from dbt_semantic_interfaces.protocols import SemanticManifest
+from dbt_semantic_interfaces.references import (
+    SemanticModelReference,
+)
+
+from metricflow_semantics.model.semantic_manifest_lookup import SemanticManifestLookup
+from metricflow_semantics.query.query_parser import MetricFlowQueryParser
+from metricflow_semantics.specs.query_param_implementations import SavedQueryParameter
+
+
+@dataclass(frozen=True)
+class SavedQueryDependencySet:
+    """The dependencies of a saved query.
+
+    The primary use case is to handle creation of the cache item associated with the saved query. The dependencies
+    listed in this class must be up-to-date before the cache associated with the saved query can be created. Otherwise,
+    running the export / creating the cache may create a cache item that is out-of-date / unusable.
+    """
+
+    # The semantic models that the saved query depends on.
+    semantic_model_references: Tuple[SemanticModelReference, ...]
+
+
+class SavedQueryDependencyResolver:
+    """Resolves the dependencies of a saved query. Also see `SavedQueryDependencySet`."""
+
+    def __init__(self, semantic_manifest: SemanticManifest) -> None:  # noqa: D107
+        self._query_parser = MetricFlowQueryParser(SemanticManifestLookup(semantic_manifest))
+
+    def resolve_dependencies(self, saved_query_name: str) -> SavedQueryDependencySet:
+        """Return the dependencies of the given saved query in the manifest."""
+        parse_result = self._query_parser.parse_and_validate_saved_query(
+            saved_query_parameter=SavedQueryParameter(saved_query_name),
+            where_filter=None,
+            limit=None,
+            time_constraint_start=None,
+            time_constraint_end=None,
+            order_by_names=None,
+            order_by_parameters=None,
+        )
+
+        return SavedQueryDependencySet(
+            semantic_model_references=tuple(
+                sorted(
+                    parse_result.queried_semantic_models,
+                    key=lambda reference: reference.semantic_model_name,
+                )
+            ),
+        )
diff --git a/metricflow-semantics/metricflow_semantics/query/saved_query_dependency_resolver.py b/metricflow-semantics/metricflow_semantics/query/saved_query_dependency_resolver.py
deleted file mode 100644
index e69de29bb2..0000000000
diff --git a/metricflow-semantics/tests_metricflow_semantics/query/test_query_dependency_resolver.py b/metricflow-semantics/tests_metricflow_semantics/query/test_query_dependency_resolver.py
new file mode 100644
index 0000000000..b11af4fb44
--- /dev/null
+++ b/metricflow-semantics/tests_metricflow_semantics/query/test_query_dependency_resolver.py
@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import pytest
+from dbt_semantic_interfaces.implementations.semantic_manifest import (
+    PydanticSemanticManifest,
+)
+from dbt_semantic_interfaces.references import (
+    SemanticModelReference,
+)
+from metricflow_semantics.query.query_dependency_resolver import SavedQueryDependencyResolver
+
+
+@pytest.fixture(scope="session")
+def resolver(  # noqa: D103
+    simple_semantic_manifest: PydanticSemanticManifest,
+) -> SavedQueryDependencyResolver:
+    return SavedQueryDependencyResolver(simple_semantic_manifest)
+
+
+def test_saved_query_dependency_resolver(resolver: SavedQueryDependencyResolver) -> None:  # noqa: D103
+    dependency_set = resolver.resolve_dependencies("p0_booking")
+    assert tuple(dependency_set.semantic_model_references) == (
+        SemanticModelReference(semantic_model_name="bookings_source"),
+        SemanticModelReference(semantic_model_name="listings_latest"),
+    )