Repository: dbt-labs/dbt-core

Issue:
serialize inferred primary key in manifest.json
### Housekeeping

- [X] I am a maintainer of dbt-core

### Short description

In #9650, we added a method for models to infer the primary key based on configured tests and/or constraints. We should serialize this as additional metadata in the `manifest.json` artifact

### Acceptance criteria

output of `model.infer_unique_key()` is written to the `manifest.json` artifact for each relevant `node`



### Suggested Tests

We already have unit testing on this method! Additional testing will be required for the artifact itself. 

### Impact to Other Teams

None

### Will backports be required?

I don't think so

### Context / Additional details

The primary key inference method on ModelNode requires a list of `data_tests` that apply to the model, which means data_tests must also be parsed by the point that we are able to populate this field on the `ModelNode` for serialization.


Discussion:

Resolving patch:
diff --git a/core/dbt/artifacts/resources/v1/model.py b/core/dbt/artifacts/resources/v1/model.py
--- a/core/dbt/artifacts/resources/v1/model.py
+++ b/core/dbt/artifacts/resources/v1/model.py
@@ -31,6 +31,7 @@ class Model(CompiledResource):
     latest_version: Optional[NodeVersion] = None
     deprecation_date: Optional[datetime] = None
     defer_relation: Optional[DeferRelation] = None
+    primary_key: List[str] = field(default_factory=list)
 
     def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
         dct = super().__post_serialize__(dct, context)
