Quality Gates: встраивание row count, null check и drift-проверок

def assert_min_row_count(
    df: DataFrame,
    min_rows: int,
    stage: str = "unnamed",
) -> DataFrame:
    """
    Fail-fast барьер: пайплайн падает, если строк меньше min_rows.
    Возвращает исходный DataFrame, чтобы встраиваться в цепочку трансформаций.
    """
    actual = df.count()
    if actual < min_rows:
        raise ValueError(
            f"[QualityGate/{stage}] Row count too low: "
            f"expected >= {min_rows:,}, got {actual:,}"
        )
    return df

bronze_df = spark.read.parquet(bronze_path)
bronze_df = assert_min_row_count(bronze_df, min_rows=1000, stage="bronze_orders")

def assert_relative_row_count(
    current_df: DataFrame,
    baseline_count: int,
    max_drop_ratio: float = 0.30,
    stage: str = "unnamed",
) -> DataFrame:
    """
    Проверяем, что текущий батч не упал более чем на max_drop_ratio
    относительно baseline (например, вчерашнего батча).
    """
    actual = current_df.count()
    if baseline_count > 0:
        drop_ratio = (baseline_count - actual) / baseline_count
        if drop_ratio > max_drop_ratio:
            raise ValueError(
                f"[QualityGate/{stage}] Row count dropped by "
                f"{drop_ratio:.1%}: baseline={baseline_count:,}, "
                f"actual={actual:,}, threshold={max_drop_ratio:.1%}"
            )
    return current_df

def get_historical_avg_count(
    spark: SparkSession,
    table: str,
    date_col: str,
    days: int = 7,
    reference_date: str = None,
) -> float:
    """
    Считаем среднее количество строк за последние N дней.
    Используем это как динамический baseline для row count gate.
    """
    ref = reference_date or "current_date()"
    query = f"""
        SELECT AVG(daily_count) as avg_count
        FROM (
            SELECT {date_col}, COUNT(*) as daily_count
            FROM {table}
            WHERE {date_col} >= date_sub({ref}, {days})
              AND {date_col} < {ref}
            GROUP BY {date_col}
        )
    """
    result = spark.sql(query).collect()
    return float(result[0]["avg_count"] or 0)

# Null propagation через арифметику
# 100 + null = null - выручка по заказу будет null!
df = df.withColumn("total", col("price") * col("quantity") + col("discount"))

# Null в GROUP BY создаёт отдельную группу
# Если user_id бывает null - это отдельная "группа" в агрегате
df.groupBy("user_id").agg(sum("amount"))

# Null в JOIN - строки не находят пары (null != null по SQL семантике)
orders.join(users, "user_id")  # строки с null user_id потеряются

# АНТИПАТТЕРН: N отдельных Action-операций = N сканов датасета
for col_name in critical_cols:
    null_count = df.filter(col(col_name).isNull()).count()
    if null_count > 0:
        raise ValueError(f"Nulls found in {col_name}: {null_count}")

from pyspark.sql import functions as F

def check_nulls_single_pass(
    df: DataFrame,
    critical_cols: list,
    nullable_cols: dict = None,
    stage: str = "unnamed",
) -> dict:
    """
    Подсчёт null-значений за один Action (один скан DataFrame).

    critical_cols - колонки, где любой null = ошибка.
    nullable_cols - {col_name: max_null_ratio}, допустимый процент null.

    Возвращает словарь {col_name: null_count} для логирования.
    """
    nullable_cols = nullable_cols or {}

    # Формируем список агрегатов за один проход
    agg_exprs = [
        F.count(F.when(F.col(c).isNull(), 1)).alias(c)
        for c in (critical_cols + list(nullable_cols.keys()))
    ]

    # Один Action = один Spark Job = один скан
    null_counts_row = df.agg(*agg_exprs).collect()[0]
    null_counts = {c: null_counts_row[c] for c in null_counts_row.__fields__}

    total_rows = df.count()  # можно совместить в один agg с COUNT(*)

    errors = []

    # Проверка критических колонок
    for col_name in critical_cols:
        count = null_counts.get(col_name, 0)
        if count > 0:
            errors.append(
                f"  {col_name}: {count:,} nulls (CRITICAL - must be 0)"
            )

    # Проверка допустимого null-ratio
    for col_name, max_ratio in nullable_cols.items():
        count = null_counts.get(col_name, 0)
        actual_ratio = count / total_rows if total_rows > 0 else 0
        if actual_ratio > max_ratio:
            errors.append(
                f"  {col_name}: null ratio {actual_ratio:.2%} "
                f"> threshold {max_ratio:.2%} ({count:,}/{total_rows:,})"
            )

    if errors:
        raise ValueError(
            f"[QualityGate/{stage}] Null check failed:\n" + "\n".join(errors)
        )

    return null_counts

def check_nulls_and_count_single_pass(
    df: DataFrame,
    critical_cols: list,
) -> tuple:
    """
    За один Action считаем и общее количество строк, и null по каждой колонке.
    Ключевая оптимизация: COUNT(*) + все COUNT(WHEN isNull) в одном select.
    """
    agg_exprs = [F.count(F.lit(1)).alias("__total__")] + [
        F.count(F.when(F.col(c).isNull(), 1)).alias(c)
        for c in critical_cols
    ]

    row = df.agg(*agg_exprs).collect()[0]
    total = row["__total__"]
    null_counts = {c: row[c] for c in critical_cols}

    return total, null_counts

