db_import.py

#!/usr/bin/env powerscript --run-level=1 --nologin
# -*- mode: python; coding: utf-8 -*-
"""
Wraps the platform's cdbimp script to run in parallel. We store the dumps of each table in separate files, so we can
import one table per subprocess, allowing for concurrent table imports based on the number of workers.

The script can also restart imports from scratch or continue from the last successful state.
"""

import argparse
import logging
import logging.handlers
import multiprocessing
import os.path
import re
import sys
import threading
import traceback
from datetime import datetime

from cdb.scripts import cdbimp

# matches the line that contains the count of all imported rows in cdb export files
EXPFILE_COUNT_LINE_RE = re.compile(r"^C(\d+)$")

# The tables to be ignored during parallel import (--workers > 1).
# cdbimp is not suited well for parallel imports, see README for more information.
# These files will stay in c_cdbexp folder and should be imported separately with
# just 1 single worker in the next run. In this case the import should be done in the
# same order as the files are listed here, therefore we use an ordered list.
FILES_EXCLUDED_FROM_PARALLEL_IMPORT = [
    "mq_system_posting.exp",
    "cdbblog_posting.exp",
    "cdbblog_topic2posting.exp",
    "cdbmq_registry.exp",
    "cdb_counter.exp",
    "cdb_audittrail.exp",
    "cdb_audittrail_detail.exp",
    "cdb_audittrail_objects.exp",
    "cdbes_jobs.exp",
    "cdbes_jobs_txt.exp",
    "cdb_object.exp",
]

LICENSE_FILES = [
    "lstatistics.exp",
    "ldbserver.exp",
    "lmnames.exp",
    "lmodules.exp",
    "lusage.exp",
]


def logger_thread(queue: multiprocessing.Queue, logfile: str):
    """
    Should be run in a thread of the parent process. All other subprocesses will write their logs to the queue and only
    this thread will actually write to the log file. This way, the subprocesses don't have to fight for a file lock on
    the log file.
    Also described in the Python Logging Cookbook:
    https://docs.python.org/3/howto/logging-cookbook.html#logging-to-a-single-file-from-multiple-processes
    """
    logger = logging.getLogger("db_import")
    logger.propagate = False
    formatter = logging.Formatter(
        "[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s"
    )
    handler = logging.FileHandler(logfile, encoding="utf-8")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    while True:
        record = queue.get()
        if record is None:
            break
        logger.handle(record)


class CDBImportOptions:
    """
    This class can be used to emulate the arguments that `cdbimp` would normally receive via the command line.
    """

    ignore = False  # because we delete everything before import, we want to know if there are any errors during import
    show_errors = False  # errors will still be logged, but we don't want them in stderr
    create_tables = True
    no_warnings = False
    no_info = False
    tablefile = None
    charset = "utf-8"
    show_comments = False
    show_progress = False
    no_change = False
    autocommit = False
    replace_tables = False

    def __init__(self, target_file: str):
        self.expfile = [target_file]


def ensure_data_directory(instancedir, reset_state=False):
    """
    Checks whether either table dump data or an archive that should contain this data is present. If the exports
    directory is present and not empty, it will be used. Else, the presence of the archive will be checked and the
    archive will be extracted. This bevahior can be overwritten by passing --reset-state, in which case the
    previous data directory will always be deleted and extracted anew.

    If neither is found, there is nothing we can do here and the program will quit.
    This behavior is taken from the previous db_mig script.
    """
    exports_dir = os.path.join(instancedir, "c_cdbexp")
    if os.path.isdir(exports_dir) and len(os.listdir(exports_dir)) > 0:
        if not reset_state:
            print(
                f"Exports directory exists and is not empty. Attempting import from {exports_dir}."
            )
            return

        print(
            f"Exports directory {exports_dir} is not empty but --reset-state was passed. Clearing old state."
        )
        import shutil

        shutil.rmtree(exports_dir)
        os.mkdir(exports_dir)

    exports_archive = os.path.join(instancedir, "c_cdbexp.tar.gz")
    if not os.path.isfile(exports_archive):
        sys.stderr.write(
            "No export data found in exports directory and archive file does not exist. Aborting.\n"
        )
        sys.exit(1)

    import tarfile

    with tarfile.open(exports_archive) as exports_archive:

        def _exp_members():
            return [
                member
                for member in exports_archive.getmembers()
                if member.name.endswith(".exp")
            ]

        print("Found exports archive. Extracting...")
        exports_archive.extractall(
            members=_exp_members(), filter="data", path=instancedir
        )


