Apache Airflow - Connection, Hook

x-airflow-common: # 각 서비스에 공통 적용될 항목들
services: # 컨테이너로 실행할 서비스를 정의
volumns: # 컨테이너에 연결할 볼륨을 정의
networks: # 컨테이너에 연결할 네트워크를 정의

x-airflow-common:
  &airflow-common
  image: ${AIRFLOW_IMAGE_NAME:-apache/airflow:3.0.1}
  environment:
    &airflow-common-env
    AIRFLOW__CORE__EXECUTOR: CeleryExecutor

services:
  airflow-apiserver:
    <<: *airflow-common
    command: api-server
    ports:
      - "8080:8080"
    healthcheck:
      test: ["CMD", "curl", "--fail", "http://localhost:8080/api/v2/version"]
      interval: 30s
      timeout: 10s
      retries: 5
      start_period: 30s
    restart: always
    depends_on:
      <<: *airflow-common-depends-on
      airflow-init:
        condition: service_completed_successfully

volumes:
  postgres-db-volume:

services:
  postgres:
    volumes:
      - postgres-db-volume:/var/lib/postgresql/data

networks:
  network_custom:
    driver: bridge
    ipam:
      driver: default
      config:
        - subnet: 172.28.0.0/16
          gateway: 172.28.0.1

% docker network inspect {network_id}
[
    {
        "Name": "airflow_default",
        ...
        "IPAM": {
            "Driver": "default",
            "Options": null,
            "Config": [
                {
                    "Subnet": "172.18.0.0/16",
                    "Gateway": "172.18.0.1"
                }
            ]
        },
        ...
    }
]

services:
  postgres_custom:
    image: postgres:13
    environment:
      POSTGRES_USER: minyeamer
      POSTGRES_PASSWORD: minyeamer
      POSGRES_DB: minyeamer
      TZ: Asia/Seoul
    volumes:
      - postgres-custom-db-volume:/var/lib/postgresql/data
    ports:
      - 5432:5432
    networks:
      network_custom:
        ipv4_address: 172.28.0.3

  postgres:
    image: postgres:13
    ...
    ports:
      - 5431:5432
    networks:
      network_custom:
        ipv4_address: 172.28.0.4

  ...

volumes:
  postgres-db-volume:
  postgres-custom-db-volume:

networks:
  network_custom:
    ...

% docker ps
CONTAINER ID   IMAGE                  COMMAND                  CREATED          STATUS                    PORTS                    NAMES
72541e89bff4   apache/airflow:3.0.1   "/usr/bin/dumb-init …"   57 seconds ago   Up 20 seconds (healthy)   8080/tcp                 airflow-airflow-worker-1
1735bc475bff   apache/airflow:3.0.1   "/usr/bin/dumb-init …"   57 seconds ago   Up 36 seconds (healthy)   0.0.0.0:8080->8080/tcp   airflow-airflow-apiserver-1
d2b680f273e3   apache/airflow:3.0.1   "/usr/bin/dumb-init …"   57 seconds ago   Up 36 seconds (healthy)   8080/tcp                 airflow-airflow-scheduler-1
f23c6d30d22c   apache/airflow:3.0.1   "/usr/bin/dumb-init …"   57 seconds ago   Up 36 seconds (healthy)   8080/tcp                 airflow-airflow-dag-processor-1
09992f09ac5e   apache/airflow:3.0.1   "/usr/bin/dumb-init …"   57 seconds ago   Up 36 seconds (healthy)   8080/tcp                 airflow-airflow-triggerer-1
5c1b13f33229   postgres:13            "docker-entrypoint.s…"   58 seconds ago   Up 56 seconds             0.0.0.0:5432->5432/tcp   airflow-postgres_custom-1
e23d4eb919fc   postgres:13            "docker-entrypoint.s…"   58 seconds ago   Up 56 seconds (healthy)   0.0.0.0:5431->5432/tcp   airflow-postgres-1
86fc82a16953   redis:7.2-bookworm     "docker-entrypoint.s…"   58 seconds ago   Up 56 seconds (healthy)   6379/tcp                 airflow-redis-1

CREATE TABLE public.dag_run (
	dag_id varchar(100) NULL,
	task_id varchar(100) NULL,
	run_id varchar(100) NULL,
	msg text NULL
);

# dags/python_with_postgres.py

from airflow.sdk import DAG
from airflow.providers.standard.operators.python import PythonOperator
import pendulum

