Skip to content

Commit

Permalink
Fix #9005: Allow singular tests to be documented
Browse files Browse the repository at this point in the history
  • Loading branch information
aranke committed Sep 19, 2024
1 parent 3308a43 commit 38aa9d5
Show file tree
Hide file tree
Showing 5 changed files with 111 additions and 2 deletions.
50 changes: 49 additions & 1 deletion core/dbt/contracts/graph/manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
SavedQuery,
SeedNode,
SemanticModel,
SingularTestNode,
SourceDefinition,
UnitTestDefinition,
UnitTestFileFixture,
Expand Down Expand Up @@ -89,7 +90,7 @@
RefName = str


def find_unique_id_for_package(storage, key, package: Optional[PackageName]):
def find_unique_id_for_package(storage, key, package: Optional[PackageName]) -> Optional[UniqueID]:
if key not in storage:
return None

Expand Down Expand Up @@ -470,6 +471,43 @@ class AnalysisLookup(RefableLookup):
_versioned_types: ClassVar[set] = set()


class SingularTestLookup(dbtClassMixin):
def __init__(self, manifest: "Manifest") -> None:
self.storage: Dict[str, Dict[PackageName, UniqueID]] = {}
self.populate(manifest)

Check warning on line 477 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L476-L477

Added lines #L476 - L477 were not covered by tests

def get_unique_id(self, search_name, package: Optional[PackageName]) -> Optional[UniqueID]:
return find_unique_id_for_package(self.storage, search_name, package)

Check warning on line 480 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L480

Added line #L480 was not covered by tests

def find(
self, search_name, package: Optional[PackageName], manifest: "Manifest"
) -> Optional[SingularTestNode]:
unique_id = self.get_unique_id(search_name, package)
if unique_id is not None:
return self.perform_lookup(unique_id, manifest)
return None

Check warning on line 488 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L485-L488

Added lines #L485 - L488 were not covered by tests

def add_singular_test(self, source: SingularTestNode) -> None:
if source.search_name not in self.storage:
self.storage[source.search_name] = {}

Check warning on line 492 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L491-L492

Added lines #L491 - L492 were not covered by tests

self.storage[source.search_name][source.package_name] = source.unique_id

Check warning on line 494 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L494

Added line #L494 was not covered by tests

def populate(self, manifest: "Manifest") -> None:
for node in manifest.nodes.values():
if isinstance(node, SingularTestNode):
self.add_singular_test(node)

Check warning on line 499 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L497-L499

Added lines #L497 - L499 were not covered by tests

def perform_lookup(self, unique_id: UniqueID, manifest: "Manifest") -> SingularTestNode:
if unique_id not in manifest.nodes:
raise dbt_common.exceptions.DbtInternalError(

Check warning on line 503 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L502-L503

Added lines #L502 - L503 were not covered by tests
f"Singular test {unique_id} found in cache but not found in manifest"
)
node = manifest.nodes[unique_id]
assert isinstance(node, SingularTestNode)
return node

Check warning on line 508 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L506-L508

Added lines #L506 - L508 were not covered by tests


def _packages_to_search(
current_project: str,
node_package: str,
Expand Down Expand Up @@ -869,6 +907,9 @@ class Manifest(MacroMethods, dbtClassMixin):
_analysis_lookup: Optional[AnalysisLookup] = field(
default=None, metadata={"serialize": lambda x: None, "deserialize": lambda x: None}
)
_singular_test_lookup: Optional[SingularTestLookup] = field(
default=None, metadata={"serialize": lambda x: None, "deserialize": lambda x: None}
)
_parsing_info: ParsingInfo = field(
default_factory=ParsingInfo,
metadata={"serialize": lambda x: None, "deserialize": lambda x: None},
Expand Down Expand Up @@ -1264,6 +1305,12 @@ def analysis_lookup(self) -> AnalysisLookup:
self._analysis_lookup = AnalysisLookup(self)
return self._analysis_lookup

@property
def singular_test_lookup(self) -> SingularTestLookup:
if self._singular_test_lookup is None:
self._singular_test_lookup = SingularTestLookup(self)
return self._singular_test_lookup

Check warning on line 1312 in core/dbt/contracts/graph/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/contracts/graph/manifest.py#L1310-L1312

Added lines #L1310 - L1312 were not covered by tests

@property
def external_node_unique_ids(self):
return [node.unique_id for node in self.nodes.values() if node.is_external_node]
Expand Down Expand Up @@ -1708,6 +1755,7 @@ def __reduce_ex__(self, protocol):
self._semantic_model_by_measure_lookup,
self._disabled_lookup,
self._analysis_lookup,
self._singular_test_lookup,
)
return self.__class__, args

Expand Down
5 changes: 5 additions & 0 deletions core/dbt/contracts/graph/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1627,6 +1627,11 @@ class ParsedMacroPatch(ParsedPatch):
arguments: List[MacroArgument] = field(default_factory=list)


@dataclass
class ParsedSingularTestPatch(ParsedPatch):
pass


# ====================================
# Node unions/categories
# ====================================
Expand Down
5 changes: 5 additions & 0 deletions core/dbt/contracts/graph/unparsed.py
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,11 @@ class UnparsedAnalysisUpdate(HasConfig, HasColumnDocs, HasColumnProps, HasYamlMe
access: Optional[str] = None


@dataclass
class UnparsedSingularTestUpdate(HasConfig, HasColumnProps, HasYamlMetadata):
pass


@dataclass
class UnparsedNodeUpdate(HasConfig, HasColumnTests, HasColumnAndTestProps, HasYamlMetadata):
quote_columns: Optional[bool] = None
Expand Down
2 changes: 2 additions & 0 deletions core/dbt/parser/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
UnparsedMacroUpdate,
UnparsedModelUpdate,
UnparsedNodeUpdate,
UnparsedSingularTestUpdate,
)
from dbt.exceptions import ParsingError
from dbt.parser.search import FileBlock
Expand All @@ -38,6 +39,7 @@ def trimmed(inp: str) -> str:
UnpatchedSourceDefinition,
UnparsedExposure,
UnparsedModelUpdate,
UnparsedSingularTestUpdate,
)


