Merge pull request #1785 from fishtown-analytics/feature/hologram-use…

…-restrict-fields

Use the restrict hologram metadata field instead of multiple enum types

core/dbt/contracts/graph/compiled.py

@@ -11,16 +11,7 @@
     TestConfig,
     PARSED_TYPES,
 )
-from dbt.node_types import (
-    NodeType,
-    AnalysisType,
-    ModelType,
-    OperationType,
-    RPCCallType,
-    SeedType,
-    SnapshotType,
-    TestType,
-)
+from dbt.node_types import NodeType
 from dbt.contracts.util import Replaceable
 
 from hologram import JsonSchemaMixin
@@ -72,28 +63,30 @@ def set_cte(self, cte_id: str, sql: str):
 
 @dataclass
 class CompiledAnalysisNode(CompiledNode):
-    resource_type: AnalysisType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Analysis]})
 
 
 @dataclass
 class CompiledHookNode(CompiledNode):
-    resource_type: OperationType
+    resource_type: NodeType = field(
+        metadata={'restrict': [NodeType.Operation]}
+    )
     index: Optional[int] = None
 
 
 @dataclass
 class CompiledModelNode(CompiledNode):
-    resource_type: ModelType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Model]})
 
 
 @dataclass
 class CompiledRPCNode(CompiledNode):
-    resource_type: RPCCallType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.RPCCall]})
 
 
 @dataclass
 class CompiledSeedNode(CompiledNode):
-    resource_type: SeedType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Seed]})
 
     @property
     def empty(self):
@@ -103,12 +96,12 @@ def empty(self):
 
 @dataclass
 class CompiledSnapshotNode(CompiledNode):
-    resource_type: SnapshotType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Snapshot]})
 
 
 @dataclass
 class CompiledTestNode(CompiledNode):
-    resource_type: TestType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Test]})
     column_name: Optional[str] = None
     config: TestConfig = field(default_factory=TestConfig)
 

core/dbt/contracts/graph/parsed.py

@@ -16,17 +16,11 @@
 )
 from dbt.contracts.util import Replaceable
 from dbt.logger import GLOBAL_LOGGER as logger  # noqa
-from dbt.node_types import (
-    NodeType, SourceType, SnapshotType, MacroType, TestType, OperationType,
-    SeedType, ModelType, AnalysisType, RPCCallType
-)
+from dbt.node_types import NodeType
 
 
-class TimestampStrategy(StrEnum):
+class SnapshotStrategy(StrEnum):
     Timestamp = 'timestamp'
-
-
-class CheckStrategy(StrEnum):
     Check = 'check'
 
 
@@ -245,28 +239,30 @@ class ParsedNode(ParsedNodeDefaults, ParsedNodeMixins):
 
 @dataclass
 class ParsedAnalysisNode(ParsedNode):
-    resource_type: AnalysisType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Analysis]})
 
 
 @dataclass
 class ParsedHookNode(ParsedNode):
-    resource_type: OperationType
+    resource_type: NodeType = field(
+        metadata={'restrict': [NodeType.Operation]}
+    )
     index: Optional[int] = None
 
 
 @dataclass
 class ParsedModelNode(ParsedNode):
-    resource_type: ModelType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Model]})
 
 
 @dataclass
 class ParsedRPCNode(ParsedNode):
-    resource_type: RPCCallType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.RPCCall]})
 
 
 @dataclass
 class ParsedSeedNode(ParsedNode):
-    resource_type: SeedType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Seed]})
 
     @property
     def empty(self):
@@ -281,7 +277,7 @@ class TestConfig(NodeConfig):
 
 @dataclass
 class ParsedTestNode(ParsedNode):
-    resource_type: TestType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Test]})
     column_name: Optional[str] = None
     config: TestConfig = field(default_factory=TestConfig)
 
