How to Integrate Soda Core into Your Airflow DAGs for Proactive Data Quality Checks

The Data Quality Gap in ML Pipelines

ML pipelines are only as good as their data. A model trained on stale data produces stale predictions. A feature pipeline that silently drops null values can introduce bias. A data source that changes its schema without notice can cause cascading failures across the entire ML stack. Yet most ML pipelines have minimal data quality checks—maybe a row count assertion or a null check on the target variable. When data quality issues slip through, they manifest as model performance degradation that's difficult to diagnose because the root cause is upstream in the data pipeline, not in the model itself.

Soda Core is an open-source data quality framework that lets you define quality checks as YAML configuration, execute them against any data source, and integrate results into orchestration tools like Airflow. This article provides a complete integration guide for embedding Soda Core checks into Airflow DAGs, covering check definition, alerting, automatic pipeline blocking, and quality trend monitoring.

Phase 1: Soda Core Configuration

Data Source Configuration

# soda/configuration.yml — Soda Core data source configuration
data_sources:
  warehouse:
    type: snowflake
    connection:
      account: ${SNOWFLAKE_ACCOUNT}
      username: ${SNOWFLAKE_USER}
      password: ${SNOWFLAKE_PASSWORD}
      warehouse: ML_ANALYTICS_WH
      database: PRODUCTION_DB
      schema: ML_FEATURES
  
  feature_store:
    type: postgres
    connection:
      host: ${FEAST_DB_HOST}
      port: 5432
      username: ${FEAST_DB_USER}
      password: ${FEAST_DB_PASSWORD}
      database: feast_online
  
  raw_data:
    type: s3
    connection:
      access_key_id: ${AWS_ACCESS_KEY_ID}
      secret_access_key: ${AWS_SECRET_ACCESS_KEY}
      region: us-east-1

# Soda Cloud connection for dashboards and alerting
soda_cloud:
  host: cloud.soda.io
  api_key_id: ${SODA_API_KEY}
  api_key_secret: ${SODA_API_SECRET}

Quality Check Definitions

# soda/checks/raw_customer_data.yml — Quality checks for raw customer data
# These checks run on the raw data BEFORE any processing

checks for raw_customers:
  # Schema checks
  - schema:
      name: "Raw customer data schema validation"
      fail:
        when required column missing:
          - customer_id
          - email
          - signup_date
          - plan_type
        when wrong column type:
          customer_id: varchar
          signup_date: timestamp
  
  # Row count checks
  - row_count:
      name: "Minimum row count"
      warn: when < 100000
      fail: when < 50000
  
  - row_count:
      name: "Row count change from yesterday"
      warn: when change > 20%
      fail: when change > 50%
  
  # Freshness check
  - freshness(signup_date):
      name: "Data freshness"
      warn: when > 24h
      fail: when > 48h
  
  # Null checks
  - missing_count(customer_id):
      name: "No null customer IDs"
      fail: when > 0
  
  - missing_percent(email):
      name: "Email completeness"
      warn: when > 5%
      fail: when > 15%
  
  # Uniqueness checks
  - duplicate_count(customer_id):
      name: "Customer ID uniqueness"
      fail: when > 0
  
  # Validity checks
  - invalid_count(email):
      name: "Email format validation"
      fail: when > 100
      valid format: email
  
  - valid_values(plan_type):
      name: "Plan type validity"
      fail: when > 0
      valid values:
        - basic
        - pro
        - enterprise
  
  # Distribution checks (detect data drift)
  - distribution_difference(age):
      name: "Age distribution stability"
      fail: when > 0.2  # KS statistic threshold
      reference_distribution:
        source: warehouse
        dataset: data_quality_reference.age_distribution
        method: ks
  
  # Anomaly detection
  - anomaly score for row_count:
      name: "Row count anomaly detection"
      warn: when > 0.7
      fail: when > 0.9
  
  - anomaly score for avg(transaction_amount):
      name: "Transaction amount anomaly detection"
      warn: when > 0.7
      fail: when > 0.9

# Checks for raw_transactions
checks for raw_transactions:
  - row_count:
      name: "Transaction volume"
      warn: when < 500000
      fail: when < 100000
  
  - freshness(transaction_date):
      name: "Transaction freshness"
      warn: when > 6h
      fail: when > 24h
  
  - min(transaction_amount):
      name: "No negative transaction amounts"
      fail: when < 0
  
  - max(transaction_amount):
      name: "Transaction amount upper bound"
      warn: when > 50000
      fail: when > 100000
  
  - missing_count(transaction_id):
      fail: when > 0
  
  - duplicate_count(transaction_id):
      fail: when > 0

