Первые наброски airflow-gp

2025-09-16 23:03:35 +03:00
parent c23512920b
commit e7e40d18f7
7 changed files with 344 additions and 0 deletions
@@ -0,0 +1,97 @@
+from __future__ import annotations
+import os, json, random
+from datetime import datetime, timedelta
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+import psycopg2
+from confluent_kafka import Producer, Consumer, KafkaException
+
+# 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")),
+}
+
+KAFKA_BOOTSTRAP = os.getenv("KAFKA_BOOTSTRAP", "kafka:9092")
+TOPIC = os.getenv("KAFKA_TOPIC", "orders")
+
+def _create_table():
+    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 (appendonly=true, orientation=column, compresstype=zlib)
+    DISTRIBUTED BY (order_id);
+    """
+    with psycopg2.connect(**GP_DSN) 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):
+        payload = {
+            "order_id": i + 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()
+
+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,
+    })
+    consumer.subscribe([TOPIC])
+
+    with psycopg2.connect(**GP_DSN) as conn, conn.cursor() as cur:
+        inserted = 0
+        while inserted < max_messages:
+            msg = consumer.poll(timeout_s)
+            if msg is None:
+                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()
+
+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,
+    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