Source code for operators.base_transform

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.postgres.hooks.postgres import PostgresHook




class BaseTransformOperator(BaseOperator):
    """
    Base class to transform data from one connection to another. Uses `sqlalchemy` and `pandas`. Defaults to SQL Server.
    :param source_conn_id: the source  onnection ID
    :type source_conn_id: str
    :param source_sql: the source SQL statement. Should return the data to be
        saved to the destination. (templated)
    :type source_sql: str
    :param dest_conn_id: the destination connection ID
    :type dest_conn_id: str
    :param dest_schema_name: the name of the destination schema
    :type dest_schema_name: str
    :param dest_table_name: the name of the destination table
    :type dest_table_name: str
    :param truncate: whether to truncate the destination table
    :type truncate: bool
    :param if_table_exists: what to do if destination table exists. Can be ``fail``, ``append``, ``replace``. Default ``fail``.
    :type if_table_exists: str
    :param query_limit: how many rows to pass in each chunk
    :type query_limit: int
    """

    # override if necessary
    template_fields = ("source_sql", "dest_table_name", "dest_schema_name")

    CHECK_TABLE_EXIST = "SELECT EXISTS (SELECT FROM pg_tables WHERE schemaname = '{}' AND tablename = '{}');"
    TRUNCATE_QUERY = "TRUNCATE TABLE {};"
    DROP_QUERY = "DROP TABLE {};"
    DEFAULT_QUERY = "SELECT * FROM {}.{}"

    @apply_defaults
    def __init__(
        self,
        source_conn_id,
        dest_conn_id,
        dest_schema_name,
        dest_table_name,
        truncate=False,
        source_sql="",
        if_table_exists="replace",
        query_limit=0,
        *args,
        **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.source_conn_id = source_conn_id
        self.source_sql = (
            source_sql
            if source_sql
            else self.DEFAULT_QUERY.format(dest_schema_name, dest_table_name)
        )
        self.dest_conn_id = dest_conn_id
        self.dest_schema_name = dest_schema_name
        self.dest_table_name = dest_table_name
        self.truncate = truncate
        self.if_table_exists = if_table_exists
        self.query_limit = query_limit



    @staticmethod
    def check_result_exist(result):
        """wrapper for db differences - depending on what the CHECK_TABLE_EXIST query does, may need to do this differently"""
        return result.fetchone()[0]




    def check_truncate_table(self, dest_conn, dest_schema_table):
        """delete all contents of the given table before writing if necessary"""
        result_exist = dest_conn.execute(
            self.CHECK_TABLE_EXIST.format(self.dest_schema_name, self.dest_table_name)
        )

        if self.check_result_exist(result_exist):
            if self.truncate:
                self.log.info("Truncating {}".format(dest_schema_table))
                dest_conn.execute(self.TRUNCATE_QUERY.format(dest_schema_table))

            # wrapping pandas functionality for chunking edge cases
            if self.if_table_exists == "replace":
                self.log.info("Replacing existing table: {}".format(dest_schema_table))
                dest_conn.execute(self.DROP_QUERY.format(dest_schema_table))

            elif self.if_table_exists == "fail":
                raise




    def df_to_sql(self, df, dest_conn, dest_schema_table, if_exists=None):
        """helper for repeated code - write as a txn"""
        with dest_conn.begin() as trans:
            try:
                df.to_sql(
                    self.dest_table_name,
                    dest_conn,
                    schema=self.dest_schema_name,
                    if_exists=if_exists or self.if_table_exists,
                    index=False,
                    method="multi",
                )
                trans.commit()

            except:
                trans.rollback()
                raise




    def chunk_query(self, source_conn, dest_conn, dest_schema_table):
        """chunk the query result to obey memory limits"""
        import pandas as pd

        self.log.info("Running query....")
        result = source_conn.execution_options(stream_results=True).execute(
            self.source_sql
        )
        self.log.info("Chunking results...")
        counter = 0

        while True:
            self.log.info("Reading chunk {}...".format(counter))
            chunk = result.fetchmany(self.query_limit)

            if not chunk:
                break

            df = pd.DataFrame(data=chunk)
            df.columns = result.keys()
            self.log.info(
                "Writing chunk {} to {}...".format(counter, dest_schema_table)
            )
            self.df_to_sql(df, dest_conn, dest_schema_table, "append")
            counter += 1




    def transform_query(self, source_conn, dest_conn, dest_schema_table):
        """read the entire query in a df then write it to a db"""
        import pandas as pd

        self.log.info("Reading data...")
        df = pd.read_sql(self.source_sql, source_conn)
        self.log.info("Writing data to {}".format(dest_schema_table))
        self.df_to_sql(df, dest_conn, dest_schema_table)




    def execute(self, context):
        """
        main call method - creates two engines and feeds them to pandas
        """
        source_engine = self.get_source_hook().get_sqlalchemy_engine()
        source_conn = source_engine.connect()

        dest_engine = self.get_dest_hook().get_sqlalchemy_engine()
        dest_conn = dest_engine.connect()

        dest_schema_table = "{}.{}".format(self.dest_schema_name, self.dest_table_name)

        self.log.info(self.source_sql)
        try:
            self.check_truncate_table(dest_conn, dest_schema_table)

            if self.query_limit:
                self.chunk_query(source_conn, dest_conn, dest_schema_table)
            else:
                self.transform_query(source_conn, dest_conn, dest_schema_table)

        finally:
            source_conn.close()
            dest_conn.close()




    def get_source_hook(self):
        """
        add src hook that is able to make a sqlalchemy connection
        """
        return PostgresHook(postgres_conn_id=self.source_conn_id)




    def get_dest_hook(self):
        """
        add dest hook that is able to make a sqlalchemy connection
        """
        return PostgresHook(postgres_conn_id=self.dest_conn_id)