@@ -307,7 +303,7 @@ def __init__(
 
 @dataclass(init=False)
 class GenericSnapshotConfig(_SnapshotConfig):
-    strategy: str
+    strategy: SnapshotStrategy
 
     def __init__(self, strategy: str, **kwargs) -> None:
         self.strategy = strategy
@@ -316,11 +312,13 @@ def __init__(self, strategy: str, **kwargs) -> None:
 
 @dataclass(init=False)
 class TimestampSnapshotConfig(_SnapshotConfig):
-    strategy: TimestampStrategy
+    strategy: SnapshotStrategy = field(metadata={
+        'restrict': [SnapshotStrategy.Timestamp]
+    })
     updated_at: str
 
     def __init__(
-        self, strategy: TimestampStrategy, updated_at: str, **kwargs
+        self, strategy: SnapshotStrategy, updated_at: str, **kwargs
     ) -> None:
         self.strategy = strategy
         self.updated_at = updated_at
@@ -329,7 +327,9 @@ def __init__(
 
 @dataclass(init=False)
 class CheckSnapshotConfig(_SnapshotConfig):
-    strategy: CheckStrategy
+    strategy: SnapshotStrategy = field(metadata={
+        'restrict': [SnapshotStrategy.Check]
+    })
     # TODO: is there a way to get this to accept tuples of strings? Adding
     # `Tuple[str, ...]` to the list of types results in this:
     # ['email'] is valid under each of {'type': 'array', 'items':
@@ -340,7 +340,7 @@ class CheckSnapshotConfig(_SnapshotConfig):
     check_cols: Union[All, List[str]]
 
     def __init__(
-        self, strategy: CheckStrategy, check_cols: Union[All, List[str]],
+        self, strategy: SnapshotStrategy, check_cols: Union[All, List[str]],
         **kwargs
     ) -> None:
         self.strategy = strategy
@@ -356,7 +356,7 @@ class IntermediateSnapshotNode(ParsedNode):
     # defined in config blocks. To fix that, we have an intermediate type that
     # uses a regular node config, which the snapshot parser will then convert
     # into a full ParsedSnapshotNode after rendering.
-    resource_type: SnapshotType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Snapshot]})
 
 
 def _create_if_else_chain(
@@ -384,7 +384,7 @@ def _create_if_else_chain(
 
 @dataclass
 class ParsedSnapshotNode(ParsedNode):
-    resource_type: SnapshotType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Snapshot]})
     config: Union[
         CheckSnapshotConfig,
         TimestampSnapshotConfig,
@@ -397,8 +397,8 @@ def json_schema(cls, embeddable=False):
 
         # mess with config
         configs = [
-            (str(CheckStrategy.Check), CheckSnapshotConfig),
-            (str(TimestampStrategy.Timestamp), TimestampSnapshotConfig),
+            (str(SnapshotStrategy.Check), CheckSnapshotConfig),
+            (str(SnapshotStrategy.Timestamp), TimestampSnapshotConfig),
         ]
 
         if embeddable:
@@ -431,7 +431,7 @@ class MacroDependsOn(JsonSchemaMixin, Replaceable):
 @dataclass
 class ParsedMacro(UnparsedMacro, HasUniqueID):
     name: str
-    resource_type: MacroType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Macro]})
     # TODO: can macros even have tags?
     tags: List[str] = field(default_factory=list)
     # TODO: is this ever populated?
@@ -465,7 +465,7 @@ class ParsedSourceDefinition(
     source_description: str
     loader: str
     identifier: str
-    resource_type: SourceType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Source]})
     quoting: Quoting = field(default_factory=Quoting)
     loaded_at_field: Optional[str] = None
     freshness: Optional[FreshnessThreshold] = None

core/dbt/contracts/graph/unparsed.py

@@ -1,4 +1,4 @@
-from dbt.node_types import UnparsedNodeType, NodeType, OperationType, MacroType
+from dbt.node_types import NodeType
 from dbt.contracts.util import Replaceable, Mergeable
 
 from hologram import JsonSchemaMixin
@@ -28,18 +28,28 @@ def empty(self):
 
 @dataclass
 class UnparsedMacro(UnparsedBaseNode, HasSQL):
-    resource_type: MacroType
+    resource_type: NodeType = field(metadata={'restrict': [NodeType.Macro]})
 
 
 @dataclass
 class UnparsedNode(UnparsedBaseNode, HasSQL):
     name: str
-    resource_type: UnparsedNodeType
+    resource_type: NodeType = field(metadata={'restrict': [
+        NodeType.Model,
+        NodeType.Analysis,
+        NodeType.Test,
+        NodeType.Snapshot,
+        NodeType.Operation,
+        NodeType.Seed,
+        NodeType.RPCCall,
+    ]})
 
 
 @dataclass
 class UnparsedRunHook(UnparsedNode):
-    resource_type: OperationType
+    resource_type: NodeType = field(
+        metadata={'restrict': [NodeType.Operation]}
+    )
     index: Optional[int] = None
 
 

core/dbt/node_types.py

@@ -35,59 +35,6 @@ def refable(cls):
         ]]
 
 
