Конвейеры ETL с Apache Airflow и Python

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
 
default_args = {
    'owner': 'data_engineer',
    'depends_on_past': False,
    'start_date': datetime(2023, 1, 1),
    'email_on_failure': True,
    'email': ['alerts@company.com'],
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}
 
def extract_data():
    # Логика извлечения данных
    return {'data': [1, 2, 3, 4, 5]}
 
def transform_data(ti):
    # Получаем данные из предыдущей задачи
    data = ti.xcom_pull(task_ids='extract_task')
    transformed_data = [x * 2 for x in data['data']]
    return {'transformed_data': transformed_data}
 
def load_data(ti):
    # Загружаем трансформированные данные
    data = ti.xcom_pull(task_ids='transform_task')
    print(f"Loading data: {data['transformed_data']}")
    # Логика загрузки в целевую систему
    
with DAG('simple_etl_pipeline',
         default_args=default_args,
         schedule_interval='@daily') as dag:
    
    extract_task = PythonOperator(
        task_id='extract_task',
        python_callable=extract_data,
    )
    
    transform_task = PythonOperator(
        task_id='transform_task',
        python_callable=transform_data,
    )
    
    load_task = PythonOperator(
        task_id='load_task',
        python_callable=load_data,
    )
    
    # Определяем последовательность выполнения
    extract_task >> transform_task >> load_task

transform_task.set_upstream(extract_task)
load_task.set_upstream(transform_task)

extract_task >> [process_task_1, process_task_2] >> merge_task

def process_daily_data(ds, **kwargs):
    file_path = f"/data/{ds}/transactions.csv"
    # Дальнейшая логика обработки файла

from airflow.hooks.postgres_hook import PostgresHook
from airflow.hooks.S3_hook import S3Hook
 
def transfer_data():
    # Получаем данные из PostgreSQL
    pg_hook = PostgresHook(postgres_conn_id='postgres_default')
    data = pg_hook.get_records("SELECT * FROM users WHERE created_at = %s", parameters=[yesterday_ds])
    
    # Загружаем в S3
    s3_hook = S3Hook(aws_conn_id='aws_default')
    s3_hook.load_string(
        string_data=json.dumps(data),
        key=f"user_data/{ds}.json",
        bucket_name="my-data-lake"
    )

from airflow.sensors.external_task_sensor import ExternalTaskSensor
from airflow.sensors.filesystem import FileSensor
from airflow.sensors.sql import SqlSensor
 
# Ждём завершения другого DAG
wait_for_upstream = ExternalTaskSensor(
    task_id='wait_for_data_processing',
    external_dag_id='upstream_dag',
    external_task_id='final_task',
    timeout=3600,
    mode='reschedule',  # Важно для длительного ожидания
    dag=dag,
)
 
# Проверяем появление файла
wait_for_file = FileSensor(
    task_id='wait_for_report',
    filepath='/path/to/expected/file.csv',
    poke_interval=300,  # Проверяем каждые 5 минут
    dag=dag,
)
 
# Ждём появления данных в БД
wait_for_record = SqlSensor(
    task_id='check_data_available',
    conn_id='postgres_dwh',
    sql="SELECT COUNT(*) FROM events WHERE date = '{{ ds }}' AND status = 'processed'",
    success=lambda cnt: cnt[0][0] > 0,
    dag=dag,
)

def process_with_explicit_xcom(ti=None, **context):
    # Какая-то логика обработки
    result = {"status": "success", "count": 42}
    # Явно сохраняем только нужные данные
    ti.xcom_push(key="processing_status", value=result["status"])
    ti.xcom_push(key="item_count", value=result["count"])
    # Возвращаемое значение тоже будет сохранено в XCom
    # под ключем "return_value"
    return "Completed"
 
process_task = PythonOperator(
    task_id='process_data',
    python_callable=process_with_explicit_xcom,
    provide_context=True,
    dag=dag,
)

def save_large_data(**context):
    large_data = generate_massive_dataset()
    s3_hook = S3Hook(aws_conn_id='aws_default')
    key = f"temp/data_for_dag_run_{context['run_id']}.parquet"
    s3_hook.load_string(
        string_data=large_data.to_parquet(),
        key=key,
        bucket_name="my-data-bucket"
    )
    # Сохраняем только ссылку в XCom
    return {"data_location": f"s3://my-data-bucket/{key}"}
 
