[Week10 Airflow] TIL 44일차 Airflow DAG, task 실습과 고도화

오늘은 앞서 작성했던 ETL코드를 airflow로 관리할 수 있도록 변경해 보자.

 

해당 실습에서는 AWS redshift를 사용한다.

AWS S3의 csv 파일을 extract, transform, load 함수를 이용해 airflow에 적재하는 데이터 파이프라인을 만든다.

Full refresh 방식을 사용한다.

Ⅰ. 실습1 : csv 파일 이용한 데이터 파이프라인 만들기

1. 초기 코드

AWS S3의 파일을 load, transform, extract해서 AWS Redshift에 적재

import psycopg2

# 1. Redshift 연결
def get_Redshift_connection():
    host = "redshift주소"
    redshift_user = "사용자이름"
    redshift_pass = "비밀번호"
    port = 5439
    dbname = "dev"
    conn = psycopg2.connect("dbname={dbname} user={user} host={host} password={password} port={port}".format(
        dbname=dbname,
        user=redshift_user,
        password=redshift_pass,
        host=host,
        port=port
    ))
    conn.set_session(autocommit=True)
    return conn.cursor()
    
# 2. ETL 함수
import requests
def extract(url):
    f = requests.get(url)
    return (f.text)
def transform(text):
    lines = text.strip().split("\n")[1:] #첫번째 라인(헤더)는 제외하고 처리
    records = []
    for l in lines:
      (name, gender) = l.split(",") #l = "이름,성별" -> [ '이름', '성별' ]
      records.append([name, gender])
    return records
def load(records):
    # BEGIN과 END를 사용해서 SQL 결과를 트랜잭션으로 만들어주는 것이 좋음
    schema="스키마이름"
    cur = get_Redshift_connection()
    # SQL 트랜잭션으로 수행
    try:
      cur.execute("BEGIN")
      # DELETE FROM을 먼저 수행 -> FULL REFRESH을 하는 형태
      cur.execute(f"DELETE FROM {schema}.name_gender;")
      for r in records:
          name = r[0]
          gender = r[1]
          print(name, "-", gender)
          sql = f"INSERT INTO {schema}.name_gender VALUES ('{name}', '{gender}')"
          cur.execute(sql)
      cur.execute("COMMIT;")
    except(Exception, psycopg2.DatabaseError) as error:
      print(error)
      cur.execute("ROLLBACK;")

 

실행

link = "AWS s3의 csv 주소"
data = extract(link)
lines = transform(data)
load(lines)

 

2. 개선1-중요한 정보는 노출되지 않게 따로 저장한다.

: airflow-admin-connections 이용

 

airflow-admin-connections에 저장할 정보는 다음과 같다.

Connection-id : 아이디 지정

Connection-type : Amazon Redshift

Database : schema이름

Port : 5439

from airflow.providers.postgres.hooks.postgres import PostgresHook

def get_Redshift_connection(autocommit=True):
    hook = PostgresHook(postgres_conn_id='redshift_dev_db')
    conn = hook.get_conn()
    conn.autocommit = autocommit
    return conn.cursor()

 

2. 개선2-각각을 하나의 task로 만들어(총 3개의 task) DAG에 지정한다.

보통 하나의 task는 하나의 operator을 사용한다.

3가지 함수를 하나의 task로 만들때는 함수의 결과를 다음 함수의 입력으로 직접 지정할 수가 있다.

하지만, 각각의 함수를 하나의 task로 만들때에는Xcom기능을 사용하여 한 함수의 출력을 다음 함수의 입력으로 지정한다.

 

Xcom이란?

각 task들이 서로 데이터를 주고받는 방식으로,  보통 Operator의 출력값을 다른 Operator가 읽어가는 형태이다.

Airflow의 경우 task의 return값을 id와 함께 저장하기 때문에 가능한 방법이다.

다만, 이 값은 메타데이터 DB에 저장이 되기 때문에 큰 데이터를 주고받는데 사용하기는 어렵다.(따라서 보통 큰 데이터라면 S3에 로드한 뒤, 해당 위치를 주고받는 형태로 데이터를 주고받는다.)

 

context["task_instance"].xcom_pull(key="return_value", task_ids="함수이름")

 

3. 개선3-operator의 params를 이용하여 매개변수로 정보를 전달하자.

airflow-admin-variables 이용 (**context)

 

보통 open_api_key, csv_url 등의 정보를 저장한다.

from airflow.models import Variable

• get : key를 이용해 airflow에서 값을 가져옴
• set : key, value를 airflow에 저장

def extract(**context):
    link = context["params"]["url"]
    task_instance = context['task_instance']
    execution_date = context['execution_date']

    logging.info(execution_date)
    f = requests.get(link)
    return (f.text)


