Убрана kafka, переделан dag на работу с csv

2025-10-14 22:43:20 +03:00
parent 50865db5e6
commit d17c1da542
8 changed files with 182 additions and 220 deletions
@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import logging
+import os
+import random
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import List
+
+import pandas as pd
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+from helpers.greenplum import get_gp_conn
+
+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,
+        order_ts TIMESTAMP NOT NULL,
+        customer_id BIGINT NOT NULL,
+        amount NUMERIC(12,2) NOT NULL
+    )
+    WITH (appendonly=true, orientation=row, compresstype=zlib, compresslevel=1)
+    DISTRIBUTED BY (order_id);
+    """
+    with get_gp_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.utcnow().strftime("%Y%m%d_%H%M%S")
+    csv_path = csv_dir / f"orders_{timestamp}.csv"
+
+    base_order_id = int(datetime.utcnow().timestamp())
+    order_ids: List[int] = list(range(base_order_id * 1_000, base_order_id * 1_000 + rows))
+    now = datetime.utcnow()
+    order_ts = [(now - timedelta(seconds=i)).isoformat() for i in range(rows)]
+    customer_ids = [random.randint(1, 1_000) for _ in range(rows)]
+    amounts = [round(random.uniform(10, 500), 2) for _ in range(rows)]
+
+    df = pd.DataFrame(
+        {
+            "order_id": order_ids,
+            "order_ts": order_ts,
+            "customer_id": customer_ids,
+            "amount": amounts,
+        }
+    )
+    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 в Greenplum через временную таблицу и anti-join."""
+    csv_file = Path(csv_path)
+    if not csv_file.exists():
+        raise FileNotFoundError(f"CSV не найден: {csv_file}")
+
+    with get_gp_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_greenplum",
+    start_date=datetime(2024, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args=default_args,
+    tags=["demo", "greenplum", "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_greenplum",
+        python_callable=_load_csv,
+        op_kwargs={
+            "csv_path": "{{ ti.xcom_pull(task_ids='generate_csv') }}",
+        },
+    )
+
+    create_table >> generate_csv >> preview_csv >> load_csv
@@ -50,7 +50,7 @@ def assert_orders_table_exists(conn) -> None:
            """
        )
        if cur.fetchone() is None:
-            raise ValueError("Таблица public.orders не найдена; запусти DAG kafka_to_greenplum.")
+            raise ValueError("Таблица public.orders не найдена; запусти DAG csv_to_greenplum.")


 def fetch_orders_schema(conn) -> Sequence[Tuple[str, str]]:
@@ -82,7 +82,7 @@ def fetch_orders_count(conn) -> int:
 def assert_orders_have_rows(conn) -> None:
    """Проверяет, что таблица orders не пустая."""
    if fetch_orders_count(conn) <= 0:
-        raise ValueError("Таблица public.orders пустая — запусти DAG kafka_to_greenplum перед проверкой.")
+        raise ValueError("Таблица public.orders пустая — запусти DAG csv_to_greenplum перед проверкой.")


 def fetch_orders_duplicates(conn) -> int:
@@ -1,146 +0,0 @@
-from __future__ import annotations
-
-import json
-import os
-import random
-from datetime import datetime, timedelta
-from typing import List, Tuple, Optional
-
-from airflow import DAG
-from airflow.operators.python import PythonOperator
-from confluent_kafka import Consumer, KafkaException, Producer
-from psycopg2.extras import execute_values
-
-from helpers.greenplum import get_gp_conn
-
-KAFKA_BOOTSTRAP = os.getenv("KAFKA_BOOTSTRAP", "kafka:29092")
-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"))
-MAX_EMPTY_POLLS = int(os.getenv("KAFKA_MAX_EMPTY_POLLS", "3"))
-def _create_table():
-    ddl = """
-    CREATE TABLE IF NOT EXISTS public.orders (
-        order_id BIGINT,
-        order_ts TIMESTAMP NOT NULL,
-        customer_id BIGINT NOT NULL,
-        amount NUMERIC(12,2) NOT NULL
-    )
-    WITH (appendonly=true, orientation=column, compresstype=zlib)
-    DISTRIBUTED BY (order_id);
-    """
-    with get_gp_conn() as conn, conn.cursor() as cur:
-        cur.execute(ddl)
-        conn.commit()
-
-
-def _produce(n=1000):
-    producer = Producer({"bootstrap.servers": KAFKA_BOOTSTRAP})
-    for idx in range(n):
-        payload = {
-            "order_id": idx + 1,
-            "order_ts": datetime.utcnow().isoformat(),
-            "customer_id": random.randint(1, 100),
-            "amount": round(random.uniform(10, 500), 2),
-        }
-        producer.produce(TOPIC, json.dumps(payload).encode("utf-8"))
-    producer.flush()
-
-
-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)
-    by_id = {int(r[0]): r for r in rows}
-    unique_rows = list(by_id.values())
-
-    # Вставка через VALUES + anti-join для GP6 (без ON CONFLICT)
-    execute_values(
-        cur,
-        """
-        INSERT INTO public.orders(order_id, order_ts, customer_id, amount)
-        SELECT v.order_id, v.order_ts, v.customer_id, v.amount
-        FROM (VALUES %s) AS v(order_id, order_ts, customer_id, amount)
-        LEFT JOIN public.orders o ON o.order_id = v.order_id
-        WHERE o.order_id IS NULL
-        """,
-        unique_rows,
-        template="(%s,%s,%s,%s)",
-    )
-
-
-def _consume_and_load(max_messages=1000, timeout_s: Optional[int] = None):
-    if timeout_s is None:
-        timeout_s = POLL_TIMEOUT_S
-
-    consumer = Consumer(
-        {
-            "bootstrap.servers": KAFKA_BOOTSTRAP,
-            "group.id": "airflow-loader-gp",
-            "auto.offset.reset": "earliest",
-            "enable.auto.commit": False,
-        }
-    )
-    consumer.subscribe([TOPIC])
-
-    with get_gp_conn() as conn, conn.cursor() as cur:
-        batch: List[Tuple] = []
-        consumed = 0
-        empty_polls = 0
-        while consumed < max_messages and empty_polls < MAX_EMPTY_POLLS:
-            msg = consumer.poll(timeout_s)
-            if msg is None:
-                empty_polls += 1
-                continue
-            empty_polls = 0
-            if msg.error():
-                raise KafkaException(msg.error())
-
-            data = json.loads(msg.value().decode("utf-8"))
-            batch.append(
-                (
-                    int(data["order_id"]),
-                    datetime.fromisoformat(data["order_ts"]),
-                    int(data["customer_id"]),
-                    float(data["amount"]),
-                )
-            )
-            consumed += 1
-
-            if len(batch) >= BATCH_SIZE:
-                _flush_batch(cur, batch)
-                conn.commit()
-                batch.clear()
-
-        # Финальный сброс, если вышли по лимиту сообщений
-        if batch:
-            _flush_batch(cur, batch)
-            conn.commit()
-
-    # Фиксируем оффсеты после успешной загрузки
-    consumer.commit()
-    consumer.close()
-
-
-default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
-
-with DAG(
-    dag_id="kafka_to_greenplum",
-    start_date=datetime(2024, 1, 1),
-    schedule=None,
-    catchup=False,
-    default_args=default_args,
-    tags=["demo", "kafka", "greenplum"],
-) as dag:
-    create_table = PythonOperator(task_id="create_table", python_callable=_create_table)
-    produce = PythonOperator(
-        task_id="produce_messages", python_callable=_produce, op_kwargs={"n": 1000}
-    )
-    consume_and_load = PythonOperator(
-        task_id="consume_and_load",
-        python_callable=_consume_and_load,
-        op_kwargs={"max_messages": 1000},
-    )
-
-    create_table >> produce >> consume_and_load
@@ -1,2 +1,2 @@
-confluent-kafka==2.3.0
 psycopg2-binary==2.9.9
+pandas==2.1.4