with DAG(
    dag_id="python_with_postgres",
    schedule=None,
    start_date=pendulum.datetime(2025, 1, 1, tz="Asia/Seoul"),
    catchup=False,
    tags=["example", "hook"],
) as dag:
    def insert_into_postgres(ip: str, port: str, dbname: str, user: str, passwd: str, **kwargs):
        import psycopg2
        from contextlib import closing

        with closing(psycopg2.connect(host=ip, dbname=dbname, user=user, password=passwd, port=int(port))) as conn:
            with closing(conn.cursor()) as cursor:
                dag_id = kwargs.get("ti").dag_id
                task_id = kwargs.get("ti").task_id
                run_id = kwargs.get("ti").run_id
                msg = "INSERT INTO 수행"
                sql = "INSERT INTO dag_run VALUES (%s,%s,%s,%s);"
                cursor.execute(sql,(dag_id,task_id,run_id,msg))
                conn.commit()

    postgres_task = PythonOperator(
        task_id="postgres_task",
        python_callable=insert_into_postgres,
        op_args=["172.28.0.3", "5432", "minyeamer", "minyeamer", "minyeamer"]
    )

def get_conn(self) -> connection:
    """Establish a connection to a postgres database."""
    conn = deepcopy(self.connection)

    # check for authentication via AWS IAM
    if conn.extra_dejson.get("iam", False):
        conn.login, conn.password, conn.port = self.get_iam_token(conn)

    conn_args = {
        "host": conn.host,
        "user": conn.login,
        "password": conn.password,
        "dbname": self.database or conn.schema,
        "port": conn.port,
    }
    raw_cursor = conn.extra_dejson.get("cursor", False)
    if raw_cursor:
        conn_args["cursor_factory"] = self._get_cursor(raw_cursor)

    if self.options:
        conn_args["options"] = self.options

    for arg_name, arg_val in conn.extra_dejson.items():
        if arg_name not in self.ignored_extra_options:
            conn_args[arg_name] = arg_val

    self.conn = psycopg2.connect(**conn_args)
    return self.conn

pip install apache-airflow-providers-postgres

# dags/python_with_postgres_hook.py

from airflow.sdk import DAG
from airflow.providers.standard.operators.python import PythonOperator
import pendulum

with DAG(
    dag_id="python_with_postgres_hook",
    schedule=None,
    start_date=pendulum.datetime(2025, 1, 1, tz="Asia/Seoul"),
    catchup=False,
    tags=["example", "hook"],
) as dag:
    def insert_into_postgres(postgres_conn_id: str, **kwargs):
        from airflow.providers.postgres.hooks.postgres import PostgresHook
        from contextlib import closing
        postgres_hook = PostgresHook(postgres_conn_id)
        with closing(postgres_hook.get_conn()) as conn:
            with closing(conn.cursor()) as cursor:
                dag_id = kwargs.get("ti").dag_id
                task_id = kwargs.get("ti").task_id
                run_id = kwargs.get("ti").run_id
                msg = "INSERT INTO 수행"
                sql = "INSERT INTO dag_run VALUES (%s,%s,%s,%s);"
                cursor.execute(sql,(dag_id,task_id,run_id,msg))
                conn.commit()

    postgres_task = PythonOperator(
        task_id="postgres_task",
        python_callable=insert_into_postgres,
        op_kwargs={"postgres_conn_id":"conn-db-postgres-custom"}
    )

def bulk_load(self, table: str, tmp_file: str) -> None:
    """Load a tab-delimited file into a database table."""
    self.copy_expert(f"COPY {table} FROM STDIN", tmp_file)

def copy_expert(self, sql: str, filename: str) -> None:
    self.log.info("Running copy expert: %s, filename: %s", sql, filename)
    if not os.path.isfile(filename):
        with open(filename, "w"):
            pass

    with open(filename, "r+") as file, closing(self.get_conn()) as conn, closing(conn.cursor()) as cur:
        cur.copy_expert(sql, file)
        file.truncate(file.tell())
        conn.commit()

copy_expert(sql, file, size=8192)

>>> cur.copy_expert("COPY test TO STDOUT WITH CSV HEADER", sys.stdout)
id,num,data
1,100,abc'def
2,,dada
...

CREATE TABLE nshopping_search(
    rank        text,
    title       text,
    link        text,
    image       text,
    lprice      text,
    hprice      text,
    mallName    text,
    productId   text,
    productType text,
    brand       text,
    maker       text,
    category1   text,
    category2   text,
    category3   text,
    category4   text
);

from airflow.sdk import DAG
from airflow.providers.standard.operators.python import PythonOperator
import pendulum

