Добавлен контроль качества данных

2025-10-12 17:24:50 +03:00
parent bf6b424a5e
commit dfc1f2ca84
6 changed files with 177 additions and 44 deletions
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+
+from helpers.greenplum import (
+    assert_orders_have_rows,
+    assert_orders_no_duplicates,
+    assert_orders_schema,
+    assert_orders_table_exists,
+    get_gp_conn,
+)
+
+
+def _run_check(check_callable):
+    """Оборачиваем проверку в контекст подключения."""
+    with get_gp_conn() as conn:
+        check_callable(conn)
+
+
+default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
+
+with DAG(
+    dag_id="greenplum_data_quality",
+    start_date=datetime(2024, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args=default_args,
+    tags=["demo", "greenplum", "quality"],
+) as dag:
+    check_exists = PythonOperator(
+        task_id="check_orders_table_exists",
+        python_callable=_run_check,
+        op_args=[assert_orders_table_exists],
+    )
+    check_schema = PythonOperator(
+        task_id="check_orders_schema",
+        python_callable=_run_check,
+        op_args=[assert_orders_schema],
+    )
+    check_has_rows = PythonOperator(
+        task_id="check_orders_has_rows",
+        python_callable=_run_check,
+        op_args=[assert_orders_have_rows],
+    )
+    check_no_duplicates = PythonOperator(
+        task_id="check_order_duplicates",
+        python_callable=_run_check,
+        op_args=[assert_orders_no_duplicates],
+    )
+
+    check_exists >> check_schema >> check_has_rows >> check_no_duplicates
@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import os
+from typing import List, Sequence, Tuple
+
+import psycopg2
+
+# Настройки для подключения к Greenplum. По умолчанию используем Airflow Connection,
+# но при проблемах можно переключиться на ENV-подключение, установив GP_USE_AIRFLOW_CONN=false.
+GP_CONN_ID = os.getenv("GP_CONN_ID", "greenplum_conn")
+GP_USE_AIRFLOW_CONN = os.getenv("GP_USE_AIRFLOW_CONN", "true").lower() in ("1", "true", "yes")
+
+EXPECTED_ORDERS_SCHEMA: List[Tuple[str, str]] = [
+    ("order_id", "bigint"),
+    ("order_ts", "timestamp without time zone"),
+    ("customer_id", "bigint"),
+    ("amount", "numeric"),
+]
+
+
+def get_gp_conn():
+    """Возвращает psycopg2 connection к Greenplum (через Airflow Connection или напрямую по ENV)."""
+    if GP_USE_AIRFLOW_CONN:
+        try:
+            from airflow.providers.postgres.hooks.postgres import PostgresHook
+
+            hook = PostgresHook(postgres_conn_id=GP_CONN_ID)
+            return hook.get_conn()
+        except Exception:
+            # Фоллбек на прямое подключение по переменным окружения.
+            pass
+
+    return psycopg2.connect(
+        dbname=os.getenv("GP_DB", "gpadmin"),
+        user=os.getenv("GP_USER", "gpadmin"),
+        password=os.getenv("GP_PASSWORD", ""),
+        host=os.getenv("GP_HOST", "greenplum"),
+        port=int(os.getenv("GP_PORT", "5432")),
+    )
+
+
+def assert_orders_table_exists(conn) -> None:
+    """Проверяет наличие таблицы orders в схеме public."""
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            SELECT 1
+            FROM pg_catalog.pg_tables
+            WHERE schemaname = 'public' AND tablename = 'orders'
+            """
+        )
+        if cur.fetchone() is None:
+            raise ValueError("Таблица public.orders не найдена; запусти DAG kafka_to_greenplum.")
+
+
+def fetch_orders_schema(conn) -> Sequence[Tuple[str, str]]:
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            SELECT column_name, data_type
+            FROM information_schema.columns
+            WHERE table_schema = 'public' AND table_name = 'orders'
+            ORDER BY ordinal_position
+            """
+        )
+        return cur.fetchall()
+
+
+def assert_orders_schema(conn) -> None:
+    """Проверяет, что схема таблицы orders соответствует ожидаемой."""
+    schema = fetch_orders_schema(conn)
+    if list(schema) != EXPECTED_ORDERS_SCHEMA:
+        raise ValueError(f"Неожиданная схема orders: {schema}. Ожидали {EXPECTED_ORDERS_SCHEMA}.")
+
+
+def fetch_orders_count(conn) -> int:
+    with conn.cursor() as cur:
+        cur.execute("SELECT COUNT(*) FROM public.orders")
+        return cur.fetchone()[0]
+
+
+def assert_orders_have_rows(conn) -> None:
+    """Проверяет, что таблица orders не пустая."""
+    if fetch_orders_count(conn) <= 0:
+        raise ValueError("Таблица public.orders пустая — запусти DAG kafka_to_greenplum перед проверкой.")
+
+
+def fetch_orders_duplicates(conn) -> int:
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            SELECT COUNT(*) FROM (
+                SELECT order_id
+                FROM public.orders
+                GROUP BY order_id
+                HAVING COUNT(*) > 1
+            ) d
+            """
+        )
+        return cur.fetchone()[0]
+
+
+def assert_orders_no_duplicates(conn) -> None:
+    """Проверяет, что в таблице нет дублей по order_id."""
+    duplicates = fetch_orders_duplicates(conn)
+    if duplicates:
+        raise ValueError(f"Обнаружены дубли по order_id ({duplicates} шт.) — проверь загрузку данных.")
@@ -9,50 +9,18 @@ from typing import List, Tuple, Optional
 from airflow import DAG
 from airflow.operators.python import PythonOperator
 from confluent_kafka import Consumer, KafkaException, Producer
-
-# Пакеты для прямого подключения и батч-загрузки
-import psycopg2
 from psycopg2.extras import execute_values

-# Airflow Connection ID для Greenplum (создаётся в UI или CLI).
-GP_CONN_ID = os.getenv("GP_CONN_ID", "greenplum_conn")
-GP_USE_AIRFLOW_CONN = os.getenv("GP_USE_AIRFLOW_CONN", "true").lower() in ("1", "true", "yes")
+from helpers.greenplum import get_gp_conn

 KAFKA_BOOTSTRAP = os.getenv("KAFKA_BOOTSTRAP", "kafka:9092")
 TOPIC = os.getenv("KAFKA_TOPIC", "orders")
 BATCH_SIZE = int(os.getenv("KAFKA_BATCH_SIZE", "500"))
 POLL_TIMEOUT_S = int(os.getenv("KAFKA_POLL_TIMEOUT", "10"))
-
-
-def _get_gp_conn():
-    """Получаем соединение с Greenplum через Airflow Connection или через ENV-DSN.
-
-    Если переменная `GP_USE_AIRFLOW_CONN` = false или отсутствует провайдер Postgres,
-    используем ENV-подключение напрямую (psycopg2).
-    """
-    if GP_USE_AIRFLOW_CONN:
-        try:
-            from airflow.providers.postgres.hooks.postgres import PostgresHook  # импорт при необходимости
-
-            hook = PostgresHook(postgres_conn_id=GP_CONN_ID)
-            return hook.get_conn()
-        except Exception:
-            # Фоллбек на прямое подключение
-            pass
-
-    return psycopg2.connect(
-        dbname=os.getenv("GP_DB", "gpadmin"),
-        user=os.getenv("GP_USER", "gpadmin"),
-        password=os.getenv("GP_PASSWORD", ""),
-        host=os.getenv("GP_HOST", "greenplum"),
-        port=int(os.getenv("GP_PORT", "5432")),
-    )
-
-
 def _create_table():
    ddl = """
    CREATE TABLE IF NOT EXISTS public.orders (
-        order_id BIGINT PRIMARY KEY,
+        order_id BIGINT,
        order_ts TIMESTAMP NOT NULL,
        customer_id BIGINT NOT NULL,
        amount NUMERIC(12,2) NOT NULL
@@ -60,7 +28,7 @@ def _create_table():
    WITH (appendonly=true, orientation=column, compresstype=zlib)
    DISTRIBUTED BY (order_id);
    """
-    with _get_gp_conn() as conn, conn.cursor() as cur:
+    with get_gp_conn() as conn, conn.cursor() as cur:
        cur.execute(ddl)
        conn.commit()

@@ -79,6 +47,7 @@ def _produce(n=1000):


 def _flush_batch(cur, rows: List[Tuple]):
+    """Insert deduplicated batch of rows into public.orders for GP6 (no PK support)."""
    if not rows:
        return
    # Дедупликация внутри батча по первичному ключу (order_id)
@@ -114,7 +83,7 @@ def _consume_and_load(max_messages=1000, timeout_s: Optional[int] = None):
    )
    consumer.subscribe([TOPIC])

-    with _get_gp_conn() as conn, conn.cursor() as cur:
+    with get_gp_conn() as conn, conn.cursor() as cur:
        batch: List[Tuple] = []
        consumed = 0
        while consumed < max_messages: