세상의 모든 데이터

# 필요한 모듈과 클래스 임포트
from airflow import DAG
from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowException
from airflow.hooks.postgres_hook import PostgresHook
from airflow.operators.python_operator import PythonOperator
import logging
from glob import glob

# DAG에서 사용할 함수: 주어진 디렉토리의 JSON 파일을 읽어들이는 함수
def load_all_jsons_into_list(path_to_json):
    # configs 리스트 초기화
    configs = []
    
    # 주어진 디렉토리에 있는 모든 .py 파일에 대해 반복
    for f_name in glob(path_to_json + '/*.py'):
        # 파일 열기
        with open(f_name) as f:
            # 파일 내용을 텍스트로 읽어들임
            dict_text = f.read()
            
            try:
                # 읽어들인 텍스트를 eval 함수를 사용하여 파이썬 딕셔너리로 변환
                dict = eval(dict_text)
            except Exception as e:
                # 변환 중 에러가 발생하면 로그에 에러 메시지 기록하고 예외 전파
                logging.info(str(e))
                raise
            else:
                # 성공적으로 변환된 딕셔너리를 configs 리스트에 추가
                configs.append(dict)
    
    # 최종적으로 변환된 딕셔너리들을 담고 있는 configs 리스트 반환
    return configs


# 주어진 테이블 이름과 테이블 설정에서 테이블을 찾는 함수
def find(table_name, table_confs):
    # table_confs 리스트를 순회하며 각 테이블 설정 딕셔너리를 확인
    for table in table_confs:
        # 현재 테이블 설정 딕셔너리의 "table" 키 값이 주어진 table_name과 일치하는지 확인
        if table.get("table") == table_name:
            # 일치하는 경우 해당 테이블 설정 딕셔너리를 반환
            return table
    
    # 모든 테이블 설정을 확인한 후에도 일치하는 것이 없으면 None 반환
    return None

# Redshift에서 요약 테이블을 생성하는 함수
def build_summary_table(dag_root_path, dag, tables_load, redshift_conn_id, start_task=None):
    # 로그에 DAG 루트 경로 출력
    logging.info(dag_root_path)
    
    # JSON 형식의 테이블 설정 파일들을 읽어들임
    table_confs = load_all_jsons_into_list(dag_root_path + "/config/")

    # DAG의 시작 작업이 지정되었으면 해당 작업을 기준으로 설정
    if start_task is not None:
        prev_task = start_task
    else:
        prev_task = None

    # tables_load에 지정된 테이블들에 대해 작업 생성 및 DAG에 추가
    for table_name in tables_load:
        # 지정된 테이블명에 해당하는 테이블 설정 찾기
        table = find(table_name, table_confs)
        
        # RedshiftSummaryOperator를 생성하고 설정값 전달
        summarizer = RedshiftSummaryOperator(
            table=table["table"],
            schema=table["schema"],
            redshift_conn_id=redshift_conn_id,
            input_check=table["input_check"],
            main_sql=table["main_sql"],
            output_check=table["output_check"],
            overwrite=table.get("overwrite", True),
            after_sql=table.get("after_sql"),
            pre_sql=table.get("pre_sql"),
            attributes=table.get("attributes", ""),
            dag=dag,
            task_id="analytics"+"__"+table["table"]
        )
        
        # 이전 작업과 현재 작업 간의 의존성 설정
        if prev_task is not None:
            prev_task >> summarizer
        prev_task = summarizer
    
    # 마지막으로 추가된 작업을 반환 (나중에 DAG 작성 시 필요)
    return prev_task

# Redshift에서 SQL을 실행하는 함수
def redshift_sql_function(**context):
    sql=context["params"]["sql"]
    print(sql)
    hook = PostgresHook(postgres_conn_id=context["params"]["redshift_conn_id"])
    hook.run(sql, True)

