InProgress: Add tags to SavedQuery

dbt-labs · Nov 13, 2024 · 10794b9 · 10794b9
1 parent aa3d8b5
commit 10794b9
Show file tree

Hide file tree

Showing 6 changed files with 165 additions and 3 deletions.
diff --git a/.changes/unreleased/Features-20241113-110648.yaml b/.changes/unreleased/Features-20241113-110648.yaml
@@ -0,0 +1,6 @@
+kind: Features
+body: Add "tags" to SavedQuery nodes, similar to existing nodes' tags.
+time: 2024-11-13T11:06:48.562566-08:00
+custom:
+    Author: theyostalservice
+    Issue: "369"
diff --git a/dbt_semantic_interfaces/implementations/saved_query.py b/dbt_semantic_interfaces/implementations/saved_query.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import Any, List, Optional, Union
 
-from typing_extensions import override
+from typing_extensions import Self, override
 
 from dbt_semantic_interfaces.implementations.base import (
     HashableBaseModel,
@@ -35,7 +35,11 @@ def _implements_protocol(self) -> SavedQueryQueryParams:
     where: Optional[PydanticWhereFilterIntersection] = None
 
 
-class PydanticSavedQuery(HashableBaseModel, ModelWithMetadataParsing, ProtocolHint[SavedQuery]):
+class PydanticSavedQuery(
+    HashableBaseModel,
+    ModelWithMetadataParsing,
+    ProtocolHint[SavedQuery],
+):
     """Pydantic implementation of SavedQuery."""
 
     @override
@@ -48,3 +52,15 @@ def _implements_protocol(self) -> SavedQuery:
     metadata: Optional[PydanticMetadata] = None
     label: Optional[str] = None
     exports: List[PydanticExport] = Field(default_factory=list)
+    tags: Union[str, List[str]] = Field(
+        default_factory=list,
+    )
+
+    @classmethod
+    def parse_obj(cls, input: Any) -> Self:  # noqa
+        if isinstance(input, dict):
+            if isinstance(input.get("tags"), str):
+                input["tags"] = [input["tags"]]
+            if isinstance(input.get("tags"), list):
+                input["tags"].sort()
+        return super(HashableBaseModel, cls).parse_obj(input)
diff --git a/dbt_semantic_interfaces/parsing/generated_json_schemas/default_explicit_schema.json b/dbt_semantic_interfaces/parsing/generated_json_schemas/default_explicit_schema.json
@@ -705,6 +705,19 @@
                 },
                 "query_params": {
                     "$ref": "#/definitions/saved_query_query_params_schema"
+                },
+                "tags": {
+                    "oneOf": [
+                        {
+                            "type": "string"
+                        },
+                        {
+                            "items": {
+                                "type": "string"
+                            },
+                            "type": "array"
+                        }
+                    ]
                 }
             },
             "required": [

diff --git a/dbt_semantic_interfaces/parsing/schemas.py b/dbt_semantic_interfaces/parsing/schemas.py
@@ -489,6 +489,15 @@
         "query_params": {"$ref": "saved_query_query_params_schema"},
         "label": {"type": "string"},
         "exports": {"type": "array", "items": {"$ref": "export_schema"}},
+        "tags": {
+            "oneOf": [
+                {"type": "string"},
+                {
+                    "type": "array",
+                    "items": {"type": "string"},
+                },
+            ],
+        },
     },
     "required": ["name", "query_params"],
     "additionalProperties": False,

diff --git a/dbt_semantic_interfaces/protocols/saved_query.py b/dbt_semantic_interfaces/protocols/saved_query.py
@@ -73,3 +73,9 @@ def label(self) -> Optional[str]:
     def exports(self) -> Sequence[Export]:
         """Exports that can run using this saved query."""
         pass
+
+    @property
+    @abstractmethod
+    def tags(self) -> Sequence[str]:
+        """List of tags to be used as part of resource selection in dbt."""
+        pass
diff --git a/tests/parsing/test_saved_query_parsing.py b/tests/parsing/test_saved_query_parsing.py
@@ -173,6 +173,118 @@ def test_saved_query_where() -> None:
     assert where == saved_query.query_params.where.where_filters[0].where_sql_template
 
 