Feature Store Quality Checks

# soda/checks/feature_store.yml — Quality checks for computed features
checks for customer_features:
  # Feature completeness
  - missing_percent(total_spend_30d):
      name: "Feature completeness: total_spend_30d"
      warn: when > 2%
      fail: when > 10%
  
  - missing_percent(dau_7d):
      name: "Feature completeness: dau_7d"
      warn: when > 5%
      fail: when > 15%
  
  # Feature value ranges
  - min(total_spend_30d):
      name: "total_spend_30d non-negative"
      fail: when < 0
  
  - max(dau_7d):
      name: "dau_7d reasonable upper bound"
      warn: when > 10000
      fail: when > 50000
  
  # Feature distribution stability (drift detection)
  - distribution_difference(total_spend_30d):
      name: "Spend feature drift"
      fail: when > 0.2
      reference_distribution:
        source: warehouse
        dataset: data_quality_reference.spend_distribution
  
  - distribution_difference(session_duration_avg):
      name: "Session duration drift"
      fail: when > 0.25
      reference_distribution:
        source: warehouse
        dataset: data_quality_reference.session_distribution
  
  # Feature correlation checks (detect broken joins)
  - failed rows:
      name: "Features join integrity"
      fail: when > 0
      query: |
        SELECT customer_id
        FROM customer_features
        WHERE total_spend_30d > 0 AND transaction_count_30d = 0
  
  # Point-in-time correctness
  - failed rows:
      name: "No future data leakage"
      fail: when > 0
      query: |
        SELECT customer_id, feature_date, last_login
        FROM customer_features
        WHERE last_login > feature_date

Phase 2: Airflow Integration

Soda Core Airflow Operator

# airflow/operators/soda_operator.py — Custom Airflow operator for Soda Core
from airflow.models import BaseOperator
from airflow.exceptions import AirflowFailException, AirflowSkipException
from soda.scan import Scan
from typing import List, Dict, Optional
import json

