Apache Airflow для оркестрации и автоматизации рабочих процессов

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.bash import BashOperator
 
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}
 
with DAG(
    'hello_world',
    default_args=default_args,
    description='Простой пример DAG',
    schedule_interval=timedelta(days=1),
    start_date=datetime(2023, 10, 1),
    catchup=False,
) as dag:
    
    t1 = BashOperator(
        task_id='print_date',
        bash_command='date',
    )
    
    t2 = BashOperator(
        task_id='print_hello',
        bash_command='echo "Привет, мир!"',
    )
    
    t1 >> t2  # Определяем последовательность задач

# Пример DAG с использованием TaskFlow API (новинка Airflow 2.0)
from airflow.decorators import dag, task
from datetime import datetime
 
@dag(
    schedule_interval='@daily',
    start_date=datetime(2023, 1, 1),
    catchup=False
)
def etl_workflow():
    
    @task()
    def extract():
        # Здесь логика извлечения данных
        return {"data": [1, 2, 3]}
    
    @task()
    def transform(data):
        # Преобразование данных
        return {"transformed": [x + 10 for x in data["data"]]}
    
    @task()
    def load(transformed_data):
        # Загрузка трансформированных данных
        print(f"Загружаем данные: {transformed_data}")
    
    # Определение потока данных
    extracted_data = extract()
    transformed_data = transform(extracted_data)
    load(transformed_data)
 
etl_dag = etl_workflow()

# Задача, которая отправляет значение через XCom
def push_function(**context):
    context['task_instance'].xcom_push(key='sample_key', value='sample_value')
    
push_task = PythonOperator(
    task_id='push_task',
    python_callable=push_function,
    provide_context=True,
    dag=dag
)
 
# Задача, которая получает значение из XCom
def pull_function(**context):
    value = context['task_instance'].xcom_pull(task_ids='push_task', key='sample_key')
    print(f"Получено значение: {value}")
    
pull_task = PythonOperator(
    task_id='pull_task',
    python_callable=pull_function,
    provide_context=True,
    dag=dag
)
 
push_task >> pull_task

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
 
# Аргументы по умолчанию для всех задач в DAG
default_args = {
    'owner': 'me',
    'retries': 2,
    'retry_delay': timedelta(minutes=1),
    'email_on_failure': False,
}
 
def extract_data():
    """Имитация извлечения данных"""
    print("Данные успешно извлечены")
    return {"успех": True, "записей": 1000}
 
def transform_data(**context):
    """Преобразование данных с использованием значения из предыдущей задачи"""
    ti = context['ti']
    extracted_data = ti.xcom_pull(task_ids='extract')
    records = extracted_data['записей']
    transformed_records = records * 1.2  # Какая-то трансформация
    print(f"Преобразовано {transformed_records} записей")
    return {"успех": True, "записей": transformed_records}
 
def load_data(**context):
    """Загрузка преобразованных данных"""
    ti = context['ti']
    transformed_data = ti.xcom_pull(task_ids='transform')
    print(f"Загружено {transformed_data['записей']} записей")
 
# Создаем DAG
with DAG(
    'simple_etl',
    default_args=default_args,
    description='Простой ETL процесс',
    schedule_interval='@daily',
    start_date=datetime(2023, 1, 1),
    catchup=False,
) as dag:
 
    # Определяем задачи
    extract = PythonOperator(
        task_id='extract',
        python_callable=extract_data,
    )
 
    transform = PythonOperator(
        task_id='transform',
        python_callable=transform_data,
        provide_context=True,
    )
 
    load = PythonOperator(
        task_id='load',
        python_callable=load_data,
        provide_context=True,
    )
 
    # Задаем последовательность выполнения
    extract >> transform >> load

import json
import pandas as pd
import requests
from datetime import datetime, timedelta
 
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
 
default_args = {
    'owner': 'data_team',
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}
 
def fetch_weather_data(**context):
    """Получаем данные о погоде из API"""
    api_key = context['params']['api_key']
    cities = context['params']['cities']
    
    weather_data = []
    for city in cities:
        url = f"https://api.openweathermap.org/data/2.5/weather?q={city}&appid={api_key}&units=metric"
        response = requests.get(url)
        if response.status_code == 200:
            weather_data.append(response.json())
    
    # Сохраняем результат в XCom
    context['ti'].xcom_push(key='weather_data', value=json.dumps(weather_data))
    return len(weather_data)
 