def run_import(
    queue: multiprocessing.Queue, fpath: str, table_name: str, skip_delete: bool
) -> None:
    """
    The main function for a subprocess job that will import the given table dump. Errors are caught and logged; if the
    import finishes without errors, the dump file will be deleted.
    """

    def _msg(log_level, msg):
        """
        One does not simply pass a Python builtin logger into the cdbimp script, so do a bit of wrapping here
        """
        match log_level:
            case cdbimp.Logger.INFO:
                logger.info(msg)
            case cdbimp.Logger.WARN:
                logger.warning(msg)
            case cdbimp.Logger.ERROR:
                logger.error(msg)
            case cdbimp.Logger.FATAL:
                logger.critical(msg)
            case _:
                raise ValueError(f"Invalid log level: {log_level}")

    logger = logging.getLogger(table_name)
    logger.propagate = False
    queue_handler = logging.handlers.QueueHandler(queue)
    logger.setLevel(logging.INFO)
    logger.addHandler(queue_handler)
    logger.msg = _msg

    try:
        from cdb import rte

        rte.ensure_run_level(rte.USER_IMPERSONATED)

        from cdb import sqlapi
        from cdb.dberrors import DBConstraintViolation, DBUndefinedObject

        imp_args = CDBImportOptions(fpath)
        if not skip_delete:
            try:
                deleted_rows: int = sqlapi.SQLdelete(f"FROM {table_name}")
                logger.info(
                    "Deleted %s rows from %s before import", deleted_rows, table_name
                )
            except (DBConstraintViolation, DBUndefinedObject):
                logger.warning(
                    "Caught DBConstraintViolation (sqlite) or DBUndefinedObject (postgres) when deleting %s, meaning the table did not exist. This is likely due to a deprecated table, and can usually be ignored. The table file is kept but will not be imported.",
                    table_name,
                )
                return
        else:
            logger.info(f"Skipping deletion of contents in table: {table_name}")

        cdbimp._process(imp_args, logger)  # pylint: disable=protected-access
        imported_rows = int(
            sqlapi.SQLnumber(sqlapi.SQLselect(f"COUNT(*) FROM {table_name}"), 0, 0)
        )
        expected_rows = count_expected_rows(fpath)
        if imported_rows != expected_rows:
            logger.fatal(
                "Expected %s lines in %s but found %s. This might be caused by an unstable database cluster; aborting the import and investigating further is strongly recommended.",
                expected_rows,
                table_name,
                imported_rows,
            )
            return
    except Exception as e:
        logger.error("[%s] Caught exception: %s", table_name, e)
        for line in traceback.format_exception(e):
            logger.error(line)
        return

    logger.info("[%s] No errors while importing file %s", table_name, fpath)
    os.remove(fpath)


def count_expected_rows(exp_file_path: str) -> int:
    """
    Each cdbexp file contains a line that counts the lines in the file, formatted like "C1234". This function matches
    this line and returns the amount of rows that should have been imported into the database.
    """
    with open(exp_file_path, "r", encoding="utf-8") as exp_file:
        for line in exp_file:
            match = re.match(EXPFILE_COUNT_LINE_RE, line)
            if not match:
                continue
            return int(match.group(1))
    raise RuntimeError(f"Did not find count of lines in cdbexp file {exp_file_path}")


def arg_parser() -> argparse.ArgumentParser:
    """
    Sets up the command line arguments for instance directory and worker count
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        "--instancedir",
        # This is the default path for our cloud images; can be overridden e.g. for testing
        default=os.environ.get("INSTANCE", "/home/contact/instance"),
        help="Instance directory."
        + " Should contain a subdirectory called 'c_cdbexp' with the database dumps to be imported.",
    )
    parser.add_argument(
        "-w",
        "--workers",
        default=multiprocessing.cpu_count(),
        type=int,
        choices=range(
            1, multiprocessing.cpu_count() + 1
        ),  # Ensure the upper bound is inclusive
    )
    parser.add_argument(
        "-r",
        "--reset-state",
        help="Clean previous state before running import.",
        action="store_true",
    )
    parser.add_argument(
        "-t",
        "--table-name",
        help="Import only this one table."
        + " The exp file is expected in 'c_cdbexp' subdirectory of the instance.",
        type=str,
    )
    parser.add_argument(
        "-d",
        "--skip-delete",
        help="Skip deleting the table contents prior to import",
        action="store_true",
        default=False,
    )
    return parser


def main():
    args = arg_parser().parse_args()

    # ensure logs directory is present
    logs_dir = os.path.join(args.instancedir, "db_import")
    try:
        os.mkdir(logs_dir)
    except FileExistsError:
        pass
    os.chmod(logs_dir, mode=0o700)

    print("Starting logging thread...")
    # set up thread for the log listener
    mp_manager = multiprocessing.Manager()
    logging_queue = mp_manager.Queue(-1)
    log_listener = threading.Thread(
        target=logger_thread,
        args=(
            logging_queue,
            os.path.join(
                logs_dir,
                f"db_import.{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.log",
            ),
        ),
        daemon=True,
    )
    log_listener.start()

    print("Starting import tasks...")
    try:
        ensure_data_directory(args.instancedir, args.reset_state)
        with multiprocessing.Pool(args.workers) as pool:
            print(f"Starting {args.workers} workers")
            files = [
                file
                for file in os.listdir(os.path.join(args.instancedir, "c_cdbexp"))
                if os.path.isfile(os.path.join(args.instancedir, "c_cdbexp", file))
                and file.endswith(".exp")
                and file not in LICENSE_FILES  # exclude license files
                # if table_name is set, filter for files with that exact name
                and (args.table_name is None or file == f"{args.table_name}.exp")
            ]

            # If parallel import, exclude certain tables.
            if args.workers > 1:
                files = [
                    f
                    for f in files
                    if f.lower() not in FILES_EXCLUDED_FROM_PARALLEL_IMPORT
                ]
            else:
                files = FILES_EXCLUDED_FROM_PARALLEL_IMPORT

            _ = [
                pool.apply_async(
                    run_import,
                    (
                        logging_queue,
                        os.path.join(args.instancedir, "c_cdbexp", file),
                        file[0:-4],
                        args.skip_delete,
                    ),
                )
                for file in files
            ]
            pool.close()
            pool.join()
            print("All tasks done; shutting down...")
    finally:
        logging_queue.put(None)
        log_listener.join()
        print("Logging thread shut down successfully. Goodbye!")


if __name__ == "__main__":
    main()