# Конфигурация проверок для Bronze -> Silver

NULL_CHECK_CONFIG = {
    # Критические поля: любой null = Fail-Fast
    "critical": [
        "transaction_id",
        "user_id",
        "created_at",
        "amount",
    ],
    # Допустимые null с порогами
    "nullable": {
        "promo_code": 0.95,       # до 95% строк могут не иметь промокода
        "delivery_address": 0.30,  # до 30% - самовывоз
        "phone": 0.15,             # до 15% - не указали телефон
        "utm_source": 0.40,        # до 40% - прямые заходы
    },
}

from pyspark.sql.types import StructType
import json

def extract_schema_snapshot(df: DataFrame) -> dict:
    """
    Преобразуем schema в сериализуемый словарь для хранения и сравнения.
    Используем json() для получения стабильного текстового представления.
    """
    return json.loads(df.schema.json())


def compare_schemas(
    current_df: DataFrame,
    expected_schema: StructType,
    strict: bool = True,
    stage: str = "unnamed",
) -> dict:
    """
    Сравниваем текущую схему DataFrame с эталонной.

    strict=True  - любое расхождение = ошибка (Fail-Fast)
    strict=False - только предупреждение, если новые колонки добавились;
                   ошибка - если существующие колонки изменились/пропали
    """
    current = {f.name: f for f in current_df.schema.fields}
    expected = {f.name: f for f in expected_schema.fields}

    missing_cols = []     # колонки исчезли из источника
    type_changes = []     # тип колонки изменился
    new_cols = []         # появились новые колонки

    for name, exp_field in expected.items():
        if name not in current:
            missing_cols.append(name)
        elif current[name].dataType != exp_field.dataType:
            type_changes.append(
                f"  {name}: expected {exp_field.dataType}, "
                f"got {current[name].dataType}"
            )

    for name in current:
        if name not in expected:
            new_cols.append(name)

    report = {
        "missing_columns": missing_cols,
        "type_changes": type_changes,
        "new_columns": new_cols,
        "is_valid": not missing_cols and not type_changes,
    }

    # Критические нарушения - всегда Fail-Fast
    errors = []
    if missing_cols:
        errors.append(f"Missing columns: {missing_cols}")
    if type_changes:
        errors.append("Type changes detected:\n" + "\n".join(type_changes))

    if errors:
        raise ValueError(
            f"[QualityGate/{stage}] Schema drift detected:\n"
            + "\n".join(errors)
        )

    # Новые колонки - зависит от режима
    if new_cols and strict:
        raise ValueError(
            f"[QualityGate/{stage}] Unexpected new columns: {new_cols}. "
            f"Update expected schema or set strict=False."
        )
    elif new_cols:
        print(
            f"[QualityGate/{stage}] WARNING: New columns detected: {new_cols}. "
            f"Notify data producers."
        )

    return report

# schemas/silver_orders_v2.json - хранится в git-репозитории

SILVER_ORDERS_SCHEMA_JSON = """
{
  "type": "struct",
  "fields": [
    {"name": "order_id", "type": "string", "nullable": false},
    {"name": "user_id", "type": "long", "nullable": false},
    {"name": "amount", "type": "double", "nullable": false},
    {"name": "status", "type": "string", "nullable": false},
    {"name": "created_at", "type": "timestamp", "nullable": false},
    {"name": "updated_at", "type": "timestamp", "nullable": true}
  ]
}
"""

def load_expected_schema(schema_json: str) -> StructType:
    return StructType.fromJson(json.loads(schema_json))

# В пайплайне:
expected = load_expected_schema(SILVER_ORDERS_SCHEMA_JSON)
compare_schemas(bronze_df, expected, strict=False, stage="bronze_to_silver")

from pyspark.sql.functions import (
    min, max, avg, stddev, count, countDistinct, percentile_approx
)

def compute_numerical_stats(
    df: DataFrame,
    numerical_cols: list,
) -> dict:
    """
    Вычисляем статистику по числовым колонкам за один проход.
    Используем approx функции там, где точность не критична.
    """
    agg_exprs = []
    for c in numerical_cols:
        agg_exprs.extend([
            F.min(c).alias(f"{c}_min"),
            F.max(c).alias(f"{c}_max"),
            F.avg(c).alias(f"{c}_avg"),
            F.stddev(c).alias(f"{c}_stddev"),
            F.percentile_approx(c, 0.5).alias(f"{c}_median"),
            F.percentile_approx(c, 0.99).alias(f"{c}_p99"),
        ])

    row = df.agg(*agg_exprs).collect()[0]

    stats = {}
    for c in numerical_cols:
        stats[c] = {
            "min": row[f"{c}_min"],
            "max": row[f"{c}_max"],
            "avg": row[f"{c}_avg"],
            "stddev": row[f"{c}_stddev"],
            "median": row[f"{c}_median"],
            "p99": row[f"{c}_p99"],
        }
    return stats

