Skip to content

Checks

Jump to bottom
Ivan Zhang edited this page Aug 29, 2023 · 5 revisions

Data Check

The library exposes the DataCheck class (defined in data_checks/classes/data_check.py) which you can use to write your checks. Begin by subclassing the DataCheck class:

from data_checks.classes.data_check import DataCheck

class MyFirstDataCheck(DataCheck):
    pass

Then define methods that check the data and will be executed when the check is run. These methods are called rules and are prefixed with rule_. For example:

from data_checks.classes.data_check import DataCheck

class MyFirstDataCheck(DataCheck):
    def rule_my_first_successful_rule(self, data="Hello World"):
        # Call functions to check the data
        assert data == "Hello World"
        # Return anything if the rule passes
    
    def rule_my_first_failed_rule(self):
        # Call functions to check the data
        # Throw an exception if the rule fails
        assert False, "This rule failed"

    def my_first_helper_function(self):
        # This function will not be run as a rule
        raise Exception("This function will not be run as a rule")

defines two rules: rule_my_first_successful_rule and rule_my_first_failed_rule. The my_first_helper_function function will not be run as a rule.

A rule is considered a failure when an exception is raised. Exceptions and tracebacks are stored in the database for failures. Successful rules should throw no exceptions and can return anything (this can be accessed in Pre- and Post- Steps). Logs are stored for all runs.

Pre- and Post- Steps

Like unit tests, you can define a series of optional pre- and post- steps that are run for setup and teardown of a check and before and after each rule. These steps are useful for doing actions based off a check and/or rule and whether or not they succeed. These methods are:

  • setup(self): This method is run before any rules are run. Use this method to load data, initialize models, etc. For example self.df = pd.read_csv("data.csv") sets a dataframe for all the rules to use. This method is run only once per check.
  • before(self, context): This method is run before each rule.
  • on_success(self, context): This method is run when a rule succeeds.
  • on_failure(self, context): This method is run when a rule fails.
  • after(self, context): This method is run after each rule. This runs regardless of success or failure of the rule.
  • teardown(self): This method is run after all rules are run. Use this method to close connections, etc. This method is run only once per check.

The context parameter is a dictionary that is passed between the four pre- and post- rule steps. You can access a few system defined keys by calling context.get_sys(key). These keys are:

  • context.get_sys("rule"): The name of the rule being run
  • context.get_sys("params"): The parameters of the rule being run as a dictionary. For rule_my_first_successful_rule in the example above, context.get_sys("params") would return: 
    {
    "args": (),
    "kwargs": {"data": "Hello World"}
}
  • If the rule succeeds, context.get_sys("result") (when called inside on_success and after) will contain the return value of the rule. If the rule fails, context.get_sys("exception") (when called inside on_failure and after) will return the exception that was raised.

This context is shared between these four functions. Thus you can use it to pass information between these steps by setting and getting like you would a dictionary. For example:

from data_checks.classes.data_check import DataCheck


class MyFirstDataCheck(DataCheck):
    def setup(self):
        """
        Setup the check. Use this to load data, initialize models, etc.
        For example `self.df = pd.read_csv("data.csv")` sets a dataframe
        for all the rules to use.
        """
        super().setup()  # DON'T FORGET TO CALL SUPER
        print("Setup")
        self.content = "Apple"

    def before(self, context):
        """
        Run before each rule. If None, the rule will not be run
        """
        super().before(context)  # DON'T FORGET TO CALL SUPER
        print("Running before each rule")
        print(context.get_sys("rule"))  # Prints the name of the rule being run
        print(context.get_sys("params"))  # Prints the parameters of the rule being run
        context["favorite_fruit"] = ["Apple", "Banana", "Orange"]

    def rule_my_first_successful_rule(self, data="Hello World"):
        # Call functions to check the data
        assert data == "Hello World"
        assert self.content == "Apple"
        # Return nothing if the rule passes
        return "Big Success!"

    def rule_my_first_failed_rule(self):
        # Call functions to check the data
        # Throw an exception if the rule fails
        assert False, "This rule failed"

    def on_success(self, context):
        """
        Called when a rule succeeds
        """
        super().on_success(context)  # DON'T FORGET TO CALL SUPER
        print("Rule succeeded")
        print(context.get_sys("result"))  # Prints "Big Success!"
        print(context["favorite_fruit"])  # Prints ["Apple", "Banana", "Orange"]

    def on_failure(self, context):
        """
        Called when a rule fails
        """
        super().on_failure(context)  # DON'T FORGET TO CALL SUPER
        print("Rule failed")
        print(context.get_sys("exception"))  # Prints the exception that was raised

    def after(self, context):
        """
        Runs after each rule
        """
        super().after(context)  # DON'T FORGET TO CALL SUPER
        print("Running after each rule")
        print(context)

    def teardown(self):
        """
        Teardown the check. Use this to close connections, etc.
        """
        super().teardown()  # DON'T FORGET TO CALL SUPER
        print("Teardown")

    def my_first_helper_function(self):
        # This function will not be run as a rule
        raise Exception("This function will not be run as a rule")

