feat: collect rich user input for tests with subprocess (#238)

* feat: collect rich user input for tests with subprocess

* fix: order of the combine list

Make nicer the logs

Secure variable

* add: feature.md

.isort.cfg

@@ -1,2 +1,2 @@
 [settings]
-known_third_party =luddite,packaging,pretty_errors,pydantic,pyfiglet,pytest,questionary,rich,snowflake,sqlalchemy,yaml,yamlloader
+known_third_party =click,luddite,packaging,pretty_errors,pydantic,pyfiglet,pytest,questionary,rich,snowflake,sqlalchemy,yaml,yamlloader

changelog/238.feature.md

@@ -0,0 +1 @@
+dbt-sugar can now add your custom tests via the console, you only need to write the test as you want to look in the schema.yml. dbt-sugar will check if the test PASSES and if it does will add the custom test to your schema.yml.

dbt_sugar/core/task/doc.py

@@ -1,12 +1,13 @@
 """Document Task module."""
 import copy
 import re
+import subprocess
 from collections import OrderedDict
 from pathlib import Path
+from shlex import quote
 from typing import Any, Dict, List, Mapping, Optional, Sequence
 
 from rich.console import Console
-from rich.progress import BarColumn, Progress
 
 from dbt_sugar.core.clients.dbt import DbtProfile
 from dbt_sugar.core.clients.yaml_helpers import open_yaml, save_yaml
@@ -52,7 +53,6 @@ def run(self) -> int:
         model = self._flags.model
         schema = self._dbt_profile.profile.get("target_schema", "")
 
-        connector = self.get_connector()
         dbt_credentials = self._dbt_profile.profile
         connector = DB_CONNECTORS.get(dbt_credentials.get("type", ""))
         if not connector:
@@ -258,51 +258,82 @@ def orchestrate_model_documentation(
         except KeyboardInterrupt:
             logger.info("The user has exited the doc task, all changes have been discarded.")
             return 0
+
         save_yaml(schema_file_path, self.order_schema_yml(content))
         self.add_primary_key_tests(schema_content=content, model_name=model_name)
-        self.check_tests(schema, model_name)
+
+        # The copy is here because it was modifying the tests.
         self.update_model_description_test_tags(
-            schema_file_path, model_name, self.column_update_payload
+            schema_file_path, model_name, copy.deepcopy(self.column_update_payload)
         )
+        self.check_tests(schema_file_path, model_name)
         # Method to update the descriptions in all the schemas.yml
         self.update_column_descriptions(self.column_update_payload)
 
         return 0
 
-    def check_tests(self, schema: str, model_name: str) -> None:
+    def delete_failed_tests_from_schema(
+        self, path_file: Path, model_name: str, tests_to_delete: Dict[str, List[str]]
+    ):
+        """
+        Method to delete the failing tests from the schema.yml.
+
+        Args:
+            path_file (Path): Path of the schema.yml file to update.
+            model_name (str): Name of the model to document.
+            tests_to_delete (Dict[str, List[str]]): with the tests that have failed.
+        """
+        content = open_yaml(path_file)
+        for model in content["models"]:
+            if model["name"] == model_name:
+                for column in model.get("columns", []):
+                    tests_to_delete_from_column = tests_to_delete.get(column["name"], [])
+                    tests_from_column = column.get("tests", [])
+                    tests_pass = [
+                        x for x in tests_from_column if x not in tests_to_delete_from_column
+                    ]
+                    if not tests_pass and tests_from_column:
+                        del column["tests"]
+                    elif tests_pass:
+                        column["tests"] = tests_pass
+        save_yaml(path_file, content)
+
+    def check_tests(self, path_file: Path, model_name: str) -> None:
         """
         Method to run and add test into a schema.yml, this method will:
 
         Run the tests and if they have been successful it will add them into the schema.yml.
 
         Args:
-            schema (str): Name of the schema where the model lives.
+            path_file (Path): Path of the schema.yml file to update.
             model_name (str): Name of the model to document.
         """
-        with Progress(
-            "[progress.description]{task.description}",
-            BarColumn(),
-            "[progress.percentage]{task.percentage:>3.0f}%",
-            transient=True,
-        ) as progress:
-            test_checking_task = progress.add_task(
-                "[bold] checking your tests...", total=len(self.column_update_payload.keys())
+        dbt_command = f"dbt test --models {quote(model_name)}".split()
+        dbt_result_command = subprocess.run(dbt_command, capture_output=True, text=True).stdout
+        tests_to_delete: Dict[str, List[str]] = {}
+
+        if "Compilation Error" in dbt_result_command:
+            logger.info(
+                "dbt encountered a compilation error in one or more of your custom tests.\n"
+                "Not able to check if the tests that you have added have PASSED.\n"
+                f"This is what dbt's compilation error says:\n{dbt_result_command}"
             )
-            for column in self.column_update_payload.keys():
-                tests = self.column_update_payload[column].get("tests", [])
-                tests_ = copy.deepcopy(tests)
-                for test in tests_:
-                    has_passed = self.connector.run_test(
-                        test,
-                        schema,
-                        model_name,
-                        column,
+
+        for column in self.column_update_payload.keys():
+            tests = self.column_update_payload[column].get("tests", [])
+            for test in tests:
+                test_name = test if type(test) == str else list(test.keys())[0]
+                test_passed_pattern = f"PASS {test_name}_{model_name}_{column}"
+                if re.search(test_passed_pattern, dbt_result_command):
+                    logger.info(f"The test {test} in the column {column} has PASSED.")
+                else:
+                    logger.info(
+                        f"The test {test} in the column {column} has FAILED to execute."
+                        "The test won't be added to your schema.yml file"
                     )
-                    message = self._generate_test_success_message(test, column, has_passed)
-                    progress.console.log(message)
-                    if not has_passed:
-                        tests.remove(test)
-                progress.advance(test_checking_task)
+                    tests_to_delete[column] = tests_to_delete.get(column, []) + [test]
+        if tests_to_delete:
+            self.delete_failed_tests_from_schema(path_file, model_name, tests_to_delete)
 
     @staticmethod
     def _generate_test_success_message(test_name: str, column_name: str, has_passed: bool):
@@ -396,12 +427,10 @@ def update_model(
 
                     columns = model.get("columns", [])
                     columns_names = [column["name"] for column in columns]
-
                     if column not in columns_names:
                         description = self.get_column_description_from_dbt_definitions(column)
                         logger.info(f"Updating column '{column.lower()}'")
                         columns.append({"name": column, "description": description})
-
         return content
 
     def create_new_model(

dbt_sugar/core/ui/cli_ui.py

@@ -3,7 +3,9 @@
 import copy
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Union, cast
 
+import click
 import questionary
+import yaml
 from pydantic import BaseModel, validator
 
 DESCRIPTION_PROMPT_MESSAGE = "Please write down your description:"
@@ -174,7 +176,7 @@ def _validate_question_payload(self) -> None:
             for payload_element_index, payload_element in enumerate(self._question_payload):
                 if payload_element_index == 0:
                     ConfirmModelDoc(**payload_element)
-                if payload_element_index == 1:
+                elif payload_element_index == 1:
                     DescriptionTextInput(**payload_element)
 
         elif self._question_type == "undocumented_columns":
@@ -188,6 +190,18 @@ def _validate_question_payload(self) -> None:
 
         self._is_valid_question_payload = True
 
+    def collect_rich_user_input(self) -> str:
+        """Uses click to open up a text editor to collect rich user input.
+
+        Returns:
+            str: string version of the text input which will then be loaded.
+        """
+        tests = click.edit(extension=".yml")
+        if tests:
+            tests = tests.replace("\t", "    ")
+            tests = yaml.safe_load(tests)
+        return tests
+
     def _iterate_through_columns(
         self, cols: List[str]
     ) -> Mapping[str, Mapping[str, Union[str, List[str]]]]:
@@ -236,10 +250,17 @@ def _iterate_through_columns(
                     message="Would you like to add any tests?"
                 ).unsafe_ask()
                 if wants_to_add_tests:
-                    tests = questionary.checkbox(
-                        message="Please select one or more tests from the list below",
-                        choices=AVAILABLE_TESTS,
+                    wants_to_pop_editor = questionary.confirm(
+                        message="Do you want to add a complex or custom tests? If, so we'll open a test editor for "
+                        f"you, otherwise you can choose from the following builtins: {AVAILABLE_TESTS}"
                     ).unsafe_ask()
+                    if wants_to_pop_editor:
+                        tests = self.collect_rich_user_input()
+                    else:
+                        tests = questionary.checkbox(
+                            message="Please select one or more tests from the list below",
+                            choices=AVAILABLE_TESTS,
+                        ).unsafe_ask()
                     if tests:
                         results[column]["tests"] = tests
 
@@ -280,8 +301,8 @@ def _document_model(
 
             if not results.get("model_description"):
                 # we return an empty dict if user decided to not enter a description in the end.
-                results = dict()
-            if collect_model_description is False:
+                results = {}
+            if not collect_model_description:
                 # if the user doesnt want to document the model we exit early even if the payload
                 # has a second entry which would trigger the description collection
                 break
@@ -299,7 +320,6 @@ def _document_undocumented_cols(
         question_payload: Sequence[Mapping[str, Any]],
     ) -> Mapping[str, Mapping[str, Union[str, List[str]]]]:
 
-        results: Mapping[str, Mapping[str, Union[str, List[str]]]] = dict()
         columns_to_document = question_payload[0].get("choices", list())
         quantifier_word = self._set_quantifier_word()
         # check if user wants to document all columns
@@ -312,14 +332,12 @@ def _document_undocumented_cols(
         ).unsafe_ask()
 
         if document_all_cols:
-            results = self._iterate_through_columns(cols=columns_to_document)
-        else:
-            # get the list of columns from user
-            columns_to_document = questionary.prompt(question_payload)
-            results = self._iterate_through_columns(
-                cols=columns_to_document["cols_to_document"],
-            )
-        return results
+            return self._iterate_through_columns(cols=columns_to_document)
+        # get the list of columns from user
+        columns_to_document = questionary.prompt(question_payload)
+        return self._iterate_through_columns(
+            cols=columns_to_document["cols_to_document"],
+        )
 
     def _document_already_documented_cols(
         self,
@@ -329,13 +347,14 @@ def _document_already_documented_cols(
         mutable_payload = cast(Sequence[Dict[str, Any]], mutable_payload)
 
         # massage the question payload
-        choices = []
-        for col, desc in mutable_payload[0].get("choices", dict()).items():
-            choices.append(f"{col} | {desc}")
+        choices = [
+            f"{col} | {desc}" for col, desc in mutable_payload[0].get("choices", dict()).items()
+        ]
+
         mutable_payload[0].update({"choices": choices})
 
         # ask user if they want to see any of the documented columns?
-        results = dict()
+        results = {}
         document_any_columns = questionary.confirm(
             message="Do you want to document any of the already documented columns in this model?",
             auto_enter=True,

tests/cli_ui_test.py

@@ -212,12 +212,12 @@ def test__document_already_documented_cols(
             {
                 "column_a": {
                     "description": "Dummy description",
-                    "tests": ["unique"],
+                    "tests": [{"accepted_values": ["a"]}],
                     "tags": ["Dummy description"],
                 },
                 "column_b": {
                     "description": "Dummy description",
-                    "tests": ["unique"],
+                    "tests": [{"accepted_values": ["a"]}],
                     "tags": ["Dummy description"],
                 },
             },
@@ -229,6 +229,7 @@ def test__iterate_through_columns(mocker, question_payload, expected_results):
     mocker.patch("questionary.text", return_value=Question("Dummy description"))
     mocker.patch("questionary.checkbox", return_value=Question(["unique"]))
     mocker.patch("questionary.confirm", return_value=Question(question_payload["ask_for_tests"]))
+    mocker.patch("click.edit", return_value="- accepted_values:\n\t- a")
     results = UserInputCollector(
         "undocumented_columns", question_payload=[], ask_for_tests=question_payload["ask_for_tests"]
     )._iterate_through_columns(cols=question_payload["col_list"])

tests/doc_task_test.py

@@ -1148,3 +1148,66 @@ def test_get_primary_key_from_sql(mocker, content, result):
     read_file = mocker.patch("dbt_sugar.core.task.doc.DocumentationTask.read_file")
     read_file.return_value = content
     assert doc_task.get_primary_key_from_sql("path") == result
+
+
+@pytest.mark.parametrize(
+    "content, model_name, tests_to_delete, result",
+    [
+        pytest.param(
+            {
+                "models": [
+                    {
+                        "name": "testmodel",
+                        "columns": [
+                            {
+                                "name": "columnA",
+                                "description": "descriptionA",
+                                "tests": ["unique", "not_null"],
+                                "tags": ["hi", "hey"],
+                            },
+                            {
+                                "name": "columnF",
+                                "description": "descriptionF",
+                                "tests": ["unique", "not_null"],
+                            },
+                        ],
+                    }
+                ]
+            },
+            "testmodel",
+            {"columnA": ["unique", "not_null"], "columnF": ["not_null"]},
+            [
+                call(
+                    PosixPath("."),
+                    {
+                        "models": [
+                            {
+                                "columns": [
+                                    {
+                                        "description": "descriptionA",
+                                        "name": "columnA",
+                                        "tags": ["hi", "hey"],
+                                    },
+                                    {
+                                        "description": "descriptionF",
+                                        "name": "columnF",
+                                        "tests": ["unique"],
+                                    },
+                                ],
+                                "name": "testmodel",
+                            }
+                        ]
+                    },
+                )
+            ],
+            id="delete_failed_test_from_schema",
+        ),
+    ],
+)
+def test_delete_failed_tests_from_schema(mocker, content, model_name, tests_to_delete, result):
+    open_yaml = mocker.patch("dbt_sugar.core.task.doc.open_yaml")
+    save_yaml = mocker.patch("dbt_sugar.core.task.doc.save_yaml")
+    open_yaml.return_value = content
+    doc_task = __init_descriptions()
+    doc_task.delete_failed_tests_from_schema(Path("."), model_name, tests_to_delete)
+    save_yaml.assert_has_calls(result)

tests/test_dbt_project/dbt_sugar_test/models/folder_to_exclude/schema.yml

@@ -1 +1 @@
-version: 1
+version: 2