def assert_no_numerical_drift(
    current_stats: dict,
    baseline_stats: dict,
    col_name: str,
    max_avg_deviation: float = 2.0,
    stage: str = "unnamed",
) -> None:
    """
    Проверяем, что среднее значение не отклонилось более чем в max_avg_deviation раз
    от исторического baseline.

    Пример: max_avg_deviation=2.0 означает, что avg может вырасти
    не более чем в 2 раза или упасть не более чем в 2 раза.
    """
    current_avg = current_stats[col_name]["avg"]
    baseline_avg = baseline_stats[col_name]["avg"]

    if baseline_avg is None or baseline_avg == 0:
        return  # нет исторических данных - пропускаем

    ratio = current_avg / baseline_avg

    if ratio > max_avg_deviation or ratio < (1.0 / max_avg_deviation):
        raise ValueError(
            f"[QualityGate/{stage}] Data drift in '{col_name}': "
            f"avg changed {ratio:.1f}x "
            f"(baseline={baseline_avg:.2f}, current={current_avg:.2f})"
        )

def compute_categorical_distribution(
    df: DataFrame,
    col_name: str,
    top_n: int = 50,
) -> dict:
    """
    Вычисляем доли каждого значения в колонке.
    Возвращает словарь {value: ratio}.
    """
    total = df.count()
    dist = (
        df.groupBy(col_name)
        .count()
        .withColumn("ratio", F.col("count") / total)
        .orderBy(F.col("count").desc())
        .limit(top_n)
        .collect()
    )
    return {str(row[col_name]): row["ratio"] for row in dist}


def assert_no_categorical_drift(
    current_dist: dict,
    baseline_dist: dict,
    col_name: str,
    max_ratio_change: float = 0.50,
    stage: str = "unnamed",
) -> None:
    """
    Проверяем, что доля каждого значения не изменилась более чем на max_ratio_change.

    Пример: если baseline[paid] = 0.70, то current[paid] должен быть в диапазоне
    [0.35, 1.0] при max_ratio_change=0.50.
    """
    alerts = []
    for value, baseline_ratio in baseline_dist.items():
        current_ratio = current_dist.get(value, 0.0)
        if baseline_ratio > 0:
            change = abs(current_ratio - baseline_ratio) / baseline_ratio
            if change > max_ratio_change:
                alerts.append(
                    f"  '{value}': {baseline_ratio:.1%} -> {current_ratio:.1%} "
                    f"(delta {change:.0%})"
                )

    # Новые неожиданные значения
    new_values = set(current_dist.keys()) - set(baseline_dist.keys())
    if new_values:
        alerts.append(f"  New values appeared: {new_values}")

    if alerts:
        raise ValueError(
            f"[QualityGate/{stage}] Categorical drift in '{col_name}':\n"
            + "\n".join(alerts)
        )

from dataclasses import dataclass

@dataclass
class QualityRule:
    name: str               # Идентификатор правила
    expression: str         # SQL-выражение, True = строка валидна
    severity: str           # "critical" | "warning"
    description: str        # Человекочитаемое описание

# Правила для Silver Orders
SILVER_ORDERS_RULES = [
    QualityRule(
        name="positive_amount",
        expression="amount > 0",
        severity="critical",
        description="Сумма заказа должна быть положительной",
    ),
    QualityRule(
        name="valid_status",
        expression="status IN ('paid', 'refunded', 'pending', 'cancelled')",
        severity="critical",
        description="Статус должен быть из допустимого справочника",
    ),
    QualityRule(
        name="delivery_after_order",
        expression="delivery_date IS NULL OR delivery_date >= created_at",
        severity="critical",
        description="Дата доставки не может быть раньше даты заказа",
    ),
    QualityRule(
        name="reasonable_amount",
        expression="amount < 10000000",
        severity="warning",
        description="Сумма заказа > 10M - возможная ошибка или мошенничество",
    ),
    QualityRule(
        name="not_future_order",
        expression="created_at <= current_timestamp()",
        severity="critical",
        description="Заказ не может быть из будущего",
    ),
]

from pyspark.sql import DataFrame
from pyspark.sql.functions import expr, count, when, lit