# Redshift에서 요약 테이블을 생성하는 PythonOperator를 확장한 사용자 정의 연산자 클래스
class RedshiftSummaryOperator(PythonOperator):
    @apply_defaults
    def __init__(self, schema, table, redshift_conn_id, input_check, main_sql, output_check, overwrite, params={}, pre_sql="", after_sql="", attributes="", *args, **kwargs):
        # 사용자가 정의한 RedshiftSummaryOperator 클래스의 초기화 메소드입니다.
        
        # 생성할 요약 테이블의 Redshift 스키마 이름
        self.schema = schema
        
        # 생성할 요약 테이블의 이름
        self.table = table
        
        # Airflow에서 사용하는 Redshift 연결 ID
        self.redshift_conn_id = redshift_conn_id
        
        # 입력 유효성을 검사하기 위한 SQL 쿼리 및 최소 레코드 수로 구성된 목록
        self.input_check = input_check
        
        # 요약 테이블을 생성하기 위한 주요 SQL 쿼리
        self.main_sql = main_sql
        
        # 출력 유효성을 검사하기 위한 SQL 쿼리 및 최소 레코드 수로 구성된 목록
        self.output_check = output_check
        
        # True인 경우 기존 테이블을 덮어쓰고, False인 경우 덮어쓰지 않고 추가합니다.
        self.overwrite = overwrite
        
        # PythonOperator에 전달할 추가 매개변수
        self.params = params
        
        # main_sql 실행 전에 실행할 SQL 쿼리
        self.pre_sql = pre_sql if pre_sql else ""
        
        # main_sql 실행 후에 실행할 SQL 쿼리
        self.after_sql = after_sql.format(schema=self.schema, table=self.table) if after_sql else ""
        
        # 생성할 테이블의 추가 속성
        self.attributes = attributes

        # temp 테이블 생성 및 데이터 적재에 사용될 SQL 문 생성
        if pre_sql:
            self.main_sql = pre_sql
            if not self.main_sql.endswith(";"):
                self.main_sql += ";"
        else:
            self.main_sql = ""
        # 임시테이블이 있으면 삭제, 
        self.main_sql += "DROP TABLE IF EXISTS {schema}.temp_{table};".format(
            schema=self.schema,
            table=self.table
        )
        # CREATE TABLE 문을 만들어서 self.main_sql에 추가합니다.
        self.main_sql += "CREATE TABLE {schema}.temp_{table} {attributes} AS ".format(
            schema=self.schema,
            table=self.table,
            attributes=self.attributes
        ) + self.main_sql

        # 상위 클래스인 PythonOperator를 호출하여 초기화
        # RedshiftSummaryOperator 클래스의 초기화 메소드에서는 두 번의 상위 클래스 초기화(super)가 이루어집니다.

        # 첫 번째 super 호출:
        super(RedshiftSummaryOperator, self).__init__(
            python_callable=redshift_sql_function,  # Python callable로 사용될 함수
            params={
                "sql": self.main_sql,  # 생성한 SQL 문
                "overwrite": self.overwrite,  # 덮어쓰기 여부
                "redshift_conn_id": self.redshift_conn_id  # Redshift 연결 ID
            },
            provide_context=True,  # Airflow 컨텍스트 제공 여부
            *args,
            **kwargs
        )

        # 두 번째 super 호출:
        # after_sql이 정의되어 있다면 해당 값을 사용하고, 그렇지 않으면 빈 문자열("")을 사용합니다.
        if after_sql:
            self.after_sql = after_sql.format(
                schema=self.schema,
                table=self.table
            )
        else:
            self.after_sql = ""

        super(RedshiftSummaryOperator, self).__init__(
            python_callable=redshift_sql_function,  # Python callable로 사용될 함수
            params={
                "sql": main_sql,  # 생성한 SQL 문
                "overwrite": overwrite,  # 덮어쓰기 여부
                "redshift_conn_id": self.redshift_conn_id  # Redshift 연결 ID
            },
            provide_context=True,  # Airflow 컨텍스트 제공 여부
            *args,
            **kwargs
        )


    # temp 테이블과 본 테이블을 스왑하는 함수
    def swap(self):
        # 원본 테이블 삭제
        # 임시테이블 원본테이블 이름으로 바꿈
        # {schema}와 {table}을 사용자가 정의한 값으로 대체하여 
        # SELECT 권한을 부여하는 SQL 문
        sql = """BEGIN;
        DROP TABLE IF EXISTS {schema}.{table} CASCADE;
        ALTER TABLE {schema}.temp_{table} RENAME TO {table};   
        GRANT SELECT ON TABLE {schema}.{table} TO GROUP analytics_users;
        END
        """.format(schema=self.schema,table=self.table)
        self.hook.run(sql, True)

    def execute(self, context):
        """
        RedshiftSummaryOperator의 execute 메소드입니다.

        1. Input_check 먼저 수행
           - input_check는 "sql"과 "count"를 포함하는 딕셔너리의 목록이어야 함
        """
        self.hook = PostgresHook(postgres_conn_id=self.redshift_conn_id)
        for item in self.input_check:
            (cnt,) = self.hook.get_first(item["sql"])
            if cnt < item["count"]:
                raise AirflowException(
                    "Input Validation Failed for " + str(item["sql"]))

        """
        2. temp 테이블 생성 및 데이터 적재 수행
        """
        return_value = super(RedshiftSummaryOperator, self).execute(context)

        """
        3. Output_check은 self.output_check 사용
        """
        for item in self.output_check:
            (cnt,) = self.hook.get_first(item["sql"].format(schema=self.schema, table=self.table))
            if item.get("op") == 'eq':
                if int(cnt) != int(item["count"]):
                    raise AirflowException(
                        "Output Validation of 'eq' Failed for " + str(item["sql"]) + ": " + str(cnt) + " vs. " + str(item["count"])
                    )
            else:
                if cnt < item["count"]:
                    raise AirflowException(
                        "Output Validation Failed for " + str(item["sql"]) + ": " + str(cnt) + " vs. " + str(item["count"])
                    )
        
        """
        4. 이제 temp 테이블 이름을 스왑
        """
        self.swap()

        """
        5. after_sql이 정의되어 있다면 실행
        """
        if self.after_sql:
            self.hook.run(self.after_sql, True)

        return return_value