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:/fishtown-analytics/dbt/issues/1030), [#2555](https:/fishtown-analytics/dbt/pull/2555/))
+- Added new node selector methods (`config`, `test_type`, `test_name`, `package`) ([#2425](https:/fishtown-analytics/dbt/issues/2425), [#2629](https:/fishtown-analytics/dbt/pull/2629))
 - Added option to specify profile when connecting to Redshift via IAM ([#2437](https:/fishtown-analytics/dbt/issues/2437), [#2581](https:/fishtown-analytics/dbt/pull/2581))
 
 ### 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)