def run_quality_rules(
    df: DataFrame,
    rules: list,
    mode: str = "fail_fast",
    stage: str = "unnamed",
) -> tuple:
    """
    Запускает все правила за один Action.

    mode="fail_fast"   - поднимает exception при первом критическом нарушении
    mode="quarantine"  - разделяет DataFrame на valid и invalid

    Возвращает (valid_df, invalid_df, quality_report).
    """
    # Строим флаги нарушений для каждого правила
    # Добавляем булевую колонку per_rule_violated
    flagged_df = df
    for rule in rules:
        flagged_df = flagged_df.withColumn(
            f"_violated_{rule.name}",
            ~expr(rule.expression),  # True если правило НАРУШЕНО
        )

    # Считаем нарушения за один Action
    agg_exprs = [
        count(when(F.col(f"_violated_{rule.name}"), lit(1))).alias(rule.name)
        for rule in rules
    ] + [count(lit(1)).alias("__total__")]

    stats_row = flagged_df.agg(*agg_exprs).collect()[0]
    total = stats_row["__total__"]

    quality_report = {
        "total_rows": total,
        "violations": {},
        "is_valid": True,
        "stage": stage,
    }

    critical_violations = []
    warning_violations = []

    for rule in rules:
        violation_count = stats_row[rule.name]
        violation_ratio = violation_count / total if total > 0 else 0

        quality_report["violations"][rule.name] = {
            "count": violation_count,
            "ratio": violation_ratio,
            "severity": rule.severity,
            "description": rule.description,
        }

        if violation_count > 0:
            msg = (
                f"  [{rule.severity.upper()}] {rule.name}: "
                f"{violation_count:,} rows ({violation_ratio:.2%}) "
                f"— {rule.description}"
            )
            if rule.severity == "critical":
                critical_violations.append(msg)
                quality_report["is_valid"] = False
            else:
                warning_violations.append(msg)

    # Логируем предупреждения
    if warning_violations:
        print(
            f"[QualityGate/{stage}] Warnings:\n"
            + "\n".join(warning_violations)
        )

    if mode == "fail_fast" and critical_violations:
        raise ValueError(
            f"[QualityGate/{stage}] Critical violations found:\n"
            + "\n".join(critical_violations)
        )

    if mode == "quarantine":
        # Разделяем: строка плохая, если нарушает хотя бы одно critical правило
        violation_condition = None
        for rule in rules:
            if rule.severity == "critical":
                cond = F.col(f"_violated_{rule.name}")
                violation_condition = (
                    cond if violation_condition is None
                    else violation_condition | cond
                )

        # Удаляем служебные колонки
        flag_cols = [f"_violated_{rule.name}" for rule in rules]

        if violation_condition is not None:
            valid_df = flagged_df.filter(~violation_condition).drop(*flag_cols)
            invalid_df = flagged_df.filter(violation_condition).drop(*flag_cols)
        else:
            valid_df = flagged_df.drop(*flag_cols)
            invalid_df = df.limit(0)  # пустой DataFrame

        return valid_df, invalid_df, quality_report

    # Режим fail_fast: возвращаем исходный DataFrame без флагов
    flag_cols = [f"_violated_{rule.name}" for rule in rules]
    return df, None, quality_report

def assert_quality_gate(
    df: DataFrame,
    rules: list,
    stage: str,
) -> DataFrame:
    """
    Fail-Fast Quality Gate.
    Бросает ValueError при первом критическом нарушении.
    Возвращает исходный DataFrame для чейнинга.
    """
    _, _, report = run_quality_rules(df, rules, mode="fail_fast", stage=stage)

    print(
        f"[QualityGate/{stage}] PASSED: "
        f"{report['total_rows']:,} rows, "
        f"{len([r for r in report['violations'].values() if r['count'] == 0])} "
        f"rules OK"
    )
    return df

from pyspark.sql.functions import current_timestamp, lit

def apply_quarantine_gate(
    df: DataFrame,
    rules: list,
    quarantine_path: str,
    stage: str,
    batch_date: str,
) -> DataFrame:
    """
    Quarantine Quality Gate.

    Разделяет DataFrame на valid/invalid.
    Invalid строки записываются в карантинную таблицу с метаданными.
    Возвращает только valid строки для дальнейшей обработки.
    """
    valid_df, invalid_df, report = run_quality_rules(
        df, rules, mode="quarantine", stage=stage
    )

    total = report["total_rows"]

    if invalid_df is not None:
        # Обогащаем карантинные строки метаданными для последующего разбора
        quarantine_df = invalid_df.withColumns({
            "_quarantine_stage": lit(stage),
            "_quarantine_date": lit(batch_date),
            "_quarantine_ts": current_timestamp(),
            "_quarantine_rules": lit(
                str([r.name for r in rules if r.severity == "critical"])
            ),
        })

        quarantine_count = quarantine_df.count()
        if quarantine_count > 0:
            (
                quarantine_df
                .write
                .mode("append")
                .partitionBy("_quarantine_date", "_quarantine_stage")
                .parquet(quarantine_path)
            )
            print(
                f"[QualityGate/{stage}] Quarantine: "
                f"{quarantine_count:,}/{total:,} rows isolated "
                f"({quarantine_count/total:.1%})"
            )

    valid_count = valid_df.count()
    print(
        f"[QualityGate/{stage}] Passing to next stage: "
        f"{valid_count:,}/{total:,} rows ({valid_count/total:.1%})"
    )

    return valid_df

-- Добавляем ограничение на существующую таблицу
ALTER TABLE silver.orders
ADD CONSTRAINT positive_amount
CHECK (amount > 0);

ALTER TABLE silver.orders
ADD CONSTRAINT valid_status
CHECK (status IN ('paid', 'refunded', 'pending', 'cancelled'));

ALTER TABLE silver.orders
ADD CONSTRAINT delivery_after_order
CHECK (delivery_date IS NULL OR delivery_date >= created_at);

