from multiprocessing import Process
import os

import config
from connectors import mysql

logger = config.logger


def main() -> None:
    """
    Get the partition count of TEMP_dig_sales_statements.
    Then for each partition, select into outfile to S3.
    """
    database = "accountingflat"
    table = "TEMP_dig_sales_statements"
    partitions = mysql.get_table_partition_count(database, table)
    batch_size = config.MYSQL_CONCURRENT_LOADS
    for batch_start in range(0, len(partitions), batch_size):
        logger.info(f"Starting from partition {batch_start}")
        partition_batch = partitions[batch_start : batch_start + batch_size]
        processes = []
        for partition in partition_batch:
            with open(
                os.path.realpath(os.path.dirname(__name__))
                + "/accounting_run_utils/queries/select_temp_dig_sales_statements_into_outfile.sql"
            ) as sql_file:
                prepared_statement = sql_file.read()
            """
            Partition names in queries cannot contain quotes, so replace 
            the partition name in the prepared statement
            with the actual partition to work around that limitation.
            """
            sql = prepared_statement.replace("PARTITION_NAME", partition)
            object_key = f"s3://{config.S3_ABACUS_OUTFILE_BUCKET}/{config.S3_ABACUS_OUTFILE_BUCKET_PREFIX}/{partition}.txt"
            keywords = {
                "sql_args": {
                    "period_id": config.PERIOD_ID,
                    "s3_object_path": object_key,
                }
            }
            processes.append(
                Process(
                    target=mysql.select_data_into_outfile,
                    args=(sql, object_key),
                    kwargs=keywords,
                )
            )
        for process in processes:
            process.start()
        for process in processes:
            process.join()


if __name__ == "__main__":
    main()