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
This commit is contained in:
2026-03-08 19:51:15 +03:00
parent f703d8e96d
commit 8028711f0c
9 changed files with 119 additions and 90 deletions
@@ -9,7 +9,7 @@ from typing import List
import pandas as pd
from airflow.operators.python import PythonOperator
from helpers.greenplum import get_gp_conn
from helpers.postgres import get_postgres_conn
from airflow import DAG
@@ -21,15 +21,13 @@ def _create_table() -> None:
"""Создаёт таблицу public.orders, если она ещё не существует."""
ddl = """
CREATE TABLE IF NOT EXISTS public.orders (
order_id BIGINT,
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=row, compresstype=zstd, compresslevel=1)
DISTRIBUTED BY (order_id);
);
"""
with get_gp_conn() as conn, conn.cursor() as cur:
with get_postgres_conn() as conn, conn.cursor() as cur:
cur.execute(ddl)
conn.commit()
@@ -90,13 +88,13 @@ def _preview_csv(csv_path: str, sample_rows: int = 5) -> None:
def _load_csv(csv_path: str) -> None:
"""Загружает CSV в Greenplum через временную таблицу и anti-join."""
"""Загружает CSV в Postgres через временную таблицу и anti-join."""
csv_file = Path(csv_path)
if not csv_file.exists():
raise FileNotFoundError(f"CSV не найден: {csv_file}")
with (
get_gp_conn() as conn,
get_postgres_conn() as conn,
conn.cursor() as cur,
csv_file.open("r", encoding="utf-8") as f,
):
@@ -129,12 +127,12 @@ def _load_csv(csv_path: str) -> None:
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
with DAG(
dag_id="csv_to_greenplum",
dag_id="csv_to_postgres",
start_date=datetime(2017, 1, 1),
schedule=None,
catchup=False,
default_args=default_args,
tags=["demo", "greenplum", "csv"],
tags=["demo", "postgres", "csv"],
) as dag:
create_table = PythonOperator(
task_id="create_orders_table",
@@ -157,11 +155,11 @@ with DAG(
)
load_csv = PythonOperator(
task_id="load_csv_to_greenplum",
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
create_table >> generate_csv >> preview_csv >> load_csv
@@ -4,12 +4,12 @@ import logging
from datetime import datetime, timedelta
from airflow.operators.python import PythonOperator
from helpers.greenplum import (
from helpers.postgres import (
assert_orders_have_rows,
assert_orders_no_duplicates,
assert_orders_schema,
assert_orders_table_exists,
get_gp_conn,
get_postgres_conn,
)
from airflow import DAG
@@ -17,7 +17,7 @@ from airflow import DAG
def _run_check(check_callable):
"""
Оборачивает проверку качества данных в контекст подключения к Greenplum.
Оборачивает проверку качества данных в контекст подключения к Postgres.
Этот DAG предназначен для автоматической проверки качества данных
после CSV-пайплайна в таблице public.orders:
@@ -33,7 +33,7 @@ def _run_check(check_callable):
check_name = check_callable.__name__.replace("assert_", "")
logging.info("🚀 Запуск проверки: %s", check_name)
with get_gp_conn() as conn:
with get_postgres_conn() as conn:
check_callable(conn)
logging.info("✅ Проверка пройдена: %s", check_name)
@@ -51,13 +51,13 @@ def _log_dq_summary():
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
with DAG(
dag_id="csv_to_greenplum_dq",
dag_id="csv_to_postgres_dq",
start_date=datetime(2017, 1, 1),
schedule=None,
catchup=False,
default_args=default_args,
tags=["demo", "greenplum", "quality", "csv", "dq"],
description="Проверки качества данных после CSV → public.orders в Greenplum",
tags=["demo", "postgres", "quality", "csv", "dq"],
description="Проверки качества данных после CSV → public.orders в Postgres",
) as dag:
# Задача 1: Проверка существования таблицы
check_exists = PythonOperator(
@@ -94,4 +94,4 @@ with DAG(
)
# Определяем последовательность выполнения задач
check_exists >> check_schema >> check_has_rows >> check_no_duplicates >> dq_summary
check_exists >> check_schema >> check_has_rows >> check_no_duplicates >> dq_summary
@@ -1,7 +1,7 @@
from __future__ import annotations
"""
Учебный DAG: применяет DDL для базовой таблицы orders в Greenplum.
Учебный DAG: применяет DDL для базовой таблицы orders в Postgres.
Запускается вручную перед CSVпайплайном или после изменения схемы.
"""
@@ -11,7 +11,7 @@ from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow import DAG
GREENPLUM_CONN_ID = "greenplum_conn"
POSTGRES_CONN_ID = "postgres_training"
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
@@ -20,13 +20,13 @@ with DAG(
start_date=datetime(2017, 1, 1),
schedule=None,
catchup=False,
template_searchpath="/sql",
template_searchpath="/opt/airflow/sql",
default_args=default_args,
tags=["demo", "greenplum", "ddl", "orders"],
tags=["demo", "postgres", "ddl", "orders"],
description="Создаёт/обновляет базовую таблицу orders в схеме public",
) as dag:
apply_orders_ddl = PostgresOperator(
task_id="apply_orders_ddl",
postgres_conn_id=GREENPLUM_CONN_ID,
postgres_conn_id=POSTGRES_CONN_ID,
sql="base/orders_ddl.sql",
)
)
@@ -1,7 +1,7 @@
from __future__ import annotations
"""
LEGACY: Вспомогательные функции для прямого подключения к Greenplum через psycopg2.
LEGACY: Вспомогательные функции для прямого подключения к Postgres через psycopg2.
Внимание: этот модуль оставлен только для поддержки базового CSV-пайплайна.
В новых DAG (ODS/DDS/DM) используйте встроенный в Airflow PostgresOperator
и штатные механизмы XCom.
@@ -13,10 +13,10 @@ 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 (
# Настройки для подключения к Postgres. По умолчанию используем Airflow Connection,
# но при проблемах можно переключиться на ENV-подключение, установив POSTGRES_USE_AIRFLOW_CONN=false.
POSTGRES_CONN_ID = os.getenv("POSTGRES_CONN_ID", "postgres_training")
POSTGRES_USE_AIRFLOW_CONN = os.getenv("POSTGRES_USE_AIRFLOW_CONN", "true").lower() in (
"1",
"true",
"yes",
@@ -31,9 +31,9 @@ EXPECTED_ORDERS_SCHEMA: List[Tuple[str, str]] = [
]
def get_gp_conn():
def get_postgres_conn():
"""
Возвращает psycopg2 connection к Greenplum.
Возвращает psycopg2 connection к Postgres.
Приоритет подключения:
1. Через Airflow Connection (если настроено и доступно)
@@ -42,11 +42,11 @@ def get_gp_conn():
Returns:
psycopg2 connection object
"""
if GP_USE_AIRFLOW_CONN:
if POSTGRES_USE_AIRFLOW_CONN:
try:
from airflow.providers.postgres.hooks.postgres import PostgresHook
hook = PostgresHook(postgres_conn_id=GP_CONN_ID)
hook = PostgresHook(postgres_conn_id=POSTGRES_CONN_ID)
conn = hook.get_conn()
logging.info("✅ Подключение через Airflow Connection успешно")
return conn
@@ -57,14 +57,14 @@ def get_gp_conn():
# Прямое подключение по переменным окружения
conn_params = {
"dbname": os.getenv("GP_DB", "gp_dwh"),
"user": os.getenv("GP_USER", "gpadmin"),
"password": os.getenv("GP_PASSWORD", ""),
"host": os.getenv("GP_HOST", "greenplum"),
"port": int(os.getenv("GP_PORT", "5432")),
"dbname": os.getenv("POSTGRES_DB", "training"),
"user": os.getenv("POSTGRES_USER", "student"),
"password": os.getenv("POSTGRES_PASSWORD", "student"),
"host": os.getenv("POSTGRES_HOST", "postgres-training"),
"port": int(os.getenv("POSTGRES_PORT", "5432")),
}
logging.info(
"🔗 Подключение к Greenplum: %s:%s/%s",
"🔗 Подключение к Postgres: %s:%s/%s",
conn_params["host"],
conn_params["port"],
conn_params["dbname"],
@@ -77,7 +77,7 @@ def assert_orders_table_exists(conn) -> None:
Проверяет наличие таблицы orders в схеме public.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Raises:
ValueError: Если таблица не найдена
@@ -93,7 +93,7 @@ def assert_orders_table_exists(conn) -> None:
)
if cur.fetchone() is None:
raise ValueError(
"❌ Таблица public.orders не найдена; запусти DAG csv_to_greenplum."
"❌ Таблица public.orders не найдена; запусти DAG csv_to_postgres."
)
logging.info("✅ Таблица public.orders существует")
@@ -103,7 +103,7 @@ def fetch_orders_schema(conn) -> Sequence[Tuple[str, str]]:
Получает схему таблицы orders из information_schema.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Returns:
Список кортежей (имя_колонки, тип_данных)
@@ -125,7 +125,7 @@ def assert_orders_schema(conn) -> None:
Проверяет, что схема таблицы orders соответствует ожидаемой.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Raises:
ValueError: Если схема не соответствует ожидаемой
@@ -147,7 +147,7 @@ def fetch_orders_count(conn) -> int:
Получает количество строк в таблице orders.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Returns:
Количество строк в таблице
@@ -162,7 +162,7 @@ def assert_orders_have_rows(conn) -> None:
Проверяет, что таблица orders не пустая.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Raises:
ValueError: Если таблица пустая
@@ -173,7 +173,7 @@ def assert_orders_have_rows(conn) -> None:
if row_count <= 0:
raise ValueError(
"❌ Таблица public.orders пустая — запусти DAG csv_to_greenplum перед проверкой."
"❌ Таблица public.orders пустая — запусти DAG csv_to_postgres перед проверкой."
)
logging.info("✅ Таблица orders содержит данные (%s строк)", row_count)
@@ -183,7 +183,7 @@ def fetch_orders_duplicates(conn) -> int:
Подсчитывает количество дубликатов по order_id.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Returns:
Количество дублирующихся order_id
@@ -207,7 +207,7 @@ def assert_orders_no_duplicates(conn) -> None:
Проверяет, что в таблице нет дублей по order_id.
Args:
conn: Подключение к Greenplum
conn: Подключение к Postgres
Raises:
ValueError: Если обнаружены дубликаты
@@ -220,4 +220,4 @@ def assert_orders_no_duplicates(conn) -> None:
raise ValueError(
f"❌ Обнаружены дубли по order_id ({duplicates} шт.) — проверь загрузку данных."
)
logging.info("✅ Дубликаты не обнаружены")
logging.info("✅ Дубликаты не обнаружены")