def add_delta_constraints(
    spark: SparkSession,
    table_name: str,
    constraints: dict,
) -> None:
    """
    Добавляем Delta CHECK constraints программно.
    constraints = {"constraint_name": "sql_expression"}
    """
    for name, expression in constraints.items():
        spark.sql(
            f"ALTER TABLE {table_name} ADD CONSTRAINT {name} CHECK ({expression})"
        )
        print(f"Added constraint: {name} -> {expression}")


add_delta_constraints(
    spark,
    "silver.orders",
    {
        "positive_amount": "amount > 0",
        "valid_status": "status IN ('paid', 'refunded', 'pending', 'cancelled')",
        "not_future_created": "created_at <= current_timestamp()",
    },
)

org.apache.spark.sql.delta.schema.DeltaInvariantViolationException:
CHECK constraint positive_amount (amount > 0) violated by row with values:
 - amount : -100.0

# Все constraints таблицы видны в SHOW TBLPROPERTIES
spark.sql("SHOW TBLPROPERTIES silver.orders")

# Или через describe detail
spark.sql("DESCRIBE DETAIL silver.orders").select("properties").show(truncate=False)

from airflow import DAG
from airflow.operators.python import BranchPythonOperator, PythonOperator
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
from airflow.utils.dates import days_ago
from datetime import timedelta

with DAG(
    dag_id="silver_orders_pipeline",
    schedule_interval="0 6 * * *",
    start_date=days_ago(1),
    catchup=False,
    default_args={"retries": 2, "retry_delay": timedelta(minutes=5)},
    tags=["silver", "orders"],
) as dag:

    # 1. Ingestion Bronze
    ingest_bronze = SparkSubmitOperator(
        task_id="ingest_bronze_orders",
        application="jobs/ingest_bronze_orders.py",
        application_args=["--date", "{{ ds }}"],
        conf={"spark.executor.memory": "4g"},
        pool="spark_heavy_pool",
    )

    # 2. Quality Gate: Bronze проверки (лёгкий job)
    bronze_quality_gate = SparkSubmitOperator(
        task_id="bronze_quality_gate",
        application="jobs/quality_gates/bronze_orders_gate.py",
        application_args=["--date", "{{ ds }}"],
        conf={"spark.executor.memory": "2g"},
        pool="spark_light_pool",
    )

    # 3. Branching: что делать при нарушении
    def route_on_quality_result(**context):
        """
        Читаем результат качества из XCom и решаем: продолжать или нет.
        В продакшне качество пишется в Metadata-таблицу.
        """
        ti = context["ti"]
        quality_passed = ti.xcom_pull(
            task_ids="bronze_quality_gate",
            key="quality_passed",
        )
        if quality_passed:
            return "transform_silver"
        else:
            return "notify_quality_failure"

    quality_branch = BranchPythonOperator(
        task_id="quality_branch",
        python_callable=route_on_quality_result,
        provide_context=True,
    )

    # 4a. Путь успеха: трансформация в Silver
    transform_silver = SparkSubmitOperator(
        task_id="transform_silver",
        application="jobs/transform_silver_orders.py",
        application_args=["--date", "{{ ds }}"],
        conf={"spark.executor.memory": "8g"},
        pool="spark_heavy_pool",
    )

    # 4b. Путь ошибки: уведомление команды
    notify_failure = PythonOperator(
        task_id="notify_quality_failure",
        python_callable=lambda: print("Alerting quality failure..."),
    )

    # 5. Silver Quality Gate
    silver_quality_gate = SparkSubmitOperator(
        task_id="silver_quality_gate",
        application="jobs/quality_gates/silver_orders_gate.py",
        application_args=["--date", "{{ ds }}"],
        pool="spark_light_pool",
    )

    # Граф зависимостей
    ingest_bronze >> bronze_quality_gate >> quality_branch
    quality_branch >> [transform_silver, notify_failure]
    transform_silver >> silver_quality_gate

# jobs/quality_gates/bronze_orders_gate.py

import sys
import json
from pyspark.sql import SparkSession


def main():
    spark = SparkSession.builder.appName("BronzeOrdersQualityGate").getOrCreate()

    date = sys.argv[sys.argv.index("--date") + 1]
    bronze_path = f"s3a://datalake/bronze/orders/date={date}/"

    df = spark.read.parquet(bronze_path)

    quality_report = {
        "date": date,
        "stage": "bronze_orders",
        "checks": {},
        "passed": True,
    }

    # Row Count Gate
    row_count = df.count()
    quality_report["checks"]["row_count"] = {
        "value": row_count,
        "threshold": 1000,
        "passed": row_count >= 1000,
    }
    if row_count < 1000:
        quality_report["passed"] = False

    # Null Gate - single pass
    _, null_counts = check_nulls_and_count_single_pass(
        df, ["order_id", "user_id", "amount"]
    )
    for col_name, null_count in null_counts.items():
        quality_report["checks"][f"null_{col_name}"] = {
            "value": null_count,
            "threshold": 0,
            "passed": null_count == 0,
        }
        if null_count > 0:
            quality_report["passed"] = False

    # Запись в Metadata-таблицу (для дашборда)
    quality_df = spark.createDataFrame([quality_report])
    (
        quality_df
        .write
        .mode("append")
        .format("delta")
        .saveAsTable("metadata.quality_reports")
    )

    # Для Airflow: non-zero exit = failure
    if not quality_report["passed"]:
        print(f"QUALITY_RESULT:FAILED:{json.dumps(quality_report)}")
        sys.exit(1)  # Airflow поймает это как task failure
    else:
        print(f"QUALITY_RESULT:PASSED:{json.dumps(quality_report)}")
        sys.exit(0)