def transform(**context):
    logging.info("Transform started")    
    text = context["task_instance"].xcom_pull(key="return_value", task_ids="extract")
    lines = text.strip().split("\n")[1:] # 첫 번째 라인을 제외하고 처리
    records = []
    for l in lines:
      (name, gender) = l.split(",") # l = "이름, 성별" -> [ '이름', '성별' ]
      records.append([name, gender])
    logging.info("Transform ended")
    return records


def load(**context):
    logging.info("load started")    
    schema = context["params"]["schema"]
    table = context["params"]["table"]

    lines = context["task_instance"].xcom_pull(key="return_value", task_ids="transform")

    # BEGIN과 END를 사용해서 SQL 결과를 트랜잭션으로 만들어주는 것이 좋음
    cur = get_Redshift_connection()
    try:
        cur.execute("BEGIN;")
        # DELETE FROM을 먼저 수행 -> FULL REFRESH을 하는 형태
        cur.execute(f"DELETE FROM {schema}.name_gender;") 
        for r in records:
            name = r[0]
            gender = r[1]
            print(name, "-", gender)
            sql = f"INSERT INTO {schema}.name_gender VALUES ('{name}', '{gender}')"
            cur.execute(sql)
        cur.execute("COMMIT;")   # cur.execute("END;") 
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
        cur.execute("ROLLBACK;")   
    logging.info("load done")

 

 

dag = DAG(
    dag_id = '아이디',
    start_date = datetime(2023,4,6), # 날짜가 미래인 경우 실행이 안됨
    schedule = '0 2 * * *',  # 적당히 조절
    catchup = False,
    max_active_runs = 1,
    default_args = {
        'retries': 1,
        'retry_delay': timedelta(minutes=3),
    }
)
extract = PythonOperator(
    task_id = 'extract',
    python_callable = extract,
    params = {
        'url':  Variable.get("csv_url")
    },
    dag = dag)

transform = PythonOperator(
    task_id = 'transform',
    python_callable = transform,
    params = { 
    },  
    dag = dag)

load = PythonOperator(
    task_id = 'load',
    python_callable = load,
    params = {
        'schema': '스키마이름',
        'table': '테이블이름'
    },
    dag = dag)

extract >> transform >> load

 

위의 개선사항들에 test decorator를 사용해 보자.

해당 데코레이터를 사용하면, 가독성 있게 코드를 만들 수 있고 Xcom을 사용하여 데이터를 주고받지 않아도 된다.

 

4. 개선4 - 가독성 있는 코드를 만들어 보자.

task decorator를 사용

def get_Redshift_connection(autocommit=True):
	#1. AIRFLOW connection 활용
    hook = PostgresHook(postgres_conn_id='redshift_dev_db')
    conn = hook.get_conn()
    conn.autocommit = autocommit
    return conn.cursor()

#2. Decorator를 이용해 각각의 함수를 task로 지정 + 파라미터 대신 매개변수 이용 
@task
def extract(url):
    logging.info(datetime.utcnow())
    f = requests.get(url)
    return f.text


@task
def transform(text):
    lines = text.strip().split("\n")[1:] # 첫 번째 라인을 제외하고 처리
    records = []
    for l in lines:
      (name, gender) = l.split(",") # l = "Keeyong,M" -> [ 'keeyong', 'M' ]
      records.append([name, gender])
    logging.info("Transform ended")
    return records


@task
def load(schema, table, records):
    logging.info("load started")    
    cur = get_Redshift_connection()   
    # BEGIN과 END를 사용해서 SQL 결과를 트랜잭션으로 만들어주는 것이 좋음
    try:
        cur.execute("BEGIN;")
        # DELETE FROM을 먼저 수행 -> FULL REFRESH을 하는 형태
        cur.execute(f"DELETE FROM {schema}.name_gender;") 
        for r in records:
            name = r[0]
            gender = r[1]
            print(name, "-", gender)
            sql = f"INSERT INTO {schema}.name_gender VALUES ('{name}', '{gender}')"
            cur.execute(sql)
        cur.execute("COMMIT;")   # cur.execute("END;") 
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
        cur.execute("ROLLBACK;")   
    logging.info("load done")


with DAG(
    dag_id='namegender_v5',
    start_date=datetime(2022, 10, 6),  # 날짜가 미래인 경우 실행이 안됨
    schedule='0 2 * * *',  # 적당히 조절
    max_active_runs=1,
    catchup=False,
    default_args={
        'retries': 1,
        'retry_delay': timedelta(minutes=3),
        # 'on_failure_callback': slack.on_failure_callback,
    }
) as dag:
	#3. Airflow variable 이용
    url = Variable.get("csv_url")
    schema = '스키마명'   ## 자신의 스키마로 변경
    table = 'name_gender' ## 테이블명
	#2. 데코데이터를 사용하여 이렇게 함수형태로 전달O
    lines = transform(extract(url))
    load(schema, table, lines)

 

 