with DAG(
    dag_id="python_with_postgres_load",
    schedule=None,
    start_date=pendulum.datetime(2025, 1, 1, tz="Asia/Seoul"),
    catchup=False,
    tags=["example", "hook"],
) as dag:
    def bulk_load_postgres(postgres_conn_id: str, table_name: str, file_path: str, **kwargs):
        from airflow.providers.postgres.hooks.postgres import PostgresHook
        postgres_hook = PostgresHook(postgres_conn_id)
        postgres_hook.bulk_load(table_name, file_path)

    postgres_task = PythonOperator(
        task_id="postgres_task",
        python_callable=bulk_load_postgres,
        op_kwargs={"postgres_conn_id":"conn-db-postgres-custom",
                    "table_name":"nshopping.search",
                    "file_name":"/opt/airflow/files/naverSearch/{{data_interval_end.in_timezone(\"Asia/Seoul\") | ds_nodash }}/shop_with_tab.csv"}
    )

@classmethod
def get_connection(cls, conn_id: str) -> Connection:
    """
    Get connection, given connection id.

    :param conn_id: connection id
    :return: connection
    """
    from airflow.models.connection import Connection

    conn = Connection.get_connection_from_secrets(conn_id)
    log.info("Connection Retrieved '%s'", conn.conn_id)
    return conn

@classmethod
def get_hook(cls, conn_id: str, hook_params: dict | None = None) -> BaseHook:
    """
    Return default hook for this connection id.

    :param conn_id: connection id
    :param hook_params: hook parameters
    :return: default hook for this connection
    """
    connection = cls.get_connection(conn_id)
    return connection.get_hook(hook_params=hook_params)

def get_conn(self) -> Any:
    """Return connection for the hook."""
    raise NotImplementedError()

# plugins/hooks/postgres.py

from airflow.hooks.base import BaseHook
from typing import Literal
import psycopg2

class CustomPostgresHook(BaseHook):
    def __init__(self, postgres_conn_id: str, **kwargs):
        self.postgres_conn_id = postgres_conn_id
        self.conn = None
        self.database = kwargs.get("database")

    def get_conn(self) -> psycopg2.extensions.connection:
        conn = BaseHook.get_connection(self.postgres_conn_id)
        conn_args = {
            "host": conn.host,
            "user": conn.login,
            "password": conn.password,
            "dbname": self.database or conn.schema,
            "port": conn.port,
        }
        self.conn = psycopg2.connect(**conn_args)
        return self.conn

    def bulk_load(self, table: str, filename: str, encoding="utf-8",
                if_exists: Literal["append","replace"]="append", sep=',', with_header=True):
        create = self._create_table_sql(table, filename, encoding, sep, with_header)
        replace = "TRUNCATE TABLE {};".format(table) if if_exists == "replace" else str()
        copy = "COPY {} FROM STDIN DELIMITER '{}' {};".format(table, sep, ("CSV HEADER" if with_header else "CSV"))
        sql = ''.join([create, replace, copy])
        self.copy_expert(sql, filename, encoding)

    def _create_table_sql(self, table: str, filename: str, encoding="utf-8", sep=',', with_header=True) -> str:
        if with_header:
            column_list = self._read_csv_column_list(filename, encoding, sep)
            return "CREATE TABLE IF NOT EXISTS {}({});".format(table, column_list)
        else:
            return str()

    def _read_csv_column_list(self, filename: str, encoding="utf-8", sep=',') -> str:
        import csv
        def is_int4_type(value: str) -> bool:
            return (not value) or (value.isdigit() and (-2147483648 <= int(value) <= 2147483647))
        with open(filename, "r+", encoding=encoding) as file:
            reader = csv.reader(file, delimiter=sep)
            header = next(reader)
            dtypes = [all(map(is_int4_type, values)) for values in zip(*[next(reader) for _ in range(5)])]
            return ", ".join(["{} {}".format(col, ("int4" if is_int4 else "text")) for col, is_int4 in zip(header, dtypes)])

    def copy_expert(self, sql: str, filename: str, encoding="utf-8") -> None:
        from contextlib import closing
        self.log.info("Running copy expert: %s, filename: %s", sql, filename)
        with open(filename, "r+", encoding=encoding) as file, closing(self.get_conn()) as conn, closing(conn.cursor()) as cur:
            cur.copy_expert(sql, file)
            file.truncate(file.tell())
            conn.commit()