diff --git a/core/dbt/parser/manifest.py b/core/dbt/parser/manifest.py
--- a/core/dbt/parser/manifest.py
+++ b/core/dbt/parser/manifest.py
@@ -50,6 +50,7 @@
 )
 from dbt.contracts.graph.nodes import (
     Exposure,
+    GenericTestNode,
     Macro,
     ManifestNode,
     Metric,
@@ -466,6 +467,7 @@ def load(self) -> Manifest:
             self.process_docs(self.root_project)
             self.process_metrics(self.root_project)
             self.process_saved_queries(self.root_project)
+            self.process_model_inferred_primary_keys()
             self.check_valid_group_config()
             self.check_valid_access_property()
 
@@ -1149,6 +1151,15 @@ def process_saved_queries(self, config: RuntimeConfig):
             # 2. process `group_by` of SavedQuery for `depends_on``
             _process_metrics_for_node(self.manifest, current_project, saved_query)
 
+    def process_model_inferred_primary_keys(self):
+        """Processes Model nodes to populate their `primary_key`."""
+        for node in self.manifest.nodes.values():
+            if not isinstance(node, ModelNode):
+                continue
+            generic_tests = self._get_generic_tests_for_model(node)
+            primary_key = node.infer_primary_key(generic_tests)
+            node.primary_key = sorted(primary_key)
+
     def update_semantic_model(self, semantic_model) -> None:
         # This has to be done at the end of parsing because the referenced model
         # might have alias/schema/database fields that are updated by yaml config.
@@ -1344,6 +1355,24 @@ def write_perf_info(self, target_path: str):
         write_file(path, json.dumps(self._perf_info, cls=dbt.utils.JSONEncoder, indent=4))
         fire_event(ParsePerfInfoPath(path=path))
 
+    def _get_generic_tests_for_model(
+        self,
+        model: ModelNode,
+    ) -> List[GenericTestNode]:
+        """Return a list of generic tests that are attached to the given model, including disabled tests"""
+        tests = []
+        for _, node in self.manifest.nodes.items():
+            if isinstance(node, GenericTestNode) and node.attached_node == model.unique_id:
+                tests.append(node)
+        for _, nodes in self.manifest.disabled.items():
+            for disabled_node in nodes:
+                if (
+                    isinstance(disabled_node, GenericTestNode)
+                    and disabled_node.attached_node == model.unique_id
+                ):
+                    tests.append(disabled_node)
+        return tests
+
 
 def invalid_target_fail_unless_test(
     node,


Test patch:
diff --git a/tests/functional/artifacts/expected_manifest.py b/tests/functional/artifacts/expected_manifest.py
--- a/tests/functional/artifacts/expected_manifest.py
+++ b/tests/functional/artifacts/expected_manifest.py
@@ -282,6 +282,7 @@ def expected_seeded_manifest(project, model_database=None, quote_model=False):
                 "database": model_database,
                 "alias": "model",
                 "description": "The test model",
+                "primary_key": ["id"],
                 "columns": {
                     "id": {
                         "name": "id",
@@ -374,6 +375,7 @@ def expected_seeded_manifest(project, model_database=None, quote_model=False):
                 "database": project.database,
                 "alias": "second_model",
                 "description": "The second test model",
+                "primary_key": [],
                 "columns": {
                     "id": {
                         "name": "id",
@@ -924,6 +926,7 @@ def expected_references_manifest(project):
                 },
                 "deprecation_date": None,
                 "description": "",
+                "primary_key": [],
                 "docs": {"node_color": None, "show": True},
                 "fqn": ["test", "ephemeral_copy"],
                 "group": None,
@@ -989,6 +992,7 @@ def expected_references_manifest(project):
                 },
                 "deprecation_date": None,
                 "description": "A summmary table of the ephemeral copy of the seed data",
+                "primary_key": [],
                 "docs": {"node_color": None, "show": True},
                 "fqn": ["test", "ephemeral_summary"],
                 "group": "test_group",
@@ -1057,6 +1061,7 @@ def expected_references_manifest(project):
                 },
                 "deprecation_date": None,
                 "description": "A view of the summary of the ephemeral copy of the seed data",
+                "primary_key": [],
                 "docs": {"node_color": None, "show": True},
                 "fqn": ["test", "view_summary"],
                 "group": None,
@@ -1502,6 +1507,7 @@ def expected_versions_manifest(project):
                 "sources": [],
                 "depends_on": {"macros": [], "nodes": []},
                 "description": "A versioned model",
+                "primary_key": ["count", "first_name"],
                 "deprecation_date": ANY,
                 "docs": {"node_color": None, "show": True},
                 "fqn": ["test", "versioned_model", "v1"],
@@ -1572,6 +1578,7 @@ def expected_versions_manifest(project):
                 "sources": [],
                 "depends_on": {"macros": [], "nodes": []},
                 "description": "A versioned model",
+                "primary_key": ["first_name"],
                 "deprecation_date": None,
                 "docs": {"node_color": None, "show": True},
                 "fqn": ["test", "versioned_model", "v2"],
@@ -1625,6 +1632,7 @@ def expected_versions_manifest(project):
                 },
                 "deprecation_date": None,
                 "description": "",
+                "primary_key": [],
                 "docs": {"node_color": None, "show": True},
                 "fqn": ["test", "ref_versioned_model"],
                 "group": None,
diff --git a/tests/functional/primary_keys/fixtures.py b/tests/functional/primary_keys/fixtures.py
new file mode 100644
--- /dev/null
+++ b/tests/functional/primary_keys/fixtures.py
@@ -0,0 +1,88 @@
+simple_model_sql = """
+select 1 as id, 'blue' as color
+"""
+
+simple_model_unique_test = """
+models:
+  - name: simple_model
+    columns:
+      - name: id
+        tests:
+          - unique
+"""
+
+simple_model_disabled_unique_test = """
+models:
+  - name: simple_model
+    columns:
+      - name: id
+        tests:
+          - unique:
+              enabled: false
+
+"""
+
+simple_model_unique_not_null_tests = """
+models:
+  - name: simple_model
+    columns:
+      - name: id
+        tests:
+          - unique
+          - not_null
+"""
+
+simple_model_unique_combo_of_columns = """
+models:
+  - name: simple_model
+    tests:
+      - dbt_utils.unique_combination_of_columns:
+          combination_of_columns: [id, color]
+"""
+
+simple_model_constraints = """
+models:
+  - name: simple_model
+    config:
+      contract:
+        enforced: true
+    columns:
+      - name: id
+        data_type: int
+        constraints:
+          - type: not_null
+          - type: primary_key
+      - name: color
+        data_type: text
+"""
+
+simple_model_two_versions_both_configured = """
+models:
+  - name: simple_model
+    latest_version: 1
+    columns:
+      - name: id
+        tests:
+          - unique
+          - not_null
+    versions:
+      - v: 1
+      - v: 2
+"""
+
+simple_model_two_versions_exclude_col = """
+models:
+  - name: simple_model
+    latest_version: 1
+    columns:
+      - name: id
+        tests:
+          - unique
+          - not_null
+    versions:
+      - v: 1
+      - v: 2
+        columns:
+          - include: all
+            exclude: [id]
+"""
diff --git a/tests/functional/primary_keys/test_primary_keys.py b/tests/functional/primary_keys/test_primary_keys.py
new file mode 100644
--- /dev/null
+++ b/tests/functional/primary_keys/test_primary_keys.py
@@ -0,0 +1,157 @@
+import pytest
+
+from dbt.tests.util import get_manifest, run_dbt
+from tests.functional.primary_keys.fixtures import (
+    simple_model_constraints,
+    simple_model_disabled_unique_test,
+    simple_model_sql,
+    simple_model_two_versions_both_configured,
+    simple_model_two_versions_exclude_col,
+    simple_model_unique_combo_of_columns,
+    simple_model_unique_not_null_tests,
+    simple_model_unique_test,
+)
+
+
+class TestSimpleModelNoYml:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model.sql": simple_model_sql,
+        }
+
+    def test_simple_model_no_yml(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node = manifest.nodes["model.test.simple_model"]
+        assert node.primary_key == []
+
+
+class TestSimpleModelConstraints:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model.sql": simple_model_sql,
+            "schema.yml": simple_model_constraints,
+        }
+
+    def test_simple_model_constraints(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node = manifest.nodes["model.test.simple_model"]
+        assert node.primary_key == ["id"]
+
+
+class TestSimpleModelUniqueNotNullTests:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model.sql": simple_model_sql,
+            "schema.yml": simple_model_unique_not_null_tests,
+        }
+
+    def test_simple_model_unique_not_null_tests(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node = manifest.nodes["model.test.simple_model"]
+        assert node.primary_key == ["id"]
+
+
+class TestSimpleModelUniqueTests:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model.sql": simple_model_sql,
+            "schema.yml": simple_model_unique_test,
+        }
+
+    def test_simple_model_unique_test(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node = manifest.nodes["model.test.simple_model"]
+        assert node.primary_key == ["id"]
+
+
+class TestSimpleModelDisabledUniqueTests:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model.sql": simple_model_sql,
+            "schema.yml": simple_model_disabled_unique_test,
+        }
+
+    def test_simple_model_disabled_unique_test(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node = manifest.nodes["model.test.simple_model"]
+        assert node.primary_key == ["id"]
+
+
+class TestVersionedSimpleModel:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model_v1.sql": simple_model_sql,
+            "simple_model_v2.sql": simple_model_sql,
+            "schema.yml": simple_model_two_versions_both_configured,
+        }
+
+    def test_versioned_simple_model(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node_v1 = manifest.nodes["model.test.simple_model.v1"]
+        node_v2 = manifest.nodes["model.test.simple_model.v2"]
+        assert node_v1.primary_key == ["id"]
+        assert node_v2.primary_key == ["id"]
+
+
+class TestVersionedSimpleModelExcludeTests:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model_v1.sql": simple_model_sql,
+            "simple_model_v2.sql": simple_model_sql,
+            "schema.yml": simple_model_two_versions_exclude_col,
+        }
+
+    def test_versioned_simple_model_exclude_col(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node_v1 = manifest.nodes["model.test.simple_model.v1"]
+        node_v2 = manifest.nodes["model.test.simple_model.v2"]
+        assert node_v1.primary_key == ["id"]
+        assert node_v2.primary_key == []
+
+
+class TestSimpleModelCombinationOfColumns:
+    @pytest.fixture(scope="class")
+    def packages(self):
+        return {
+            "packages": [
+                {
+                    "git": "https://github.com/dbt-labs/dbt-utils.git",
+                    "revision": "1.1.0",
+                },
+            ]
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_model.sql": simple_model_sql,
+            "schema.yml": simple_model_unique_combo_of_columns,
+        }
+
+    def test_versioned_simple_combo_of_columns(self, project):
+        run_dbt(["deps"])
+        run_dbt(["run"])
+        manifest = get_manifest(project.project_root)
+        node = manifest.nodes["model.test.simple_model"]
+        assert node.primary_key == ["color", "id"]
diff --git a/tests/unit/contracts/graph/test_manifest.py b/tests/unit/contracts/graph/test_manifest.py
--- a/tests/unit/contracts/graph/test_manifest.py
+++ b/tests/unit/contracts/graph/test_manifest.py
@@ -74,6 +74,7 @@
         "raw_code",
         "language",
         "description",
+        "primary_key",
         "columns",
         "fqn",
         "build_path",
diff --git a/tests/unit/fixtures.py b/tests/unit/fixtures.py
--- a/tests/unit/fixtures.py
+++ b/tests/unit/fixtures.py
@@ -26,6 +26,7 @@ def model_node():
         metrics=[],
         depends_on=DependsOn(),
         description="",
+        primary_key=[],
         database="test_db",
         schema="test_schema",
         alias="bar",
diff --git a/tests/unit/test_contracts_graph_compiled.py b/tests/unit/test_contracts_graph_compiled.py
--- a/tests/unit/test_contracts_graph_compiled.py
+++ b/tests/unit/test_contracts_graph_compiled.py
@@ -149,6 +149,7 @@ def basic_compiled_dict():
         "depends_on": {"macros": [], "nodes": []},
         "database": "test_db",
         "description": "",
+        "primary_key": [],
         "schema": "test_schema",
         "alias": "bar",
         "tags": [],
diff --git a/tests/unit/test_contracts_graph_parsed.py b/tests/unit/test_contracts_graph_parsed.py
--- a/tests/unit/test_contracts_graph_parsed.py
+++ b/tests/unit/test_contracts_graph_parsed.py
@@ -167,6 +167,7 @@ def base_parsed_model_dict():
         "depends_on": {"macros": [], "nodes": []},
         "database": "test_db",
         "description": "",
+        "primary_key": [],
         "schema": "test_schema",
         "alias": "bar",
         "tags": [],
@@ -220,6 +221,7 @@ def basic_parsed_model_object():
         metrics=[],
         depends_on=DependsOn(),
         description="",
+        primary_key=[],
         database="test_db",
         schema="test_schema",
         alias="bar",
@@ -274,6 +276,7 @@ def complex_parsed_model_dict():
         "depends_on": {"macros": [], "nodes": ["model.test.bar"]},
         "database": "test_db",
         "description": "My parsed node",
+        "primary_key": [],
         "schema": "test_schema",
         "alias": "bar",
         "tags": ["tag"],