transfer_operator.py

class TransferPsqlToMysql(BaseOperator):

    template_fields = ["sql"]
    template_ext = [".sql"]
    template_fields_renderers = {"sql": "sql"}

    # pylint: disable=too-many-arguments
    def __init__(
        self,
        postgres_conn_id,
        mysql_conn_id,
        sql,
        dst_table_name,
        truncate_table,
        **kwargs,
    ):
        self.postgres_conn_id = postgres_conn_id
        self.mysql_conn_id = mysql_conn_id
        self.sql = sql
        self.dst_table_name = dst_table_name
        self.truncate_table = truncate_table
        super().__init__(**kwargs)

    def execute(self, context: Context) -> Any:
        postgres_hook = PostgresHook(self.postgres_conn_id)
        mysql_hook = MySqlHook(self.mysql_conn_id)

        with NamedTemporaryFile("w+b") as file:
            logging.info("Exporting data to csv file")
            postgres_hook.copy_expert(
                f"COPY ({self.sql}) TO STDOUT WITH (FORMAT csv, DELIMITER '\t')",
                file.name,
            )
            file.flush()
            file.seek(0)

            if self.truncate_table:
                mysql_hook.run(f"TRUNCATE TABLE {self.dst_table_name}")

            mysql_hook.bulk_load(self.dst_table_name, file.name)