def get_conn(self) -> psycopg2.extensions.connection:
    conn = BaseHook.get_connection(self.postgres_conn_id)
    conn_args = {
        "host": conn.host,
        "user": conn.login,
        "password": conn.password,
        "dbname": self.database or conn.schema,
        "port": conn.port,
    }
    self.conn = psycopg2.connect(**conn_args)
    return self.conn

def bulk_load(self, table: str, filename: str, encoding="utf-8",
            if_exists: Literal["append","replace"]="append", sep=',', with_header=True):
    create = self._create_table_sql(table, filename, encoding, sep, with_header)
    replace = "TRUNCATE TABLE {};".format(table) if if_exists == "replace" else str()
    copy = "COPY {} FROM STDIN DELIMITER '{}' {};".format(table, sep, ("CSV HEADER" if with_header else "CSV"))
    sql = ''.join([create, replace, copy])
    self.copy_expert(sql, filename, encoding)

def _create_table_sql(self, table: str, filename: str, encoding="utf-8", sep=',', with_header=True) -> str:
    if with_header:
        column_list = self._read_csv_column_list(filename, encoding, sep)
        return "CREATE TABLE IF NOT EXISTS {}({});".format(table, column_list)
    else:
        return str()

def _read_csv_column_list(self, filename: str, encoding="utf-8", sep=',') -> str:
    import csv
    def is_int4_type(value: str) -> bool:
        return (not value) or (value.isdigit() and (-2147483648 <= int(value) <= 2147483647))
    with open(filename, "r+", encoding=encoding) as file:
        reader = csv.reader(file, delimiter=sep)
        header = next(reader)
        dtypes = [all(map(is_int4_type, values)) for values in zip(*[next(reader) for _ in range(5)])]
        return ", ".join(["{} {}".format(col, ("int4" if is_int4 else "text")) for col, is_int4 in zip(header, dtypes)])

def copy_expert(self, sql: str, filename: str, encoding="utf-8") -> None:
    from contextlib import closing
    self.log.info("Running copy expert: %s, filename: %s", sql, filename)
    with open(filename, "r+", encoding=encoding) as file, closing(self.get_conn()) as conn, closing(conn.cursor()) as cur:
        cur.copy_expert(sql, file)
        file.truncate(file.tell())
        conn.commit()

# dags/python_with_postgres_custom.py

from airflow.sdk import DAG
from airflow.providers.standard.operators.python import PythonOperator
from hooks.postgres import CustomPostgresHook
import pendulum

with DAG(
        dag_id="python_with_postgres_custom",
        schedule=None,
        start_date=pendulum.datetime(2025, 1, 1, tz="Asia/Seoul"),
        catchup=False,
        tags=["example", "hook"],
) as dag:
    def bulk_load_postgres(postgres_conn_id: str, table: str, filename: str, **kwargs):
        custom_postgres_hook = CustomPostgresHook(postgres_conn_id=postgres_conn_id)
        custom_postgres_hook.bulk_load(table=table, filename=filename, if_exists="replace", sep=",", with_header=True)

    bulk_load_postgres = PythonOperator(
        task_id="bulk_load_postgres",
        python_callable=bulk_load_postgres,
        op_kwargs={"postgres_conn_id": "conn-db-postgres-custom",
                    "table":"nshopping.search2",
                    "filename":"/opt/airflow/files/naverSearch/{{data_interval_end.in_timezone(\"Asia/Seoul\") | ds_nodash }}/shop.csv"}
    )

[2025-06-11, 01:24:16] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-06-11, 01:24:16] INFO - Filling up the DagBag from /opt/airflow/dags/python_with_postgres_custom.py: source="airflow.models.dagbag.DagBag"
[2025-06-11, 01:24:16] INFO - Running copy expert: CREATE TABLE IF NOT EXISTS nshopping.search2(rank int4, title text, link text, image text, lprice int4, hprice int4, mallName text, productId text, productType int4, brand text, maker text, category1 text, category2 text, category3 int4, category4 int4);TRUNCATE TABLE nshopping.search2;COPY nshopping.search2 FROM STDIN DELIMITER ',' CSV HEADER;, filename: /opt/airflow/files/naverSearch/20250611/shop.csv: source="hooks.postgres.CustomPostgresHook"
[2025-06-11, 01:24:16] INFO - Secrets backends loaded for worker: count=1: backend_classes=["EnvironmentVariablesBackend"]: source="supervisor"
[2025-06-11, 01:24:16] INFO - Connection Retrieved 'conn-db-postgres-custom': source="airflow.hooks.base"
[2025-06-11, 01:24:16] INFO - Done. Returned value was: None: source="airflow.task.operators.airflow.providers.standard.operators.python.PythonOperator"