def process_large_data(**context):
    ti = context['ti']
    data_info = ti.xcom_pull(task_ids='save_large_data')
    s3_path = data_info["data_location"]
    # Загружаем данные из S3
    # ...

# Наша функция трансформации
def normalize_customer_data(row):
    return {
        'customer_id': str(row.get('id', '')).strip(),
        'full_name': f"{row.get('first_name', '')} {row.get('last_name', '')}".strip(),
        'email': row.get('email', '').lower(),
        'active': bool(row.get('status') == 'active')
    }
 
# Тест для функции трансформации
def test_normalize_customer_data():
    input_data = {
        'id': 123,
        'first_name': 'John ',
        'last_name': 'Doe',
        'email': 'JOHN.DOE@example.com',
        'status': 'active'
    }
    expected = {
        'customer_id': '123',
        'full_name': 'John Doe',
        'email': 'john.doe@example.com',
        'active': True
    }
    assert normalize_customer_data(input_data) == expected

from airflow.models import DagBag
 
def test_dag_loads():
    """Проверяем, что DAG загружается без ошыбок."""
    dagbag = DagBag()
    dag = dagbag.get_dag('my_etl_dag')
    assert dagbag.import_errors == {}
    assert dag is not None
 
def test_dag_structure():
    """Проверяем структуру DAG (задачи и их зависимости)."""
    dagbag = DagBag()
    dag = dagbag.get_dag('my_etl_dag')
    # Проверяем наличие ожидаемых задач
    task_ids = [t.task_id for t in dag.tasks]
    assert 'extract_data' in task_ids
    assert 'transform_data' in task_ids
    assert 'load_data' in task_ids
    # Проверяем зависимости
    extract_task = dag.get_task('extract_data')
    transform_task = dag.get_task('transform_data')
    assert transform_task.upstream_task_ids == {'extract_data'}

def validate_data_quality(**context):
    from great_expectations.data_context import DataContext
    
    # Инициализируем контекст Great Expectations
    data_context = DataContext('/path/to/great_expectations')
    
    # Запускаем набор проверок качества
    result = data_context.run_checkpoint(
        checkpoint_name="my_checkpoint",
        batch_kwargs={
            "path": context['ti'].xcom_pull(task_ids='extract_transform_task'),
            "datasource": "my_datasource"
        }
    )
    
    if not result["success"]:
        raise Exception("Проверка качества данных не пройдена!")
    
    return result
 
validate_task = PythonOperator(
    task_id='validate_data',
    python_callable=validate_data_quality,
    provide_context=True,
    dag=dag
)

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.postgres_operator import PostgresOperator
from airflow.hooks.postgres_hook import PostgresHook
from airflow.hooks.http_hook import HttpHook
import pandas as pd
import json
 
default_args = {
'owner': 'data_team',
'depends_on_past': False,
'start_date': datetime(2023, 1, 1),
'email_on_failure': True,
'email': ['data-alerts@example.com'],
'retries': 2,
'retry_delay': timedelta(minutes=5),
}
 
dag = DAG(
'sales_etl_pipeline',
default_args=default_args,
description='ETL для агрегации данных о продажах',
schedule_interval='0 3 * * *',  # Каждый день в 3:00
catchup=False,
max_active_runs=1
)
 
# Функции для извлечения данных из разных источников
def extract_web_sales(**context):
http_hook = HttpHook(http_conn_id='web_api', method='GET')
response = http_hook.run(f"/api/sales/daily/{context['ds']}")
sales_data = json.loads(response.content)
# Сохраняем во временную таблицу
pg_hook = PostgresHook(postgres_conn_id='dwh_connection')
df = pd.DataFrame(sales_data)
df.to_sql('temp_web_sales', pg_hook.get_sqlalchemy_engine(), if_exists='replace', index=False)
return len(sales_data)
 
def extract_mobile_sales([B]context):
# Похожая логика для мобильных продаж
pass
 
def extract_pos_sales([/B]context):
# Логика для POS-терминалов
pass
 