if __name__ == "__main__":
    main()

from pydeequ.analyzers import AnalysisRunner, Size, Completeness, Uniqueness
from pydeequ.checks import Check, CheckLevel
from pydeequ.verification import VerificationSuite

# Определяем проверки
check = (
    Check(spark, CheckLevel.Error, "Order Quality")
    .hasSize(lambda x: x >= 1000, "Too few rows")
    .isComplete("order_id")
    .isComplete("user_id")
    .isUnique("order_id")
    .isContainedIn("status", ["paid", "refunded", "pending", "cancelled"])
    .isPositive("amount")
    .isNonNegative("discount")
)

# Запускаем верификацию
result = (
    VerificationSuite(spark)
    .onData(df)
    .addCheck(check)
    .run()
)

# Анализируем результат
from pydeequ.verification import VerificationResult
result_df = VerificationResult.checkResultsAsDataFrame(spark, result)
result_df.show()

import great_expectations as gx

# Контекст - центральный объект GX
context = gx.get_context()

# Создаём Data Source для Spark DataFrame
data_source = context.sources.add_spark("spark_source")
data_asset = data_source.add_dataframe_asset("orders")

# Определяем Expectation Suite
suite = context.add_expectation_suite("silver_orders")

# Добавляем ожидания
suite.expect_column_values_to_not_be_null("order_id")
suite.expect_column_values_to_be_in_set(
    "status", ["paid", "refunded", "pending", "cancelled"]
)
suite.expect_column_values_to_be_between("amount", min_value=0.01)
suite.expect_table_row_count_to_be_between(min_value=1000)

# Валидируем
batch_request = data_asset.build_batch_request(dataframe=df)
results = context.run_validation_operator(
    "action_list_operator",
    assets_to_validate=[batch_request],
    run_name="daily_run",
)

# АНТИПАТТЕРН: магические числа без объяснения
if df.count() < 50000:
    raise ValueError("Too few rows")

# ПРАВИЛЬНО: именованные константы с документацией
# Исторический минимум в non-holiday weekday
# Источник: analytics_report_2024-01.pdf, page 12
# Обновить при существенном изменении бизнеса
MIN_DAILY_ORDERS = 50_000

if df.count() < MIN_DAILY_ORDERS:
    raise ValueError(
        f"Row count below historical minimum: {MIN_DAILY_ORDERS:,}"
    )

# АНТИПАТТЕРН: бинарный threshold без учёта дня недели
if current_count < 50000:
    raise ValueError()  # срабатывает каждый понедельник из-за выходных

# ПРАВИЛЬНО: динамический baseline с учётом дня недели
def get_dow_adjusted_baseline(table: str, day_of_week: int) -> int:
    """
    Разные пороги для будней/выходных.
    Понедельник: исторически на 20% меньше из-за выходных.
    Берём исторический avg только за те же дни недели.
    """
    base = get_historical_avg_count(
        table, days=28, same_dow=True
    )
    return int(base * 0.70)  # допускаем падение до 30% от DOW-avg

# АНТИПАТТЕРН: строгая валидация бизнес-правил на Bronze
# Bronze должен принимать всё, включая «плохие» данные
bronze_df = assert_quality_gate(
    df,
    [
        QualityRule(
            "valid_status",
            "status IN ('paid', 'refunded')",
            "critical",
            "...",
        ),
    ],
    stage="bronze",
)
# Результат: Bronze не принял данные из источника-нарушителя.
# Теперь мы потеряли сырые данные и не можем расследовать проблему.

# ПРАВИЛЬНО: Bronze принимает всё, Silver фильтрует
# Bronze gate: только техническая валидность (схема, non-null PK)
# Silver gate: бизнес-правила + quarantine

# АНТИПАТТЕРН: бессмысленная проверка
def assert_not_empty(df):
    assert df.count() >= 0  # Всегда True!

# АНТИПАТТЕРН: проверка без assert
def check_nulls(df, cols):
    null_counts = {c: df.filter(col(c).isNull()).count() for c in cols}
    return null_counts  # Считает, но никогда не падает!

# ПРАВИЛЬНО: явные assertions с условиями отказа
def assert_no_critical_nulls(df, cols):
    _, null_counts = check_nulls_and_count_single_pass(df, cols)
    for col_name, count in null_counts.items():
        assert count == 0, f"Found {count} nulls in critical column '{col_name}'"

# АНТИПАТТЕРН: качество перемешано с трансформацией
def process_orders(df):
    df = df.filter(col("amount") > 0)   # Молча удаляем строки с amount <= 0
    df = df.withColumn("total", col("amount") * col("qty"))
    df = df.dropDuplicates(["order_id"])  # Молча убиваем дубликаты
    return df