+def test_saved_query_with_single_tag_string() -> None:
+    """Test for parsing a single string (not a list) tag in a saved query."""
+    yaml_contents = textwrap.dedent(
+        """\
+        saved_query:
+          name: test_saved_query_group_bys
+          tags: "tag_1"
+          query_params:
+              metrics:
+                - test_metric_a
+        """
+    )
+    file = YamlConfigFile(filepath="test_dir/inline_for_test", contents=yaml_contents)
+
+    build_result = parse_yaml_files_to_semantic_manifest(files=[file, EXAMPLE_PROJECT_CONFIGURATION_YAML_CONFIG_FILE])
+    assert len(build_result.semantic_manifest.saved_queries) == 1
+    saved_query = build_result.semantic_manifest.saved_queries[0]
+    assert saved_query.tags is not None
+    assert len(saved_query.tags) == 1
+    assert saved_query.tags == ["tag_1"]
+
+
+def test_saved_query_with_multiline_list_of_tags() -> None:
+    """Test for parsing a multiline list of tags in a saved query."""
+    yaml_contents = textwrap.dedent(
+        """\
+        saved_query:
+          name: test_saved_query_group_bys
+          tags: ["tag_1", "tag_2"]
+          query_params:
+              metrics:
+                - test_metric_a
+        """
+    )
+    file = YamlConfigFile(filepath="test_dir/inline_for_test", contents=yaml_contents)
+
+    build_result = parse_yaml_files_to_semantic_manifest(files=[file, EXAMPLE_PROJECT_CONFIGURATION_YAML_CONFIG_FILE])
+    assert len(build_result.semantic_manifest.saved_queries) == 1
+    saved_query = build_result.semantic_manifest.saved_queries[0]
+    assert saved_query.tags is not None
+    assert len(saved_query.tags) == 2
+    assert saved_query.tags == ["tag_1", "tag_2"]
+
+
+def test_saved_query_with_single_line_list_of_tags() -> None:
+    """Test for parsing a single-line list of tags in a saved query."""
+    yaml_contents = textwrap.dedent(
+        """\
+        saved_query:
+          name: test_saved_query_group_bys
+          tags:
+            - "tag_1"
+            - "tag_2"
+          query_params:
+              metrics:
+                - test_metric_a
+        """
+    )
+    file = YamlConfigFile(filepath="test_dir/inline_for_test", contents=yaml_contents)
+
+    build_result = parse_yaml_files_to_semantic_manifest(files=[file, EXAMPLE_PROJECT_CONFIGURATION_YAML_CONFIG_FILE])
+    assert len(build_result.semantic_manifest.saved_queries) == 1
+    saved_query = build_result.semantic_manifest.saved_queries[0]
+    assert saved_query.tags is not None
+    assert len(saved_query.tags) == 2
+    assert saved_query.tags == ["tag_1", "tag_2"]
+
+
+def test_saved_query_tags_are_sorted() -> None:
+    """Test tags in a saved query are SORTED after parsing."""
+    yaml_contents = textwrap.dedent(
+        """\
+        saved_query:
+          name: test_saved_query_group_bys
+          tags:
+            - "tag_2"
+            - "tag_1"
+          query_params:
+              metrics:
+                - test_metric_a
+        """
+    )
+    file = YamlConfigFile(filepath="test_dir/inline_for_test", contents=yaml_contents)
+
+    build_result = parse_yaml_files_to_semantic_manifest(files=[file, EXAMPLE_PROJECT_CONFIGURATION_YAML_CONFIG_FILE])
+    assert len(build_result.semantic_manifest.saved_queries) == 1
+    saved_query = build_result.semantic_manifest.saved_queries[0]
+    assert saved_query.tags is not None
+    assert len(saved_query.tags) == 2
+    assert saved_query.tags == ["tag_1", "tag_2"]
+
+
+def test_saved_query_with_no_tags_defaults_to_empty_list() -> None:
+    """Test tags in a saved query will default to empty list if missing."""
+    yaml_contents = textwrap.dedent(
+        """\
+        saved_query:
+          name: test_saved_query_group_bys
+          query_params:
+              metrics:
+                - test_metric_a
+        """
+    )
+    file = YamlConfigFile(filepath="test_dir/inline_for_test", contents=yaml_contents)
+
+    build_result = parse_yaml_files_to_semantic_manifest(files=[file, EXAMPLE_PROJECT_CONFIGURATION_YAML_CONFIG_FILE])
+    assert len(build_result.semantic_manifest.saved_queries) == 1
+    saved_query = build_result.semantic_manifest.saved_queries[0]
+    assert saved_query.tags is not None
+    assert saved_query.tags == []
+
+
 def test_saved_query_exports() -> None:
     """Test for parsing exports referenced in a saved query."""
     yaml_contents = textwrap.dedent(