dvt-core 1.11.0b4__py3-none-any.whl

dvt/graph/queue.py

@@ -0,0 +1,213 @@
+import threading
+from queue import PriorityQueue
+from typing import Dict, Generator, List, Optional, Set
+
+import networkx as nx  # type: ignore
+from dvt.contracts.graph.manifest import Manifest
+from dvt.contracts.graph.nodes import (
+    Exposure,
+    GraphMemberNode,
+    Metric,
+    SourceDefinition,
+)
+from dvt.node_types import NodeType
+
+from .graph import UniqueId
+
+
+class GraphQueue:
+    """A fancy queue that is backed by the dependency graph.
+    Note: this will mutate input!
+
+    This queue is thread-safe for `mark_done` calls, though you must ensure
+    that separate threads do not call `.empty()` or `__len__()` and `.get()` at
+    the same time, as there is an unlocked race!
+    """
+
+    def __init__(
+        self,
+        graph: nx.DiGraph,
+        manifest: Manifest,
+        selected: Set[UniqueId],
+        preserve_edges: bool = True,
+    ) -> None:
+        # 'create_empty_copy' returns a copy of the graph G with all of the edges removed, and leaves nodes intact.
+        self.graph = graph if preserve_edges else nx.classes.function.create_empty_copy(graph)
+        self.manifest = manifest
+        self._selected = selected
+        # store the queue as a priority queue.
+        self.inner: PriorityQueue = PriorityQueue()
+        # things that have been popped off the queue but not finished
+        # and worker thread reservations
+        self.in_progress: Set[UniqueId] = set()
+        # things that are in the queue
+        self.queued: Set[UniqueId] = set()
+        # this lock controls most things
+        self.lock = threading.Lock()
+        # store the 'score' of each node as a number. Lower is higher priority.
+        self._scores = self._get_scores(self.graph)
+        # populate the initial queue
+        self._find_new_additions(list(self.graph.nodes()))
+        # awaits after task end
+        self.some_task_done = threading.Condition(self.lock)
+
+    def get_selected_nodes(self) -> Set[UniqueId]:
+        return self._selected.copy()
+
+    def _include_in_cost(self, node_id: UniqueId) -> bool:
+        node = self.manifest.expect(node_id)
+        if node.resource_type != NodeType.Model:
+            return False
+        # must be a Model - tell mypy this won't be a Source or Exposure or Metric
+        assert not isinstance(node, (SourceDefinition, Exposure, Metric))
+        if node.is_ephemeral:
+            return False
+        return True
+
+    @staticmethod
+    def _grouped_topological_sort(
+        graph: nx.DiGraph,
+    ) -> Generator[List[str], None, None]:
+        """Topological sort of given graph that groups ties.
+
+        Adapted from `nx.topological_sort`, this function returns a topo sort of a graph however
+        instead of arbitrarily ordering ties in the sort order, ties are grouped together in
+        lists.
+
+        Args:
+            graph: The graph to be sorted.
+
+        Returns:
+            A generator that yields lists of nodes, one list per graph depth level.
+        """
+        indegree_map = {v: d for v, d in graph.in_degree() if d > 0}
+        zero_indegree = [v for v, d in graph.in_degree() if d == 0]
+
+        while zero_indegree:
+            yield zero_indegree
+            new_zero_indegree = []
+            for v in zero_indegree:
+                for _, child in graph.edges(v):
+                    indegree_map[child] -= 1
+                    if not indegree_map[child]:
+                        new_zero_indegree.append(child)
+            zero_indegree = new_zero_indegree
+
+    def _get_scores(self, graph: nx.DiGraph) -> Dict[str, int]:
+        """Scoring nodes for processing order.
+
+        Scores are calculated by the graph depth level. Lowest score (0) should be processed first.
+
+        Args:
+            graph: The graph to be scored.
+
+        Returns:
+            A dictionary consisting of `node name`:`score` pairs.
+        """
+        # split graph by connected subgraphs
+        subgraphs = (graph.subgraph(x) for x in nx.connected_components(nx.Graph(graph)))
+
+        # score all nodes in all subgraphs
+        scores = {}
+        for subgraph in subgraphs:
+            grouped_nodes = self._grouped_topological_sort(subgraph)
+            for level, group in enumerate(grouped_nodes):
+                for node in group:
+                    scores[node] = level
+
+        return scores
+
+    def get(self, block: bool = True, timeout: Optional[float] = None) -> GraphMemberNode:
+        """Get a node off the inner priority queue. By default, this blocks.
+
+        This takes the lock, but only for part of it.
+
+        :param block: If True, block until the inner queue has data
+        :param timeout: If set, block for timeout seconds waiting for data.
+        :return: The node as present in the manifest.
+
+        See `queue.PriorityQueue` for more information on `get()` behavior and
+        exceptions.
+        """
+        _, node_id = self.inner.get(block=block, timeout=timeout)
+        with self.lock:
+            self._mark_in_progress(node_id)
+        return self.manifest.expect(node_id)
+
+    def __len__(self) -> int:
+        """The length of the queue is the number of tasks left for the queue to
+        give out, regardless of where they are. Incomplete tasks are not part
+        of the length.
+
+        This takes the lock.
+        """
+        with self.lock:
+            return len(self.graph) - len(self.in_progress)
+
+    def empty(self) -> bool:
+        """The graph queue is 'empty' if it all remaining nodes in the graph
+        are in progress.
+
+        This takes the lock.
+        """
+        return len(self) == 0
+
+    def _already_known(self, node: UniqueId) -> bool:
+        """Decide if a node is already known (either handed out as a task, or
+        in the queue).
+
+        Callers must hold the lock.
+
+        :param str node: The node ID to check
+        :returns bool: If the node is in progress/queued.
+        """
+        return node in self.in_progress or node in self.queued
+
+    def _find_new_additions(self, candidates) -> None:
+        """Find any nodes in the graph that need to be added to the internal
+        queue and add them.
+        """
+        for node in candidates:
+            if self.graph.in_degree(node) == 0 and not self._already_known(node):
+                self.inner.put((self._scores[node], node))
+                self.queued.add(node)
+
+    def mark_done(self, node_id: UniqueId) -> None:
+        """Given a node's unique ID, mark it as done.
+
+        This method takes the lock.
+
+        :param str node_id: The node ID to mark as complete.
+        """
+        with self.lock:
+            self.in_progress.remove(node_id)
+            successors = list(self.graph.successors(node_id))
+            self.graph.remove_node(node_id)
+            self._find_new_additions(successors)
+            self.inner.task_done()
+            self.some_task_done.notify_all()
+
+    def _mark_in_progress(self, node_id: UniqueId) -> None:
+        """Mark the node as 'in progress'.
+
+        Callers must hold the lock.
+
+        :param str node_id: The node ID to mark as in progress.
+        """
+        self.queued.remove(node_id)
+        self.in_progress.add(node_id)
+
+    def join(self) -> None:
+        """Join the queue. Blocks until all tasks are marked as done.
+
+        Make sure not to call this before the queue reports that it is empty.
+        """
+        self.inner.join()
+
+    def wait_until_something_was_done(self) -> int:
+        """Block until a task is done, then return the number of unfinished
+        tasks.
+        """
+        with self.lock:
+            self.some_task_done.wait()
+            return self.inner.unfinished_tasks

