main.py

import asyncio
import logging
import os
import os.path as path
import sys
import time
from contextlib import asynccontextmanager as AsyncContextManager
from datetime import datetime
from pathlib import Path
from typing import List

from fastlogging import LogInit
from lark import logger as lark_logger
from tqdm.asyncio import tqdm_asyncio

from utils import (
    DataBase,
    Deid_Pipeline,
    PipeConfig,
    build_db_link,
    check_schema,
)


@AsyncContextManager
async def setup():
    working_dir = Path(path.abspath("."))
    # configure logger
    logfile = f"{datetime.now().strftime('%Y-%m-%d')}.log"
    logfile = working_dir / "logs" / logfile
    if not logfile.parent.exists():
        os.makedirs(logfile.parent)
    with open(logfile, "w") as f:
        f.write(f"# Log file {logfile}\n")
    logger = LogInit("deid_ppl", console=True, colors=True, pathName=str(logfile))
    lark_logger.setLevel(logging.WARNING)
    # read config
    config_path = PipeConfig.find_config_file(working_dir)
    config = PipeConfig(config_path, working_dir, logger)

    source = {
        table: getattr(config, table)["note_column_names"]
        for table in config.table_names
    }
    await check_schema(
        logger,
        config.db["host"],
        int(config.db["port"]),
        config.db["user"],
        config.db["password"],
        config.db["database"],
        config.db["charset"],
        source,
        config.write_back,
        config.write_to_col,
        config.db["deid_suffix"],
    )

    # prepare database connection
    async with DataBase(
        build_db_link(
            config.db["host"],
            int(config.db["port"]),
            config.db["user"],
            config.db["password"],
            config.db["database"],
            config.db["charset"],
        ),
        logger,
        len(config.table_names) * 10,
    ) as db:
        logger.info("Setup phase is finished.")
        yield config, db, logger


async def main():
    start = time.monotonic()
    async with setup() as (config, db, logger):
        # temp_file_type = config.deid["temp_file_type"]
        table_names: List[str] = config.table_names
        logger.debug(f"Working with tables: {', '.join(table_names)}")

        tasks = []
        for i, table_name in enumerate(table_names):
            table_conf = getattr(config, table_name)
            tasks.append(
                asyncio.create_task(
                    Deid_Pipeline(
                        db,
                        table_name,
                        table_conf["note_column_names"],
                        config.batch_size,
                        config.write_to_col,
                        config.db["deid_suffix"],
                        i,
                        config.write_back,
                        Path(config.deid["temp_dir"]),
                        config.deid["script"],
                        config.deid["config"],
                        config.autoclean,
                    ).run(
                        config.worker_size,
                        table_conf["start_row"],
                        table_conf["end_row"],
                    ),
                    name=f"table_{table_name}",
                )
            )

        await tqdm_asyncio.gather(
            *tasks,
            desc="Deid pipeline",
            unit="table",
            position=len(table_names),
            leave=False,
        )
        # await asyncio.gather(*tasks)
        end = time.monotonic()
        logger.info("All deid tasks are done, total time: %.2fs", end - start)
        logger.join()


if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    asyncio.run(main(), debug=False)
    sys.exit(0)
    # try:
    #     loop = asyncio.new_event_loop()
    #     asyncio.set_event_loop(loop)
    #     asyncio.run(main(), debug=False)
    #     sys.exit(0)
    # except Exception:
    #     sys.exit(255)