def transform_sales_data(**context):
pg_hook = PostgresHook(postgres_conn_id='dwh_connection')
# Объединяем, чистим и трансформируем данные
transform_sql = """
CREATE TABLE IF NOT EXISTS temp_transformed_sales AS
SELECT 
  coalesce(w.transaction_id, m.transaction_id, p.transaction_id) as transaction_id,
  coalesce(w.sale_date, m.sale_date, p.sale_date) as sale_date,
  coalesce(w.product_id, m.product_id, p.product_id) as product_id,
  coalesce(w.customer_id, m.customer_id, p.customer_id) as customer_id,
  coalesce(w.quantity, m.quantity, p.quantity) as quantity,
  coalesce(w.price, m.price, p.price) as price,
  CASE 
    WHEN w.transaction_id IS NOT NULL THEN 'web'
    WHEN m.transaction_id IS NOT NULL THEN 'mobile'
    ELSE 'pos'
  END as channel
FROM temp_web_sales w
FULL OUTER JOIN temp_mobile_sales m ON w.transaction_id = m.transaction_id
FULL OUTER JOIN temp_pos_sales p ON w.transaction_id = p.transaction_id OR m.transaction_id = p.transaction_id;
"""
pg_hook.run(transform_sql)
return "Данные трансформированы"
 
def load_data_to_dwh(**context):
pg_hook = PostgresHook(postgres_conn_id='dwh_connection')
load_sql = """
INSERT INTO sales_fact (transaction_id, sale_date, product_id, customer_id, quantity, price, channel)
SELECT transaction_id, sale_date, product_id, customer_id, quantity, price, channel
FROM temp_transformed_sales
ON CONFLICT (transaction_id) DO UPDATE 
SET quantity = EXCLUDED.quantity, price = EXCLUDED.price;
 
-- Обновляем агрегаты
INSERT INTO daily_sales_agg (sale_date, channel, total_sales, total_quantity)
SELECT sale_date, channel, SUM(price * quantity), SUM(quantity)
FROM temp_transformed_sales
GROUP BY sale_date, channel
ON CONFLICT (sale_date, channel) DO UPDATE 
SET total_sales = EXCLUDED.total_sales, total_quantity = EXCLUDED.total_quantity;
"""
pg_hook.run(load_sql)
return "Данные загружены в хранилище"
 
# Определяем задачи
start = DummyOperator(task_id='start', dag=dag)
end = DummyOperator(task_id='end', dag=dag)
 
extract_web = PythonOperator(
  task_id='extract_web_sales',
  python_callable=extract_web_sales,
  provide_context=True,
  dag=dag
)
 
extract_mobile = PythonOperator(
  task_id='extract_mobile_sales',
  python_callable=extract_mobile_sales,
  provide_context=True,
  dag=dag
)
 
extract_pos = PythonOperator(
  task_id='extract_pos_sales',
  python_callable=extract_pos_sales,
  provide_context=True,
  dag=dag
)
 
transform = PythonOperator(
  task_id='transform_sales_data',
  python_callable=transform_sales_data,
  provide_context=True,
  dag=dag
)
 
load = PythonOperator(
  task_id='load_data_to_dwh',
  python_callable=load_data_to_dwh,
  provide_context=True,
  dag=dag
)
 
cleanup = PostgresOperator(
  task_id='cleanup_temp_tables',
  postgres_conn_id='dwh_connection',
  sql="""
  DROP TABLE IF EXISTS temp_web_sales;
  DROP TABLE IF EXISTS temp_mobile_sales;
  DROP TABLE IF EXISTS temp_pos_sales;
  DROP TABLE IF EXISTS temp_transformed_sales;
  """,
  dag=dag
)
 
# Определяем последовательность
start >> [extract_web, extract_mobile, extract_pos] >> transform >> load >> cleanup >> end

def log_duration(func):
    def wrapper(*args, **kwargs):
        import time
        start_time = time.time()
        result = func(*args, [B]kwargs)
        duration = time.time() - start_time
        print(f"Задача {func.__name__} выполнялась {duration:.2f} секунд")
        return result
    return wrapper
 
@log_duration
def transform_sales_data([/B]context):
    # Логика трансформации