-class UnparsedNodeType(StrEnum):
-    Model = str(NodeType.Model)
-    Analysis = str(NodeType.Analysis)
-    Test = str(NodeType.Test)
-    Snapshot = str(NodeType.Snapshot)
-    Operation = str(NodeType.Operation)
-    Seed = str(NodeType.Seed)
-    RPCCall = str(NodeType.RPCCall)
-
-
 class RunHookType(StrEnum):
     Start = 'on-run-start'
     End = 'on-run-end'
-
-# It would be nice to use hologram.StrLiteral for these, but it results in
-# un-pickleable types :(
-
-
-class AnalysisType(StrEnum):
-    Analysis = str(NodeType.Analysis)
-
-
-class DocumentationType(StrEnum):
-    Documentation = str(NodeType.Documentation)
-
-
-class MacroType(StrEnum):
-    Macro = str(NodeType.Macro)
-
-
-class ModelType(StrEnum):
-    Model = str(NodeType.Model)
-
-
-class OperationType(StrEnum):
-    Operation = str(NodeType.Operation)
-
-
-class RPCCallType(StrEnum):
-    RPCCall = str(NodeType.RPCCall)
-
-
-class SeedType(StrEnum):
-    Seed = str(NodeType.Seed)
-
-
-class SnapshotType(StrEnum):
-    Snapshot = str(NodeType.Snapshot)
-
-
-class SourceType(StrEnum):
-    Source = str(NodeType.Source)
-
-
-class TestType(StrEnum):
-    Test = str(NodeType.Test)

core/dbt/parser/base.py

@@ -21,7 +21,7 @@
     CompilationException, validator_error_message
 )
 from dbt.include.global_project import PROJECT_NAME as GLOBAL_PROJECT_NAME
-from dbt.node_types import NodeType, UnparsedNodeType
+from dbt.node_types import NodeType
 from dbt.source_config import SourceConfig
 from dbt.parser.results import ParseResult, ManifestNodes
 from dbt.parser.search import FileBlock
@@ -236,7 +236,7 @@ def _create_error_node(
         # message reasons
         return UnparsedNode(
             name=name,
-            resource_type=UnparsedNodeType(self.resource_type),
+            resource_type=self.resource_type,
             path=path,
             original_file_path=original_file_path,
             root_path=self.project.project_root,

core/dbt/parser/macros.py

@@ -7,7 +7,7 @@
 from dbt.contracts.graph.parsed import ParsedMacro
 from dbt.exceptions import CompilationException
 from dbt.logger import GLOBAL_LOGGER as logger
-from dbt.node_types import NodeType, MacroType
+from dbt.node_types import NodeType
 from dbt.parser.base import BaseParser
 from dbt.parser.search import FileBlock, FilesystemSearcher
 from dbt.utils import MACRO_PREFIX
@@ -78,7 +78,7 @@ def parse_file(self, block: FileBlock):
             package_name=self.project.project_name,
             raw_sql=source_file.contents,
             root_path=self.project.project_root,
-            resource_type=MacroType(NodeType.Macro),
+            resource_type=NodeType.Macro,
         )
 
         for node in self.parse_unparsed_macros(base_node):

core/dbt/parser/rpc.py

@@ -6,7 +6,7 @@
 from dbt.contracts.graph.parsed import ParsedRPCNode, ParsedMacro
 from dbt.contracts.graph.unparsed import UnparsedMacro
 from dbt.exceptions import InternalException
-from dbt.node_types import NodeType, MacroType
+from dbt.node_types import NodeType
 from dbt.parser.base import SimpleSQLParser
 from dbt.parser.macros import MacroParser
 from dbt.parser.search import FileBlock
@@ -56,7 +56,7 @@ def parse_remote(self, contents) -> Iterable[ParsedMacro]:
             package_name=self.project.project_name,
             raw_sql=contents,
             root_path=self.project.project_root,
-            resource_type=MacroType(NodeType.Macro),
+            resource_type=NodeType.Macro,
         )
         for node in self.parse_unparsed_macros(base):
             yield node

core/dbt/parser/schemas.py

@@ -21,7 +21,7 @@
     validator_error_message, JSONValidationException,
     raise_invalid_schema_yml_version, ValidationException, CompilationException
 )
-from dbt.node_types import NodeType, SourceType
+from dbt.node_types import NodeType
 from dbt.parser.base import SimpleParser
 from dbt.parser.search import FileBlock, FilesystemSearcher
 from dbt.parser.schema_test_builders import (
@@ -355,7 +355,7 @@ def generate_source_node(
             loaded_at_field=loaded_at_field,
             freshness=freshness,
             quoting=quoting,
-            resource_type=SourceType(NodeType.Source),
+            resource_type=NodeType.Source,
             fqn=[self.project.project_name, source.name, table.name],
         )