# Никто не знает, сколько строк потерялось. Quality "gate" невидим.

# ПРАВИЛЬНО: явный gate до трансформации
def process_orders(df, quarantine_path, batch_date):
    # Сначала - явный gate с отчётом
    df = apply_quarantine_gate(
        df,
        rules=ORDERS_RULES,
        quarantine_path=quarantine_path,
        stage="silver_orders",
        batch_date=batch_date,
    )
    # Потом - трансформация только над чистыми данными
    df = df.withColumn("total", col("amount") * col("qty"))
    return df

# pipeline_config.py

from pyspark.sql.types import (
    StructType, StructField, StringType, LongType, DoubleType,
    TimestampType, IntegerType,
)

BRONZE_ORDERS_SCHEMA = StructType([
    StructField("order_id", StringType(), nullable=False),
    StructField("user_id", LongType(), nullable=False),
    StructField("amount", DoubleType(), nullable=False),
    StructField("currency", StringType(), nullable=True),
    StructField("status", StringType(), nullable=True),
    StructField("created_at", TimestampType(), nullable=False),
    StructField("updated_at", TimestampType(), nullable=True),
    StructField("items_count", IntegerType(), nullable=True),
    StructField("delivery_date", TimestampType(), nullable=True),
    StructField("promo_code", StringType(), nullable=True),
    StructField("_source_file", StringType(), nullable=True),
    StructField("_ingested_at", TimestampType(), nullable=True),
])

SILVER_ORDERS_RULES = [
    QualityRule(
        "positive_amount",
        "amount > 0",
        "critical",
        "Сумма заказа должна быть положительной",
    ),
    QualityRule(
        "valid_status",
        "status IN ('paid', 'refunded', 'pending', 'cancelled')",
        "critical",
        "Статус из допустимого справочника",
    ),
    QualityRule(
        "valid_currency",
        "currency IN ('RUB', 'USD', 'EUR')",
        "warning",
        "Валюта из поддерживаемых",
    ),
    QualityRule(
        "reasonable_amount",
        "amount < 10000000",
        "warning",
        "Сумма > 10M требует проверки",
    ),
    QualityRule(
        "delivery_after_order",
        "delivery_date IS NULL OR delivery_date >= created_at",
        "critical",
        "Дата доставки >= дата заказа",
    ),
    QualityRule(
        "not_future_created",
        "created_at <= current_timestamp()",
        "critical",
        "Заказ не из будущего",
    ),
]

# jobs/daily_orders_pipeline.py

import sys
from datetime import datetime
from pyspark.sql import SparkSession
from pyspark.sql import functions as F


