Доработки по удобству

2025-10-12 16:41:32 +03:00
parent 9341983f26
commit bf6b424a5e
5 changed files with 211 additions and 49 deletions
@@ -1,22 +1,53 @@
 from __future__ import annotations
-import os, json, random
+
+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
-import psycopg2
-from confluent_kafka import Producer, Consumer, KafkaException
+from confluent_kafka import Consumer, KafkaException, Producer

-# Greenplum (Postgres wire protocol)
-GP_DSN = {
-    "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")),
-}
+# Пакеты для прямого подключения и батч-загрузки
+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")

 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 = """
@@ -29,69 +60,119 @@ def _create_table():
    WITH (appendonly=true, orientation=column, compresstype=zlib)
    DISTRIBUTED BY (order_id);
    """
-    with psycopg2.connect(**GP_DSN) as conn, conn.cursor() as cur:
+    with _get_gp_conn() as conn, conn.cursor() as cur:
        cur.execute(ddl)
        conn.commit()

+
 def _produce(n=1000):
-    p = Producer({"bootstrap.servers": KAFKA_BOOTSTRAP})
-    for i in range(n):
+    producer = Producer({"bootstrap.servers": KAFKA_BOOTSTRAP})
+    for idx in range(n):
        payload = {
-            "order_id": i + 1,
+            "order_id": idx + 1,
            "order_ts": datetime.utcnow().isoformat(),
            "customer_id": random.randint(1, 100),
            "amount": round(random.uniform(10, 500), 2),
        }
-        p.produce(TOPIC, json.dumps(payload).encode("utf-8"))
-    p.flush()
+        producer.produce(TOPIC, json.dumps(payload).encode("utf-8"))
+    producer.flush()

-def _consume_and_load(max_messages=1000, timeout_s=10):
-    consumer = Consumer({
-        "bootstrap.servers": KAFKA_BOOTSTRAP,
-        "group.id": "airflow-loader-gp",
-        "auto.offset.reset": "earliest",
-        "enable.auto.commit": False,
-    })
+
+def _flush_batch(cur, rows: List[Tuple]):
+    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 psycopg2.connect(**GP_DSN) as conn, conn.cursor() as cur:
-        inserted = 0
-        while inserted < max_messages:
+    with _get_gp_conn() as conn, conn.cursor() as cur:
+        batch: List[Tuple] = []
+        consumed = 0
+        while consumed < max_messages:
            msg = consumer.poll(timeout_s)
            if msg is None:
+                # Нет новых сообщений — сбрасываем остаток батча и выходим
+                if batch:
+                    _flush_batch(cur, batch)
+                    conn.commit()
+                    batch.clear()
                break
            if msg.error():
                raise KafkaException(msg.error())
-            d = json.loads(msg.value().decode("utf-8"))
-            # GPDB6 не поддерживает ON CONFLICT — используем WHERE NOT EXISTS
-            cur.execute(
-                """
-                INSERT INTO public.orders(order_id, order_ts, customer_id, amount)
-                SELECT %s, %s, %s, %s
-                WHERE NOT EXISTS (
-                    SELECT 1 FROM public.orders WHERE order_id = %s
-                );
-                """,
-                (d["order_id"], d["order_ts"], d["customer_id"], d["amount"], d["order_id"]),
-            )
-            inserted += 1
-        conn.commit()

-    consumer.commit(); consumer.close()
+            data = json.loads(msg.value().decode("utf-8"))
+            batch.append(
+                (
+                    int(data["order_id"]),
+                    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_interval=None,
+    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})
+    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