Merge pull request #2629 from fishtown-analytics/feature/new-node-sel…

…ectors

new node selectors (#2425)

CHANGELOG.md

@@ -5,6 +5,7 @@
 
 ### Features
 - Added support for Snowflake query tags at the connection and model level ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030), [#2555](https://github.com/fishtown-analytics/dbt/pull/2555/))
+- Added new node selector methods (`config`, `test_type`, `test_name`, `package`) ([#2425](https://github.com/fishtown-analytics/dbt/issues/2425), [#2629](https://github.com/fishtown-analytics/dbt/pull/2629))
 - Added option to specify profile when connecting to Redshift via IAM ([#2437](https://github.com/fishtown-analytics/dbt/issues/2437), [#2581](https://github.com/fishtown-analytics/dbt/pull/2581))
 - Add more helpful error message for misconfiguration in profiles.yml ([#2569](https://github.com/fishtown-analytics/dbt/issues/2569), [#2627](https://github.com/fishtown-analytics/dbt/pull/2627))
 ### Fixes

core/dbt/contracts/graph/model_config.py

@@ -120,6 +120,7 @@ def insensitive_patterns(*patterns: str):
 
 
 Severity = NewType('Severity', str)
+
 register_pattern(Severity, insensitive_patterns('warn', 'error'))
 
 

core/dbt/exceptions.py

@@ -393,6 +393,12 @@ def __init__(self, thread_id, known, node=None):
         )
 
 
+class InvalidSelectorException(RuntimeException):
+    def __init__(self, name: str):
+        self.name = name
+        super().__init__(name)
+
+
 def raise_compiler_error(msg, node=None) -> NoReturn:
     raise CompilationException(msg, node)
 

core/dbt/graph/__init__.py

@@ -9,6 +9,9 @@
     ResourceTypeSelector,
     NodeSelector,
 )
-from .cli import parse_difference  # noqa: F401
+from .cli import (  # noqa: F401
+    parse_difference,
+    parse_test_selectors,
+)
 from .queue import GraphQueue  # noqa: F401
 from .graph import Graph, UniqueId  # noqa: F401

core/dbt/graph/cli.py

@@ -17,6 +17,8 @@
 
 DEFAULT_INCLUDES: List[str] = ['fqn:*', 'source:*']
 DEFAULT_EXCLUDES: List[str] = []
+DATA_TEST_SELECTOR: str = 'test_type:data'
+SCHEMA_TEST_SELECTOR: str = 'test_type:schema'
 
 
 def parse_union(
@@ -64,3 +66,34 @@ def parse_difference(
     included = parse_union_from_default(include, DEFAULT_INCLUDES)
     excluded = parse_union_from_default(exclude, DEFAULT_EXCLUDES)
     return SelectionDifference(components=[included, excluded])
+
+
+def parse_test_selectors(
+    data: bool, schema: bool, base: SelectionSpec
+) -> SelectionSpec:
+    union_components = []
+
+    if data:
+        union_components.append(
+            SelectionCriteria.from_single_spec(DATA_TEST_SELECTOR)
+        )
+    if schema:
+        union_components.append(
+            SelectionCriteria.from_single_spec(SCHEMA_TEST_SELECTOR)
+        )
+
+    intersect_with: SelectionSpec
+    if not union_components:
+        return base
+    elif len(union_components) == 1:
+        intersect_with = union_components[0]
+    else:  # data and schema tests
+        intersect_with = SelectionUnion(
+            components=union_components,
+            expect_exists=True,
+            raw=[DATA_TEST_SELECTOR, SCHEMA_TEST_SELECTOR],
+        )
+
+    return SelectionIntersection(
+        components=[base, intersect_with], expect_exists=True
+    )

core/dbt/graph/selector.py

@@ -1,23 +1,18 @@
 
-from typing import (
-    Set, List, Dict, Union, Type
-)
+from typing import Set, List, Union
 
 from .graph import Graph, UniqueId
 from .queue import GraphQueue
-from .selector_methods import (
-    MethodName,
-    SelectorMethod,
-    QualifiedNameSelectorMethod,
-    TagSelectorMethod,
-    SourceSelectorMethod,
-    PathSelectorMethod,
-)
+from .selector_methods import MethodManager
 from .selector_spec import SelectionCriteria, SelectionSpec
 
 from dbt.logger import GLOBAL_LOGGER as logger
 from dbt.node_types import NodeType
-from dbt.exceptions import InternalException, warn_or_error
+from dbt.exceptions import (
+    InternalException,
+    InvalidSelectorException,
+    warn_or_error,
+)
 from dbt.contracts.graph.compiled import NonSourceNode, CompileResultNode
 from dbt.contracts.graph.manifest import Manifest
 from dbt.contracts.graph.parsed import ParsedSourceDefinition
@@ -35,16 +30,9 @@ def alert_non_existence(raw_spec, nodes):
         )
 
 
-class InvalidSelectorError(Exception):
-    # this internal exception should never escape the module.
-    pass
-
-
-class NodeSelector:
+class NodeSelector(MethodManager):
     """The node selector is aware of the graph and manifest,
     """
-    SELECTOR_METHODS: Dict[MethodName, Type[SelectorMethod]] = {}
-
     def __init__(
         self,
         graph: Graph,
@@ -53,52 +41,51 @@ def __init__(
         self.full_graph = graph
         self.manifest = manifest
 
-    @classmethod
-    def register_method(cls, name: MethodName, method: Type[SelectorMethod]):
-        cls.SELECTOR_METHODS[name] = method
-
-    def get_method(self, method: MethodName) -> SelectorMethod:
-        if method in self.SELECTOR_METHODS:
-            cls: Type[SelectorMethod] = self.SELECTOR_METHODS[method]
-            return cls(self.manifest)
-        else:
-            raise InvalidSelectorError(method)
+        # 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)
+        method = self.get_method(spec.method, spec.method_arguments)
         return set(method.search(included_nodes, spec.value))
 
     def get_nodes_from_criteria(
-        self, graph: Graph, spec: SelectionCriteria
+        self,
+        spec: SelectionCriteria,
     ) -> Set[UniqueId]:
-        """Given a Graph, get all nodes specified by the spec.
+        """Get all nodes specified by the single selection criteria.
 
         - collect the directly included nodes
         - find their specified relatives
         - perform any selector-specific expansion
         """
-        nodes = graph.nodes()
+
+        nodes = self.graph.nodes()
         try:
             collected = self.select_included(nodes, spec)
-        except InvalidSelectorError:
+        except InvalidSelectorException:
             valid_selectors = ", ".join(self.SELECTOR_METHODS)
             logger.info(
                 f"The '{spec.method}' selector specified in {spec.raw} is "
                 f"invalid. Must be one of [{valid_selectors}]"
             )
             return set()
 
-        extras = self.collect_specified_neighbors(spec, graph, collected)
-        result = self.expand_selection(graph, collected | extras)
+        extras = self.collect_specified_neighbors(spec, collected)
+        result = self.expand_selection(collected | extras)
         return result
 
     def collect_specified_neighbors(
-        self, spec: SelectionCriteria, graph: Graph, selected: Set[UniqueId]
+        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 ("+"/"@").
@@ -107,35 +94,32 @@ def collect_specified_neighbors(
         """
         additional: Set[UniqueId] = set()
         if spec.select_childrens_parents:
-            additional.update(graph.select_childrens_parents(selected))
+            additional.update(self.graph.select_childrens_parents(selected))
+
         if spec.select_parents:
-            additional.update(
-                graph.select_parents(selected, spec.select_parents_max_depth)
-            )
+            depth = spec.select_parents_max_depth
+            additional.update(self.graph.select_parents(selected, depth))
+
         if spec.select_children:
-            additional.update(
-                graph.select_children(selected, spec.select_children_max_depth)
-            )
+            depth = spec.select_children_max_depth
+            additional.update(self.graph.select_children(selected, depth))
         return additional
 
-    def select_nodes(self, graph: Graph, spec: SelectionSpec) -> Set[UniqueId]:
+    def select_nodes(self, spec: SelectionSpec) -> Set[UniqueId]:
         """Select the nodes in the graph according to the spec.
 
         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):
-            result = self.get_nodes_from_criteria(graph, spec)
+            result = self.get_nodes_from_criteria(spec)
         else:
             node_selections = [
-                self.select_nodes(graph, component)
+                self.select_nodes(component)
                 for component in spec
             ]
-            if node_selections:
-                result = spec.combine_selections(node_selections)
-            else:
-                result = set()
+            result = spec.combined(node_selections)
             if spec.expect_exists:
                 alert_non_existence(spec.raw, result)
         return result
@@ -168,16 +152,6 @@ def _is_match(self, unique_id: UniqueId) -> bool:
             )
         return self.node_is_match(node)
 
-    def build_graph_member_subgraph(self) -> Graph:
-        """Build a subgraph of all enabled, non-empty nodes based on the full
-        graph.
-        """
-        graph_members = {
-            unique_id for unique_id in self.full_graph.nodes()
-            if self._is_graph_member(unique_id)
-        }
-        return self.full_graph.subgraph(graph_members)
-
     def filter_selection(self, selected: Set[UniqueId]) -> Set[UniqueId]:
         """Return the subset of selected nodes that is a match for this
         selector.
@@ -186,17 +160,13 @@ def filter_selection(self, selected: Set[UniqueId]) -> Set[UniqueId]:
             unique_id for unique_id in selected if self._is_match(unique_id)
         }
 
-    def expand_selection(
-        self, filtered_graph: Graph, selected: Set[UniqueId]
-    ) -> Set[UniqueId]:
+    def expand_selection(self, selected: Set[UniqueId]) -> Set[UniqueId]:
         """Perform selector-specific expansion."""
         return selected
 
     def get_selected(self, spec: SelectionSpec) -> Set[UniqueId]:
         """get_selected runs trhough the node selection process:
 
-            - build a subgraph containing only non-empty, enabled nodes and
-                enabled sources.
             - node selection. Based on the include/exclude sets, the set
                 of matched unique IDs is returned
                 - expand the graph at each leaf node, before combination
@@ -206,8 +176,7 @@ def get_selected(self, spec: SelectionSpec) -> Set[UniqueId]:
                 - selectors can filter the nodes after all of them have been
                   selected
         """
-        filtered_graph = self.build_graph_member_subgraph()
-        selected_nodes = self.select_nodes(filtered_graph, spec)
+        selected_nodes = self.select_nodes(spec)
         filtered_nodes = self.filter_selection(selected_nodes)
         return filtered_nodes
 
@@ -236,9 +205,3 @@ def __init__(
 
     def node_is_match(self, node):
         return node.resource_type in self.resource_types
-
-
-NodeSelector.register_method(MethodName.FQN, QualifiedNameSelectorMethod)
-NodeSelector.register_method(MethodName.Tag, TagSelectorMethod)
-NodeSelector.register_method(MethodName.Source, SourceSelectorMethod)
-NodeSelector.register_method(MethodName.Path, PathSelectorMethod)