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:
@@ -9,7 +9,7 @@ from typing import List
|
|||||||
|
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
from airflow.operators.python import PythonOperator
|
from airflow.operators.python import PythonOperator
|
||||||
from helpers.greenplum import get_gp_conn
|
from helpers.postgres import get_postgres_conn
|
||||||
|
|
||||||
from airflow import DAG
|
from airflow import DAG
|
||||||
|
|
||||||
@@ -21,15 +21,13 @@ def _create_table() -> None:
|
|||||||
"""Создаёт таблицу public.orders, если она ещё не существует."""
|
"""Создаёт таблицу public.orders, если она ещё не существует."""
|
||||||
ddl = """
|
ddl = """
|
||||||
CREATE TABLE IF NOT EXISTS public.orders (
|
CREATE TABLE IF NOT EXISTS public.orders (
|
||||||
order_id BIGINT,
|
order_id BIGINT PRIMARY KEY,
|
||||||
order_ts TIMESTAMP NOT NULL,
|
order_ts TIMESTAMP NOT NULL,
|
||||||
customer_id BIGINT NOT NULL,
|
customer_id BIGINT NOT NULL,
|
||||||
amount NUMERIC(12,2) 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)
|
cur.execute(ddl)
|
||||||
conn.commit()
|
conn.commit()
|
||||||
|
|
||||||
@@ -90,13 +88,13 @@ def _preview_csv(csv_path: str, sample_rows: int = 5) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def _load_csv(csv_path: str) -> None:
|
def _load_csv(csv_path: str) -> None:
|
||||||
"""Загружает CSV в Greenplum через временную таблицу и anti-join."""
|
"""Загружает CSV в Postgres через временную таблицу и anti-join."""
|
||||||
csv_file = Path(csv_path)
|
csv_file = Path(csv_path)
|
||||||
if not csv_file.exists():
|
if not csv_file.exists():
|
||||||
raise FileNotFoundError(f"CSV не найден: {csv_file}")
|
raise FileNotFoundError(f"CSV не найден: {csv_file}")
|
||||||
|
|
||||||
with (
|
with (
|
||||||
get_gp_conn() as conn,
|
get_postgres_conn() as conn,
|
||||||
conn.cursor() as cur,
|
conn.cursor() as cur,
|
||||||
csv_file.open("r", encoding="utf-8") as f,
|
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)}
|
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
|
||||||
|
|
||||||
with DAG(
|
with DAG(
|
||||||
dag_id="csv_to_greenplum",
|
dag_id="csv_to_postgres",
|
||||||
start_date=datetime(2017, 1, 1),
|
start_date=datetime(2017, 1, 1),
|
||||||
schedule=None,
|
schedule=None,
|
||||||
catchup=False,
|
catchup=False,
|
||||||
default_args=default_args,
|
default_args=default_args,
|
||||||
tags=["demo", "greenplum", "csv"],
|
tags=["demo", "postgres", "csv"],
|
||||||
) as dag:
|
) as dag:
|
||||||
create_table = PythonOperator(
|
create_table = PythonOperator(
|
||||||
task_id="create_orders_table",
|
task_id="create_orders_table",
|
||||||
@@ -157,11 +155,11 @@ with DAG(
|
|||||||
)
|
)
|
||||||
|
|
||||||
load_csv = PythonOperator(
|
load_csv = PythonOperator(
|
||||||
task_id="load_csv_to_greenplum",
|
task_id="load_csv_to_postgres",
|
||||||
python_callable=_load_csv,
|
python_callable=_load_csv,
|
||||||
op_kwargs={
|
op_kwargs={
|
||||||
"csv_path": "{{ ti.xcom_pull(task_ids='generate_csv') }}",
|
"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
|
||||||
+8
-8
@@ -4,12 +4,12 @@ import logging
|
|||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
from airflow.operators.python import PythonOperator
|
from airflow.operators.python import PythonOperator
|
||||||
from helpers.greenplum import (
|
from helpers.postgres import (
|
||||||
assert_orders_have_rows,
|
assert_orders_have_rows,
|
||||||
assert_orders_no_duplicates,
|
assert_orders_no_duplicates,
|
||||||
assert_orders_schema,
|
assert_orders_schema,
|
||||||
assert_orders_table_exists,
|
assert_orders_table_exists,
|
||||||
get_gp_conn,
|
get_postgres_conn,
|
||||||
)
|
)
|
||||||
|
|
||||||
from airflow import DAG
|
from airflow import DAG
|
||||||
@@ -17,7 +17,7 @@ from airflow import DAG
|
|||||||
|
|
||||||
def _run_check(check_callable):
|
def _run_check(check_callable):
|
||||||
"""
|
"""
|
||||||
Оборачивает проверку качества данных в контекст подключения к Greenplum.
|
Оборачивает проверку качества данных в контекст подключения к Postgres.
|
||||||
|
|
||||||
Этот DAG предназначен для автоматической проверки качества данных
|
Этот DAG предназначен для автоматической проверки качества данных
|
||||||
после CSV-пайплайна в таблице public.orders:
|
после CSV-пайплайна в таблице public.orders:
|
||||||
@@ -33,7 +33,7 @@ def _run_check(check_callable):
|
|||||||
check_name = check_callable.__name__.replace("assert_", "")
|
check_name = check_callable.__name__.replace("assert_", "")
|
||||||
logging.info("🚀 Запуск проверки: %s", check_name)
|
logging.info("🚀 Запуск проверки: %s", check_name)
|
||||||
|
|
||||||
with get_gp_conn() as conn:
|
with get_postgres_conn() as conn:
|
||||||
check_callable(conn)
|
check_callable(conn)
|
||||||
|
|
||||||
logging.info("✅ Проверка пройдена: %s", check_name)
|
logging.info("✅ Проверка пройдена: %s", check_name)
|
||||||
@@ -51,13 +51,13 @@ def _log_dq_summary():
|
|||||||
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
|
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
|
||||||
|
|
||||||
with DAG(
|
with DAG(
|
||||||
dag_id="csv_to_greenplum_dq",
|
dag_id="csv_to_postgres_dq",
|
||||||
start_date=datetime(2017, 1, 1),
|
start_date=datetime(2017, 1, 1),
|
||||||
schedule=None,
|
schedule=None,
|
||||||
catchup=False,
|
catchup=False,
|
||||||
default_args=default_args,
|
default_args=default_args,
|
||||||
tags=["demo", "greenplum", "quality", "csv", "dq"],
|
tags=["demo", "postgres", "quality", "csv", "dq"],
|
||||||
description="Проверки качества данных после CSV → public.orders в Greenplum",
|
description="Проверки качества данных после CSV → public.orders в Postgres",
|
||||||
) as dag:
|
) as dag:
|
||||||
# Задача 1: Проверка существования таблицы
|
# Задача 1: Проверка существования таблицы
|
||||||
check_exists = PythonOperator(
|
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
|
||||||
+6
-6
@@ -1,7 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
"""
|
"""
|
||||||
Учебный DAG: применяет DDL для базовой таблицы orders в Greenplum.
|
Учебный DAG: применяет DDL для базовой таблицы orders в Postgres.
|
||||||
Запускается вручную перед CSV‑пайплайном или после изменения схемы.
|
Запускается вручную перед CSV‑пайплайном или после изменения схемы.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -11,7 +11,7 @@ from airflow.providers.postgres.operators.postgres import PostgresOperator
|
|||||||
|
|
||||||
from airflow import DAG
|
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)}
|
default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
|
||||||
|
|
||||||
@@ -20,13 +20,13 @@ with DAG(
|
|||||||
start_date=datetime(2017, 1, 1),
|
start_date=datetime(2017, 1, 1),
|
||||||
schedule=None,
|
schedule=None,
|
||||||
catchup=False,
|
catchup=False,
|
||||||
template_searchpath="/sql",
|
template_searchpath="/opt/airflow/sql",
|
||||||
default_args=default_args,
|
default_args=default_args,
|
||||||
tags=["demo", "greenplum", "ddl", "orders"],
|
tags=["demo", "postgres", "ddl", "orders"],
|
||||||
description="Создаёт/обновляет базовую таблицу orders в схеме public",
|
description="Создаёт/обновляет базовую таблицу orders в схеме public",
|
||||||
) as dag:
|
) as dag:
|
||||||
apply_orders_ddl = PostgresOperator(
|
apply_orders_ddl = PostgresOperator(
|
||||||
task_id="apply_orders_ddl",
|
task_id="apply_orders_ddl",
|
||||||
postgres_conn_id=GREENPLUM_CONN_ID,
|
postgres_conn_id=POSTGRES_CONN_ID,
|
||||||
sql="base/orders_ddl.sql",
|
sql="base/orders_ddl.sql",
|
||||||
)
|
)
|
||||||
+25
-25
@@ -1,7 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
"""
|
"""
|
||||||
LEGACY: Вспомогательные функции для прямого подключения к Greenplum через psycopg2.
|
LEGACY: Вспомогательные функции для прямого подключения к Postgres через psycopg2.
|
||||||
Внимание: этот модуль оставлен только для поддержки базового CSV-пайплайна.
|
Внимание: этот модуль оставлен только для поддержки базового CSV-пайплайна.
|
||||||
В новых DAG (ODS/DDS/DM) используйте встроенный в Airflow PostgresOperator
|
В новых DAG (ODS/DDS/DM) используйте встроенный в Airflow PostgresOperator
|
||||||
и штатные механизмы XCom.
|
и штатные механизмы XCom.
|
||||||
@@ -13,10 +13,10 @@ from typing import List, Sequence, Tuple
|
|||||||
|
|
||||||
import psycopg2
|
import psycopg2
|
||||||
|
|
||||||
# Настройки для подключения к Greenplum. По умолчанию используем Airflow Connection,
|
# Настройки для подключения к Postgres. По умолчанию используем Airflow Connection,
|
||||||
# но при проблемах можно переключиться на ENV-подключение, установив GP_USE_AIRFLOW_CONN=false.
|
# но при проблемах можно переключиться на ENV-подключение, установив POSTGRES_USE_AIRFLOW_CONN=false.
|
||||||
GP_CONN_ID = os.getenv("GP_CONN_ID", "greenplum_conn")
|
POSTGRES_CONN_ID = os.getenv("POSTGRES_CONN_ID", "postgres_training")
|
||||||
GP_USE_AIRFLOW_CONN = os.getenv("GP_USE_AIRFLOW_CONN", "true").lower() in (
|
POSTGRES_USE_AIRFLOW_CONN = os.getenv("POSTGRES_USE_AIRFLOW_CONN", "true").lower() in (
|
||||||
"1",
|
"1",
|
||||||
"true",
|
"true",
|
||||||
"yes",
|
"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 (если настроено и доступно)
|
1. Через Airflow Connection (если настроено и доступно)
|
||||||
@@ -42,11 +42,11 @@ def get_gp_conn():
|
|||||||
Returns:
|
Returns:
|
||||||
psycopg2 connection object
|
psycopg2 connection object
|
||||||
"""
|
"""
|
||||||
if GP_USE_AIRFLOW_CONN:
|
if POSTGRES_USE_AIRFLOW_CONN:
|
||||||
try:
|
try:
|
||||||
from airflow.providers.postgres.hooks.postgres import PostgresHook
|
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()
|
conn = hook.get_conn()
|
||||||
logging.info("✅ Подключение через Airflow Connection успешно")
|
logging.info("✅ Подключение через Airflow Connection успешно")
|
||||||
return conn
|
return conn
|
||||||
@@ -57,14 +57,14 @@ def get_gp_conn():
|
|||||||
|
|
||||||
# Прямое подключение по переменным окружения
|
# Прямое подключение по переменным окружения
|
||||||
conn_params = {
|
conn_params = {
|
||||||
"dbname": os.getenv("GP_DB", "gp_dwh"),
|
"dbname": os.getenv("POSTGRES_DB", "training"),
|
||||||
"user": os.getenv("GP_USER", "gpadmin"),
|
"user": os.getenv("POSTGRES_USER", "student"),
|
||||||
"password": os.getenv("GP_PASSWORD", ""),
|
"password": os.getenv("POSTGRES_PASSWORD", "student"),
|
||||||
"host": os.getenv("GP_HOST", "greenplum"),
|
"host": os.getenv("POSTGRES_HOST", "postgres-training"),
|
||||||
"port": int(os.getenv("GP_PORT", "5432")),
|
"port": int(os.getenv("POSTGRES_PORT", "5432")),
|
||||||
}
|
}
|
||||||
logging.info(
|
logging.info(
|
||||||
"🔗 Подключение к Greenplum: %s:%s/%s",
|
"🔗 Подключение к Postgres: %s:%s/%s",
|
||||||
conn_params["host"],
|
conn_params["host"],
|
||||||
conn_params["port"],
|
conn_params["port"],
|
||||||
conn_params["dbname"],
|
conn_params["dbname"],
|
||||||
@@ -77,7 +77,7 @@ def assert_orders_table_exists(conn) -> None:
|
|||||||
Проверяет наличие таблицы orders в схеме public.
|
Проверяет наличие таблицы orders в схеме public.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: Если таблица не найдена
|
ValueError: Если таблица не найдена
|
||||||
@@ -93,7 +93,7 @@ def assert_orders_table_exists(conn) -> None:
|
|||||||
)
|
)
|
||||||
if cur.fetchone() is None:
|
if cur.fetchone() is None:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"❌ Таблица public.orders не найдена; запусти DAG csv_to_greenplum."
|
"❌ Таблица public.orders не найдена; запусти DAG csv_to_postgres."
|
||||||
)
|
)
|
||||||
logging.info("✅ Таблица public.orders существует")
|
logging.info("✅ Таблица public.orders существует")
|
||||||
|
|
||||||
@@ -103,7 +103,7 @@ def fetch_orders_schema(conn) -> Sequence[Tuple[str, str]]:
|
|||||||
Получает схему таблицы orders из information_schema.
|
Получает схему таблицы orders из information_schema.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Список кортежей (имя_колонки, тип_данных)
|
Список кортежей (имя_колонки, тип_данных)
|
||||||
@@ -125,7 +125,7 @@ def assert_orders_schema(conn) -> None:
|
|||||||
Проверяет, что схема таблицы orders соответствует ожидаемой.
|
Проверяет, что схема таблицы orders соответствует ожидаемой.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: Если схема не соответствует ожидаемой
|
ValueError: Если схема не соответствует ожидаемой
|
||||||
@@ -147,7 +147,7 @@ def fetch_orders_count(conn) -> int:
|
|||||||
Получает количество строк в таблице orders.
|
Получает количество строк в таблице orders.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Количество строк в таблице
|
Количество строк в таблице
|
||||||
@@ -162,7 +162,7 @@ def assert_orders_have_rows(conn) -> None:
|
|||||||
Проверяет, что таблица orders не пустая.
|
Проверяет, что таблица orders не пустая.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: Если таблица пустая
|
ValueError: Если таблица пустая
|
||||||
@@ -173,7 +173,7 @@ def assert_orders_have_rows(conn) -> None:
|
|||||||
|
|
||||||
if row_count <= 0:
|
if row_count <= 0:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"❌ Таблица public.orders пустая — запусти DAG csv_to_greenplum перед проверкой."
|
"❌ Таблица public.orders пустая — запусти DAG csv_to_postgres перед проверкой."
|
||||||
)
|
)
|
||||||
logging.info("✅ Таблица orders содержит данные (%s строк)", row_count)
|
logging.info("✅ Таблица orders содержит данные (%s строк)", row_count)
|
||||||
|
|
||||||
@@ -183,7 +183,7 @@ def fetch_orders_duplicates(conn) -> int:
|
|||||||
Подсчитывает количество дубликатов по order_id.
|
Подсчитывает количество дубликатов по order_id.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Количество дублирующихся order_id
|
Количество дублирующихся order_id
|
||||||
@@ -207,7 +207,7 @@ def assert_orders_no_duplicates(conn) -> None:
|
|||||||
Проверяет, что в таблице нет дублей по order_id.
|
Проверяет, что в таблице нет дублей по order_id.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
conn: Подключение к Greenplum
|
conn: Подключение к Postgres
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: Если обнаружены дубликаты
|
ValueError: Если обнаружены дубликаты
|
||||||
@@ -220,4 +220,4 @@ def assert_orders_no_duplicates(conn) -> None:
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"❌ Обнаружены дубли по order_id ({duplicates} шт.) — проверь загрузку данных."
|
f"❌ Обнаружены дубли по order_id ({duplicates} шт.) — проверь загрузку данных."
|
||||||
)
|
)
|
||||||
logging.info("✅ Дубликаты не обнаружены")
|
logging.info("✅ Дубликаты не обнаружены")
|
||||||
@@ -61,6 +61,7 @@ services:
|
|||||||
volumes:
|
volumes:
|
||||||
- ./dags:/opt/airflow/dags
|
- ./dags:/opt/airflow/dags
|
||||||
- ./data:/opt/airflow/data
|
- ./data:/opt/airflow/data
|
||||||
|
- ./tests:/opt/airflow/tests
|
||||||
depends_on:
|
depends_on:
|
||||||
airflow-init:
|
airflow-init:
|
||||||
condition: service_completed_successfully
|
condition: service_completed_successfully
|
||||||
@@ -82,6 +83,7 @@ services:
|
|||||||
volumes:
|
volumes:
|
||||||
- ./dags:/opt/airflow/dags
|
- ./dags:/opt/airflow/dags
|
||||||
- ./data:/opt/airflow/data
|
- ./data:/opt/airflow/data
|
||||||
|
- ./tests:/opt/airflow/tests
|
||||||
depends_on:
|
depends_on:
|
||||||
airflow-init:
|
airflow-init:
|
||||||
condition: service_completed_successfully
|
condition: service_completed_successfully
|
||||||
@@ -100,6 +102,7 @@ services:
|
|||||||
volumes:
|
volumes:
|
||||||
- ./dags:/opt/airflow/dags
|
- ./dags:/opt/airflow/dags
|
||||||
- ./data:/opt/airflow/data
|
- ./data:/opt/airflow/data
|
||||||
|
- ./tests:/opt/airflow/tests
|
||||||
command: >
|
command: >
|
||||||
bash -ceuo pipefail "
|
bash -ceuo pipefail "
|
||||||
mkdir -p /opt/airflow/data &&
|
mkdir -p /opt/airflow/data &&
|
||||||
|
|||||||
@@ -12,3 +12,6 @@ mimesis==15.1.0
|
|||||||
|
|
||||||
# Airflow PostgreSQL provider (used in sql_basic_dag.py and data_processing_dag.py)
|
# Airflow PostgreSQL provider (used in sql_basic_dag.py and data_processing_dag.py)
|
||||||
apache-airflow-providers-postgres==5.11.1
|
apache-airflow-providers-postgres==5.11.1
|
||||||
|
|
||||||
|
# Testing framework
|
||||||
|
pytest==7.4.4
|
||||||
|
|||||||
@@ -2,10 +2,8 @@
|
|||||||
-- Выполняется идемпотентно: таблица создаётся, если ещё не существует.
|
-- Выполняется идемпотентно: таблица создаётся, если ещё не существует.
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS public.orders (
|
CREATE TABLE IF NOT EXISTS public.orders (
|
||||||
order_id BIGINT,
|
order_id BIGINT PRIMARY KEY,
|
||||||
order_ts TIMESTAMP NOT NULL,
|
order_ts TIMESTAMP NOT NULL,
|
||||||
customer_id BIGINT NOT NULL,
|
customer_id BIGINT NOT NULL,
|
||||||
amount NUMERIC(12,2) NOT NULL
|
amount NUMERIC(12,2) NOT NULL
|
||||||
)
|
);
|
||||||
WITH (appendonly=true, orientation=row, compresstype=zstd, compresslevel=1)
|
|
||||||
DISTRIBUTED BY (order_id);
|
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def patch_postgres_hook(monkeypatch, fake_hook_class) -> None:
|
||||||
|
"""
|
||||||
|
Патчит PostgresHook для тестирования.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
monkeypatch: pytest monkeypatch fixture
|
||||||
|
fake_hook_class: Класс-имитация PostgresHook
|
||||||
|
"""
|
||||||
|
# Патчим PostgresHook на уровне airflow.providers.postgres.hooks.postgres
|
||||||
|
# Это нужно, так как get_postgres_conn() импортирует его оттуда
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"airflow.providers.postgres.hooks.postgres.PostgresHook",
|
||||||
|
fake_hook_class,
|
||||||
|
raising=False,
|
||||||
|
)
|
||||||
+42
-35
@@ -1,12 +1,19 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sys
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
from typing import Any, List, Sequence
|
from typing import Any, List, Sequence
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
import airflow.dags.helpers.greenplum as greenplum
|
# Добавляем путь к dags в sys.path для импорта модулей
|
||||||
from tests.conftest import patch_postgres_hook
|
dags_path = str(Path(__file__).parent.parent / "dags")
|
||||||
|
if dags_path not in sys.path:
|
||||||
|
sys.path.insert(0, dags_path)
|
||||||
|
|
||||||
|
import helpers.postgres as postgres_helpers
|
||||||
|
from conftest import patch_postgres_hook
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -49,7 +56,7 @@ class FakeConn:
|
|||||||
self.commits += 1
|
self.commits += 1
|
||||||
|
|
||||||
|
|
||||||
def test_get_gp_conn_uses_airflow_hook(monkeypatch) -> None:
|
def test_get_postgres_conn_uses_airflow_hook(monkeypatch) -> None:
|
||||||
class FakeHook:
|
class FakeHook:
|
||||||
def __init__(self, postgres_conn_id: str) -> None:
|
def __init__(self, postgres_conn_id: str) -> None:
|
||||||
self.postgres_conn_id = postgres_conn_id
|
self.postgres_conn_id = postgres_conn_id
|
||||||
@@ -58,15 +65,15 @@ def test_get_gp_conn_uses_airflow_hook(monkeypatch) -> None:
|
|||||||
return "hook_connection"
|
return "hook_connection"
|
||||||
|
|
||||||
patch_postgres_hook(monkeypatch, FakeHook)
|
patch_postgres_hook(monkeypatch, FakeHook)
|
||||||
monkeypatch.setattr(greenplum, "GP_CONN_ID", "demo_conn", raising=False)
|
monkeypatch.setattr(postgres_helpers, "POSTGRES_CONN_ID", "demo_conn", raising=False)
|
||||||
monkeypatch.setattr(greenplum, "GP_USE_AIRFLOW_CONN", True, raising=False)
|
monkeypatch.setattr(postgres_helpers, "POSTGRES_USE_AIRFLOW_CONN", True, raising=False)
|
||||||
|
|
||||||
conn = greenplum.get_gp_conn()
|
conn = postgres_helpers.get_postgres_conn()
|
||||||
|
|
||||||
assert conn == "hook_connection"
|
assert conn == "hook_connection"
|
||||||
|
|
||||||
|
|
||||||
def test_get_gp_conn_fallback_to_psycopg(monkeypatch) -> None:
|
def test_get_postgres_conn_fallback_to_psycopg(monkeypatch) -> None:
|
||||||
class BrokenHook:
|
class BrokenHook:
|
||||||
def __init__(self, postgres_conn_id: str) -> None:
|
def __init__(self, postgres_conn_id: str) -> None:
|
||||||
self.postgres_conn_id = postgres_conn_id
|
self.postgres_conn_id = postgres_conn_id
|
||||||
@@ -75,13 +82,13 @@ def test_get_gp_conn_fallback_to_psycopg(monkeypatch) -> None:
|
|||||||
raise RuntimeError("boom")
|
raise RuntimeError("boom")
|
||||||
|
|
||||||
patch_postgres_hook(monkeypatch, BrokenHook)
|
patch_postgres_hook(monkeypatch, BrokenHook)
|
||||||
monkeypatch.setattr(greenplum, "GP_USE_AIRFLOW_CONN", True, raising=False)
|
monkeypatch.setattr(postgres_helpers, "POSTGRES_USE_AIRFLOW_CONN", True, raising=False)
|
||||||
monkeypatch.setattr(greenplum, "GP_CONN_ID", "demo_conn", raising=False)
|
monkeypatch.setattr(postgres_helpers, "POSTGRES_CONN_ID", "demo_conn", raising=False)
|
||||||
monkeypatch.setenv("GP_DB", "demo_db")
|
monkeypatch.setenv("POSTGRES_DB", "demo_db")
|
||||||
monkeypatch.setenv("GP_USER", "demo_user")
|
monkeypatch.setenv("POSTGRES_USER", "demo_user")
|
||||||
monkeypatch.setenv("GP_PASSWORD", "secret")
|
monkeypatch.setenv("POSTGRES_PASSWORD", "secret")
|
||||||
monkeypatch.setenv("GP_HOST", "greenplum-host")
|
monkeypatch.setenv("POSTGRES_HOST", "postgres-host")
|
||||||
monkeypatch.setenv("GP_PORT", "5434")
|
monkeypatch.setenv("POSTGRES_PORT", "5434")
|
||||||
|
|
||||||
captured_kwargs = {}
|
captured_kwargs = {}
|
||||||
|
|
||||||
@@ -89,27 +96,27 @@ def test_get_gp_conn_fallback_to_psycopg(monkeypatch) -> None:
|
|||||||
captured_kwargs.update(kwargs)
|
captured_kwargs.update(kwargs)
|
||||||
return "psycopg_connection"
|
return "psycopg_connection"
|
||||||
|
|
||||||
monkeypatch.setattr(greenplum.psycopg2, "connect", fake_connect)
|
monkeypatch.setattr(postgres_helpers.psycopg2, "connect", fake_connect)
|
||||||
|
|
||||||
conn = greenplum.get_gp_conn()
|
conn = postgres_helpers.get_postgres_conn()
|
||||||
|
|
||||||
assert conn == "psycopg_connection"
|
assert conn == "psycopg_connection"
|
||||||
assert captured_kwargs == {
|
assert captured_kwargs == {
|
||||||
"dbname": "demo_db",
|
"dbname": "demo_db",
|
||||||
"user": "demo_user",
|
"user": "demo_user",
|
||||||
"password": "secret",
|
"password": "secret",
|
||||||
"host": "greenplum-host",
|
"host": "postgres-host",
|
||||||
"port": 5434,
|
"port": 5434,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
def test_get_gp_conn_without_airflow(monkeypatch) -> None:
|
def test_get_postgres_conn_without_airflow(monkeypatch) -> None:
|
||||||
monkeypatch.setattr(greenplum, "GP_USE_AIRFLOW_CONN", False, raising=False)
|
monkeypatch.setattr(postgres_helpers, "POSTGRES_USE_AIRFLOW_CONN", False, raising=False)
|
||||||
monkeypatch.setenv("GP_DB", "demo_db")
|
monkeypatch.setenv("POSTGRES_DB", "demo_db")
|
||||||
monkeypatch.setenv("GP_USER", "demo_user")
|
monkeypatch.setenv("POSTGRES_USER", "demo_user")
|
||||||
monkeypatch.setenv("GP_PASSWORD", "secret")
|
monkeypatch.setenv("POSTGRES_PASSWORD", "secret")
|
||||||
monkeypatch.setenv("GP_HOST", "greenplum-host")
|
monkeypatch.setenv("POSTGRES_HOST", "postgres-host")
|
||||||
monkeypatch.setenv("GP_PORT", "5435")
|
monkeypatch.setenv("POSTGRES_PORT", "5435")
|
||||||
|
|
||||||
captured_kwargs = {}
|
captured_kwargs = {}
|
||||||
|
|
||||||
@@ -117,9 +124,9 @@ def test_get_gp_conn_without_airflow(monkeypatch) -> None:
|
|||||||
captured_kwargs.update(kwargs)
|
captured_kwargs.update(kwargs)
|
||||||
return "direct_psycopg"
|
return "direct_psycopg"
|
||||||
|
|
||||||
monkeypatch.setattr(greenplum.psycopg2, "connect", fake_connect)
|
monkeypatch.setattr(postgres_helpers.psycopg2, "connect", fake_connect)
|
||||||
|
|
||||||
conn = greenplum.get_gp_conn()
|
conn = postgres_helpers.get_postgres_conn()
|
||||||
|
|
||||||
assert conn == "direct_psycopg"
|
assert conn == "direct_psycopg"
|
||||||
assert captured_kwargs["port"] == 5435
|
assert captured_kwargs["port"] == 5435
|
||||||
@@ -128,51 +135,51 @@ def test_get_gp_conn_without_airflow(monkeypatch) -> None:
|
|||||||
def test_assert_orders_table_exists_ok() -> None:
|
def test_assert_orders_table_exists_ok() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchone_value=(1,))])
|
conn = FakeConn([FakeCursor(fetchone_value=(1,))])
|
||||||
|
|
||||||
greenplum.assert_orders_table_exists(conn)
|
postgres_helpers.assert_orders_table_exists(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_table_exists_missing() -> None:
|
def test_assert_orders_table_exists_missing() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchone_value=None)])
|
conn = FakeConn([FakeCursor(fetchone_value=None)])
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
greenplum.assert_orders_table_exists(conn)
|
postgres_helpers.assert_orders_table_exists(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_schema_ok() -> None:
|
def test_assert_orders_schema_ok() -> None:
|
||||||
expected = list(greenplum.EXPECTED_ORDERS_SCHEMA)
|
expected = list(postgres_helpers.EXPECTED_ORDERS_SCHEMA)
|
||||||
conn = FakeConn([FakeCursor(fetchall_value=expected)])
|
conn = FakeConn([FakeCursor(fetchall_value=expected)])
|
||||||
|
|
||||||
greenplum.assert_orders_schema(conn)
|
postgres_helpers.assert_orders_schema(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_schema_mismatch() -> None:
|
def test_assert_orders_schema_mismatch() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchall_value=[("order_id", "bigint")])])
|
conn = FakeConn([FakeCursor(fetchall_value=[("order_id", "bigint")])])
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
greenplum.assert_orders_schema(conn)
|
postgres_helpers.assert_orders_schema(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_have_rows_ok() -> None:
|
def test_assert_orders_have_rows_ok() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchone_value=(5,))])
|
conn = FakeConn([FakeCursor(fetchone_value=(5,))])
|
||||||
|
|
||||||
greenplum.assert_orders_have_rows(conn)
|
postgres_helpers.assert_orders_have_rows(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_have_rows_empty() -> None:
|
def test_assert_orders_have_rows_empty() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchone_value=(0,))])
|
conn = FakeConn([FakeCursor(fetchone_value=(0,))])
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
greenplum.assert_orders_have_rows(conn)
|
postgres_helpers.assert_orders_have_rows(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_no_duplicates_ok() -> None:
|
def test_assert_orders_no_duplicates_ok() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchone_value=(0,))])
|
conn = FakeConn([FakeCursor(fetchone_value=(0,))])
|
||||||
|
|
||||||
greenplum.assert_orders_no_duplicates(conn)
|
postgres_helpers.assert_orders_no_duplicates(conn)
|
||||||
|
|
||||||
|
|
||||||
def test_assert_orders_no_duplicates_detected() -> None:
|
def test_assert_orders_no_duplicates_detected() -> None:
|
||||||
conn = FakeConn([FakeCursor(fetchone_value=(3,))])
|
conn = FakeConn([FakeCursor(fetchone_value=(3,))])
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
greenplum.assert_orders_no_duplicates(conn)
|
postgres_helpers.assert_orders_no_duplicates(conn)
|
||||||
Reference in New Issue
Block a user