def run_pipeline(spark: SparkSession, batch_date: str) -> None:
    """Full Quality-Gated Pipeline: Bronze -> Silver -> Gold."""

    print(f"=== Pipeline start: {batch_date} ===")

    bronze_path = f"s3a://datalake/bronze/orders/date={batch_date}/"
    quarantine_path = "s3a://datalake/quarantine/orders/"
    silver_path = "s3a://datalake/silver/orders/"
    gold_path = "s3a://datalake/gold/daily_kpi/"

    # ─────────────────────────────────────────────────────────────────
    # BRONZE: Ingestion + техническая валидация
    # ─────────────────────────────────────────────────────────────────
    print("Step 1: Reading Bronze...")
    bronze_df = spark.read.schema(BRONZE_ORDERS_SCHEMA).json(bronze_path)

    # Gate 1a: Schema Drift - ломаем при критических расхождениях
    print("Gate 1a: Schema drift check...")
    compare_schemas(
        bronze_df,
        BRONZE_ORDERS_SCHEMA,
        strict=False,
        stage="bronze_ingestion",
    )

    # Gate 1b: Row Count - пустой батч = сломан источник
    print("Gate 1b: Row count check...")
    bronze_count = bronze_df.count()
    if bronze_count == 0:
        raise ValueError(
            f"[CRITICAL] Bronze is empty for {batch_date}. "
            "Upstream failure suspected."
        )
    print(f"  Bronze rows: {bronze_count:,}")

    # Gate 1c: Null Check на критических полях (single-pass)
    print("Gate 1c: Critical null check...")
    _, null_counts = check_nulls_and_count_single_pass(
        bronze_df,
        ["order_id", "user_id", "amount", "created_at"],
    )
    critical_nulls = {k: v for k, v in null_counts.items() if v > 0}
    if critical_nulls:
        raise ValueError(
            f"[CRITICAL] Null in PK/mandatory fields: {critical_nulls}. "
            "Upstream data is poisoned."
        )

    # ─────────────────────────────────────────────────────────────────
    # SILVER: Нормализация + бизнес-валидация с Quarantine
    # ─────────────────────────────────────────────────────────────────
    print("\nStep 2: Transforming to Silver...")
    silver_df = (
        bronze_df
        .dropDuplicates(["order_id"])
        .withColumn("processing_date", F.lit(batch_date))
    )

    # Gate 2: Бизнес-правила с Quarantine
    print("Gate 2: Business rules (quarantine mode)...")
    silver_valid, silver_invalid, silver_report = run_quality_rules(
        silver_df,
        SILVER_ORDERS_RULES,
        mode="quarantine",
        stage="silver_orders",
    )

    # Сохраняем карантин
    if silver_invalid is not None:
        quarantine_count = silver_invalid.count()
        if quarantine_count > 0:
            (
                silver_invalid
                .withColumn("_quarantine_date", F.lit(batch_date))
                .withColumn("_quarantine_ts", F.current_timestamp())
                .write
                .mode("append")
                .partitionBy("_quarantine_date")
                .parquet(quarantine_path)
            )
            print(f"  Quarantined: {quarantine_count:,} rows")

            # Если карантин > 20% - алерт, но продолжаем
            quarantine_ratio = quarantine_count / bronze_count
            if quarantine_ratio > 0.20:
                print(
                    f"  [WARNING] High quarantine ratio: {quarantine_ratio:.1%}. "
                    "Notify data producers."
                )

    # Записываем Silver
    valid_count = silver_valid.count()
    print(f"  Silver valid rows: {valid_count:,}")

    (
        silver_valid
        .write
        .mode("overwrite")
        .format("delta")
        .partitionBy("processing_date")
        .save(silver_path)
    )

    # Gate 2b: Финальный контроль - если потеряли > 50% строк, это уже Fail-Fast
    if valid_count < int(bronze_count * 0.50):
        raise ValueError(
            f"[CRITICAL] Silver has only {valid_count:,}/{bronze_count:,} rows "
            f"({valid_count/bronze_count:.1%}). Too many quarantined."
        )

    # ─────────────────────────────────────────────────────────────────
    # GOLD: Агрегации + финальные проверки
    # ─────────────────────────────────────────────────────────────────
    print("\nStep 3: Computing Gold KPIs...")
    gold_df = (
        spark.read.format("delta").load(silver_path)
        .filter(F.col("processing_date") == batch_date)
        .filter(F.col("status") == "paid")
        .agg(
            F.count("order_id").alias("total_orders"),
            F.sum("amount").alias("total_revenue"),
            F.avg("amount").alias("avg_order_value"),
            F.countDistinct("user_id").alias("unique_buyers"),
        )
        .withColumn("report_date", F.lit(batch_date))
    )

    # Gate 3: Проверка разумности Gold-метрик
    print("Gate 3: Gold metrics sanity check...")
    gold_row = gold_df.collect()[0]

    EXPECTED_MIN_REVENUE = 1_000_000  # 1M RUB минимальная выручка в день
    EXPECTED_MAX_AOV = 500_000        # 500K RUB максимальный средний чек

    if gold_row["total_revenue"] < EXPECTED_MIN_REVENUE:
        raise ValueError(
            f"[CRITICAL] Daily revenue too low: "
            f"{gold_row['total_revenue']:,.0f} RUB. "
            f"Expected >= {EXPECTED_MIN_REVENUE:,}"
        )

    if gold_row["avg_order_value"] > EXPECTED_MAX_AOV:
        raise ValueError(
            f"[CRITICAL] AOV suspiciously high: "
            f"{gold_row['avg_order_value']:,.0f} RUB. "
            f"Possible currency bug."
        )

    # Записываем Gold
    (
        gold_df
        .write
        .mode("overwrite")
        .format("delta")
        .partitionBy("report_date")
        .save(gold_path)
    )

    print(f"""
=== Pipeline COMPLETE: {batch_date} ===
  Bronze rows:         {bronze_count:,}
  Silver valid:        {valid_count:,}
  Silver quarantined:  {bronze_count - valid_count:,}
  Gold revenue:        {gold_row['total_revenue']:,.0f} RUB
  Gold orders:         {gold_row['total_orders']:,}
  Gold AOV:            {gold_row['avg_order_value']:,.0f} RUB
""")


if __name__ == "__main__":
    spark = (
        SparkSession.builder
        .appName("DailyOrdersPipeline")
        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
        .getOrCreate()
    )

    batch_date = sys.argv[1] if len(sys.argv) > 1 else str(datetime.today().date())
    run_pipeline(spark, batch_date)

from dataclasses import dataclass, asdict
from datetime import datetime

@dataclass
class QualityRunRecord:
    """Одна запись результата Quality Gate."""
    run_date: str
    pipeline: str
    stage: str
    check_name: str
    check_type: str    # row_count | null_check | schema | business_rule | drift
    passed: bool
    metric_value: float
    threshold_value: float
    row_count: int
    severity: str
    recorded_at: str = None

    def __post_init__(self):
        if self.recorded_at is None:
            self.recorded_at = datetime.utcnow().isoformat()


def write_quality_metrics(
    spark: SparkSession,
    records: list,
    metrics_table: str = "metadata.quality_runs",
) -> None:
    """
    Записываем результаты всех проверок в единую метрики-таблицу.
    Используется для построения Quality Dashboard в BI-инструменте.
    """
    rows = [asdict(r) for r in records]
    metrics_df = spark.createDataFrame(rows)
    (
        metrics_df
        .write
        .mode("append")
        .format("delta")
        .partitionBy("run_date", "pipeline")
        .saveAsTable(metrics_table)
    )