fix(dags): сделать учебные сценарии воспроизводимыми

Управляемые ветки и ошибки, зависимости общих обработчиков, сохранение SQL-результата и фактическое чтение результатов Extract. Добавлен пример пулов и XCom; DQ-сводка сохраняет отказ при all_done.

Проверено на Airflow 2.9.2: импорт и запуск девяти исходных DAG, выбор веток, повтор и окончательные отказы обеих рабочих задач, отрицательный DQ-сценарий и восстановление, фильтрация Extract и ожидание слотов.

airflow-docker/dags/branching_dag.py

@@ -1,103 +1,70 @@
+"""Выбор ветки через Variable branch_format: csv или json.
+
+Обработка файлов имитируется. В основной практике изучаются граф и статусы.
 """
-DAG для демонстрации условного выполнения задач в Airflow
-Уровень: Продвинутый
-"""
-from datetime import datetime, timedelta
+import logging
+from datetime import datetime
+
 from airflow import DAG
-from airflow.operators.python import PythonOperator, BranchPythonOperator
-from airflow.operators.dummy import DummyOperator
-import random
+from airflow.models import Variable
+from airflow.operators.empty import EmptyOperator
+from airflow.operators.python import BranchPythonOperator, PythonOperator
 
-# Определение DAG
-default_args = {
-    'owner': 'student',
-    'depends_on_past': False,
-    'start_date': datetime(2023, 1, 1),
-    'email_on_failure': False,
-    'email_on_retry': False,
-    'retries': 1,
-    'retry_delay': timedelta(minutes=5)
-}
 
-dag = DAG(
-    'branching_dag',
-    default_args=default_args,
-    description='DAG для изучения условного выполнения задач в Airflow',
-    schedule_interval=None,
-    catchup=False,
-    tags=['educational', 'branching', 'advanced']
-)
+def choose_format():
+    """Возвращает task_id выбранной ветки."""
+    file_format = Variable.get("branch_format", default_var="csv")
+    branches = {"csv": "process_csv_branch", "json": "process_json_branch"}
+    if file_format not in branches:
+        raise ValueError(f"branch_format должен быть csv или json, получено: {file_format}")
+    logging.info("Выбран формат: %s", file_format)
+    return branches[file_format]
 
-def check_data_quality():
-    """Проверка качества данных - случайным образом определяет, какие данные использовать"""
-    # В реальном сценарии здесь будет проверка качества данных
-    # Для учебных целей просто случайное решение
-    quality_score = random.random()  # случайное число от 0 до 1
-    
-    if quality_score > 0.5:
-        print(f"Качество данных хорошее (оценка: {quality_score:.2f}), используем CSV")
-        return 'process_csv_branch'
-    else:
-        print(f"Качество данных требует внимания (оценка: {quality_score:.2f}), используем JSON")
-        return 'process_json_branch'
 
 def process_csv_data():
-    """Обработка CSV данных"""
-    print("Обработка CSV файла...")
-    # Здесь будет логика обработки CSV файла
-    return "CSV данные обработаны"
+    """Имитирует обработку CSV."""
+    logging.info("Учебная CSV-ветка выполнена; файлы не читались")
+    return "csv"
+
 
 def process_json_data():
-    """Обработка JSON данных"""
-    print("Обработка JSON файла...")
-    # Здесь будет логика обработки JSON файла
-    return "JSON данные обработаны"
+    """Имитирует обработку JSON."""
+    logging.info("Учебная JSON-ветка выполнена; файлы не читались")
+    return "json"
 
-def merge_results():
-    """Объединение результатов из разных веток"""
-    print("Объединение результатов из разных веток...")
-    return "Результаты объединены"
 
-# Определение задач
-start_task = DummyOperator(
-    task_id='start_task',
-    dag=dag
-)
+def merge_results(ti):
+    """Читает результат выполненной ветки текущего запуска."""
+    results = ti.xcom_pull(task_ids=["process_csv_branch", "process_json_branch"])
+    selected = [value for value in results if value is not None]
+    logging.info("Результат выбранной ветки: %s", selected)
+    return selected
 
-check_quality_task = BranchPythonOperator(
-    task_id='check_data_quality',
-    python_callable=check_data_quality,
-    dag=dag
-)
 
-process_csv_task = PythonOperator(
-    task_id='process_csv_branch',
-    python_callable=process_csv_data,
-    dag=dag
-)
+with DAG(
+    "branching_dag",
+    start_date=datetime(2023, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"owner": "student", "retries": 0},
+    tags=["educational", "branching"],
+) as dag:
+    start_task = EmptyOperator(task_id="start_task")
+    choose_task = BranchPythonOperator(
+        task_id="choose_format", python_callable=choose_format,
+    )
+    process_csv_task = PythonOperator(
+        task_id="process_csv_branch", python_callable=process_csv_data,
+    )
+    process_json_task = PythonOperator(
+        task_id="process_json_branch", python_callable=process_json_data,
+    )
+    merge_task = PythonOperator(
+        task_id="merge_results",
+        python_callable=merge_results,
+        trigger_rule="none_failed_min_one_success",
+    )
+    end_task = EmptyOperator(task_id="end_task")
 
-process_json_task = PythonOperator(
-    task_id='process_json_branch',
-    python_callable=process_json_data,
-    dag=dag
-)
-
-merge_task = PythonOperator(
-    task_id='merge_results',
-    python_callable=merge_results,
-    trigger_rule='none_failed_or_skipped',  # Выполняется, когда одна из веток завершена
-    dag=dag
-)
-
-end_task = DummyOperator(
-    task_id='end_task',
-    dag=dag
-)
-
-# Установка зависимостей
-# BranchPythonOperator автоматически пропускает (skips) задачи в невыбранной ветке
-# Например, если check_data_quality возвращает 'process_csv_branch',
-# то задача process_json_branch будет пропущена (статус skipped)
-start_task >> check_quality_task
-check_quality_task >> [process_csv_task, process_json_task]
-[process_csv_task, process_json_task] >> merge_task >> end_task
+    start_task >> choose_task >> [process_csv_task, process_json_task]
+    [process_csv_task, process_json_task] >> merge_task >> end_task

airflow-docker/dags/csv_to_postgres_dq.py

@@ -2,7 +2,7 @@
 DAG для проверки качества данных (Data Quality) после загрузки в PostgreSQL
 Уровень: Средний
 
-Этот DAG запускается после csv_to_postgres и проверяет таблицу public.orders:
+Запустите этот DAG вручную после csv_to_postgres. Он проверяет public.orders:
 1. Существование таблицы
 2. Соответствие схемы (колонки и типы данных)
 3. Наличие данных (таблица не пуста)
@@ -14,12 +14,11 @@ DAG для проверки качества данных (Data Quality) пос
 from __future__ import annotations
 
 import logging
-from datetime import datetime, timedelta
-
-from airflow.operators.python import PythonOperator
-from airflow.providers.postgres.hooks.postgres import PostgresHook
+from datetime import datetime
 
 from airflow import DAG
+from airflow.operators.python import PythonOperator
+from airflow.providers.postgres.hooks.postgres import PostgresHook
 
 POSTGRES_CONN_ID = "postgres_training"
 
@@ -106,13 +105,20 @@ def _check_no_duplicates():
         conn.close()
 
 
-def _log_dq_summary():
-    """Логирует итоговую сводку по качеству данных."""
-    logging.info("Все проверки качества данных пройдены успешно!")
-    logging.info("Качество данных в таблице orders соответствует требованиям.")
+def _log_dq_summary(dag_run, ti):
+    """Проверяет итог, в том числе при учебном переходе на all_done."""
+    states = {
+        task.task_id: task.state
+        for task in dag_run.get_task_instances()
+        if task.task_id != ti.task_id
+    }
+    logging.info("Результаты проверок: %s", states)
+    if any(state != "success" for state in states.values()):
+        raise ValueError("Не все проверки качества данных прошли успешно")
+    logging.info("Все проверки качества данных пройдены")
 
 
-default_args = {"owner": "airflow", "retries": 1, "retry_delay": timedelta(seconds=30)}
+default_args = {"owner": "student", "retries": 0}
 
 with DAG(
     dag_id="csv_to_postgres_dq",

airflow-docker/dags/data_processing_dag.py

@@ -6,7 +6,7 @@ DAG для демонстрации ETL процессов в Airflow
 1. Создание тестовых данных (клиенты + заказы) в CSV
 2. Параллельное извлечение (Extract) клиентов и заказов
 3. Трансформация (Transform) — объединение таблиц, вычисляемые поля
-4. Загрузка (Load) — подготовка SQL для вставки в БД
+4. Имитация загрузки: число подготовленных строк и пример DDL в логах
 5. Генерация текстового отчёта со статистикой
 
 Результат: файлы в /opt/airflow/data/output/ (extracted_*, transformed_data.csv, report.txt).
@@ -14,7 +14,6 @@ DAG для демонстрации ETL процессов в Airflow
 from datetime import datetime, timedelta
 from airflow import DAG
 from airflow.operators.python import PythonOperator
-from airflow.providers.postgres.operators.postgres import PostgresOperator
 
 # Определение DAG
 default_args = {
@@ -85,8 +84,8 @@ def extract_orders():
 def transform_data():
     """Преобразование данных - объединение клиентов и заказов"""
     import pandas as pd
-    customers_df = pd.read_csv('/opt/airflow/data/input/customers.csv')
-    orders_df = pd.read_csv('/opt/airflow/data/input/orders.csv')
+    customers_df = pd.read_csv('/opt/airflow/data/output/extracted_customers.csv')
+    orders_df = pd.read_csv('/opt/airflow/data/output/extracted_orders.csv')
     
     # Объединяем данные
     merged_df = pd.merge(orders_df, customers_df, on='customer_id', how='left')
@@ -101,14 +100,14 @@ def transform_data():
     
     return f"Преобразованы {len(merged_df)} записей"
 
-def load_to_database():
+def simulate_load():
     """Загрузка данных в базу данных (симуляция)"""
     import pandas as pd
     df = pd.read_csv('/opt/airflow/data/output/transformed_data.csv')
     
     # В реальном сценарии здесь был бы код для загрузки в базу данных
     # Для учебных целей просто логируем
-    print(f"Загружено в базу данных: {len(df)} записей")
+    print(f"Имитация загрузки: подготовлено {len(df)} записей. Записи в БД нет.")
     
     # Создаем SQL для создания таблицы (в реальном сценарии)
     create_table_sql = """