def transform_sales_data(**context):
    pg_hook = PostgresHook(postgres_conn_id='dwh_connection')
    
    # Разбиваем данные на партиции по первой цифре ID транзакции
    for partition in range(10):
        transform_sql = f"""
        CREATE TABLE IF NOT EXISTS temp_transformed_sales_{partition} AS
        SELECT 
            /* те же поля, что и раньше */
        FROM temp_web_sales w
        FULL OUTER JOIN temp_mobile_sales m 
            ON w.transaction_id = m.transaction_id 
            AND w.transaction_id % 10 = {partition}
        FULL OUTER JOIN temp_pos_sales p 
            ON (w.transaction_id = p.transaction_id OR m.transaction_id = p.transaction_id)
            AND p.transaction_id % 10 = {partition}
        WHERE 
            COALESCE(w.transaction_id, m.transaction_id, p.transaction_id) % 10 = {partition}
        """
        pg_hook.run(transform_sql)
    
    # Объединяем партиции
    pg_hook.run("""
    CREATE TABLE temp_transformed_sales AS
    SELECT * FROM temp_transformed_sales_0
    UNION ALL SELECT * FROM temp_transformed_sales_1
    /* и так далее для всех партиций */
    """)

def extract_web_sales(**context):
    # Получаем время последней успешной загрузки
    last_etl_time = get_last_successful_etl_time('web_sales')
    
    http_hook = HttpHook(http_conn_id='web_api', method='GET')
    # Запрашиваем только новые или изменённые записи
    response = http_hook.run(
        f"/api/sales/changes?since={last_etl_time}&until={context['execution_date']}"
    )
    sales_data = json.loads(response.content)
    
    if not sales_data:
        return "Нет новых данных"
    
    # Остальная логика загрузки
    # ...
    
    # Сохраняем время успешной загрузки
    save_successful_etl_time('web_sales', context['execution_date'])
    return f"Загружено {len(sales_data)} новых записей"

def get_last_successful_etl_time(source_name):
    pg_hook = PostgresHook(postgres_conn_id='dwh_connection')
    result = pg_hook.get_first(
        "SELECT last_successful_run FROM etl_watermarks WHERE source_name = %s",
        parameters=(source_name,)
    )
    return result[0] if result else datetime(2000, 1, 1).isoformat()
 
def save_successful_etl_time(source_name, execution_time):
    pg_hook = PostgresHook(postgres_conn_id='dwh_connection')
    pg_hook.run("""
    INSERT INTO etl_watermarks (source_name, last_successful_run, records_processed)
    VALUES (%s, %s, %s)
    ON CONFLICT (source_name) 
    DO UPDATE SET 
        last_successful_run = EXCLUDED.last_successful_run,
        records_processed = EXCLUDED.records_processed,
        updated_at = CURRENT_TIMESTAMP
    """, parameters=(source_name, execution_time, 0))

# dag_extract.py
extract_dag = DAG(
    'extract_data_sources',
    default_args=default_args,
    schedule_interval='0 1 * * *',  # Запуск в 1:00
    catchup=False
)
 
# Задачи извлечения...
 
# dag_transform.py
transform_dag = DAG(
    'transform_data',
    default_args=default_args,
    schedule_interval=None,  # Триггерится извне
    catchup=False
)
 
# Задачи трансформации...
 
# dag_load.py
load_dag = DAG(
    'load_data_warehouse',
    default_args=default_args,
    schedule_interval=None,  # Триггерится извне
    catchup=False
)
 
# Задачи загрузки...

from airflow.operators.trigger_dagrun import TriggerDagRunOperator
 
trigger_transform = TriggerDagRunOperator(
    task_id='trigger_transform',
    trigger_dag_id='transform_data',
    execution_date='{{ ds }}',
    reset_dag_run=True,
    wait_for_completion=True,  # Ждём завершения следующего DAG
    dag=extract_dag
)
 
trigger_load = TriggerDagRunOperator(
    task_id='trigger_load',
    trigger_dag_id='load_data_warehouse',
    execution_date='{{ ds }}',
    reset_dag_run=True,
    wait_for_completion=True,
    dag=transform_dag
)

from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
 
process_large_dataset = SparkSubmitOperator(
    task_id='process_large_dataset',
    application='/path/to/spark_script.py',
    conn_id='spark_default',
    application_args=['{{ ds }}', 'other_param'],
    conf={
        'spark.driver.memory': '4g',
        'spark.executor.memory': '2g',
        'spark.executor.cores': '2',
        'spark.dynamicAllocation.enabled': 'true',
        'spark.dynamicAllocation.maxExecutors': '20'
    },
    dag=dag
)

