feat(greenplum): добавлен CSV-пайплайн и проверки качества данных
- Зачем: - Необходим пример загрузки данных из CSV в Greenplum через Airflow - Добавлены проверки качества данных для валидации загруженных данных - Что: - создан DAG csv_to_greenplum для генерации и загрузки CSV в таблицу orders - создан DAG csv_to_greenplum_dq для проверок качества данных - создан DAG ddl_greenplum_base для применения DDL таблицы orders - добавлен модуль helpers/greenplum с функциями подключения и валидации - добавлен SQL-скрипт sql/base/orders_ddl.sql для создания таблицы - добавлены unit-тесты test_greenplum_helpers.py - Проверка: - запуск DAG csv_to_greenplum в Airflow UI - запуск тестов pytest tests/test_greenplum_helpers.py
This commit is contained in:
@@ -0,0 +1,223 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""
|
||||
LEGACY: Вспомогательные функции для прямого подключения к Greenplum через psycopg2.
|
||||
Внимание: этот модуль оставлен только для поддержки базового CSV-пайплайна.
|
||||
В новых DAG (ODS/DDS/DM) используйте встроенный в Airflow PostgresOperator
|
||||
и штатные механизмы XCom.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
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 (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
|
||||
# Ожидаемая схема таблицы orders для проверки качества данных
|
||||
EXPECTED_ORDERS_SCHEMA: List[Tuple[str, str]] = [
|
||||
("order_id", "bigint"),
|
||||
("order_ts", "timestamp without time zone"),
|
||||
("customer_id", "bigint"),
|
||||
("amount", "numeric"),
|
||||
]
|
||||
|
||||
|
||||
def get_gp_conn():
|
||||
"""
|
||||
Возвращает psycopg2 connection к Greenplum.
|
||||
|
||||
Приоритет подключения:
|
||||
1. Через Airflow Connection (если настроено и доступно)
|
||||
2. Прямое подключение по переменным окружения (фоллбек)
|
||||
|
||||
Returns:
|
||||
psycopg2 connection object
|
||||
"""
|
||||
if GP_USE_AIRFLOW_CONN:
|
||||
try:
|
||||
from airflow.providers.postgres.hooks.postgres import PostgresHook
|
||||
|
||||
hook = PostgresHook(postgres_conn_id=GP_CONN_ID)
|
||||
conn = hook.get_conn()
|
||||
logging.info("✅ Подключение через Airflow Connection успешно")
|
||||
return conn
|
||||
except Exception as e:
|
||||
logging.warning("⚠️ Не удалось подключиться через Airflow Connection: %s", e)
|
||||
logging.info("🔄 Переключаемся на прямое подключение по ENV переменным")
|
||||
# Фоллбек на прямое подключение по переменным окружения.
|
||||
|
||||
# Прямое подключение по переменным окружения
|
||||
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")),
|
||||
}
|
||||
logging.info(
|
||||
"🔗 Подключение к Greenplum: %s:%s/%s",
|
||||
conn_params["host"],
|
||||
conn_params["port"],
|
||||
conn_params["dbname"],
|
||||
)
|
||||
return psycopg2.connect(**conn_params)
|
||||
|
||||
|
||||
def assert_orders_table_exists(conn) -> None:
|
||||
"""
|
||||
Проверяет наличие таблицы orders в схеме public.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Raises:
|
||||
ValueError: Если таблица не найдена
|
||||
"""
|
||||
logging.info("🔍 Проверяем существование таблицы public.orders...")
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT 1
|
||||
FROM pg_catalog.pg_tables
|
||||
WHERE schemaname = 'public' AND tablename = 'orders'
|
||||
"""
|
||||
)
|
||||
if cur.fetchone() is None:
|
||||
raise ValueError(
|
||||
"❌ Таблица public.orders не найдена; запусти DAG csv_to_greenplum."
|
||||
)
|
||||
logging.info("✅ Таблица public.orders существует")
|
||||
|
||||
|
||||
def fetch_orders_schema(conn) -> Sequence[Tuple[str, str]]:
|
||||
"""
|
||||
Получает схему таблицы orders из information_schema.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Returns:
|
||||
Список кортежей (имя_колонки, тип_данных)
|
||||
"""
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT column_name, data_type
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public' AND table_name = 'orders'
|
||||
ORDER BY ordinal_position
|
||||
"""
|
||||
)
|
||||
return cur.fetchall()
|
||||
|
||||
|
||||
def assert_orders_schema(conn) -> None:
|
||||
"""
|
||||
Проверяет, что схема таблицы orders соответствует ожидаемой.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Raises:
|
||||
ValueError: Если схема не соответствует ожидаемой
|
||||
"""
|
||||
logging.info("📋 Проверяем схему таблицы orders...")
|
||||
schema = fetch_orders_schema(conn)
|
||||
logging.info("📊 Фактическая схема: %s", list(schema))
|
||||
logging.info("📊 Ожидаемая схема: %s", EXPECTED_ORDERS_SCHEMA)
|
||||
|
||||
if list(schema) != EXPECTED_ORDERS_SCHEMA:
|
||||
raise ValueError(
|
||||
f"❌ Неожиданная схема orders: {schema}. Ожидали {EXPECTED_ORDERS_SCHEMA}."
|
||||
)
|
||||
logging.info("✅ Схема таблицы orders соответствует ожиданиям")
|
||||
|
||||
|
||||
def fetch_orders_count(conn) -> int:
|
||||
"""
|
||||
Получает количество строк в таблице orders.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Returns:
|
||||
Количество строк в таблице
|
||||
"""
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("SELECT COUNT(*) FROM public.orders")
|
||||
return cur.fetchone()[0]
|
||||
|
||||
|
||||
def assert_orders_have_rows(conn) -> None:
|
||||
"""
|
||||
Проверяет, что таблица orders не пустая.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Raises:
|
||||
ValueError: Если таблица пустая
|
||||
"""
|
||||
logging.info("📊 Проверяем наличие данных в таблице orders...")
|
||||
row_count = fetch_orders_count(conn)
|
||||
logging.info("📈 Количество строк в orders: %s", row_count)
|
||||
|
||||
if row_count <= 0:
|
||||
raise ValueError(
|
||||
"❌ Таблица public.orders пустая — запусти DAG csv_to_greenplum перед проверкой."
|
||||
)
|
||||
logging.info("✅ Таблица orders содержит данные (%s строк)", row_count)
|
||||
|
||||
|
||||
def fetch_orders_duplicates(conn) -> int:
|
||||
"""
|
||||
Подсчитывает количество дубликатов по order_id.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Returns:
|
||||
Количество дублирующихся order_id
|
||||
"""
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT COUNT(*) FROM (
|
||||
SELECT order_id
|
||||
FROM public.orders
|
||||
GROUP BY order_id
|
||||
HAVING COUNT(*) > 1
|
||||
) d
|
||||
"""
|
||||
)
|
||||
return cur.fetchone()[0]
|
||||
|
||||
|
||||
def assert_orders_no_duplicates(conn) -> None:
|
||||
"""
|
||||
Проверяет, что в таблице нет дублей по order_id.
|
||||
|
||||
Args:
|
||||
conn: Подключение к Greenplum
|
||||
|
||||
Raises:
|
||||
ValueError: Если обнаружены дубликаты
|
||||
"""
|
||||
logging.info("🔍 Проверяем отсутствие дубликатов по order_id...")
|
||||
duplicates = fetch_orders_duplicates(conn)
|
||||
logging.info("📊 Найдено дубликатов: %s", duplicates)
|
||||
|
||||
if duplicates:
|
||||
raise ValueError(
|
||||
f"❌ Обнаружены дубли по order_id ({duplicates} шт.) — проверь загрузку данных."
|
||||
)
|
||||
logging.info("✅ Дубликаты не обнаружены")
|
||||
Reference in New Issue
Block a user