refactor(dags): миграция CSV-пайплайна с Greenplum на PostgreSQL

- Зачем: - адаптация пайплайна для работы с чистым PostgreSQL вместо Greenplum - обеспечение совместимости с учебным стендом на базе Docker Compose - Что: - переименованы файлы (greenplum -> postgres) - обновлены DAG-файлы: замена get_gp_conn на get_postgres_conn, адаптация DDL под PostgreSQL - обновлён helpers/postgres.py: замена функций и переменных окружения - добавлен pytest в requirements.txt для запуска тестов - создан tests/conftest.py с функцией patch_postgres_hook - обновлён tests/test_postgres_helpers.py: исправлены импорты - добавлено монтирование папки tests в docker-compose.yml - адаптирован sql/base/orders_ddl.sql под PostgreSQL - Проверка: - docker-compose exec airflow-webserver pytest /opt/airflow/tests/test_postgres_helpers.py
2026-03-08 19:51:15 +03:00
parent f703d8e96d
commit 8028711f0c
9 changed files with 119 additions and 90 deletions
@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import logging
+import os
+import random
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import List
+
+import pandas as pd
+from airflow.operators.python import PythonOperator
+from helpers.postgres import get_postgres_conn
+
+from airflow import DAG
+
+CSV_DIR = Path(os.getenv("CSV_DIR", "/opt/airflow/data"))
+CSV_ROWS = int(os.getenv("CSV_ROWS", "1000"))
+
+
+def _create_table() -> None:
+    """Создаёт таблицу public.orders, если она ещё не существует."""
+    ddl = """
+    CREATE TABLE IF NOT EXISTS public.orders (
+        order_id BIGINT PRIMARY KEY,
+        order_ts TIMESTAMP NOT NULL,
+        customer_id BIGINT NOT NULL,
+        amount NUMERIC(12,2) NOT NULL
+    );
+    """
+    with get_postgres_conn() as conn, conn.cursor() as cur:
+        cur.execute(ddl)
+        conn.commit()
+
+
+def _generate_csv(rows: int, csv_dir: Path) -> str:
+    """Генерирует CSV c заказами с помощью pandas и сохраняет на диск."""
+    csv_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
+    csv_path = csv_dir / f"orders_{timestamp}.csv"
+
+    # Генерируем данные в pandas-стиле
+    base_order_id = int(datetime.now(UTC).timestamp() * 1_000)
+
+    # Создаём DataFrame с использованием pandas методов
+    df = pd.DataFrame(
+        {
+            # Уникальные order_id начиная с базового значения
+            "order_id": pd.Series(
+                range(base_order_id, base_order_id + rows), dtype="int64"
+            ),
+            # Временные метки с интервалом в 1 секунду в обратном порядке
+            "order_ts": pd.date_range(
+                end=datetime.now(UTC), periods=rows, freq="1S"
+            ).sort_values(ascending=False),
+            # Случайные customer_id от 1 до 1000
+            "customer_id": pd.Series(
+                random.choices(range(1, 1001), k=rows), dtype="int64"
+            ),
+            # Случайные суммы от 10 до 500 с округлением до 2 знаков
+            "amount": pd.Series(
+                [round(random.uniform(10, 500), 2) for _ in range(rows)],
+                dtype="float64",
+            ),
+        }
+    )
+
+    # Сохраняем CSV без индекса
+    df.to_csv(csv_path, index=False)
+    logging.info("CSV сохранён: %s (строк: %s)", csv_path, len(df))
+    return str(csv_path)
+
+
+def _preview_csv(csv_path: str, sample_rows: int = 5) -> None:
+    """Отображает предпросмотр CSV через pandas (head и describe)."""
+    df = pd.read_csv(csv_path)
+    df["order_ts"] = pd.to_datetime(df["order_ts"], errors="coerce")
+    logging.info(
+        "Первые %s строк:\n%s", sample_rows, df.head(sample_rows).to_string(index=False)
+    )
+    numeric_summary = df.describe(include="number")
+    logging.info("Числовая статистика:\n%s", numeric_summary.to_string())
+    if df["order_ts"].notna().any():
+        logging.info(
+            "Диапазон order_ts: %s → %s",
+            df["order_ts"].min().isoformat(),
+            df["order_ts"].max().isoformat(),
+        )
+
+
+def _load_csv(csv_path: str) -> None:
+    """Загружает CSV в Postgres через временную таблицу и anti-join."""
+    csv_file = Path(csv_path)
+    if not csv_file.exists():
+        raise FileNotFoundError(f"CSV не найден: {csv_file}")
+
+    with (
+        get_postgres_conn() as conn,
+        conn.cursor() as cur,
+        csv_file.open("r", encoding="utf-8") as f,
+    ):
+        cur.execute(
+            "CREATE TEMP TABLE tmp_orders (LIKE public.orders INCLUDING DEFAULTS) ON COMMIT DROP;"
+        )
+        cur.copy_expert(
+            "COPY tmp_orders (order_id, order_ts, customer_id, amount) FROM STDIN WITH CSV HEADER",
+            f,
+        )
+
+        cur.execute("SELECT COUNT(*) FROM tmp_orders")
+        tmp_rows = cur.fetchone()[0]
+
+        cur.execute(
+            """
+            INSERT INTO public.orders(order_id, order_ts, customer_id, amount)
+            SELECT t.order_id, t.order_ts, t.customer_id, t.amount
+            FROM tmp_orders t
+            LEFT JOIN public.orders o ON o.order_id = t.order_id
+            WHERE o.order_id IS NULL
+            """
+        )
+        inserted = cur.rowcount if cur.rowcount != -1 else 0
+        conn.commit()
+
+    logging.info("Загружено строк: %s (прочитано из CSV: %s)", inserted, tmp_rows)
+
+
+default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
+
+with DAG(
+    dag_id="csv_to_postgres",
+    start_date=datetime(2017, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args=default_args,
+    tags=["demo", "postgres", "csv"],
+) as dag:
+    create_table = PythonOperator(
+        task_id="create_orders_table",
+        python_callable=_create_table,
+    )
+
+    generate_csv = PythonOperator(
+        task_id="generate_csv",
+        python_callable=_generate_csv,
+        op_kwargs={"rows": CSV_ROWS, "csv_dir": CSV_DIR},
+    )
+
+    preview_csv = PythonOperator(
+        task_id="preview_csv",
+        python_callable=_preview_csv,
+        op_kwargs={
+            "csv_path": "{{ ti.xcom_pull(task_ids='generate_csv') }}",
+            "sample_rows": 5,
+        },
+    )
+
+    load_csv = PythonOperator(
+        task_id="load_csv_to_postgres",
+        python_callable=_load_csv,
+        op_kwargs={
+            "csv_path": "{{ ti.xcom_pull(task_ids='generate_csv') }}",
+        },
+    )
+
+    create_table >> generate_csv >> preview_csv >> load_csv