CT 2251 model yaml frontmatter

.changes/unreleased/Features-20230302-095613.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Support embedded yaml frontmatter in model nodes
+time: 2023-03-02T09:56:13.106933-05:00
+custom:
+  Author: gshank
+  Issue: "7099"

core/dbt/clients/yaml_helper.py

@@ -1,5 +1,6 @@
+import re
 import dbt.exceptions
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 import yaml
 
 # the C version is faster, but it doesn't always exist
@@ -8,6 +9,9 @@
 except ImportError:
     from yaml import Loader, SafeLoader, Dumper  # type: ignore  # noqa: F401
 
+FRONTMATTER_CHECK = ["---\n", "---\r\n"]
+FRONTMATTER_DELIMITER = re.compile(r"^---", re.MULTILINE)
+NON_WHITESPACE = re.compile(r"\S")
 
 YAML_ERROR_MESSAGE = """
 Syntax error near line {line_number}
@@ -61,3 +65,42 @@ def load_yaml_text(contents, path=None):
             error = str(e)
 
         raise dbt.exceptions.DbtValidationError(error)
+
+
+def split_yaml_frontmatter(content: str) -> Tuple[Optional[str], str]:
+    """Splits `content` into raw YAML frontmatter (as a raw string) and everything else proceeding.
+
+    Frontmatter is defined as a block of YAML appearing between two `---` tokens in an otherwise non-YAML document.
+    The frontmatter must be placed at the beginning of the file: anything other than whitespace preceding the first `---`
+    will cause the frontmatter block to be ignored, with a warning.
+    """
+    parts = FRONTMATTER_DELIMITER.split(content, 2)
+    if len(parts) != 3:
+        # Zero or one `---` token, so return the original string
+        return None, content
+    elif NON_WHITESPACE.search(parts[0]) is not None:
+        # No frontmatter section or non-whitespace preceding the first `---`, so skip frontmatter
+        return None, content
+
+    frontmatter_content, after_footer = parts[1:]
+    return frontmatter_content, after_footer
+
+
+def parse_yaml_frontmatter(frontmatter_content: str, original_content: str):
+    try:
+        parsed_yaml = safe_load(frontmatter_content)
+    except (yaml.scanner.ScannerError, yaml.YAMLError) as e:
+        if hasattr(e, "problem_mark"):
+            error = contextualized_yaml_error(original_content, e)
+        else:
+            error = str(e)
+        error = f"Error parsing YAML frontmatter:  {error}"
+        raise dbt.exceptions.DbtValidationError(error)
+
+    return parsed_yaml
+
+
+def has_yaml_frontmatter(content: str) -> bool:
+    """Check first line for yaml frontmatter"""
+
+    return content.startswith(FRONTMATTER_CHECK[0]) or content.startswith(FRONTMATTER_CHECK[1])

core/dbt/contracts/graph/nodes.py

@@ -361,7 +361,8 @@ def patch(self, patch: "ParsedNodePatch"):
         # explicitly pick out the parts to update so we don't inadvertently
         # step on the model name or anything
         # Note: config should already be updated
-        self.patch_path: Optional[str] = patch.file_id
+        if patch.file_id != self.file_id:
+            self.patch_path: Optional[str] = patch.file_id
         # update created_at so process_docs will run in partial parsing
         self.created_at = time.time()
         self.description = patch.description
@@ -487,6 +488,7 @@ class HookNode(CompiledNode):
 class ModelNode(CompiledNode):
     resource_type: NodeType = field(metadata={"restrict": [NodeType.Model]})
     access: AccessType = AccessType.Protected
+    yaml_config_dict: Dict[str, Any] = field(default_factory=dict)
 
 
 # TODO: rm?
@@ -684,6 +686,19 @@ def same_contents(self, other) -> bool:
     def test_node_type(self):
         return "generic"
 
+    def write_node(self, target_path: str, subdirectory: str, payload: str):
+        # If this test came from yaml frontmatter, it will try to write to a directory
+        # with the same name as the model file, which won't work. So change the
+        # directory to model file name + .yaml.
+        file_path = self.original_file_path
+        if file_path.endswith(".sql") or file_path.endswith(".py"):
+            file_path = file_path + ".yaml"
+        path = os.path.join(file_path, self.path)
+        full_path = os.path.join(target_path, subdirectory, self.package_name, path)
+
+        write_file(full_path, payload)
+        return full_path
+
 
 # ====================================
 # Snapshot node

core/dbt/parser/base.py

@@ -393,9 +393,13 @@ def parse_node(self, block: ConfiguredBlockType) -> FinalNode:
             fqn=fqn,
         )
         self.render_update(node, config)
-        result = self.transform(node)
-        self.add_result_node(block, result)
-        return result
+        # Snapshot is converted from IntermediateSnapshotNode in transform
+        final_node = self.transform(node)
+        self.add_result_node(block, final_node)
+        # This handles calling the schema model parser for the model parser.
+        # It needs to be after "add_result_node" has been called.
+        self.post_parse_node(final_node)
+        return final_node
 
     def _update_node_relation_name(self, node: ManifestNode):
         # Seed and Snapshot nodes and Models that are not ephemeral,
@@ -418,6 +422,9 @@ def parse_file(self, file_block: FileBlock) -> None:
     def transform(self, node: IntermediateNode) -> FinalNode:
         pass
 
+    def post_parse_node(self, node):
+        pass
+
 
 class SimpleParser(
     ConfiguredParser[ConfiguredBlockType, FinalNode, FinalNode],

core/dbt/parser/models.py

@@ -8,7 +8,14 @@
 from dbt.node_types import NodeType, ModelLanguage
 from dbt.parser.base import SimpleSQLParser
 from dbt.parser.search import FileBlock
+from dbt.parser.schemas import SchemaParser, TestablePatchParser
+from dbt.parser.generic_test_builders import YamlBlock
 from dbt.clients.jinja import get_rendered
+from dbt.clients.yaml_helper import (
+    has_yaml_frontmatter,
+    split_yaml_frontmatter,
+    parse_yaml_frontmatter,
+)
 import dbt.tracking as tracking
 from dbt import utils
 from dbt_extractor import ExtractionError, py_extract_from_source  # type: ignore
@@ -174,10 +181,44 @@ def verify_python_model_code(node):
 
 class ModelParser(SimpleSQLParser[ModelNode]):
     def parse_from_dict(self, dct, validate=True) -> ModelNode:
+
+        # change raw_code to not have the yaml frontmatter
+        if dct["raw_code"] and has_yaml_frontmatter(dct["raw_code"]):
+            yaml_frontmatter, rest_of_contents = split_yaml_frontmatter(dct["raw_code"])
+            if yaml_frontmatter:
+                yaml_config_dict = parse_yaml_frontmatter(yaml_frontmatter, dct["raw_code"])
+                dct["yaml_config_dict"] = yaml_config_dict
+                dct["raw_code"] = rest_of_contents.strip()
+
         if validate:
             ModelNode.validate(dct)
         return ModelNode.from_dict(dct)
 
+    def post_parse_node(self, node: ModelNode):
+        # This handles parsing yaml frontmatter by calling the
+        # schema model parser
+        if node.yaml_config_dict:
+            # The generic test processing expects to be able to lookup the node, so we add it here.
+            # The node should already have been added to the manifest by "add_result_node"
+            # in the base "parse_node" method.
+            self.manifest.ref_lookup.add_node(node)
+
+            # Create the YamlBlock that the schema model parser needs.
+            source_file = self.manifest.files[node.file_id]
+            block = FileBlock(file=source_file)
+            entry = deepcopy(node.yaml_config_dict)
+            # Make a dictionary like the schema model parser expects.
+            entry["name"] = node.name
+            dct = {"models": [entry]}
+            yaml_block = YamlBlock.from_file_block(block, dct)
+
+            # Create the SchemaParser needed by the "TestablePatchParser"
+            schema_parser = SchemaParser(self.project, self.manifest, self.root_project)
+            model_parser = TestablePatchParser(schema_parser, yaml_block, "models")
+            # parse the schema model patch and create the tests
+            for test_block in model_parser.parse():
+                schema_parser.parse_tests(test_block)
+
     @property
     def resource_type(self) -> NodeType:
         return NodeType.Model

core/dbt/parser/partial.py

@@ -303,7 +303,8 @@ def update_mssat_in_saved(self, new_source_file, old_source_file):
         if self.already_scheduled_for_parsing(old_source_file):
             return
 
-        # These files only have one node except for snapshots
+        # These files only have one node except for snapshots and
+        # model files with yaml frontmatter
         unique_ids = []
         if old_source_file.nodes:
             unique_ids = old_source_file.nodes

core/dbt/parser/schemas.py

@@ -863,8 +863,8 @@ def parse_patch(self, block: TargetBlock[NodeTarget], refs: ParserRef) -> None:
             config=block.target.config,
             access=block.target.access,
         )
-        assert isinstance(self.yaml.file, SchemaSourceFile)
-        source_file: SchemaSourceFile = self.yaml.file
+        # This used to assert SchemaSourceFile, but removed to support yaml frontmatter
+        source_file = self.yaml.file
         if patch.yaml_key in ["models", "seeds", "snapshots"]:
             unique_id = self.manifest.ref_lookup.get_unique_id(patch.name, None)
             if unique_id:
@@ -907,9 +907,11 @@ def parse_patch(self, block: TargetBlock[NodeTarget], refs: ParserRef) -> None:
 
                 # all nodes in the disabled dict have the same unique_id so just grab the first one
                 # to append with the uniqe id
-                source_file.append_patch(patch.yaml_key, found_nodes[0].unique_id)
+                if isinstance(source_file, SchemaSourceFile):
+                    source_file.append_patch(patch.yaml_key, found_nodes[0].unique_id)
                 for node in found_nodes:
-                    node.patch_path = source_file.file_id
+                    if isinstance(source_file, SchemaSourceFile):
+                        node.patch_path = source_file.file_id
                     # re-calculate the node config with the patch config.  Always do this
                     # for the case when no config is set to ensure the default of true gets captured
                     if patch.config:
@@ -929,11 +931,23 @@ def parse_patch(self, block: TargetBlock[NodeTarget], refs: ParserRef) -> None:
         # patches can't be overwritten
         node = self.manifest.nodes.get(unique_id)
         if node:
-            if node.patch_path:
-                package_name, existing_file_path = node.patch_path.split("://")
+            # If patch_path is set, then a schema file has already been processed for this node.
+            # If the node has a yaml_config_dict and THIS invocation of the schema parser is from
+            # a different file, then this is a schema file duplicate of config already done in the model
+            # via yaml frontmatter.
+            if node.patch_path or (
+                hasattr(node, "yaml_config_dict")
+                and node.yaml_config_dict
+                and node.file_id != source_file.file_id
+            ):
+                if node.patch_path:
+                    _, existing_file_path = node.patch_path.split("://")
+                else:  # yaml frontmatter
+                    existing_file_path = node.original_file_path
                 raise DuplicatePatchPathError(patch, existing_file_path)
 
-            source_file.append_patch(patch.yaml_key, node.unique_id)
+            if isinstance(source_file, SchemaSourceFile):
+                source_file.append_patch(patch.yaml_key, node.unique_id)
             # re-calculate the node config with the patch config.  Always do this
             # for the case when no config is set to ensure the default of true gets captured
             if patch.config: