Orchestrate Great Expectations with Airflow

1
from great_expectations_provider.operators.validate_dataframe import GXValidateDataFrameOperator
2
3
def configure_dataframe():
4
    """Load data into a DataFrame using an Airflow hook."""
5
    from airflow.providers.common.sql.hooks.sql import DbApiHook
6
7
    hook = DbApiHook.get_hook("my_db_conn")
8
    return hook.get_pandas_df("SELECT * FROM daily_planet_report")
9
10
11
def configure_expectations(dataframe):
12
    """Define expectations for the DataFrame."""
13
    import great_expectations as gx
14
15
    suite = gx.ExpectationSuite(name="daily_report_suite")
16
    suite.add_expectation(
17
        gx.expectations.ExpectColumnValuesToNotBeNull(column="planet_name")
18
    )
19
    suite.add_expectation(
20
        gx.expectations.ExpectColumnValuesToNotBeNull(column="total_passengers")
21
    )
22
    suite.add_expectation(
23
        gx.expectations.ExpectColumnValuesToBeBetween(
24
            column="total_net_fare_usd", min_value=0
25
        )
26
    )
27
    suite.add_expectation(
28
        gx.expectations.ExpectColumnValuesToBeBetween(
29
            column="total_discounts_usd", min_value=0
30
        )
31
    )
32
    return suite
33
34
35
_validate_dataframe = GXValidateDataFrameOperator(
36
    task_id="validate_with_gx_dataframe",
37
    configure_dataframe=configure_dataframe,
38
    configure_expectations=configure_expectations,
39
    result_format="SUMMARY",
40
    context_type="ephemeral",
41
)

1
from great_expectations_provider.operators.validate_batch import GXValidateBatchOperator
2
from great_expectations_provider.common.external_connections import build_snowflake_connection_string
3
4
SNOWFLAKE_CONN_ID = "snowflake_default"
5
6
7
def configure_batch_definition(context):
8
    """Configure a batch definition for a SQL table."""
9
    connection_string = build_snowflake_connection_string(
10
        conn_id=SNOWFLAKE_CONN_ID,
11
        schema="my_schema"
12
    )
13
14
    data_source = context.data_sources.add_sql(
15
        name="snowflake_ds",
16
        connection_string=connection_string,
17
    )
18
19
    table_asset = data_source.add_table_asset(
20
        name="daily_report_table",
21
        table_name="daily_planet_report",
22
    )
23
24
    return table_asset.add_batch_definition_whole_table(name="full_table")
25
26
27
def configure_expectations(context):
28
    """Define expectations and add them to the context."""
29
    import great_expectations.expectations as gxe
30
    from great_expectations import ExpectationSuite
31
32
    return context.suites.add_or_update(
33
        ExpectationSuite(
34
            name="daily_report_batch_suite",
35
            expectations=[
36
                gxe.ExpectColumnValuesToNotBeNull(column="planet_name"),
37
                gxe.ExpectColumnValuesToNotBeNull(column="total_passengers"),
38
                gxe.ExpectColumnValuesToBeBetween(
39
                    column="total_net_fare_usd", min_value=0
40
                ),
41
                gxe.ExpectColumnValuesToBeBetween(
42
                    column="total_discounts_usd", min_value=0
43
                ),
44
            ],
45
        )
46
    )
47
48
49
_validate_batch = GXValidateBatchOperator(
50
    task_id="validate_with_gx_batch",
51
    configure_batch_definition=configure_batch_definition,
52
    configure_expectations=configure_expectations,
53
    result_format="SUMMARY",
54
    context_type="ephemeral",
55
)

1
def configure_batch_definition(context):
2
    """Configure a batch definition to read from a CSV file."""
3
    from pathlib import Path
4
5
    data_source = context.data_sources.add_pandas_filesystem(
6
        name="my_datasource",
7
        base_directory=Path("/path/to/data"),
8
    )
9
    csv_asset = data_source.add_csv_asset(name="daily_report_csv")
10
    return csv_asset.add_batch_definition_path(
11
        name="daily_report_batch",
12
        path="daily_report.csv",
13
    )

1
from great_expectations_provider.operators.validate_checkpoint import GXValidateCheckpointOperator
2
from great_expectations_provider.common.external_connections import build_snowflake_connection_string
3
4
SNOWFLAKE_CONN_ID = "snowflake_default"
5
6
7
def configure_checkpoint(context):
8
    """Configure a full GX checkpoint with validation and actions."""
9
    import great_expectations.expectations as gxe
10
    from great_expectations import Checkpoint, ExpectationSuite, ValidationDefinition
11
12
    connection_string = build_snowflake_connection_string(
13
        conn_id=SNOWFLAKE_CONN_ID,
14
        schema="my_schema"
15
    )
16
17
    batch_definition = (
18
        context.data_sources.add_sql(
19
            name="snowflake_ds",
20
            connection_string=connection_string,
21
        )
22
        .add_table_asset(
23
            name="daily_report_table",
24
            table_name="daily_planet_report",
25
        )
26
        .add_batch_definition_whole_table(name="full_table")
27
    )
28
29
    expectation_suite = context.suites.add(
30
        ExpectationSuite(
31
            name="daily_report_checkpoint_suite",
32
            expectations=[
33
                gxe.ExpectColumnValuesToNotBeNull(column="planet_name"),
34
                gxe.ExpectColumnValuesToNotBeNull(column="total_passengers"),
35
                gxe.ExpectColumnValuesToBeBetween(
36
                    column="total_net_fare_usd", min_value=0
37
                ),
38
                gxe.ExpectColumnValuesToBeBetween(
39
                    column="total_discounts_usd", min_value=0
40
                ),
41
            ],
42
        )
43
    )
44
45
    validation_definition = context.validation_definitions.add(
46
        ValidationDefinition(
47
            name="daily_report_validation",
48
            data=batch_definition,
49
            suite=expectation_suite,
50
        )
51
    )
52
53
    return context.checkpoints.add(
54
        Checkpoint(
55
            name="daily_report_checkpoint",
56
            validation_definitions=[validation_definition],
57
            actions=[],  # Add SlackNotificationAction, EmailAction, etc.
58
        )
59
    )
60
61
62
_validate_checkpoint = GXValidateCheckpointOperator(
63
    task_id="validate_with_gx_checkpoint",
64
    configure_checkpoint=configure_checkpoint,
65
    context_type="ephemeral",
66
)