Expand Down
51 changes: 50 additions & 1 deletion core/dbt/parser/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
ModelNode,
ParsedMacroPatch,
ParsedNodePatch,
ParsedSingularTestPatch,
UnpatchedSourceDefinition,
)
from dbt.contracts.graph.unparsed import (
Expand All @@ -27,6 +28,7 @@
UnparsedMacroUpdate,
UnparsedModelUpdate,
UnparsedNodeUpdate,
UnparsedSingularTestUpdate,
UnparsedSourceDefinition,
)
from dbt.events.types import (
Expand Down Expand Up @@ -221,6 +223,10 @@ def parse_file(self, block: FileBlock, dct: Optional[Dict] = None) -> None:
parser = MacroPatchParser(self, yaml_block, "macros")
parser.parse()

if "data_tests" in dct:
parser = SingularTestPatchParser(self, yaml_block, "data_tests")
parser.parse()

Check warning on line 228 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L227-L228

Added lines #L227 - L228 were not covered by tests

# PatchParser.parse() (but never test_blocks)
if "analyses" in dct:
parser = AnalysisPatchParser(self, yaml_block, "analyses")
Expand Down Expand Up @@ -266,14 +272,17 @@ def parse_file(self, block: FileBlock, dct: Optional[Dict] = None) -> None:
saved_query_parser.parse()


Parsed = TypeVar("Parsed", UnpatchedSourceDefinition, ParsedNodePatch, ParsedMacroPatch)
Parsed = TypeVar(
"Parsed", UnpatchedSourceDefinition, ParsedNodePatch, ParsedMacroPatch, ParsedSingularTestPatch
)
NodeTarget = TypeVar("NodeTarget", UnparsedNodeUpdate, UnparsedAnalysisUpdate, UnparsedModelUpdate)
NonSourceTarget = TypeVar(
"NonSourceTarget",
UnparsedNodeUpdate,
UnparsedAnalysisUpdate,
UnparsedMacroUpdate,
UnparsedModelUpdate,
UnparsedSingularTestUpdate,
)


Expand Down Expand Up @@ -1055,6 +1064,46 @@ def _target_type(self) -> Type[UnparsedAnalysisUpdate]:
return UnparsedAnalysisUpdate


class SingularTestPatchParser(PatchParser[UnparsedSingularTestUpdate, ParsedSingularTestPatch]):
def get_block(self, node: UnparsedSingularTestUpdate) -> TargetBlock:
return TargetBlock.from_yaml_block(self.yaml, node)

Check warning on line 1069 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1069

Added line #L1069 was not covered by tests

def _target_type(self) -> Type[UnparsedSingularTestUpdate]:
return UnparsedSingularTestUpdate

Check warning on line 1072 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1072

Added line #L1072 was not covered by tests

def parse_patch(self, block: TargetBlock[UnparsedSingularTestUpdate], refs: ParserRef) -> None:
patch = ParsedSingularTestPatch(

Check warning on line 1075 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1075

Added line #L1075 was not covered by tests
name=block.target.name,
description=block.target.description,
meta=block.target.meta,
docs=block.target.docs,
config=block.target.config,
original_file_path=block.target.original_file_path,
yaml_key=block.target.yaml_key,
package_name=block.target.package_name,
)

assert isinstance(self.yaml.file, SchemaSourceFile)
source_file: SchemaSourceFile = self.yaml.file

Check warning on line 1087 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1086-L1087

Added lines #L1086 - L1087 were not covered by tests

unique_id = self.manifest.singular_test_lookup.get_unique_id(

Check warning on line 1089 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1089

Added line #L1089 was not covered by tests
block.name, block.target.package_name
)
assert unique_id is not None
node = self.manifest.nodes.get(unique_id)
assert node is not None

Check warning on line 1094 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1092-L1094

Added lines #L1092 - L1094 were not covered by tests

source_file.append_patch(patch.yaml_key, unique_id)
if patch.config:
self.patch_node_config(node, patch)

Check warning on line 1098 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1096-L1098

Added lines #L1096 - L1098 were not covered by tests

node.patch_path = patch.file_id
node.description = patch.description
node.created_at = time.time()
node.meta = patch.meta
node.docs = patch.docs

Check warning on line 1104 in core/dbt/parser/schemas.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/schemas.py#L1100-L1104

Added lines #L1100 - L1104 were not covered by tests


class MacroPatchParser(PatchParser[UnparsedMacroUpdate, ParsedMacroPatch]):
def get_block(self, node: UnparsedMacroUpdate) -> TargetBlock:
return TargetBlock.from_yaml_block(self.yaml, node)
Expand Down

0 comments on commit 38aa9d5

Please sign in to comment.