def transform_weather_data(**context):
    """Преобразуем данные в структурированный формат"""
    ti = context['ti']
    weather_data_str = ti.xcom_pull(task_ids='fetch_weather', key='weather_data')
    weather_data = json.loads(weather_data_str)
    
    # Создаем DataFrame и выбираем нужные колонки
    records = []
    for data in weather_data:
        records.append({
            'city': data['name'],
            'country': data['sys']['country'],
            'temperature': data['main']['temp'],
            'humidity': data['main']['humidity'],
            'description': data['weather'][0]['description'],
            'timestamp': datetime.utcfromtimestamp(data['dt'])
        })
    
    df = pd.DataFrame(records)
    
    # Сохраняем преобразованные данные
    transformed_data = df.to_dict(orient='records')
    context['ti'].xcom_push(key='transformed_data', value=json.dumps(transformed_data))
    return len(transformed_data)
 
def load_to_postgres(**context):
    """Загружаем данные в Postgres"""
    ti = context['ti']
    transformed_data_str = ti.xcom_pull(task_ids='transform', key='transformed_data')
    transformed_data = json.loads(transformed_data_str)
    
    # Подключаемся к Postgres через hook
    pg_hook = PostgresHook(postgres_conn_id='postgres_default')
    
    # Вставляем данные
    rows = 0
    for record in transformed_data:
        pg_hook.run(
            """
            INSERT INTO weather_data
            (city, country, temperature, humidity, description, timestamp)
            VALUES (%s, %s, %s, %s, %s, %s)
            """,
            parameters=(
                record['city'], 
                record['country'],
                record['temperature'],
                record['humidity'],
                record['description'],
                record['timestamp']
            )
        )
        rows += 1
    
    return f"Загружено {rows} записей"
 
with DAG(
    'weather_etl',
    default_args=default_args,
    description='Загрузка данных о погоде из API в базу данных',
    schedule_interval='0 */3 * * *',  # Каждые 3 часа
    start_date=datetime(2023, 1, 1),
    catchup=False,
    params={
        'api_key': 'your_api_key_here',
        'cities': ['Moscow', 'London', 'New York', 'Tokyo', 'Sydney']
    }
) as dag:
 
    # Создаем таблицу, если не существует
    create_table = PostgresOperator(
        task_id='create_weather_table',
        postgres_conn_id='postgres_default',
        sql="""
        CREATE TABLE IF NOT EXISTS weather_data (
            id SERIAL PRIMARY KEY,
            city VARCHAR(50),
            country VARCHAR(10),
            temperature FLOAT,
            humidity FLOAT,
            description TEXT,
            timestamp TIMESTAMP
        );
        """
    )
    
    # Задачи ETL
    fetch_weather = PythonOperator(
        task_id='fetch_weather',
        python_callable=fetch_weather_data,
    )
    
    transform = PythonOperator(
        task_id='transform',
        python_callable=transform_weather_data,
    )
    
    load = PythonOperator(
        task_id='load',
        python_callable=load_to_postgres,
    )
    
    # Порядок выполнения
    create_table >> fetch_weather >> transform >> load

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
 
def create_dag(dag_id, schedule, default_args, source_config):
    """Фабрика DAG для разных источников данных"""
    
    def process_source():
        print(f"Обрабатываем источник: {source_config['name']}")
        # Логика обработки
    
    with DAG(dag_id, schedule_interval=schedule, default_args=default_args) as dag:
        
        process_task = PythonOperator(
            task_id='process_source',
            python_callable=process_source,
        )
        
        return dag
 
# Конфигурации источников
sources = [
    {"name": "source_1", "id": "src_1", "schedule": "0 0 * * *"},
    {"name": "source_2", "id": "src_2", "schedule": "0 12 * * *"},
    {"name": "source_3", "id": "src_3", "schedule": "0 0 * * 0"} # Еженедельно
]
 
# Создаем DAG для каждого источника
for src in sources:
    dag_id = f"process_{src['id']}"
    
    default_args = {
        'owner': 'airflow',
        'start_date': datetime(2023, 1, 1),
        'retries': 2,
    }
    
    globals()[dag_id] = create_dag(
        dag_id=dag_id,
        schedule=src['schedule'],
        default_args=default_args,
        source_config=src
    )