5. 위의 코드가 저장된 깃허브에서 파일들을 복사해서 실습

(AWS-airflow, GCP-Docker-airflow 두 경우 모두 실습)

 

1. 깃허브에서 코드 복사(클론+모든 폴더/파일 복사)

git clone 레포지토리 주소

cp -r 레포지토리파일(폴더)경로/* 복사할 폴더 경로(보통 airflow계정의 dags 폴더 안에)

 

2. dag 기본 저장 위치 확인 airflow.cfg

aws의 경우

dags_folder = /var/lib/airflow/dags

 

docker의 경우

dags_folder = /docker/airflow-setup/dags

 

해당 위치에 dag 들을 위치해야 Airflow가 자동으로 dag로 인식을 하게 된다.

 

3-1. Web으로 Airflow 접속

airflow의 DAG는 웹에서 조작할 수도 있고, shell에서 조작할 수도 있다. 

aws의 경우

ec2주소:8080

 

docker의 경우

VM외부ip주소:8080

 

4. Unpause DAG 선택 후 정상적으로 작동되는지 확인

만약, 특정 DAG 상태가 failed이면 → failed인 task에 들어가서 log 확인

 

3-2. shell에서 확인

EC2-airflow 이용하는 경우

(공통 명령어를 이용한다)

 

GCP-Docker-airflow 이용하는 경우

- docker 컨테이너 정보 확인

docker -ps

- airflow-scheduler 컨테이너 접속

docker exec -it airflow-scheduler-컨테이너id sh

(공통 명령어를 이용한다)

 

(공통 명령어)

- 모든 dags 목록 확인

airflow dags list

- 특정 dag의 task 목록 확인

airflow tasks list dagid이름

- variables 목록 확인

airflow variables list

- 특정 dag 실행

airflow task [text | run] DAG이름 task이름 날짜(YYYY-MM-DD)

• test : task의 실행이 메타데이터 DB에 기록 O
• run : task의 실행이 메타데이터 DB에 기록 O

 

 

Ⅱ. 실습2 : api 파일 이용한 데이터 파이프라인 만들기

Yahoo finance api를 활용해 주식의 정보를 airflow를 활용해 수집해오자.

yfinance api는 지난 30일의 정보를 받아올 수 있다(특정 날짜에 대한 정보는 어렵다) 

 

구현해야 하는 기능들은 다음과 같다.

1. Task1 : Extract, Transform

2. Task2-1 : Load(Full refresh) UpdateSymbol.py

• delete * from table : 모든 레코드 삭제(트랜잭션)
• drop table : 테이블을 통째로 삭제(이 방법 사용)

3. Task2-2 : Load(Incremental Update) UpdateSymbol_v2.py

1) 임시 테이블 생성(CTAS) : 기존 테이블 복사

cur.execute(f"CREATE TEMP TABLE t AS SELECT * FROM {schema}.{table};")

 

2) 임시 테이블에 finanace api 레코드 적재

for r in records:
    sql = f"INSERT INTO t VALUES ('{r[0]}', {r[1]}, {r[2]}, {r[3]}, {r[4]}, {r[5]});"
    cur.execute(sql)

 

3) 원본 테이블 삭제하고 새로 생성

_create_table(cur, schema, table, True)

 

4) 원본 테이블에 임시 테이블 복사(단 CTASelect distinc 사용하여 중복 제거)

cur.execute(f"INSERT INTO {schema}.{table} SELECT DISTINCT * FROM t;")
cur.execute("COMMIT;")   # cur.execute("END;")

 

 

Task 3. DAG 실행

1. airflow web UI 에서 실행

2. shell에서 실행

airflow dats test dagid이름 2023-05-30

 

 

Ⅱ. SQL 명령어에 대한 견해

1. FULL Refresh일때 DDL과 DML의 사용

데이터가 많은 경우 delete ~ from table, truncate ~ from table, drop table 중 어느 명령어를 사용할지 고민이 된다.

 

여기에서는 delete와 drop 명령어의 사용에 대해 비교해 보자.

 

DELETE를 사용하는 경우:

• 특정 조건에 맞는 행을 삭제할 때
• 특정 행을 삭제한 후에도 테이블의 구조를 유지하고 싶을 때
• 롤백이 필요할 때(트랜잭션을 사용할 수 있음).

DROP을 사용하는 경우:

• 테이블 전체를 삭제하고 해당 테이블의 구조까지 완전히 제거할 때
• 테이블을 완전히 재생성해야 할 때

DROP 명령어를 이용하면 테이블 전체가 삭제되기 때문에 보통 더 빠르다.

하지만, 데이터가 많은 경우 복구해야할 필요성이 있을 수도 있으므로 데이터의 쓰임에 대해 고민한 뒤, 명령어를 고르자! 

 

 

Ⅶ. 숙제2) 데이터 파이프라인 코드 작성하기