Repository: dbt-checkpoint/dbt-checkpoint (language: python)

Issue:
feature: `create a check-model-has-labels-keys` for dbt with BigQuery
**Describe the feature you'd like**
dbt-checkpoint is great and the `check-model-has-meta-keys` is optimal to enforce Data Governance labels that we need in a big dbt project such as one with a data mesh architecture. With BigQuery, dbt gives us the possibility to assign also a list of keys to a dictionary called `labels` and this keys will be passed to BigQuery and they also will be assigned to the table in the Data Warehouse.

The idea here is to create a hook to enforce the labels keys to be filled, instead of the meta keys.

**Additional context**
I've tried to implement myself the hook by changing all the occurrences of "meta" in the `check-model-has-meta-keys.py` file for "labels", but it obviously didn't work as expected. After adding it to the `setup.cfg` and `.pre-commit-hooks.yaml` files, I encounter the following error:

![image](https://github.com/dbt-checkpoint/dbt-checkpoint/assets/68338969/6a7216a2-8a65-4091-8194-e1280f0fea02)

So the conclusion is that I would like help to implement that, if it is possible somehow. I think it is a good feature for BigQuery + dbt users.
feature: `create a check-model-has-labels-keys` for dbt with BigQuery
**Describe the feature you'd like**
dbt-checkpoint is great and the `check-model-has-meta-keys` is optimal to enforce Data Governance labels that we need in a big dbt project such as one with a data mesh architecture. With BigQuery, dbt gives us the possibility to assign also a list of keys to a dictionary called `labels` and this keys will be passed to BigQuery and they also will be assigned to the table in the Data Warehouse.

The idea here is to create a hook to enforce the labels keys to be filled, instead of the meta keys.

**Additional context**
I've tried to implement myself the hook by changing all the occurrences of "meta" in the `check-model-has-meta-keys.py` file for "labels", but it obviously didn't work as expected. After adding it to the `setup.cfg` and `.pre-commit-hooks.yaml` files, I encounter the following error:

![image](https://github.com/dbt-checkpoint/dbt-checkpoint/assets/68338969/6a7216a2-8a65-4091-8194-e1280f0fea02)

So the conclusion is that I would like help to implement that, if it is possible somehow. I think it is a good feature for BigQuery + dbt users.


Discussion / hints:

Fix patch:
diff --git a/.pre-commit-hooks.yaml b/.pre-commit-hooks.yaml
index b9ee299f..a3b0ca85 100644
--- a/.pre-commit-hooks.yaml
+++ b/.pre-commit-hooks.yaml
@@ -51,6 +51,13 @@
   language: python
   types_or: [yaml, sql]
   require_serial: true # because we need to process YAML and SQL
+- id: check-model-has-labels-keys
+  name: Check the model has keys in the labels part
+  description: Ensures that the model has a list of valid labels keys.
+  entry: check-model-has-labels-keys
+  language: python
+  types_or: [yaml, sql]
+  require_serial: true # because we need to process YAML and SQL
 - id: check-model-has-properties-file
   name: Check the model has properties file
   description: Ensures that the model has properties file (schema file).
@@ -169,6 +176,12 @@
   entry: check-source-has-meta-keys
   language: python
   types_or: [yaml]
+- id: check-source-has-labels-keys
+  name: Check the source has keys in the labels part
+  description: Ensures that the source has a list of valid labels keys.
+  entry: check-source-has-labels-keys
+  language: python
+  types_or: [yaml]
 - id: check-source-has-tests-by-name
   name: Check the source tests by test name
   description: Ensures that the source has a number of tests of a certain name (e.g. data, unique).
diff --git a/HOOKS.md b/HOOKS.md
index 758e45cb..185dfcbf 100644
--- a/HOOKS.md
+++ b/HOOKS.md
@@ -12,6 +12,7 @@
 - [`check-model-has-all-columns`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-all-columns): Check the model has all columns in the properties file.
 - [`check-model-has-description`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-description): Check the model has description.
 - [`check-model-has-meta-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-meta-keys): Check the model has keys in the meta part.
+- [`check-model-has-labels-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-labels-keys): Check the model has keys in the labels part.
 - [`check-model-has-properties-file`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-properties-file): Check the model has properties file.
 - [`check-model-has-tests-by-name`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-tests-by-name): Check the model has a number of tests by test name.
 - [`check-model-has-tests-by-type`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-tests-by-type): Check the model has a number of tests by test type.
@@ -37,6 +38,7 @@
 - [`check-source-has-freshness`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-freshness): Check the source has the freshness.
 - [`check-source-has-loader`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-loader): Check the source has loader option.
 - [`check-source-has-meta-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-meta-keys): Check the source has keys in the meta part.
+- [`check-source-has-labels-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-labels-keys): Check the source has keys in the labels part.
 - [`check-source-has-tests-by-name`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-tests-by-name): Check the source has a number of tests by test name.
 - [`check-source-has-tests-by-type`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-tests-by-type): Check the source has a number of tests by test type.
 - [`check-source-has-tests`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-tests): Check the source has a number of tests.
@@ -374,6 +376,59 @@ If you `run` your model and then you delete meta keys from a properties file, th
 
 ---
 
+### `check-model-has-labels-keys`
+
+Ensures that the model has a list of valid labels keys. (usually `schema.yml`).
+
+By default, it does not allow the model to have any other labels keys other than the ones required. An optional argument can be used to allow for extra keys.
+
+#### Arguments
+
+`--manifest`: location of `manifest.json` file. Usually `target/manifest.json`. This file contains a full representation of dbt project. **Default: `target/manifest.json`**<br/>
+`--labels-keys`: list of the required keys in the labels part of the model.<br/>
+`--allow-extra-keys`: whether extra keys are allowed. **Default: `False`**.
+
+#### Example
+
+```
+repos:
+- repo: https://github.com/dbt-checkpoint/dbt-checkpoint
+ rev: v1.0.0
+ hooks:
+ - id: check-model-has-labels-keys
+   args: ['--labels-keys', 'foo', 'bar', "--"]
+```
+
+:warning: do not forget to include `--` as the last argument. Otherwise `pre-commit` would not be able to separate a list of files with args.
+
+#### When to use it
+
+If every model needs to have certain labels keys.
+
+#### Requirements
+
+| Model exists in `manifest.json` <sup id="a1">[1](#f1)</sup> | Model exists in `catalog.json` <sup id="a2">[2](#f2)</sup> |
+| :---------------------------------------------------------: | :--------------------------------------------------------: |
+|   :x: Not needed since it also validates properties files   |                       :x: Not needed                       |
+
+<sup id="f1">1</sup> It means that you need to run `dbt run`, `dbt compile` before run this hook.<br/>
+<sup id="f2">2</sup> It means that you need to run `dbt docs generate` before run this hook.
+
+#### How it works
+
+- Hook takes all changed `yml` and `SQL` files.
+- The model name is obtained from the `SQL` file name.
+- The manifest is scanned for a model.
+- Modified `yml` files are scanned for a model.
+- If any model (from a manifest or `yml` files) does not have specified labels keys, the hook fails.
+- The labels keys must be in either the yml file **or** the manifest.
+
+#### Known limitations
+
+If you `run` your model and then you delete labels keys from a properties file, the hook success since the labels keys is still present in `manifest.json`.
+
+---
+
 ### `check-model-has-properties-file`
 
 Ensures that the model has a properties file (schema file).
@@ -1168,6 +1223,48 @@ If every source needs to have certain meta keys.
 
 ---
 
+### `check-source-has-labels-keys`
+
+Ensures that the source has a list of valid labels keys. (usually `schema.yml`).
+
+#### Arguments
+
+`--labels-keys`: list of the required keys in the labels part of the model.
+
+#### Example
+
+```
+repos:
+- repo: https://github.com/dbt-checkpoint/dbt-checkpoint
+ rev: v1.0.0
+ hooks:
+ - id: check-source-has-labels-keys
+   args: ['--labels-keys', 'foo', 'bar', "--"]
+```
+
+:warning: do not forget to include `--` as the last argument. Otherwise `pre-commit` would not be able to separate a list of files with args.
+
+#### When to use it
+
+If every source needs to have certain labels keys.
+
+#### Requirements
+
+| Model exists in `manifest.json` <sup id="a1">[1](#f1)</sup> | Model exists in `catalog.json` <sup id="a2">[2](#f2)</sup> |
+| :---------------------------------------------------------: | :--------------------------------------------------------: |
+|   :x: Not needed since it also validates properties files   |                       :x: Not needed                       |
+
+<sup id="f1">1</sup> It means that you need to run `dbt run`, `dbt compile` before run this hook.<br/>
+<sup id="f2">2</sup> It means that you need to run `dbt docs generate` before run this hook.
+
+#### How it works
+
+- Hook takes all changed `yml`.
+- All sources from yml file are parsed.
+- If the source does not have the required labels keys set, the hook fails.
+
+---
+
 ### `check-source-has-tests-by-name`
 
 Ensures that the source has a number of tests of a certain name (e.g. data, unique).
diff --git a/README.md b/README.md
index 80b766c0..12a21adb 100644
--- a/README.md
+++ b/README.md
@@ -82,6 +82,7 @@ Since the root-level `exclude` statement is handled by pre-commit, when those ho
 - [`check-model-has-all-columns`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-all-columns): Check the model has all columns in the properties file.
 - [`check-model-has-description`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-description): Check the model has description.
 - [`check-model-has-meta-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-meta-keys): Check the model has keys in the meta part.
+- [`check-model-has-labels-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-labels-keys): Check the model has keys in the labels part.
 - [`check-model-has-properties-file`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-properties-file): Check the model has properties file.
 - [`check-model-has-tests-by-name`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-tests-by-name): Check the model has a number of tests by test name.
 - [`check-model-has-tests-by-type`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-model-has-tests-by-type): Check the model has a number of tests by test type.
@@ -107,6 +108,7 @@ Since the root-level `exclude` statement is handled by pre-commit, when those ho
 - [`check-source-has-freshness`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-freshness): Check the source has the freshness.
 - [`check-source-has-loader`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-loader): Check the source has loader option.
 - [`check-source-has-meta-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-meta-keys): Check the source has keys in the meta part.
+- [`check-source-has-labels-keys`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-labels-keys): Check the source has keys in the labels part.
 - [`check-source-has-tests-by-group`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-tests-by-group): Check the source has a number of tests from a group of tests.
 - [`check-source-has-tests-by-name`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-tests-by-name): Check the source has a number of tests by test name.
 - [`check-source-has-tests-by-type`](https://github.com/dbt-checkpoint/dbt-checkpoint/blob/main/HOOKS.md#check-source-has-tests-by-type): Check the source has a number of tests by test type.
diff --git a/dbt_checkpoint/check_model_has_labels_keys.py b/dbt_checkpoint/check_model_has_labels_keys.py
new file mode 100644
index 00000000..ccdf3c3c
--- /dev/null
+++ b/dbt_checkpoint/check_model_has_labels_keys.py
@@ -0,0 +1,134 @@
+import argparse
+import os
+import time
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import Optional
+from typing import Sequence
+
+from dbt_checkpoint.tracking import dbtCheckpointTracking
+from dbt_checkpoint.utils import add_default_args
+from dbt_checkpoint.utils import get_dbt_manifest
+from dbt_checkpoint.utils import get_filenames
+from dbt_checkpoint.utils import get_model_schemas
+from dbt_checkpoint.utils import get_model_sqls
+from dbt_checkpoint.utils import get_models
+from dbt_checkpoint.utils import JsonOpenError
+
+
+def validate_keys(
+    actual: Iterable[str], expected: Iterable[str], allow_extra_keys: bool
+) -> bool:
+    actual = set(actual)
+    expected = set(expected)
+
+    if allow_extra_keys:
+        return expected.issubset(actual)
+    else:
+        return expected == actual
+
+
+def has_labels_key(
+    paths: Sequence[str],
+    manifest: Dict[str, Any],
+    labels_keys: Sequence[str],
+    allow_extra_keys: bool,
+) -> int:
+    status_code = 0
+    ymls = get_filenames(paths, [".yml", ".yaml"])
+    sqls = get_model_sqls(paths, manifest)
+    filenames = set(sqls.keys())
+    models = get_models(manifest, filenames)
+    schemas = get_model_schemas(list(ymls.values()), filenames)
+
+    in_models = set()
+    for model in models:
+        model_config = model.node.get("config")
+        if model_config is not None:
+            model_labels_data = model_config.get("labels")
+        else:
+            model_labels_data = None
+        model_labels = set(model_labels_data.keys()) if model_labels_data else set()
+        
+        if validate_keys(model_labels, labels_keys, allow_extra_keys):
+            in_models.add(model.filename)
+
+
+    in_schemas = set()
+    for schema in schemas:
+        schema_config = schema.schema.get("config")
+        schema_labels_data = schema_config.get("labels")
+        schema_labels = set(schema_labels_data.keys()) if schema_labels_data else set()
+        
+        if validate_keys(schema_labels, labels_keys, allow_extra_keys):
+            in_schemas.add(schema.model_name)
+
+    missing = filenames.difference(in_models, in_schemas)
+
+    for model in missing:
+        status_code = 1
+        result = "\n- ".join(list(labels_keys))
+        print(
+            f"{sqls.get(model)}: "
+            f"does not have some of the labels keys defined:\n- {result}",
+        )
+
+    return status_code
+
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser()
+    add_default_args(parser)
+
+    parser.add_argument(
+        "--labels-keys",
+        nargs="+",
+        required=True,
+        help="List of required key in labels part of model.",
+    )
+
+    parser.add_argument(
+        "--allow-extra-keys",
+        action="store_true",
+        required=False,
+        help="Whether extra keys are allowed.",
+    )
+
+    args = parser.parse_args(argv)
+
+    try:
+        manifest = get_dbt_manifest(args)
+    except JsonOpenError as e:
+        print(f"Unable to load manifest file ({e})")
+        return 1
+
+    start_time = time.time()
+    status_code = has_labels_key(
+        paths=args.filenames,
+        manifest=manifest,
+        labels_keys=args.labels_keys,
+        allow_extra_keys=args.allow_extra_keys,
+    )
+    end_time = time.time()
+    script_args = vars(args)
+
+    tracker = dbtCheckpointTracking(script_args=script_args)
+    tracker.track_hook_event(
+        event_name="Hook Executed",
+        manifest=manifest,
+        event_properties={
+            "hook_name": os.path.basename(__file__),
+            "description": "Check model has labels keys",
+            "status": status_code,
+            "execution_time": end_time - start_time,
+            "is_pytest": script_args.get("is_test"),
+        },
+    )
+
+    return status_code
+
+
+if __name__ == "__main__":
+    exit(main())
diff --git a/dbt_checkpoint/check_source_has_labels_keys.py b/dbt_checkpoint/check_source_has_labels_keys.py
new file mode 100644
index 00000000..3e0753f0
--- /dev/null
+++ b/dbt_checkpoint/check_source_has_labels_keys.py
@@ -0,0 +1,81 @@
+import argparse
+import os
+import time
+from pathlib import Path
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Sequence
+
+from dbt_checkpoint.tracking import dbtCheckpointTracking
+from dbt_checkpoint.utils import add_default_args
+from dbt_checkpoint.utils import get_dbt_manifest
+from dbt_checkpoint.utils import get_source_schemas
+from dbt_checkpoint.utils import JsonOpenError
+
+
+def has_labels_key(paths: Sequence[str], labels_keys: Sequence[str]) -> Dict[str, Any]:
+    status_code = 0
+    ymls = [Path(path) for path in paths]
+
+    # if user added schema but did not rerun
+    schemas = get_source_schemas(ymls)
+
+    for schema in schemas:
+        schema_labels = set(schema.source_schema.get("labels", {}).keys())
+        table_labels = set(schema.table_schema.get("labels", {}).keys())
+        all_labels = schema_labels.union(table_labels)
+        diff_labels = set(labels_keys).difference(all_labels)
+
+        if diff_labels:
+            status_code = 1
+            result = "\n- ".join(list(diff_labels))
+            print(
+                f"{schema.source_name}.{schema.table_name}: "
+                f"does not have some of the labels keys defined:\n- {result}",
+            )
+    return {"status_code": status_code}
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser()
+    add_default_args(parser)
+
+    parser.add_argument(
+        "--labels-keys",
+        nargs="+",
+        required=True,
+        help="List of required key in labels part of source.",
+    )
+
+    args = parser.parse_args(argv)
+
+    try:
+        manifest = get_dbt_manifest(args)
+    except JsonOpenError as e:
+        print(f"Unable to load manifest file ({e})")
+        return 1
+
+    start_time = time.time()
+    hook_properties = has_labels_key(paths=args.filenames, labels_keys=args.labels_keys)
+    end_time = time.time()
+    script_args = vars(args)
+
+    tracker = dbtCheckpointTracking(script_args=script_args)
+    tracker.track_hook_event(
+        event_name="Hook Executed",
+        manifest=manifest,
+        event_properties={
+            "hook_name": os.path.basename(__file__),
+            "description": "Check the source has keys in the labels part.",
+            "status": hook_properties.get("status_code"),
+            "execution_time": end_time - start_time,
+            "is_pytest": script_args.get("is_test"),
+        },
+    )
+
+    return hook_properties.get("status_code")
+
+
+if __name__ == "__main__":
+    exit(main())
diff --git a/setup.cfg b/setup.cfg
index 5a6c28da..9b013c0e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -35,6 +35,7 @@ console_scripts =
     check-model-has-all-columns = dbt_checkpoint.check_model_has_all_columns:main
     check-model-has-description = dbt_checkpoint.check_model_has_description:main
     check-model-has-meta-keys = dbt_checkpoint.check_model_has_meta_keys:main
+    check-model-has-labels-keys = dbt_checkpoint.check_model_has_labels_keys:main
     check-model-has-properties-file = dbt_checkpoint.check_model_has_properties_file:main
     check-model-has-tests-by-name = dbt_checkpoint.check_model_has_tests_by_name:main
     check-model-has-tests-by-type = dbt_checkpoint.check_model_has_tests_by_type:main
@@ -55,6 +56,7 @@ console_scripts =
     check-source-has-freshness = dbt_checkpoint.check_source_has_freshness:main
     check-source-has-loader = dbt_checkpoint.check_source_has_loader:main
     check-source-has-meta-keys = dbt_checkpoint.check_source_has_meta_keys:main
+    check-source-has-labels-keys = dbt_checkpoint.check_source_has_labels_keys:main
     check-source-has-tests-by-group = dbt_checkpoint.check_source_has_tests_by_group:main
     check-source-has-tests-by-name = dbt_checkpoint.check_source_has_tests_by_name:main
     check-source-has-tests-by-type = dbt_checkpoint.check_source_has_tests_by_type:main


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index ca277209..c99f718d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -66,6 +66,28 @@
             "path": "aa/bb/without_meta.sql",
             "root_path": "/path/to/aa",
         },
+        "model.with_labels": {
+            "config": {
+                "labels": {"foo": "test", "bar": "test"},
+            },
+            "patch_path": "project://bb/with_labels.yml",
+            "path": "aa/bb/with_labels.sql",
+            "root_path": "/path/to/aa",
+        },
+        "model.with_labels_foo": {
+            "config": {
+                "labels": {"foo": "test"},
+            },
+            "patch_path": "project://bb/with_labels_foo.yml",
+            "path": "aa/bb/with_labels_foo.sql",
+            "root_path": "/path/to/aa",
+        },
+        "model.without_labels": {
+            "patch_path": "project://bb/without_labels.yml",
+            "path": "aa/bb/without_labels.sql",
+            "root_path": "/path/to/aa",
+        },
+        
         "model.with_tags": {
             "tags": ["foo", "bar"],
             "patch_path": "project://bb/with_tags.yml",
diff --git a/tests/unit/check_model_has_labels_keys.py b/tests/unit/check_model_has_labels_keys.py
new file mode 100644
index 00000000..d0a6bdd7
--- /dev/null
+++ b/tests/unit/check_model_has_labels_keys.py
@@ -0,0 +1,126 @@
+import argparse
+import os
+import time
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import Optional
+from typing import Sequence
+
+from dbt_checkpoint.tracking import dbtCheckpointTracking
+from dbt_checkpoint.utils import add_default_args
+from dbt_checkpoint.utils import get_dbt_manifest
+from dbt_checkpoint.utils import get_filenames
+from dbt_checkpoint.utils import get_model_schemas
+from dbt_checkpoint.utils import get_model_sqls
+from dbt_checkpoint.utils import get_models
+from dbt_checkpoint.utils import JsonOpenError
+
+
+def validate_keys(
+    actual: Iterable[str], expected: Iterable[str], allow_extra_keys: bool
+) -> bool:
+    actual = set(actual)
+    expected = set(expected)
+
+    if allow_extra_keys:
+        return expected.issubset(actual)
+    else:
+        return expected == actual
+
+
+def has_labels_key(
+    paths: Sequence[str],
+    manifest: Dict[str, Any],
+    labels_keys: Sequence[str],
+    allow_extra_keys: bool,
+) -> int:
+    status_code = 0
+    ymls = get_filenames(paths, [".yml", ".yaml"])
+    sqls = get_model_sqls(paths, manifest)
+    filenames = set(sqls.keys())
+    models = get_models(manifest, filenames)
+    schemas = get_model_schemas(list(ymls.values()), filenames)
+
+    in_models = set()
+    for model in models:
+        model_config = model.node.get("config", {})
+        model_labels = set(model_config.get("labels", {}).keys())
+        if validate_keys(model_labels, labels_keys, allow_extra_keys):
+            in_models.add(model.filename)
+
+    in_schemas = set()
+    for schema in schemas:
+        schema_config = schema.schema.get("config", {})
+        schema_labels = set(schema_config.get("labels", {}).keys())
+        
+        if validate_keys(schema_labels, labels_keys, allow_extra_keys):
+            in_schemas.add(schema.model_name)
+
+    missing = filenames.difference(in_models, in_schemas)
+
+    for model in missing:
+        status_code = 1
+        result = "\n- ".join(list(labels_keys))
+        print(
+            f"{sqls.get(model)}: "
+            f"does not have some of the labels keys defined:\n- {result}",
+        )
+
+    return status_code
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser()
+    add_default_args(parser)
+
+    parser.add_argument(
+        "--labels-keys",
+        nargs="+",
+        required=True,
+        help="List of required key in labels part of model.",
+    )
+
+    parser.add_argument(
+        "--allow-extra-keys",
+        action="store_true",
+        required=False,
+        help="Whether extra keys are allowed.",
+    )
+
+    args = parser.parse_args(argv)
+
+    try:
+        manifest = get_dbt_manifest(args)
+    except JsonOpenError as e:
+        print(f"Unable to load manifest file ({e})")
+        return 1
+
+    start_time = time.time()
+    status_code = has_labels_key(
+        paths=args.filenames,
+        manifest=manifest,
+        labels_keys=args.labels_keys,
+        allow_extra_keys=args.allow_extra_keys,
+    )
+    end_time = time.time()
+    script_args = vars(args)
+
+    tracker = dbtCheckpointTracking(script_args=script_args)
+    tracker.track_hook_event(
+        event_name="Hook Executed",
+        manifest=manifest,
+        event_properties={
+            "hook_name": os.path.basename(__file__),
+            "description": "Check model has labels keys",
+            "status": status_code,
+            "execution_time": end_time - start_time,
+            "is_pytest": script_args.get("is_test"),
+        },
+    )
+
+    return status_code
+
+
+if __name__ == "__main__":
+    exit(main())
diff --git a/tests/unit/check_source_has_labels_keys.py b/tests/unit/check_source_has_labels_keys.py
new file mode 100644
index 00000000..3e0753f0
--- /dev/null
+++ b/tests/unit/check_source_has_labels_keys.py
@@ -0,0 +1,81 @@
+import argparse
+import os
+import time
+from pathlib import Path
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Sequence
+
+from dbt_checkpoint.tracking import dbtCheckpointTracking
+from dbt_checkpoint.utils import add_default_args
+from dbt_checkpoint.utils import get_dbt_manifest
+from dbt_checkpoint.utils import get_source_schemas
+from dbt_checkpoint.utils import JsonOpenError
+
+
+def has_labels_key(paths: Sequence[str], labels_keys: Sequence[str]) -> Dict[str, Any]:
+    status_code = 0
+    ymls = [Path(path) for path in paths]
+
+    # if user added schema but did not rerun
+    schemas = get_source_schemas(ymls)
+
+    for schema in schemas:
+        schema_labels = set(schema.source_schema.get("labels", {}).keys())
+        table_labels = set(schema.table_schema.get("labels", {}).keys())
+        all_labels = schema_labels.union(table_labels)
+        diff_labels = set(labels_keys).difference(all_labels)
+
+        if diff_labels:
+            status_code = 1
+            result = "\n- ".join(list(diff_labels))
+            print(
+                f"{schema.source_name}.{schema.table_name}: "
+                f"does not have some of the labels keys defined:\n- {result}",
+            )
+    return {"status_code": status_code}
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser()
+    add_default_args(parser)
+
+    parser.add_argument(
+        "--labels-keys",
+        nargs="+",
+        required=True,
+        help="List of required key in labels part of source.",
+    )
+
+    args = parser.parse_args(argv)
+
+    try:
+        manifest = get_dbt_manifest(args)
+    except JsonOpenError as e:
+        print(f"Unable to load manifest file ({e})")
+        return 1
+
+    start_time = time.time()
+    hook_properties = has_labels_key(paths=args.filenames, labels_keys=args.labels_keys)
+    end_time = time.time()
+    script_args = vars(args)
+
+    tracker = dbtCheckpointTracking(script_args=script_args)
+    tracker.track_hook_event(
+        event_name="Hook Executed",
+        manifest=manifest,
+        event_properties={
+            "hook_name": os.path.basename(__file__),
+            "description": "Check the source has keys in the labels part.",
+            "status": hook_properties.get("status_code"),
+            "execution_time": end_time - start_time,
+            "is_pytest": script_args.get("is_test"),
+        },
+    )
+
+    return hook_properties.get("status_code")
+
+
+if __name__ == "__main__":
+    exit(main())
diff --git a/tests/unit/test_check_model_has_labels_keys.py b/tests/unit/test_check_model_has_labels_keys.py
new file mode 100644
index 00000000..a530d44a
--- /dev/null
+++ b/tests/unit/test_check_model_has_labels_keys.py
@@ -0,0 +1,113 @@
+from unittest.mock import mock_open
+from unittest.mock import patch
+
+import pytest
+
+from dbt_checkpoint.check_model_has_labels_keys import main
+
+# Input args, valid manifest, expected return value
+TESTS = (
+    (
+        [
+            "aa/bb/with_labels.sql",
+            "--labels-keys",
+            "foo",
+            "bar",
+        ],
+        {"models": [{"name": "with_labels", "config": {"labels": {"foo": "bar"}}}]},
+        True,
+        True,
+        0,
+    ),
+    (
+        [
+            "aa/bb/with_labels.sql",
+            "--labels-keys",
+            "foo",
+            "bar",
+        ],
+        {"models": [{"name": "with_labels", "config": {"labels": {"foo": "bar"}}}]},
+        False,
+        True,
+        1,
+    ),
+    (
+        [
+            "aa/bb/without_labels.sql",
+            "--labels-keys",
+            "foo",
+            "bar",
+        ],
+        {"models": [{"name": "without_labels"}]},
+        True,
+        True,
+        1,
+    ),
+    (
+        ["aa/bb/with_labels.sql", "--labels-keys", "foo", "--allow-extra-keys"],
+        {"models": [{"name": "with_labels", "config":{"labels": {"foo": "bar", "baz": "test"}}}]},
+        True,
+        True,
+        0,
+    ),
+    (
+        ["aa/bb/with_labels.sql", "--labels-keys", "baz"],
+        {"models": [{"name": "with_labels", "config":{"labels": {"foo": "bar", "baz": "test"}}}]},
+        True,
+        True,
+        1,
+    ),
+)
+
+
+@pytest.mark.parametrize(
+    ("input_args", "schema", "valid_manifest", "valid_config", "expected_status_code"),
+    TESTS,
+)
+def test_check_model_labels_keys(
+    input_args,
+    schema,
+    valid_manifest,
+    valid_config,
+    expected_status_code,
+    manifest_path_str,
+    config_path_str,
+):
+    if valid_manifest:
+        input_args.extend(["--manifest", manifest_path_str])
+    if valid_config:
+        input_args.extend(["--config", config_path_str])
+    with patch("builtins.open", mock_open(read_data="data")):
+        with patch("dbt_checkpoint.utils.safe_load") as mock_safe_load:
+            mock_safe_load.return_value = schema
+    status_code = main(input_args)
+    assert status_code == expected_status_code
+
+
+@pytest.mark.parametrize("extension", [("yml"), ("yaml")])
+def test_check_model_labels_keys_in_changed(extension, tmpdir, manifest_path_str):
+    schema_yml = """
+version: 2
+
+models:
+-   name: in_schema_labels
+    config:
+        labels:
+            foo: test
+            bar: test
+-   name: xxx
+    """
+    yml_file = tmpdir.join(f"schema.{extension}")
+    yml_file.write(schema_yml)
+    result = main(
+        argv=[
+            "in_schema_labels.sql",
+            str(yml_file),
+            "--labels-keys",
+            "foo",
+            "bar",
+            "--manifest",
+            manifest_path_str,
+        ],
+    )
+    assert result == 0
diff --git a/tests/unit/test_check_source_has_labels_keys.py b/tests/unit/test_check_source_has_labels_keys.py
new file mode 100644
index 00000000..d333a939
--- /dev/null
+++ b/tests/unit/test_check_source_has_labels_keys.py
@@ -0,0 +1,154 @@
+import pytest
+
+from dbt_checkpoint.check_source_has_labels_keys import main
+
+# Input schema, expected return value
+TESTS = (
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    labels:
+        foo: test
+        bar: test
+    tables:
+    -   name: test
+    """,
+        True,
+        True,
+        0,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    labels:
+        foo: test
+        bar: test
+    tables:
+    -   name: test
+    """,
+        False,
+        True,
+        1,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    labels:
+        foo: test
+    tables:
+    -   name: test
+        labels:
+            bar: test
+    """,
+        True,
+        True,
+        0,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    tables:
+    -   name: test
+        labels:
+            bar: test
+            foo: test
+    """,
+        True,
+        True,
+        0,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    tables:
+    -   name: test
+        labels:
+            bar: test
+    """,
+        True,
+        True,
+        1,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    labels:
+        bar: test
+    tables:
+    -   name: test
+    """,
+        True,
+        True,
+        1,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    tables:
+    -   name: test
+    """,
+        True,
+        True,
+        1,
+    ),
+    (
+        """
+sources:
+-   name: src
+    loader: test
+    labels:
+        foo: test
+        bar: test
+    tables:
+    -   name: test
+    """,
+        True,
+        False,
+        0,
+    ),
+)
+
+
+@pytest.mark.parametrize(
+    ("input_schema", "valid_manifest", "valid_config", "expected_status_code"), TESTS
+)
+def test_check_source_has_labels_keys(
+    input_schema,
+    valid_manifest,
+    valid_config,
+    expected_status_code,
+    tmpdir,
+    manifest_path_str,
+    config_path_str,
+):
+    input_args = [
+        "--labels-keys",
+        "foo",
+        "bar",
+        "--is_test",
+    ]
+
+    if valid_manifest:
+        input_args.extend(["--manifest", manifest_path_str])
+
+    if valid_config:
+        input_args.extend(["--config", config_path_str])
+
+    yml_file = tmpdir.join("schema.yml")
+    yml_file.write(input_schema)
+    status_code = main(argv=[str(yml_file), *input_args])
+    assert status_code == expected_status_code