Running this check will print:

[1/1 checks] MyFirstDataCheck
Setup
	[1/2 Rules] rule_my_first_successful_rule
Running before each rule
rule_my_first_successful_rule
{'args': (), 'kwargs': {}}
		rule_my_first_successful_rule took 1.1920928955078125e-06 seconds
Rule succeeded
Big Success!
['Apple', 'Banana', 'Orange']
Running after each rule
{'sys': {'rule': 'rule_my_first_successful_rule', 'params': {'args': (), 'kwargs': {}}, 'result': 'Big Success!'}, 'favorite_fruit': ['Apple', 'Banana', 'Orange']}
	[2/2 Rules] rule_my_first_failed_rule
Running before each rule
rule_my_first_failed_rule
{'args': (), 'kwargs': {}}
This rule failed
Rule failed
DataCheckException(severity=1.0, exception=This rule failed, metadata={'rule': 'rule_my_first_failed_rule', 'params': {'args': (), 'kwargs': {}}})
Running after each rule
{'sys': {'rule': 'rule_my_first_failed_rule', 'params': {'args': (), 'kwargs': {}}, 'exception': DataCheckException()}, 'favorite_fruit': ['Apple', 'Banana', 'Orange']}
Teardown

❗IMPORTANT
Make sure to call super().setup(), super().before(context), super().on_success(context), super().on_failure(context), super().after(context), and super().teardown() in your methods. Their omission may break the scheduling and alerting features of this library.

In addition to these methods, you can also define these two additional methods:

from data_checks.classes.data_check import DataCheck

class MyFirstDataCheck(DataCheck):
    ...
    @classmethod
    def defined_rules(cls) -> list[str]:
        return [
            "rule_my_first_successful_rule"
        ]

    @classmethod
    def check_config(cls) -> dict:
        return {
            "param1": value1,
            "param2": value2,
            "rules_config": { # Configuration for each rule. If not provided, rule config will be empty dictionary.
                "rule_1": {
                    "silenced_until": "2021-08-22 00:00:00.359828-00", # "%Y-%m-%d %H:%M:%S.%f%z" formatted date or datetime until which the rule will be silenced. If not provided, rule will not be silenced.
                    ...
                },
                "rule_2": {
                    ...
                },
            }
        }
    ...

  • defined_rules(cls) -> list[str]: This method returns a list of rules that should be run. If this method is not defined, all rules (with the prefix) will be run. This method is useful if you want to run only a subset of rules. For example, you can use this method to run only rules that are not silenced. You may also override this method in a subclass to exclude certain rules.

  • check_config(cls) -> dict: This method returns a dictionary of configurations for the check. These configuration options are stored in the database and can be used to define additional features for your checks. For instance you may use a description field to describe the check. Note that the rules_config are not stored with checks but instead with each respective rule (i.e. rule_1 configuration options are stored in the Rules database table). For each rule, you can define a silenced_until field that silences the rule until that point.

Defining Checks in CHECKS_MODULE

Your check should be written inside the specified CHECKS_MODULE in your settings file. For example, if you set CHECKS_MODULE = "my_checks", then you should write your check in my_checks/my_first_data_check.py. Make sure that CHECKS_MODULE and any nested modules are properly defined as directories (i.e. have an __init__.py file). Doing this allows the library to automatically find and run checks.

Instantiating and Running Checks

DataCheck also defines additional methods as well as instance variables that can are used to instantiate and run checks. The DataCheck class exposes the following instance variables:

  • name: The name of the check. This is used to identify the check in the database. If not provided, the name of the class is used.

  • rules_params: A dictionary of parameters for each rule. The keys are the names of the rules and the values are the parameters to be passed to that rule. If a rule does not have any required parameters or has default parameters values that work, you can omit it from the dictionary. If a rule has parameters, you must provide them in the dictionary. If a rule has parameters but you do not provide them, an exception will be thrown. If a rule does not have parameters and you provide them, an exception will be thrown. Assume you have the following rule:

    from data_checks.classes.data_check import DataCheck

