Skip to content
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 setup(self):
        """
        Setup the check. Use this to load data, initialize models, etc.
        """
        super().setup() # DON'T FORGET TO CALL SUPER
        self.content = "Apple"

    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
    
    def rule_my_first_failed_rule(self):
        # Call functions to check the data
        # Throw an exception if the rule fails
        assert False

    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")

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 return nothing. Logs are for all runs. In the above example, the setup(self) method is used to initialize the check. This method is called before any rules are run. The setup(self) method is optional. If you do not need to initialize anything, you can omit it. If you do need to initialize anything like datasets, models, etc. make sure to call super().setup() before doing so.

In addition to these basic 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).

Clone this wiki locally