dvt/graph/selector.py

@@ -0,0 +1,375 @@
+from typing import List, Optional, Set, Tuple
+
+from dvt import selected_resources
+from dvt.contracts.graph.manifest import Manifest
+from dvt.contracts.graph.nodes import GraphMemberNode
+from dvt.contracts.state import PreviousState
+from dvt.events.types import NoNodesForSelectionCriteria, SelectorReportInvalidSelector
+from dvt.exceptions import DbtInternalError, InvalidSelectorError
+from dvt.node_types import NodeType
+
+from dbt_common.events.functions import fire_event, warn_or_error
+
+from .graph import Graph, UniqueId
+from .queue import GraphQueue
+from .selector_methods import MethodManager
+from .selector_spec import IndirectSelection, SelectionCriteria, SelectionSpec
+
+
+def get_package_names(nodes):
+    return set([node.split(".")[1] for node in nodes])
+
+
+def can_select_indirectly(node):
+    """If a node is not selected itself, but its parent(s) are, it may qualify
+    for indirect selection.
+    Today, only Test nodes can be indirectly selected. In the future,
+    other node types or invocation flags might qualify.
+    """
+    if node.resource_type == NodeType.Test:
+        return True
+    elif node.resource_type == NodeType.Unit:
+        return True
+    else:
+        return False
+
+
+class NodeSelector(MethodManager):
+    """The node selector is aware of the graph and manifest"""
+
+    def __init__(
+        self,
+        graph: Graph,
+        manifest: Manifest,
+        previous_state: Optional[PreviousState] = None,
+        include_empty_nodes: bool = False,
+    ) -> None:
+        super().__init__(manifest, previous_state)
+        self.full_graph: Graph = graph
+        self.include_empty_nodes: bool = include_empty_nodes
+
+        # build a subgraph containing only non-empty, enabled nodes and enabled
+        # sources.
+        graph_members = {
+            unique_id for unique_id in self.full_graph.nodes() if self._is_graph_member(unique_id)
+        }
+        self.graph = self.full_graph.subgraph(graph_members)
+
+    def select_included(
+        self,
+        included_nodes: Set[UniqueId],
+        spec: SelectionCriteria,
+    ) -> Set[UniqueId]:
+        """Select the explicitly included nodes, using the given spec. Return
+        the selected set of unique IDs.
+        """
+        method = self.get_method(spec.method, spec.method_arguments)
+        return set(method.search(included_nodes, spec.value))
+
+    def get_nodes_from_criteria(
+        self, spec: SelectionCriteria
+    ) -> Tuple[Set[UniqueId], Set[UniqueId]]:
+        """Get all nodes specified by the single selection criteria.
+
+        - collect the directly included nodes
+        - find their specified relatives
+        - perform any selector-specific expansion
+        """
+
+        nodes = self.graph.nodes()
+        try:
+            collected = self.select_included(nodes, spec)
+        except InvalidSelectorError:
+            valid_selectors = ", ".join(self.SELECTOR_METHODS)
+            fire_event(
+                SelectorReportInvalidSelector(
+                    valid_selectors=valid_selectors, spec_method=spec.method, raw_spec=spec.raw
+                )
+            )
+            return set(), set()
+
+        neighbors = self.collect_specified_neighbors(spec, collected)
+        selected = collected | neighbors
+
+        # if --indirect-selection EMPTY, do not expand to adjacent tests
+        if spec.indirect_selection == IndirectSelection.Empty:
+            return selected, set()
+        else:
+            direct_nodes, indirect_nodes = self.expand_selection(
+                selected=selected, indirect_selection=spec.indirect_selection
+            )
+            return direct_nodes, indirect_nodes
+
+    def collect_specified_neighbors(
+        self, spec: SelectionCriteria, selected: Set[UniqueId]
+    ) -> Set[UniqueId]:
+        """Given the set of models selected by the explicit part of the
+        selector (like "tag:foo"), apply the modifiers on the spec ("+"/"@").
+        Return the set of additional nodes that should be collected (which may
+        overlap with the selected set).
+        """
+        additional: Set[UniqueId] = set()
+        if spec.childrens_parents:
+            additional.update(self.graph.select_childrens_parents(selected))
+
+        if spec.parents:
+            depth = spec.parents_depth
+            additional.update(self.graph.select_parents(selected, depth))
+
+        if spec.children:
+            depth = spec.children_depth
+            additional.update(self.graph.select_children(selected, depth))
+        return additional
+
+    def select_nodes_recursively(self, spec: SelectionSpec) -> Tuple[Set[UniqueId], Set[UniqueId]]:
+        """If the spec is a composite spec (a union, difference, or intersection),
+        recurse into its selections and combine them. If the spec is a concrete
+        selection criteria, resolve that using the given graph.
+        """
+        if isinstance(spec, SelectionCriteria):
+            direct_nodes, indirect_nodes = self.get_nodes_from_criteria(spec)
+        else:
+            bundles = [self.select_nodes_recursively(component) for component in spec]
+
+            direct_sets = []
+            indirect_sets = []
+
+            for direct, indirect in bundles:
+                direct_sets.append(direct)
+                indirect_sets.append(direct | indirect)
+
+            initial_direct = spec.combined(direct_sets)
+            indirect_nodes = spec.combined(indirect_sets)
+
+            direct_nodes = self.incorporate_indirect_nodes(
+                initial_direct, indirect_nodes, spec.indirect_selection
+            )
+
+            if spec.expect_exists and len(direct_nodes) == 0:
+                warn_or_error(NoNodesForSelectionCriteria(spec_raw=str(spec.raw)))
+
+        return direct_nodes, indirect_nodes
+
+    def select_nodes(self, spec: SelectionSpec) -> Tuple[Set[UniqueId], Set[UniqueId]]:
+        """Select the nodes in the graph according to the spec.
+
+        This is the main point of entry for turning a spec into a set of nodes:
+        - Recurse through spec, select by criteria, combine by set operation
+        - Return final (unfiltered) selection set
+        """
+        direct_nodes, indirect_nodes = self.select_nodes_recursively(spec)
+        indirect_only = indirect_nodes.difference(direct_nodes)
+        return direct_nodes, indirect_only
+
+    def _is_graph_member(self, unique_id: UniqueId) -> bool:
+        if unique_id in self.manifest.sources:
+            source = self.manifest.sources[unique_id]
+            return source.config.enabled
+        elif unique_id in self.manifest.exposures:
+            return True
+        elif unique_id in self.manifest.functions:
+            function = self.manifest.functions[unique_id]
+            return function.config.enabled
+        elif unique_id in self.manifest.metrics:
+            metric = self.manifest.metrics[unique_id]
+            return metric.config.enabled
+        elif unique_id in self.manifest.semantic_models:
+            semantic_model = self.manifest.semantic_models[unique_id]
+            return semantic_model.config.enabled
+        elif unique_id in self.manifest.unit_tests:
+            unit_test = self.manifest.unit_tests[unique_id]
+            return unit_test.config.enabled
+        elif unique_id in self.manifest.saved_queries:
+            saved_query = self.manifest.saved_queries[unique_id]
+            return saved_query.config.enabled
+        elif unique_id in self.manifest.exposures:
+            exposure = self.manifest.exposures[unique_id]
+            return exposure.config.enabled
+        else:
+            node = self.manifest.nodes[unique_id]
+            return node.config.enabled
+
+    def _is_empty_node(self, unique_id: UniqueId) -> bool:
+        if unique_id in self.manifest.nodes:
+            node = self.manifest.nodes[unique_id]
+            return node.empty
+        else:
+            return False
+
+    def node_is_match(self, node: GraphMemberNode) -> bool:
+        """Determine if a node is a match for the selector. Non-match nodes
+        will be excluded from results during filtering.
+        """
+        return True
+
+    def _is_match(self, unique_id: UniqueId) -> bool:
+        node: GraphMemberNode
+        if unique_id in self.manifest.nodes:
+            node = self.manifest.nodes[unique_id]
+        elif unique_id in self.manifest.sources:
+            node = self.manifest.sources[unique_id]
+        elif unique_id in self.manifest.exposures:
+            node = self.manifest.exposures[unique_id]
+        elif unique_id in self.manifest.functions:
+            node = self.manifest.functions[unique_id]
+        elif unique_id in self.manifest.metrics:
+            node = self.manifest.metrics[unique_id]
+        elif unique_id in self.manifest.semantic_models:
+            node = self.manifest.semantic_models[unique_id]
+        elif unique_id in self.manifest.unit_tests:
+            node = self.manifest.unit_tests[unique_id]
+        elif unique_id in self.manifest.saved_queries:
+            node = self.manifest.saved_queries[unique_id]
+        else:
+            raise DbtInternalError(f"Node {unique_id} not found in the manifest!")
+        return self.node_is_match(node)
+
+    def filter_selection(self, selected: Set[UniqueId]) -> Set[UniqueId]:
+        """Return the subset of selected nodes that is a match for this
+        selector.
+        """
+        return {
+            unique_id
+            for unique_id in selected
+            if self._is_match(unique_id)
+            and (self.include_empty_nodes or not self._is_empty_node(unique_id))
+        }
+
+    def expand_selection(
+        self,
+        selected: Set[UniqueId],
+        indirect_selection: IndirectSelection = IndirectSelection.Eager,
+    ) -> Tuple[Set[UniqueId], Set[UniqueId]]:
+        # Test selection by default expands to include an implicitly/indirectly selected tests.
+        # `dbt test -m model_a` also includes tests that directly depend on `model_a`.
+        # Expansion has four modes, EAGER, CAUTIOUS and BUILDABLE, EMPTY.
+        #
+        # EAGER mode: If ANY parent is selected, select the test.
+        #
+        # CAUTIOUS mode:
+        #  - If ALL parents are selected, select the test.
+        #  - If ANY parent is missing, return it separately. We'll keep it around
+        #    for later and see if its other parents show up.
+        #
+        # BUILDABLE mode:
+        #  - If ALL parents are selected, or the parents of the test are themselves parents of the selected, select the test.
+        #  - If ANY parent is missing, return it separately. We'll keep it around
+        #    for later and see if its other parents show up.
+        #
+        # EMPTY mode: Only select the given node and ignore attached nodes (i.e. ignore tests attached to a model)
+        #
+        # Users can opt out of inclusive EAGER mode by passing --indirect-selection cautious
+        # CLI argument or by specifying `indirect_selection: true` in a yaml selector
+
+        direct_nodes = set(selected)
+        indirect_nodes = set()
+        selected_and_parents = set()
+        if indirect_selection == IndirectSelection.Buildable:
+            selected_and_parents = selected.union(self.graph.select_parents(selected)).union(
+                self.manifest.sources
+            )
+
+        for unique_id in self.graph.select_successors(selected):
+            if unique_id in self.manifest.nodes or unique_id in self.manifest.unit_tests:
+                if unique_id in self.manifest.nodes:
+                    node = self.manifest.nodes[unique_id]
+                elif unique_id in self.manifest.unit_tests:
+                    node = self.manifest.unit_tests[unique_id]  # type: ignore
+                # Test nodes that are not selected themselves, but whose parents are selected.
+                # (Does not include unit tests because they can only have one parent.)
+                if can_select_indirectly(node):
+                    # should we add it in directly?
+                    if indirect_selection == IndirectSelection.Eager or set(
+                        node.depends_on_nodes
+                    ) <= set(selected):
+                        direct_nodes.add(unique_id)
+                    elif indirect_selection == IndirectSelection.Buildable and set(
+                        node.depends_on_nodes
+                    ) <= set(selected_and_parents):
+                        direct_nodes.add(unique_id)
+                    elif indirect_selection == IndirectSelection.Empty:
+                        pass
+                    else:
+                        indirect_nodes.add(unique_id)
+
+        return direct_nodes, indirect_nodes
+
+    def incorporate_indirect_nodes(
+        self,
+        direct_nodes: Set[UniqueId],
+        indirect_nodes: Set[UniqueId] = set(),
+        indirect_selection: IndirectSelection = IndirectSelection.Eager,
+    ) -> Set[UniqueId]:
+        # Check tests previously selected indirectly to see if ALL their
+        # parents are now present.
+
+        # performance: if identical, skip the processing below
+        if set(direct_nodes) == set(indirect_nodes):
+            return direct_nodes
+
+        selected = set(direct_nodes)
+
+        if indirect_selection == IndirectSelection.Cautious:
+            for unique_id in indirect_nodes:
+                if unique_id in self.manifest.nodes:
+                    node = self.manifest.nodes[unique_id]
+                    if set(node.depends_on_nodes) <= set(selected):
+                        selected.add(unique_id)
+        elif indirect_selection == IndirectSelection.Buildable:
+            selected_and_parents = selected.union(self.graph.select_parents(selected))
+            for unique_id in indirect_nodes:
+                if unique_id in self.manifest.nodes:
+                    node = self.manifest.nodes[unique_id]
+                    if set(node.depends_on_nodes) <= set(selected_and_parents):
+                        selected.add(unique_id)
+
+        return selected
+
+    def get_selected(self, spec: SelectionSpec) -> Set[UniqueId]:
+        """get_selected runs through the node selection process:
+
+        - node selection. Based on the include/exclude sets, the set
+            of matched unique IDs is returned
+            - includes direct + indirect selection (for tests)
+        - filtering:
+            - selectors can filter the nodes after all of them have been
+              selected
+        """
+        selected_nodes, indirect_only = self.select_nodes(spec)
+        filtered_nodes = self.filter_selection(selected_nodes)
+
+        return filtered_nodes
+
+    def get_graph_queue(self, spec: SelectionSpec, preserve_edges: bool = True) -> GraphQueue:
+        """Returns a queue over nodes in the graph that tracks progress of
+        dependencies.
+        """
+        # Filtering happens in get_selected
+        selected_nodes = self.get_selected(spec)
+        # Save to global variable
+        selected_resources.set_selected_resources(selected_nodes)
+        # Construct a new graph using the selected_nodes
+        new_graph = self.full_graph.get_subset_graph(selected_nodes)
+        # should we give a way here for consumers to mutate the graph?
+        return GraphQueue(new_graph.graph, self.manifest, selected_nodes, preserve_edges)
+
+
+class ResourceTypeSelector(NodeSelector):
+    def __init__(
+        self,
+        graph: Graph,
+        manifest: Manifest,
+        previous_state: Optional[PreviousState],
+        resource_types: List[NodeType],
+        include_empty_nodes: bool = False,
+    ) -> None:
+        super().__init__(
+            graph=graph,
+            manifest=manifest,
+            previous_state=previous_state,
+            include_empty_nodes=include_empty_nodes,
+        )
+        self.resource_types: Set[NodeType] = set(resource_types)
+
+    def node_is_match(self, node):
+        return node.resource_type in self.resource_types