class MyFirstDataCheck(DataCheck):
    ...
    def rule_with_required_arguments(self, data):
        assert data == "Hello World", "This rule failed"
    ...

    You can define the data parameter value of rule_with_required_arguments in several ways:

    • A dictionary of keyword arguments: 
      check = MyFirstDataCheck(
    rules_params={"rule_with_required_arguments": {"data": "Hello World"}}
)
    • A tuple of positional arguments 
      check = MyFirstDataCheck(
    rules_params={"rule_with_required_arguments": ("Hello World",)}
)
    • A dictionary with a args key for positional arguments and a kwargs key for keyword arguments 
      check = MyFirstDataCheck(
    rules_params={"rule_with_required_arguments": {"args": ("Hello World",), "kwargs": {}}} 
)
    • A list of any of the three previous options. This runs the rule multiple times with each set of arguments. 
      check = MyFirstDataCheck(
    rules_params={
        "rule_with_required_arguments": [
            {"data": "Hello World"},
            ("Hello World",),
            {"args": ("Hello World",), "kwargs": {}},
        ]
    }
)
    • A callable that returns any of the four previous options. This allows you to dynamically define the arguments. The callable can have any arguments but should return one of the four options. 
      check = MyFirstDataCheck(
    rules_params={"rule_with_required_arguments": lambda: {"data": "Hello World"}}
)

  • excluded_rules: A list of rules to exclude from the check. This is useful if you want to run all rules except for a few. For example, you can use this to exclude rules that are silenced.

  • only_run_specified_rules: A boolean that determines whether to run only the rules specified in rules_params. If True, only the rules specified in rules_params will be run. If False, all rules will be run except for those specified in excluded_rules.

  • Additional keyword arguments that are passed to the Check class become instance fields of the check.

You can pass these arguments to the DataCheck constructor. For example:

from my_checks.my_first_data_check import MyFirstDataCheck

check = MyFirstDataCheck(
    name="My First Data Check",
    rules_params={"rule_with_required_arguments": {"data": "Hello World"}},
    excluded_rules=["rule_my_first_failed_rule"],
    only_run_specified_rules=True,
    param1=value1,
    param2=value2,
)

creates a check with the name "My First Data Check" that runs only the rule_with_required_arguments rule with the data parameter set to "Hello World" and excludes the rule_my_first_failed_rule rule. The check also has the param1 and param2 instance fields set to value1 and value2 respectively that can be accessed via check.param1 and check.param2.

You can call a few built-in methods on the check. These are:

  • get_rules_to_run(): Returns a set of rules that will be run. This is useful if you want to know which rules will be run before running the check.
  • run_all(): Runs all rules in the check. This is the default method that is called when you run a check.
  • run_all_async(): Runs all rules in the check asynchronously.
  • run(rule): Runs a specific rule.
  • run_async(rule, wait_for_completion=True): Runs a specific rule asynchronously. If wait_for_completion is True, the method will wait for the rule to finish before returning. If False, the method will return a list of Process objects that can be used to wait for the rule to finish (i.e. by calling process.join()).

Rules With Arguments

You can also define rules with arguments. For example:

from data_checks.classes.data_check import DataCheck

class MyFirstDataCheck(DataCheck):
    ...
    # This rule has an argument
    def rule_with_required_arguments(self, data):
        assert data == "Hello World", "This rule failed"
    ...

To run this check, you must pass in the arguments. For example:

# run_check.py
from my_checks.my_first_data_check import MyFirstDataCheck

check = MyFirstDataCheck(
    rules_params={"rule_with_required_arguments": {"data": "Hello World1"}}
)
check.run_all() # check.run_all_async() to run all rules asynchronously

Run this check with python run_check.py. The output should be:

	[1/3 Rules] rule_my_first_successful_rule
		rule_my_first_successful_rule took 9.5367431640625e-07 seconds
	[2/3 Rules] rule_with_required_arguments
This rule failed
	[3/3 Rules] rule_my_first_failed_rule
This rule failed

A similar process applies for checks that are run within a suite (see [Advanced] Suites). Note that checks with arguments cannot be run from the command line.

More Advanced Checks

If you have existing check classes, you can still subclass DataCheck and use the library in the same manner noted above. Make sure that your class does not accidentally override any of the methods in Check (see data_checks/base/check.py)

The DataCheck class is a simplified and beginner friendly subclass of the base Check class (data_checks/base/check.py). The user can also directly subclass the Check class to create more advanced checks (see References).

Documentation

Clone this wiki locally