def produce_value(**context):
    # Генерация случайного числа и сохранение в XCom
    import random
    value = random.randint(1, 100)
    context['task_instance'].xcom_push(key='random_value', value=value)
 
def consume_value(**context):
    # Получение значения из XCom
    value = context['task_instance'].xcom_pull(
        task_ids='producer_task', 
        key='random_value'
    )
    print(f"Получено значение: {value}")
    # Дальнейшая обработка...

def notify_on_failure(**context):
    # Отправка уведомления о сбое
    failed_task = context['task_instance']
    dag_id = failed_task.dag_id
    task_id = failed_task.task_id
    execution_date = failed_task.execution_date
    
    # Отправка сообщения в Slack, почту или другую систему
    message = f"Задача {task_id} в DAG {dag_id} завершилась с ошибкой {execution_date}"
    # код отправки...
 
# В настройках DAG
default_args = {
    # ...
    'on_failure_callback': notify_on_failure,
    # также доступны: on_success_callback, on_retry_callback
}

from airflow.sensors.base import BaseSensorOperator
from airflow.utils.decorators import apply_defaults
 
class CustomThresholdSensor(BaseSensorOperator):
    @apply_defaults
    def __init__(self, threshold_value, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.threshold_value = threshold_value
    
    def poke(self, context):
        # Проверка условия - возвращает True, если сенсор должен завершиться
        current_value = self._get_current_value()  # Реализация получения значения
        return current_value > self.threshold_value
 
# Использование в DAG
threshold_check = CustomThresholdSensor(
    task_id='wait_for_threshold',
    threshold_value=100,
    poke_interval=60,  # Проверять каждую минуту
    timeout=60*60*2,   # Максимальное время ожидания - 2 часа
    mode='reschedule'  # Освобождать слот между проверками
)

task = PythonOperator(
    task_id='high_priority_task',
    python_callable=some_function,
    pool='limited_pool',           # Группировка задач в пул
    pool_slots=2,                  # Потребление 2 слотов из пула 
    priority_weight=10             # Высокий приоритет (по умолчанию 1)
)

# Работа с S3
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.operators.s3_copy_object import S3CopyObjectOperator
 
# Запуск Spark-задачи
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
 
# Выполнение HTTP-запроса
from airflow.providers.http.operators.http import SimpleHttpOperator

from airflow.models.baseoperator import BaseOperator
from airflow.utils.decorators import apply_defaults
 
class CustomDBOperator(BaseOperator):
    @apply_defaults
    def __init__(self, conn_id, sql, [B]kwargs):
        super().__init__([/B]kwargs)
        self.conn_id = conn_id
        self.sql = sql
    
    def execute(self, context):
        # Вся магия происходит в методе execute
        from mycompany.db import DBConnection
        
        conn = DBConnection(self.conn_id)
        result = conn.execute(self.sql)
        
        # Возвращаемое значение автоматически отправляется в XCom
        return result

# Пример подхода к оптимизации при большом количестве DAG
[H2]Вместо множества отдельных файлов используем динамическую генерацию[/H2]
 
def create_dag(dag_id, config):
    dag = DAG(
        dag_id,
        schedule_interval=config['schedule'],
        default_args=default_args,
        catchup=False
    )
    
    # Динамическое создание задач
    for task_config in config['tasks']:
        # Создание задач на основе конфигурации
        pass
        
    return dag
 
# Загрузка конфигурации из единого источника (БД, JSON и т.д.)
configs = load_configs_from_source()
 
# Создание DAG для каждой конфигурации
for config in configs:
    globals()[f"dag_{config['id']}"] = create_dag(f"dag_{config['id']}", config)

# Пример задачи в Luigi
import luigi
 
class FetchDataTask(luigi.Task):
    date = luigi.DateParameter()
    
    def output(self):
        return luigi.LocalTarget(f"data_{self.date}.txt")
    
    def run(self):
        # Логика получения данных
        with self.output().open('w') as f:
            f.write("данные")
 
class ProcessDataTask(luigi.Task):
    date = luigi.DateParameter()
    
    def requires(self):
        return FetchDataTask(date=self.date)
    
    def output(self):
        return luigi.LocalTarget(f"processed_{self.date}.txt")
    
    def run(self):
        # Обработка данных из входной задачи
        with self.input().open() as in_file, self.output().open('w') as out_file:
            data = in_file.read()
            out_file.write(data.upper())