@@ -185,8 +184,8 @@ transform_task = PythonOperator(
 )
 
 load_task = PythonOperator(
-    task_id='load_to_database',
-    python_callable=load_to_database,
+    task_id='simulate_load',
+    python_callable=simulate_load,
     dag=dag
 )
 
@@ -197,4 +196,4 @@ report_task = PythonOperator(
 )
 
 # Установка зависимостей
-create_data_task >> [extract_customers_task, extract_orders_task] >> transform_task >> load_task >> report_task
+create_data_task >> [extract_customers_task, extract_orders_task] >> transform_task >> load_task >> report_task

airflow-docker/dags/error_handling_dag.py

@@ -1,108 +1,89 @@
+"""Повторы и обработчики результата. Режим задается Variable error_mode.
+
+retry_once: вторая задача падает при первой попытке, затем выполняется.
+success: обе задачи выполняются сразу.
+fail_first / fail_second / fail_both: выбранные задачи падают при каждой попытке.
+Меняйте режим после завершения предыдущего запуска.
 """
-DAG для демонстрации обработки ошибок в Airflow
-Уровень: Продвинутый
-"""
+import logging
 from datetime import datetime, timedelta
+
 from airflow import DAG
+from airflow.models import Variable
+from airflow.operators.empty import EmptyOperator
 from airflow.operators.python import PythonOperator
-from airflow.operators.dummy import DummyOperator
-from airflow.models.baseoperator import chain
-import random
 
-# Определение DAG
-default_args = {
-    'owner': 'student',
-    'depends_on_past': False,
-    'start_date': datetime(2023, 1, 1),
-    'email_on_failure': False,  # Отключаем email уведомления для простоты
-    'email_on_retry': False,
-    'retries': 3,  # Количество попыток при ошибке
-    'retry_delay': timedelta(seconds=10)  # Задержка между попытками
-}
 
-dag = DAG(
-    'error_handling_dag',
-    default_args=default_args,
-    description='DAG для изучения обработки ошибок в Airflow',
-    schedule_interval=None,
+def get_error_mode():
+    """Читает режим во время выполнения задачи."""
+    mode = Variable.get("error_mode", default_var="retry_once")
+    if mode not in {"success", "retry_once", "fail_first", "fail_second", "fail_both"}:
+        raise ValueError(f"Неизвестный error_mode: {mode}")
+    return mode
+
+
+def run_first_task():
+    """Имитирует первую рабочую задачу."""
+    if get_error_mode() in {"fail_first", "fail_both"}:
+        raise ValueError("Учебная ошибка первой задачи")
+    logging.info("Первая задача выполнена")
+
+
+def run_retry_task(ti):
+    """Имитирует временную или постоянную ошибку второй задачи."""
+    mode = get_error_mode()
+    logging.info("Режим: %s; попытка: %s", mode, ti.try_number)
+    if mode in {"fail_second", "fail_both"}:
+        raise ValueError("Постоянная учебная ошибка второй задачи")
+    if mode == "retry_once" and ti.try_number == 1:
+        raise ValueError("Временная учебная ошибка: следующая попытка выполнится")
+    logging.info("Вторая задача выполнена")
+
+
+def handle_success():
+    """Сообщает об успехе обеих рабочих задач."""
+    logging.info("Обе рабочие задачи выполнены")
+
+
+def handle_failure():
+    """Сообщает об окончательном отказе хотя бы одной задачи."""
+    logging.error("Рабочая задача исчерпала повторы. Откройте ее лог")
+
+
+with DAG(
+    "error_handling_dag",
+    start_date=datetime(2023, 1, 1),
+    schedule=None,
     catchup=False,
-    tags=['educational', 'error_handling', 'advanced']
-)
-
-def unreliable_task():
-    """Задача, которая может завершиться с ошибкой"""
-    # В реальном сценарии это может быть задача, зависящая от внешних факторов
-    # Для учебных целей случайным образом генерируем ошибку
-    if random.random() < 0.3:  # 30% вероятность ошибки
-        print("Ошибка: задача не выполнена успешно!")
-        raise Exception("Случайная ошибка в задаче")
-
-    print("Задача выполнена успешно!")
-    return "Задача выполнена"
-
-def success_handler():
-    """Обработчик успешного выполнения"""
-    print("Поздравляем! Все задачи выполнены успешно!")
-    return "Успешно завершено"
-
-def failure_handler():
-    """Обработчик ошибок"""
-    print("Одна или несколько задач завершились с ошибкой!")
-    print("Проверьте логи для получения дополнительной информации")
-    return "Ошибка обработана"
-
-def retry_task():
-    """Задача с механизмом повторных попыток"""
-    # Имитируем задачу, которая может завершиться с ошибкой, но со временем исправляется
-    import time
-    time.sleep(2)  # Имитация работы
-
-    # С вероятностью 50% задача завершится с ошибкой
-    if random.random() < 0.5:
-        print("Ошибка в retry_task!")
-        raise Exception("Ошибка в задаче с повторными попытками")
-
-    print("retry_task выполнена успешно!")
-    return "retry_task завершена"
-
-# Определение задач
-start_task = DummyOperator(
-    task_id='start_task',
-    dag=dag
-)
-
-unreliable_task = PythonOperator(
-    task_id='unreliable_task',
-    python_callable=unreliable_task,
-    dag=dag
-)
-
-retry_task = PythonOperator(
-    task_id='retry_task',
-    python_callable=retry_task,
-    dag=dag
-)
-
-success_handler_task = PythonOperator(
-    task_id='success_handler',
-    python_callable=success_handler,
-    trigger_rule='all_success',  # Выполняется только если все предыдущие задачи успешны
-    dag=dag
-)
-
-failure_handler_task = PythonOperator(
-    task_id='failure_handler',
-    python_callable=failure_handler,
-    trigger_rule='one_failed',  # Выполняется если хотя бы одна предыдущая задача завершилась с ошибкой
-    dag=dag
-)
-
-end_task = DummyOperator(
-    task_id='end_task',
-    dag=dag
-)
-
-# Установка зависимостей
-# Используем chain, чтобы наглядно показать ученикам построение ветвящихся зависимостей без ручного перечисления операторов.
-chain(start_task, [unreliable_task, retry_task], [success_handler_task, failure_handler_task], end_task)
+    default_args={
+        "owner": "student",
+        "retries": 1,
+        "retry_delay": timedelta(seconds=10),
+    },
+    tags=["educational", "errors"],
+) as dag:
+    start_task = EmptyOperator(task_id="start_task")
+    unreliable_task = PythonOperator(
+        task_id="unreliable_task", python_callable=run_first_task,
+    )
+    retry_task = PythonOperator(
+        task_id="retry_task", python_callable=run_retry_task,
+    )
+    success_handler_task = PythonOperator(
+        task_id="success_handler",
+        python_callable=handle_success,
+        trigger_rule="all_success",
+        retries=0,
+    )
+    failure_handler_task = PythonOperator(
+        task_id="failure_handler",
+        python_callable=handle_failure,
+        trigger_rule="one_failed",
+        retries=0,
+    )
 
+    start_task >> [unreliable_task, retry_task]
+    [unreliable_task, retry_task] >> success_handler_task
+    [unreliable_task, retry_task] >> failure_handler_task
+    # Оба обработчика завершающие. При отказе success_handler получает
+    # upstream_failed, поэтому успешное уведомление не скрывает ошибку DAG Run.

airflow-docker/dags/resource_management_dag.py

@@ -0,0 +1,61 @@
+"""Пулы и XCom на двух задачах с наблюдаемой длительностью.
+
+До запуска создайте пул training_pool с двумя слотами.
+Чтение источников имитируется; подключений и входных файлов не требуется.
+"""
+import logging
+import time
+from datetime import datetime
+
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+
+
+def read_source(source, rows):
+    """Имитирует чтение источника в течение 20 секунд."""
+    logging.info("Начало чтения %s", source)
+    time.sleep(20)
+    logging.info("Чтение %s завершено: %s строк (имитация)", source, rows)
+    return rows
+
+
+def calculate_metrics(ti):
+    """Получает небольшие результаты двух задач через XCom."""
+    customers = ti.xcom_pull(task_ids="read_customers")
+    orders = ti.xcom_pull(task_ids="read_orders")
+    return {"customers": customers, "orders": orders}
+
+
+def log_metrics(ti):
+    """Печатает словарь, возвращенный предыдущей задачей."""
+    logging.info("Метрики: %s", ti.xcom_pull(task_ids="calculate_metrics"))
+
+
+with DAG(
+    "resource_management_dag",
+    start_date=datetime(2023, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"owner": "student", "retries": 0},
+    tags=["educational", "pools", "xcom"],
+) as dag:
+    customers_task = PythonOperator(
+        task_id="read_customers",
+        python_callable=read_source,
+        op_kwargs={"source": "customers", "rows": 20},
+        pool="training_pool",
+        pool_slots=2,
+    )
+    orders_task = PythonOperator(
+        task_id="read_orders",
+        python_callable=read_source,
+        op_kwargs={"source": "orders", "rows": 30},
+        pool="training_pool",
+        pool_slots=1,
+    )
+    metrics_task = PythonOperator(
+        task_id="calculate_metrics", python_callable=calculate_metrics,
+    )
+    log_task = PythonOperator(task_id="log_metrics", python_callable=log_metrics)
+
+    [customers_task, orders_task] >> metrics_task >> log_task

airflow-docker/dags/sql_basic_dag.py

@@ -4,7 +4,6 @@ DAG для демонстрации работы с SQL в Airflow
 """
 from datetime import datetime, timedelta
 from airflow import DAG
-from airflow.operators.python import PythonOperator
 from airflow.providers.postgres.operators.postgres import PostgresOperator
 
 # Определение DAG
@@ -76,11 +75,12 @@ query_data_task = PostgresOperator(
 )
 
 truncate_table_task = PostgresOperator(
-    task_id='drop_table',
+    task_id='truncate_table',
     postgres_conn_id='postgres_training',
     sql=truncate_table_sql,
     dag=dag
 )
 
 # Установка зависимостей
-create_table_task >> insert_data_task >> query_data_task >> truncate_table_task
+# Очистка перед вставкой: после запуска в таблице остаются три строки.
+create_table_task >> truncate_table_task >> insert_data_task >> query_data_task