# coding: utf-8

from contextlib import closing
import logging

from airflow.exceptions import AirflowException


def insert_rows(mysql_conn, table, rows, target_fields, batch_size=1000, upsert=False):
    """
    Insert a set of tuples into a table, a new transaction is created every commit_every rows

    :param mysql_conn: MySQL Connection
    :param table: Name of the target table
    :param rows: The rows to insert into the table
    :param target_fields: The names of the columns to fill in the table
    :param batch_size: The maximum number of rows to insert in one
            batch. Set to 0 to insert all rows in one batch.
    :param upsert: Whether to use INSERT ... ON DUPLICATE KEY UPDATE
    """

    if (not isinstance(target_fields, list)) or len(target_fields) < 1:
        raise AirflowException('target_fields should be a nonempty list')

    place_holder = ['%s'] * len(target_fields)
    place_holder = ', '.join(place_holder)

    insert_sql = f"insert into {table} ({', '.join(target_fields)}) values ({place_holder})"
    if upsert:
        update_clauses = [f'{col} = values({col})' for col in target_fields]
        update_clauses = ', '.join(update_clauses)

        insert_sql = insert_sql + ' on duplicate key update ' + update_clauses
    logging.info(insert_sql)

    i = 0
    with closing(mysql_conn.cursor()) as cursor:
        batch = []
        for i, row in enumerate(rows, 1):
            batch.append(row)

            if batch_size and i % batch_size == 0:
                cursor.executemany(insert_sql, batch)
                mysql_conn.commit()
                logging.info("Loaded %s rows into %s so far", i, table)
                batch = []

        if batch:
            cursor.executemany(insert_sql, batch)
            mysql_conn.commit()
            logging.info("Loaded %s rows into %s so far", i, table)

    logging.info("Done loading. Loaded a total of %s rows", i)