class SodaCoreCheckOperator(BaseOperator):
    """
    Airflow operator that runs Soda Core data quality checks.
    
    Fails the task if any checks fail (blocks downstream pipeline).
    Warns (but doesn't fail) on warning-level check failures.
    """
    
    template_fields = ["data_source", "checks_files", "variables"]
    
    def __init__(
        self,
        data_source: str,
        checks_files: List[str],
        variables: Optional[Dict[str, str]] = None,
        fail_on_warning: bool = False,
        publish_to_soda_cloud: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.data_source = data_source
        self.checks_files = checks_files
        self.variables = variables or {}
        self.fail_on_warning = fail_on_warning
        self.publish_to_soda_cloud = publish_to_soda_cloud
    
    def execute(self, context):
        """Execute Soda Core checks."""
        
        scan = Scan()
        
        # Configure data source
        scan.set_data_source_name(self.data_source)
        scan.add_configuration_yaml_files(
            paths=["soda/configuration.yml"]
        )
        
        # Add check files
        for checks_file in self.checks_files:
            scan.add_sodacl_yaml_file(checks_file)
        
        # Add variables (e.g., date partitions)
        for key, value in self.variables.items():
            scan.add_variables({key: value})
        
        # Set scan name for Soda Cloud
        scan.set_scan_definition_name(
            f"{context['dag'].dag_id}.{context['task_instance'].task_id}"
        )
        scan.set_verbose(True)
        
        # Execute scan
        scan.execute()
        
        # Process results
        scan_result = scan.get_scan_results()
        
        # Log results
        self.log.info(f"Soda Core scan complete:")
        self.log.info(f"  Checks passed: {scan_result.pass_count}")
        self.log.info(f"  Checks failed: {scan_result.fail_count}")
        self.log.info(f"  Checks warned: {scan_result.warn_count}")
        self.log.info(f"  Checks errored: {scan_result.error_count}")
        
        # Log individual check results
        for check in scan_result.checks:
            status_emoji = {
                "pass": "✅",
                "fail": "❌",
                "warn": "⚠️",
                "error": "🔥",
            }.get(check.outcome, "❓")
            
            self.log.info(
                f"  {status_emoji} {check.name}: {check.outcome}"
            )
            if check.outcome == "fail":
                self.log.info(f"      Diagnostics: {check.diagnostic}")
        
        # Push results to XCom for downstream tasks
        context['ti'].xcom_push(
            key='soda_scan_results',
            value={
                "pass_count": scan_result.pass_count,
                "fail_count": scan_result.fail_count,
                "warn_count": scan_result.warn_count,
                "error_count": scan_result.error_count,
                "checks": [
                    {
                        "name": c.name,
                        "outcome": c.outcome,
                        "diagnostic": str(c.diagnostic) if c.diagnostic else None,
                    }
                    for c in scan_result.checks
                ],
            },
        )
        
        # Fail the task if checks failed
        if scan_result.fail_count > 0 or scan_result.error_count > 0:
            raise AirflowFailException(
                f"Soda Core checks failed: {scan_result.fail_count} failures, "
                f"{scan_result.error_count} errors"
            )
        
        # Optionally fail on warnings
        if self.fail_on_warning and scan_result.warn_count > 0:
            raise AirflowFailException(
                f"Soda Core checks warned: {scan_result.warn_count} warnings"
            )
        
        return {
            "status": "passed" if scan_result.warn_count == 0 else "passed_with_warnings",
            "pass_count": scan_result.pass_count,
            "fail_count": scan_result.fail_count,
            "warn_count": scan_result.warn_count,
        }

Full DAG with Soda Core Checks

# airflow/dags/ml_pipeline_with_quality.py — ML pipeline with data quality gates
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.external_task import ExternalTaskSensor
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.operators.empty import EmptyOperator
from datetime import datetime, timedelta
from operators.soda_operator import SodaCoreCheckOperator

default_args = {
    "owner": "ml-platform",
    "retries": 1,
    "retry_delay": timedelta(minutes=15),
}

with DAG(
    "ml_pipeline_with_quality",
    default_args=default_args,
    schedule_interval="0 6 * * *",  # 6 AM daily
    start_date=datetime(2025, 1, 1),
    catchup=False,
    max_active_runs=1,
    tags=["ml", "data-quality"],
) as dag:
    
    # ================================================================
    # Stage 1: Raw Data Quality Checks
    # ================================================================
    
    raw_data_quality = SodaCoreCheckOperator(
        task_id="check_raw_data_quality",
        data_source="warehouse",
        checks_files=[
            "soda/checks/raw_customer_data.yml",
            "soda/checks/raw_transactions.yml",
        ],
        variables={
            "date": "{{ ds }}",
        },
        fail_on_warning=False,
    )
    
    # ================================================================
    # Stage 2: Data Processing
    # ================================================================
    
    process_data = PythonOperator(
        task_id="process_data",
        python_callable=run_data_processing,
    )
    
    # ================================================================
    # Stage 3: Feature Store Quality Checks
    # ================================================================
    
    feature_quality = SodaCoreCheckOperator(
        task_id="check_feature_quality",
        data_source="feature_store",
        checks_files=[
            "soda/checks/feature_store.yml",
        ],
        variables={
            "date": "{{ ds }}",
        },
        fail_on_warning=True,  # Fail on warnings for features (closer to model)
    )
    
    # ================================================================
    # Stage 4: Training Data Quality Checks
    # ================================================================
    
    training_data_quality = SodaCoreCheckOperator(
        task_id="check_training_data_quality",
        data_source="warehouse",
        checks_files=[
            "soda/checks/training_data.yml",
        ],
        variables={
            "date": "{{ ds }}",
        },
    )
    
    # ================================================================
    # Stage 5: Model Training
    # ================================================================
    
    train_model = PythonOperator(
        task_id="train_model",
        python_callable=run_model_training,
    )
    
    # ================================================================
    # Stage 6: Model Evaluation Quality Checks
    # ================================================================
    
    model_quality = SodaCoreCheckOperator(
        task_id="check_model_quality",
        data_source="warehouse",
        checks_files=[
            "soda/checks/model_predictions.yml",
        ],
    )
    
    # ================================================================
    # Stage 7: Deploy Model
    # ================================================================
    
    deploy_model = PythonOperator(
        task_id="deploy_model",
        python_callable=deploy_to_production,
    )
    
    # ================================================================
    # Quality gate failure handler
    # ================================================================
    
    def handle_quality_failure(**kwargs):
        """Handle data quality check failures."""
        ti = kwargs['ti']
        
        # Get failed check details from upstream
        raw_results = ti.xcom_pull(
            task_ids="check_raw_data_quality",
            key="soda_scan_results",
        )
        
        feature_results = ti.xcom_pull(
            task_ids="check_feature_quality",
            key="soda_scan_results",
        )
        
        # Send alert with details
        failed_checks = []
        for source, results in [("raw_data", raw_results), ("features", feature_results)]:
            if results:
                for check in results.get("checks", []):
                    if check["outcome"] in ["fail", "error"]:
                        failed_checks.append(f"[{source}] {check['name']}: {check.get('diagnostic', '')}")
        
        alert_message = "ML Pipeline blocked due to data quality failures:\n\n"
        alert_message += "\n".join(failed_checks)
        
        # Send Slack alert
        send_slack_alert(channel="#ml-alerts", message=alert_message)
        
        # Create PagerDuty incident for critical failures
        create_pagerduty_incident(
            summary="ML Pipeline Data Quality Failure",
            details=alert_message,
            severity="critical",
        )
    
    quality_failure_handler = PythonOperator(
        task_id="handle_quality_failure",
        python_callable=handle_quality_failure,
        trigger_rule="one_failed",
    )
    
    # ================================================================
    # DAG Dependencies
    # ================================================================
    
    raw_data_quality >> process_data >> feature_quality >> training_data_quality >> train_model >> model_quality >> deploy_model
    
    # All quality checks connect to failure handler
    [raw_data_quality, feature_quality, training_data_quality, model_quality] >> quality_failure_handler

Phase 3: Quality Monitoring

Data Quality Dashboard

# monitoring/quality_dashboard.py — Data quality trend monitoring
import pandas as pd
import plotly.express as px
from datetime import datetime, timedelta

class DataQualityDashboard:
    """Monitor data quality trends over time."""
    
    def __init__(self, soda_cloud_client):
        self.soda = soda_cloud_client
    
    def get_quality_trends(self, dataset: str, days: int = 30) -> pd.DataFrame:
        """Get quality check results over time for a dataset."""
        
        scans = self.soda.get_scans(
            dataset=dataset,
            start_date=datetime.utcnow() - timedelta(days=days),
        )
        
        rows = []
        for scan in scans:
            for check in scan.checks:
                rows.append({
                    "scan_date": scan.timestamp,
                    "check_name": check.name,
                    "outcome": check.outcome,
                    "diagnostic": check.diagnostic,
                    "dataset": dataset,
                })
        
        return pd.DataFrame(rows)
    
    def generate_health_score(self, dataset: str, days: int = 7) -> float:
        """
        Compute a data health score (0-100) based on recent quality checks.
        """
        trends = self.get_quality_trends(dataset, days)
        
        if trends.empty:
            return 100.0
        
        total_checks = len(trends)
        passed = len(trends[trends["outcome"] == "pass"])
        warned = len(trends[trends["outcome"] == "warn"])
        failed = len(trends[trends["outcome"].isin(["fail", "error"])])
        
        # Weighted score: pass=100%, warn=50%, fail=0%
        score = (passed * 100 + warned * 50) / total_checks
        
        return round(score, 1)
    
    def detect_quality_degradation(self, dataset: str) -> dict:
        """Detect if data quality is degrading over time."""
        
        trends = self.get_quality_trends(dataset, days=30)
        
        if len(trends) < 10:
            return {"degrading": False, "reason": "insufficient data"}
        
        # Compare last 7 days vs. previous 7 days
        recent = trends[trends["scan_date"] >= datetime.utcnow() - timedelta(days=7)]
        previous = trends[
            (trends["scan_date"] >= datetime.utcnow() - timedelta(days=14)) &
            (trends["scan_date"] < datetime.utcnow() - timedelta(days=7))
        ]
        
        recent_fail_rate = len(recent[recent["outcome"].isin(["fail", "error"])]) / max(len(recent), 1)
        previous_fail_rate = len(previous[previous["outcome"].isin(["fail", "error"])]) / max(len(previous), 1)
        
        degradation = recent_fail_rate - previous_fail_rate
        
        return {
            "degrading": degradation > 0.05,
            "recent_fail_rate": recent_fail_rate,
            "previous_fail_rate": previous_fail_rate,
            "degradation_amount": degradation,
        }

Conclusion

Integrating Soda Core into Airflow DAGs transforms data quality from a reactive afterthought into a proactive pipeline gate. Quality checks defined as YAML are version-controlled, reviewable, and reusable across pipelines. When embedded at every stage of the ML pipeline—raw data, features, training data, predictions—quality checks catch issues before they propagate downstream and corrupt models. The Airflow operator blocks pipeline execution when checks fail, preventing bad data from reaching training jobs or production models. Soda Cloud provides trend monitoring and alerting, enabling teams to detect gradual quality degradation before it causes model failures. Teams that implement systematic data quality checks reduce model performance incidents by 60-80% and spend significantly less time debugging data-related issues in production.