Yahoo finance API 사용해서 Airflow DAG 만들어보기 - TIL230607

📚 KDT WEEK 10 DAY 3 TIL

• Yahoo finance API Airflow DAG 코드 구현
  • Full Refresh 방식 구현
  • Incremental Update 방식 구현

 

---

 

Yahoo finance API를 사용하여 Airflow ETL 구현하기

 

먼저 Yahoo finance API를 사용하려면 yfinance 모듈을 다운로드받아야한다.

 

나는 Airflow를 도커 환경에서 사용하고 있으므로 도커에 설치를 해주었다.

 

docker ps # Airflow scheduler의 Container ID를 확인하자
docker exec -it SchedulerContainerID sh # 찾은 scheduler의 Container ID를 입력해서 접속한다.

# airflow 접속 후
pip install yfinance

 

🟥 Full Refresh(매번 전체를 복사해여 적재하는 방식) 구현

 

다음은 Airflow DAG에 넣을 Python 코드이다.

Extract와 Transform을 합쳐서 get_historical_prices 함수로 구현하고, Load는 load 함수로 구현하였다.

 

from airflow import DAG
from airflow.decorators import task
from airflow.providers.postgres.hooks.postgres import PostgresHook
from datetime import datetime
from pandas import Timestamp

import yfinance as yf
import pandas as pd
import logging

# Airflow에 등록된 Connection 정보를 통해 Redshift 연결
def get_Redshift_connection(autocommit=True):
    hook = PostgresHook(postgres_conn_id='redshift_dev_db')
    conn = hook.get_conn()
    conn.autocommit = autocommit
    return conn.cursor()


# Extract + Transform
@task
def get_historical_prices(symbol):
    ticket = yf.Ticker(symbol) # 심볼(특정 주식의 ID? 같은것)을 통해 yahoo finance 주식 정보 가져오기
    data = ticket.history() # 지난 30일간의 해당 주식 정보를 DataFrame 형태로 반환
    records = []

    for index, row in data.iterrows():
        date = index.strftime('%Y-%m-%d %H:%M:%S')

		# 차례로 날짜, 시작주가, 상한가, 하한가, 종료주가, 매도매수건
        records.append([date, row["Open"], row["High"], row["Low"], row["Close"], row["Volume"]])

    return records

# Load
@task
def load(schema, table, records):
    logging.info("load started")
    cur = get_Redshift_connection()
    try:
        cur.execute("BEGIN;")
        cur.execute(f"DROP TABLE IF EXISTS {schema}.{table};")
        cur.execute(f"""
CREATE TABLE {schema}.{table} (
    date date,
    "open" float,
    high float,
    low float,
    close float,
    volume bigint
);""")

        for r in records:
            sql = f"INSERT INTO {schema}.{table} VALUES ('{r[0]}', {r[1]}, {r[2]}, {r[3]}, {r[4]}, {r[5]});"
            print(sql)
            cur.execute(sql)
        cur.execute("COMMIT;")   # cur.execute("END;")
    except Exception as error:
        print(error)
        cur.execute("ROLLBACK;")
        raise

    logging.info("load done")


with DAG(
    dag_id = 'UpdateSymbol',
    start_date = datetime(2023,5,30),
    catchup=False,
    tags=['API'],
    schedule = '0 10 * * *' # 매일 10시 정각마다
) as dag:

    results = get_historical_prices("AAPL") # 애플 주식의 심볼 AAPL
    load("sunhokim_public", "stock_info", results) # 순서대로 스키마이름, 테이블이름, transform 함수의 결과

 

Airflow의 Connection 정보를 연결하는 방법은 바로 이전 글을 참고하면 알 수 있다.

 

Airflow Decorator인 @task를 사용했기 때문에 간단하게 구현이 된 모습이다.

Decorator을 사용하면 함수에 인자도 원하는 이름으로 지어 넣을 수 있고, 가독성도 높아진다는 점이 큰 장점인 것 같다.

 

이 코드를 airflow폴더 > dags 폴더에 넣으면 Airflow 웹 UI에 나타날 것이다. (반영에 5분정도 걸림)

 

이제 정상적으로 구동이 되는지 DAG를 실행시켜 보자.

웹 UI(로컬포트)로도 가능하지만 나는 CLI로 실행시켜보았다.

 

이제 아까 yfinance를 설치하기 위해 접속했던 airflow 환경에 다시 들어가서,

다음의 명령어를 입력한다.

cd dags # dags 폴더로 이동
ls # 안에 있는 파일 확인
airflow dags test UpdateSymbol 2023-05-30 # UpdateSymbol이 아까 만들었던 코드파일의 이름이다.

 

정상적으로 구동되었다면, 다음과 같은 화면으로 마무리될 것이다!

INSERT가 정상적으로 동작하였고, 마무리되었다.

 

---

 

🟦 Incremental Update(새로운 데이터만 계속 추가) 구현

 

Full Refresh 방식의 경우 항상 한 달치의 데이터만 테이블에 저장되어 있을 것이다. (소스 데이터가 그렇다) 

이번에는 새로운 데이터가 생기면 본래 테이블에 추가해주는 방식을 구현해보자.

 

아까 Full Refresh 코드에서 _create_table 함수를 추가하고 load 함수를 수정하였다.

# drop_first는 테이블이 이미 존재할 경우 True를 인자로 넣는다
def _create_table(cur, schema, table, drop_first):
    if drop_first:
        cur.execute(f"DROP TABLE IF EXISTS {schema}.{table};")
    cur.execute(f"""
CREATE TABLE IF NOT EXISTS {schema}.{table} (
    date date,
    "open" float,
    high float,
    low float,
    close float,
    volume bigint
);""")


@task
def load(schema, table, records):
    logging.info("load started")
    cur = get_Redshift_connection()
    try:
        cur.execute("BEGIN;")
        # 원본 테이블이 없으면 생성 - 테이블이 처음 한번 만들어질 때 필요한 코드
        _create_table(cur, schema, table, False)
        # 임시 테이블로 원본 테이블을 복사
        cur.execute(f"CREATE TEMP TABLE t AS SELECT * FROM {schema}.{table};")
        for r in records:
            sql = f"INSERT INTO t VALUES ('{r[0]}', {r[1]}, {r[2]}, {r[3]}, {r[4]}, {r[5]});"
            print(sql)
            cur.execute(sql)

        # 원본 테이블 생성
        _create_table(cur, schema, table, True)
        # 임시 테이블 내용을 원본 테이블로 복사
        cur.execute(f"INSERT INTO {schema}.{table} SELECT DISTINCT * FROM t;")
        cur.execute("COMMIT;")   # cur.execute("END;")
    except Exception as error:
        print(error)
        cur.execute("ROLLBACK;") 
        raise
    logging.info("load done")

load 함수의 작동 방식은 다음과 같이 이루어진다.

 

원본 테이블 없으면 생성 > 임시 테이블로 내용 복사 > 임시 테이블에 소스 데이터 추가

> 원본 테이블 삭제 후 재생성 > 원본 테이블에 중복없이 데이터 추가

 

임시 테이블에 원본 테이블의 내용이 들어간 후에 소스 데이터도 추가되었기 때문에 중복 데이터가 발생할 것이다.(첫 시도가 아니라면)

이때 원본 테이블에 임시 테이블의 내용을 DISTINCT를 사용하여 중복없이 데이터를 삽입하면 새로운 데이터가 추가된 형태로 남는다.