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 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
@@ -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
@@ -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",
) )
@@ -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("✅ Дубликаты не обнаружены")
+3
View File
@@ -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 &&
+3
View File
@@ -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 -4
View File
@@ -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);
+20
View File
@@ -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,
)
@@ -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)