[Airflow] 코인 데이터 배치 활용하기

오늘은 airflow를 통해 5분간격으로 데이터를 넣어보는 작업을 진행했다.

일단 upbit의 코인데이터를 5분 간격으로 데이터를 적재했다.

당연히 중복데이터를 제거하고 새로운 데이터만 받아왔다!

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.operators.dummy_operator import DummyOperator
from airflow.hooks.base_hook import BaseHook
from datetime import datetime, timedelta
import pandas as pd
import pyupbit

# SQLAlchemy 관련 라이브러리
from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, UniqueConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# 데이터베이스 연결 설정
conn = BaseHook.get_connection('postgres_test')
DATABASE_URL = f"postgresql+psycopg2://{conn.login}:{conn.password}@{conn.host}:{conn.port}/{conn.schema}"
engine = create_engine(DATABASE_URL, echo=True)
Session = sessionmaker(bind=engine)

일단 라이브러리 다 임포트 해줬다. 저번에 airflow 웹에서 커넥트 설정했었는데,

sqlalchemy쓰면 굳이 operator 안 써도 된다고 하여, postgreoperator은 사용하지 않았다.

# 기본 베이스 클래스 생성
Base = declarative_base()

# 테이블 정의
class OHLCV(Base):
    __tablename__ = 'raw_data'
    date = Column(DateTime, primary_key=True)
    ticker = Column(String, primary_key=True)
    open = Column(Float)
    high = Column(Float)
    low = Column(Float)
    close = Column(Float)
    volume = Column(Float)
    __table_args__ = (
        UniqueConstraint('date', 'ticker', name='_date_ticker_uc'),
    )

# 데이터 수집 함수
def fetch_data(**kwargs):
    # 코인 목록 가져오기
    tickers = pyupbit.get_tickers('KRW')
    df_list = []

    # 각 코인에 대해 OHLCV 데이터 수집
    for ticker in tickers:
        df = pyupbit.get_ohlcv(ticker, interval='minute5', count=10)
        if df is not None:
            df['ticker'] = ticker  # 코인 심볼 추가
            df_list.append(df)

    # 데이터프레임 합치기
    if df_list:
        merged_df = pd.concat(df_list)
        merged_df.reset_index(inplace=True)
        merged_df = merged_df.rename(columns={'index': 'date'})
        
        # 'date' 열을 문자열로 변환하여 JSON 직렬화 가능하게 하기
        merged_df['date'] = merged_df['date'].apply(lambda x: x.strftime('%Y-%m-%d %H:%M:%S'))

        # 데이터프레임을 XCom에 저장
        kwargs['ti'].xcom_push(key='ohlcv_data', value=merged_df.to_dict('records'))

# 데이터베이스 삽입 함수
def insert_data(**kwargs):
    # 데이터베이스 세션 시작
    session = Session()

    try:
        # XCom에서 데이터 가져오기
        merged_df = kwargs['ti'].xcom_pull(key='ohlcv_data', task_ids='fetch_data')

        # 데이터베이스에 데이터 삽입 또는 업데이트
        for row in merged_df:
            # 데이터베이스에 이미 있는지 확인
            existing_data = session.query(OHLCV).filter_by(date=row['date'], ticker=row['ticker']).first()
            if existing_data:
                # 데이터가 존재하면 업데이트
                existing_data.open = row['open']
                existing_data.high = row['high']
                existing_data.low = row['low']
                existing_data.close = row['close']
                existing_data.volume = row['volume']
            else:
                # 데이터가 없으면 새로 삽입
                ohlcv_data = OHLCV(
                    date=row['date'],
                    open=row['open'],
                    high=row['high'],
                    low=row['low'],
                    close=row['close'],
                    volume=row['volume'],
                    ticker=row['ticker']
                )
                session.add(ohlcv_data)

        # 변경 내용을 커밋
        session.commit()

        # 삽입된 데이터프레임을 XCom에 저장
        # df = pd.DataFrame(merged_df, columns=['date', 'open', 'high', 'low', 'close', 'volume', 'ticker'])
        #kwargs['ti'].xcom_push(key='inserted_data', value=df)

    except Exception as e:
        session.rollback()
        raise e
    finally:
        session.close()

# Airflow DAG 설정
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2023, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

dag = DAG(
    'fetch_ohlcv_data_with_postgres_operator',
    default_args=default_args,
    description='Fetch OHLCV data every 5 minutes and insert into PostgreSQL',
    schedule_interval=timedelta(minutes=5),
    catchup=False,
)

start_task = DummyOperator(
    task_id='start',
    dag=dag
)

fetch_data_task = PythonOperator(
    task_id='fetch_data',
    python_callable=fetch_data,
    provide_context=True,
    dag=dag,
)

insert_data_task = PythonOperator(
    task_id='insert_data',
    python_callable=insert_data,
    provide_context=True,
    dag=dag,
)

end_task = DummyOperator(
    task_id='end',
    dag=dag
)

start_task >> fetch_data_task >> insert_data_task >> end_task

날짜 + 종목으로 기본키를 설정해주었다. 5분 단위로 데이터를 받아 왔고, 날짜+종목이 없는 종목만 받아왔다.