Как работать Air Flow
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
import pandas as pd
import psycopg2
import os
# ============================================================
# НАСТРОЙКИ
# ============================================================
CSV_PATH = '/home/user/airflow/data/traffic_full_30min.csv' # ПУТЬ К ВАШЕМУ CSV
DB_CONFIG = {
'host': 'localhost',
'database': 'dbname',
'user': 'username',
'password': 'password'
}
# ============================================================
# ЗАДАЧА 1: СОЗДАНИЕ ТАБЛИЦ
# ============================================================
def create_tables():
conn = psycopg2.connect(**DB_CONFIG)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS raw_traffic (
timestamp TIMESTAMP PRIMARY KEY,
intensity INTEGER,
cars INTEGER,
trucks INTEGER,
busses INTEGER,
avg_speed NUMERIC,
weather_type TEXT,
load_time TIMESTAMP
);
CREATE TABLE IF NOT EXISTS clean_traffic (
timestamp TIMESTAMP PRIMARY KEY,
intensity INTEGER,
cars INTEGER,
trucks INTEGER,
busses INTEGER,
total_vehicles INTEGER,
quality_flag BOOLEAN,
update_time TIMESTAMP
);
CREATE TABLE IF NOT EXISTS etl_log (
id SERIAL PRIMARY KEY,
run_time TIMESTAMP,
rows_processed INTEGER,
status TEXT,
message TEXT
);
""")
conn.commit()
cursor.close()
conn.close()
print("✅ Таблицы созданы")
# ============================================================
# ЗАДАЧА 2: ЗАГРУЗКА CSV (EXTRACT)
# ============================================================
def extract_data():
df = pd.read_csv(CSV_PATH, parse_dates=['timestamp'])
print(f"📥 Загружено {len(df)} строк из CSV")
# Сохраняем в raw таблицу
conn = psycopg2.connect(**DB_CONFIG)
cursor = conn.cursor()
for _, row in df.iterrows():
cursor.execute("""
INSERT INTO raw_traffic (timestamp, intensity, cars, trucks, busses, avg_speed, weather_type, load_time)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (timestamp) DO NOTHING
""", (row['timestamp'], row['intensity_30min'], row['cars'],
row['trucks'], row['busses'], row['avg_speed'],
row['weather_type'], datetime.now()))
conn.commit()
cursor.close()
conn.close()
# Возвращаем df для следующих задач
return df.to_dict('records')
# ============================================================
# ЗАДАЧА 3: ОЧИСТКА И ТРАНСФОРМАЦИЯ
# ============================================================
def transform_data(**context):
# Получаем df из предыдущей задачи
ti = context['ti']
data = ti.xcom_pull(task_ids='extract_data')
df = pd.DataFrame(data)
issues = []
# Удаляем дубликаты
before = len(df)
df = df.drop_duplicates(subset=['timestamp'])
if len(df) < before:
issues.append(f"Удалено дубликатов: {before - len(df)}")
# Проверяем скорость
invalid = df[(df['avg_speed'] < 0) | (df['avg_speed'] > 150)].shape[0]
if invalid > 0:
issues.append(f"Некорректная скорость: {invalid}")
df['avg_speed'] = df['avg_speed'].clip(0, 150)
# Обогащение
df['total_vehicles'] = df['cars'] + df['trucks'] + df['busses']
df['quality_flag'] = len(issues) == 0
df['update_time'] = datetime.now()
# Логируем
conn = psycopg2.connect(**DB_CONFIG)
cursor = conn.cursor()
cursor.execute("""
INSERT INTO etl_log (run_time, rows_processed, status, message)
VALUES (%s, %s, %s, %s)
""", (datetime.now(), len(df), 'SUCCESS' if len(issues)==0 else 'WARNING',
'; '.join(issues) if issues else 'OK'))
conn.commit()
cursor.close()
conn.close()
# Возвращаем очищенный df
return df.to_dict('records')
# 5. НАСТРОЙКА DAG
# ============================================================
default_args = {
'owner': 'analyst',
'depends_on_past': False,
'start_date': datetime(2026, 4, 1),
'retries': 1,
'retry_delay': timedelta(minutes=5)
}
dag = DAG(
'traffic_etl_pipeline',
default_args=default_args,
description='ETL для транспортных данных',
schedule_interval='*/30 * * * *', # каждые 30 минут
catchup=False
)
# Определяем задачи
t1 = PythonOperator(task_id='create_tables', python_callable=create_tables, dag=dag)
t2 = PythonOperator(task_id='extract_data', python_callable=extract_data, dag=dag)
t3 = PythonOperator(task_id='transform_data', python_callable=transform_data, dag=dag, provide_context=True)
t4 = PythonOperator(task_id='load_data', python_callable=load_data, dag=dag, provide_context=True)
# Порядок выполнения
t1 >> t2 >> t3 >> t4
# ============================================================
# ЗАДАЧА 4: ЗАГРУЗКА В CLEAN_TRAFFIC (LOAD)
# ============================================================
def load_data(**context):
ti = context['ti']
data = ti.xcom_pull(task_ids='transform_data')
df = pd.DataFrame(data)
conn = psycopg2.connect(**DB_CONFIG)
cursor = conn.cursor()
for _, row in df.iterrows():
cursor.execute("""
INSERT INTO clean_traffic (timestamp, intensity, cars, trucks, busses, total_vehicles, quality_flag, update_time)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (timestamp) DO UPDATE SET
intensity = EXCLUDED.intensity,
cars = EXCLUDED.cars,
trucks = EXCLUDED.trucks,
busses = EXCLUDED.busses,
total_vehicles = EXCLUDED.total_vehicles,
quality_flag = EXCLUDED.quality_flag,
update_time = EXCLUDED.update_time
""", (row['timestamp'], row['intensity_30min'], row['cars'],
row['trucks'], row['busses'], row['total_vehicles'],
row['quality_flag'], row['update_time']))
conn.commit()
cursor.close()
conn.close()
print(f"✅ Загружено {len(df)} записей в clean_traffic")