# spark_script.py
from pyspark.sql import SparkSession
import sys
 
# Получаем аргументы из Airflow
execution_date = sys.argv[1]
 
spark = SparkSession.builder \
    .appName(f"LargeDataProcessing_{execution_date}") \
    .getOrCreate()
 
# Читаем огромный набор данных
df = spark.read.parquet(f"s3://data-lake/raw/transactions/{execution_date}/")
 
# Выполняем сложные трансформации
result = df.groupBy("customer_id", "product_category") \
    .agg({"amount": "sum", "transaction_id": "count"}) \
    .withColumnRenamed("sum(amount)", "total_spent") \
    .withColumnRenamed("count(transaction_id)", "transaction_count")
 
# Сохраняем результат
result.write.mode("overwrite").parquet(f"s3://data-lake/processed/customer_stats/{execution_date}/")

# Задача для запуска Kafka-коннектора, который считывает CDC из базы
start_kafka_connector = HttpOperator(
    task_id='start_kafka_connector',
    http_conn_id='kafka_connect_api',
    endpoint='/connectors',
    method='POST',
    data=json.dumps({
        "name": "mysql-source-connector",
        "config": {
            "connector.class": "io.debezium.connector.mysql.MySqlConnector",
            "database.hostname": "mysql-server",
            "database.port": "3306",
            "database.user": "debezium",
            "database.password": "dbz",
            "database.server.id": "42",
            "database.server.name": "mysql-server-1",
            "table.whitelist": "inventory.customers",
            "database.history.kafka.bootstrap.servers": "kafka:9092",
            "database.history.kafka.topic": "schema-changes.inventory"
        }
    }),
    headers={"Content-Type": "application/json"},
    dag=dag
)
 
# Запускаем Spark Streaming джобу, которая обрабатывает поток данных
start_spark_streaming = SparkSubmitOperator(
    task_id='start_spark_streaming',
    application='/path/to/streaming_script.py',
    conn_id='spark_default',
    conf={'spark.jars.packages': 'org.apache.spark:spark-sql-kafka-0-10_2.12:3.0.0'},
    dag=dag
)

# streaming_script.py
from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import *
 
spark = SparkSession.builder \
    .appName("KafkaStreamProcessing") \
    .getOrCreate()
 
# Читаем поток из Kafka
df = spark \
    .readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "kafka:9092") \
    .option("subscribe", "mysql-server-1.inventory.customers") \
    .load()
 
# Парсим JSON из значения
schema = StructType([
    StructField("id", LongType()),
    StructField("first_name", StringType()),
    StructField("last_name", StringType()),
    StructField("email", StringType()),
    # Другие поля...
])
 
parsed = df.select(
    from_json(col("value").cast("string"), schema).alias("data")
).select("data.*")
 
# Применяем бизнес-логику
processed = parsed.withColumn(
    "full_name", 
    concat(col("first_name"), lit(" "), col("last_name"))
)
 
# Записываем результаты в хранилище
query = processed \
    .writeStream \
    .outputMode("append") \
    .format("parquet") \
    .option("path", "/data/processed/customers") \
    .option("checkpointLocation", "/tmp/checkpoints/customers") \
    .trigger(processingTime="1 minute") \
    .start()
 
query.awaitTermination()

with DAG('dynamic_tasks_dag', default_args=default_args, schedule_interval='@daily') as dag:
    # Список источников данных
    data_sources = {
        'users': {'table': 'users', 'primary_key': 'user_id'},
        'orders': {'table': 'orders', 'primary_key': 'order_id'},
        'products': {'table': 'products', 'primary_key': 'product_id'},
    }
    
    start = DummyOperator(task_id='start')
    end = DummyOperator(task_id='end')
    
    processing_tasks = {}
    
    # Динамически создаём задачи для каждого источника
    for source_name, config in data_sources.items():
        extract_task = PythonOperator(
            task_id=f'extract_{source_name}',
            python_callable=extract_data,
            op_kwargs={'table': config['table'], 'pk': config['primary_key']},
        )
        
        transform_task = PythonOperator(
            task_id=f'transform_{source_name}',
            python_callable=transform_data,
            op_kwargs={'source': source_name},
        )
        
        load_task = PythonOperator(
            task_id=f'load_{source_name}',
            python_callable=load_data,
            op_kwargs={'target_table': f'dwh_{source_name}'},
        )
        
        # Сохраняем задачи в словаре для последующего определения зависимостей
        processing_tasks[source_name] = {
            'extract': extract_task,
            'transform': transform_task,
            'load': load_task,
        }
        
        # Определяем последовательность для этого источника
        extract_task >> transform_task >> load_task
    
    # Связываем с общими начальной и конечной точками
    for source in processing_tasks.values():
        start >> source['extract']
        source['load'] >> end

def get_config_from_db():
    hook = PostgresHook(postgres_conn_id='airflow_db')
    query = "SELECT source_name, config FROM etl_sources WHERE active = true"
    sources = hook.get_records(query)
    return {source[0]: json.loads(source[1]) for source in sources}
 
data_sources = get_config_from_db()

def process_with_robust_error_handling(**context):
    try:
        # Основная логика
        result = complex_processing()
        return result
    except ConnectionError as e:
        # Сетевые ошыбки часто временные, стоит повторить
        context['ti'].xcom_push(key='retry_reason', value=str(e))
        raise Exception(f"Сетевая ошыбка: {e}")
    except ValidationError as e:
        # Проблемы с данными требуют вмешательства человека
        send_alert(f"Критическая ошибка в данных: {e}")
        context['ti'].xcom_push(key='validation_errors', value=str(e))
        raise AirflowSkipException("Задача пропущена из-за ошибок валидации")
    except Exception as e:
        # Неизвестные ошибки логируем детально
        logging.error(f"Неожиданная ошибка: {e}")
        logging.error(traceback.format_exc())
        raise

def task_failure_slack_alert(context):
    failed_alert = SlackWebhookOperator(
        task_id='slack_alert',
        http_conn_id='slack_webhook',
        webhook_token=slack_webhook_token,
        message=f"""
        :red_circle: Задача упала!
        *Даг*: {context.get('task_instance').dag_id}
        *Задача*: {context.get('task_instance').task_id}
        *Дата выполнения*: {context.get('execution_date')}
        *Ошибка*: {context.get('exception')}
        """,
        username='airflow-bot',
    )
    return failed_alert.execute(context=context)
 
default_args = {
    # Другие аргументы...
    'on_failure_callback': task_failure_slack_alert,
}

def predict_dag_issues():
    hook = PostgresHook(postgres_conn_id='airflow_db')
    # Анализируем статистику выполнения
    query = """
    SELECT 
        dag_id,
        AVG(end_date - start_date) as avg_duration,
        STDDEV(end_date - start_date) as stddev_duration,
        COUNT(*) FILTER (WHERE state = 'failed') / COUNT(*) as failure_rate
    FROM dag_run
    WHERE start_date > NOW() - INTERVAL '7 days'
    GROUP BY dag_id
    """
    stats = hook.get_pandas_df(query)
    
    # Выявляем аномалии
    for index, row in stats.iterrows():
        if row['failure_rate'] > 0.2:  # Больше 20% падений
            send_alert(f"Высокий процент сбоев у DAG {row['dag_id']}: {row['failure_rate']*100}%")
        
        if row['stddev_duration'] > row['avg_duration'] * 0.5:
            # Высокая вариативность времени выполнения - признак нестабильности
            send_alert(f"Нестабильное время выполнения у DAG {row['dag_id']}")

# Ограничение параллелизма для конкретного DAG
dag = DAG(
    'parallel_processing',
    default_args=default_args,
    concurrency=20,  # Не более 20 задач одновременно
    max_active_runs=3  # Не более 3 запусков DAG параллельно
)

def process_in_parallel(partition_keys):
    from concurrent.futures import ThreadPoolExecutor
    
    def process_partition(key):
        # Обработка конкретной партиции
        return f"Партиция {key} обработана"
    
    # Распараллеливаем внутри одной задачи
    with ThreadPoolExecutor(max_workers=10) as executor:
        results = list(executor.map(process_partition, partition_keys))
    
    return results

task = PythonOperator(
    task_id='resource_intensive_task',
    python_callable=process_function,
    pool='database_pool',  # Задачи из этого пула выполняются с ограниченным параллелизмом
    dag=dag
)