ody_migr_transaction.py

#!/usr/bin/env python
"""Main module that handles all database operations.

- Each database operation can take relevant optional parameters.
- Safe sql query generation is performed in ody_migr_db_handler.py script.

Specific handler such as `menu_insert_cu_featuremenu` must catch and
raise psycopg2 specific database errors and raise so that execute_migration
sees them and handles appropriately. A decorator `pg_crsr_hndlr_decrtr` is
used for that purpose.
"""


import logging
import psycopg2
from functools import wraps
import ody_migr_db_handler as pg_handler
from ody_migr_utils import (get_valid_json,
                            sanitize_insert_values,
                            sanitize_dict,
                            sanitize_insert_collection)

# import only necessary variables for settings
from ody_migr_config import (CREATE,
                             INSERT,
                             INSERT_ONE,
                             SELECT,
                             SELECT_RETURN,
                             SELECT_COUNT,
                             DELETE,
                             SUMMARY,
                             TABLE_EXISTS,
                             KEY_CU_PREFIXED_TBLS,
                             KEY_COMMON_TBLS,
                             ODY_MUST_EXIST_TABLES,
                             FEATURE_NOT_FOUND)

LOGGER = logging.getLogger(__name__)

# common sql and transaction related default keywords
# used in kwargs in various db operations
PG_ARG_DEFAULT = "arg_default"
PGARG_COL_SCHMA = "column_schema"
PGARG_COLL = "collection"
PGARG_CLN_BFRE_INSRT = "cleanup_before_insert"
PGARG_RTRNG_COLMN = "returning_col"
PGARG_VRBSE = "verbose"
PGARG_WHR_CONDS = "where_conditions"
PGARG_RTRN_ONLY_CNT = "return_only_count"
PGARG_SLCT_COLS = "select_columns"
PGARG_HNDLR_SPECIFIC = "handler_specific_args"

# dictionary of default values of different
# kwargs arguments being used in PgTransaction methods
PGDEF_ARGS_DICT = {
    PGARG_COL_SCHMA: None,
    PGARG_COLL: [],
    PGARG_CLN_BFRE_INSRT: False,
    PGARG_RTRNG_COLMN: "",
    PGARG_VRBSE: False,
    PGARG_WHR_CONDS: {},
    PGARG_RTRN_ONLY_CNT: False,
    PGARG_SLCT_COLS: [],
    PGARG_HNDLR_SPECIFIC: {}
}


def pg_crsr_hndlr_decrtr(original_action_handler):
    """Common psycopg2 exception handler decorator

    Catches and raise psycopg2 related errors and warnings.

    Args:
        original_action_handler: function to wrap
    """
    @wraps(original_action_handler)
    def wrapper(*args, **kwargs):
        try:
            return original_action_handler(*args, **kwargs)
        except (psycopg2.Error, psycopg2.Warning):
            raise
    return wrapper


class PgTransaction(object):
    """Main Transaction class to enwrap multiple db operations into a
    single transaction.

    Currently, a transaction is on single connection and uses
    a single cursor. This class is the central object to handle
    all database/sql operations with the help of ody_migr_db_handler.py
    module being used for safe sql generation.
    """

    def __init__(self, conn, cur):
        """PgTransaction constructor

        Args:
            conn: psycopg2 connection object
            cur: psycopg2 cursor object, must be a cursor of conn
        """
        self.conn = conn
        self.cur = cur

    def __call__(self, db_operation, arg0, *args, **kwargs):
        """Central transaction executing calling function of PgTransaction

        Performs appropriate method call based on the db_operation

        Args:
            db_operation: well defined database operation keyword
                         (CREATE, DELETE, etc)
           arg0:         table name for all db operation except TABLE_EXISTS,
                         data_category for TABLE_EXISTS and SUMMARY
                         in which case it is data_category
           args:         list of positional arguments to be forwarded to
                         a specified database operation handling method.
           kwargs:       list of keyword arguments to be forwarded to
                         a specified database operation handling method, mainly
                         used to generate sql scripts
        """

        try:
            # default data strucure for insertion is dict
            insert_data_str = kwargs.pop("data_struct", list)
            assert insert_data_str in [list, dict]
            tbl_name = ""

            # if required table exist
            if db_operation == TABLE_EXISTS:
                data_ctgry = arg0
                cu_tbl_prefix = args[0]
                return self.table_exists(data_ctgry, cu_tbl_prefix, **kwargs)

            # logs summary of records count for all the tables associated with
            # the specified data category
            elif db_operation == SUMMARY:
                data_ctgry = arg0
                cu_tbl_prefix = args[0]
                return self.display_summary(
                    data_ctgry,
                    cu_tbl_prefix,
                    **kwargs
                )

            else:
                tbl_name = arg0

                # create database table
                if db_operation == CREATE:
                    return self.create(tbl_name, **kwargs)

                # insert one record
                elif db_operation == INSERT_ONE:
                    if insert_data_str == list:
                        assert len(args) == 2
                        assert type(args[0]) in [list, tuple]
                        assert type(args[1]) in [list, tuple]
                        cols = args[0]
                        values = args[1]
                        return self.insert_one_record(
                            tbl_name,
                            cols,
                            values,
                            **kwargs
                        )
                    else:
                        assert len(args) == 1
                        assert args[0] == dict
                        _data_dict = args[0]
                        return self.insert_one_record_dict(
                            tbl_name,
                            _data_dict,
                            **kwargs
                        )

                # insert many records
                elif db_operation == INSERT:
                    return self.insert(tbl_name, **kwargs)

                # select records from a table and log output
                elif db_operation == SELECT:
                    return self.select_and_print(tbl_name, **kwargs)

                # select records from a table and return list
                elif db_operation == SELECT_RETURN:
                    return self.select_and_return(tbl_name, **kwargs)

                # select count of records
                elif db_operation == SELECT_COUNT:
                    return self.select_count(tbl_name, **kwargs)

                # delete records from a table
                elif db_operation == DELETE:
                    return self.delete(tbl_name, **kwargs)

                # exit if not a valid db operation
                else:
                    raise SystemExit(
                        "'{}' is not a valid sql operator.".format(
                            db_operation))

        # some database error caught, rollback the transaction
        except psycopg2.Error as err:
            err_dict = {
                "PSYCOPG_DETAIL": err.diag.message_primary,
                "PSYCOPG_PGCODE": err.pgcode,
                "PSYCOPG DIAGNOSTICS":
                pg_handler.pg_prepare_diagnostic_error(err.diag),
                "PSYCOPG_{}".format(err.diag.severity): err.pgerror
            }
            LOGGER.error(get_valid_json(err_dict))
            # rollback on any error while building a transaction and exit
            self.conn.rollback()
            # any failure that will cause a rollback of a
            # transaction halts the script processing
            raise SystemExit(
                "[Error] Rolling back transaction: {} ({})"
                .format(tbl_name, db_operation))

        # postgres warning, possibly due to data truncations during insertions
        except psycopg2.Warning as warning:
            LOGGER.warning(get_valid_json(warning))
            # we rollback for any warning while buildinga transaction and exit
            self.conn.rollback()
            raise SystemExit(
                "[Warning] Rolling back transaction: {} ({})"
                .format(tbl_name, db_operation))

        # Following will be executed if the operation in the try clause
        # does not raise an exception.
        # Even though we are good to commit for the operation recenly
        # executed on this transaction so far, but we enforce to
        # commit explicitly after all the operations are done executed
        else:
            # this is a typical place to commit the transactions, but we choose
            # to commit by explicitly calling the commit method when ready to
            # have more control over the transaction.
            LOGGER.info(
                "Successfully executed sql operation: `{}` on table `{}`."
                .format(tbl_name, db_operation))

    def _single_table_exists(self, _tbl_name):
        """Check if a single table exists in the database

        Args:
            _tbl_name:  name of the table to check the existence of
        """
        (check_script,
         check_values) = pg_handler.pg_check_table_exists(_tbl_name)
        self.cur.execute(check_script, check_values)
        return self.cur.fetchone()[0]

    def table_exists(self, _data_category, cu_tbl_prefix, **kwargs):
        """Check if required tables for each data category exist in the db

        Args:
            _data_category:     one of {memdata, memhist, admin, settings}
            cu_tbl_prefix:      current credit union code
            kwargs:             optional keyword arguments

        Returns:
            all_tbls_exist:     flag to notify if all required tables exist
            do_not_exist_list:  list of tables that donot exist if
                                all_tbls_exist is False
        """
        all_tbls_exist = True
        do_not_exist_list = []

        for cu_prefixed_tbl in ODY_MUST_EXIST_TABLES[_data_category][
                KEY_CU_PREFIXED_TBLS]:
            tbl = cu_prefixed_tbl.format(cu_tbl_prefix)
            if not self._single_table_exists(tbl):
                do_not_exist_list.append(tbl)
                all_tbls_exist = False

        for common_tbl in ODY_MUST_EXIST_TABLES[_data_category][
                KEY_COMMON_TBLS]:
            if not self._single_table_exists(common_tbl):
                do_not_exist_list.append(common_tbl)
                all_tbls_exist = False
        return all_tbls_exist, do_not_exist_list

    def display_summary(self, data_ctgry, _cu, **kwargs):
        """Log tables' records count summary for each data category

        Args:
            data_ctgry:         one of {memdata, memhist, admin, settings}
            _cu:                current credit union code
            kwargs:             optional keyword arguments
        """
        msg_prefix = kwargs.get("msg_prefix", "")
        settings_summary = {}
        for tbl in ODY_MUST_EXIST_TABLES[data_ctgry][KEY_CU_PREFIXED_TBLS]:
            cu_tbl = tbl.format(_cu.lower())
            this_count = self.select_count(cu_tbl)
            settings_summary[cu_tbl] = this_count

        for tbl in ODY_MUST_EXIST_TABLES[data_ctgry][KEY_COMMON_TBLS]:
            # handle special case tables that do not condition on <cu> code
            if tbl in ["cuadmquestselect", "cucontact", "lnappschemadetail",
                       "lnappuser_questselect", "lnappuserresponse"]:
                this_count = self.select_count(tbl)
                settings_summary[tbl] = this_count
                continue
            this_count = self.select_count(tbl, where_conditions={"cu": _cu})
            settings_summary[tbl] = this_count

        if msg_prefix == "":
            msg = "Summary of table records count:"
        else:
            msg = "[{}] Summary of table records count:".format(msg_prefix)

        LOGGER.info("{} {}".format(msg, get_valid_json(settings_summary)))

    def create(self, _tbl_name, **kwargs):
        """Handles table creation operation

        Args:
            _tbl_name: name of table to create
            kwargs:    optional keyword arguments
        """
        (table_exists_script,
         table_exists_values) = pg_handler.pg_check_table_exists(_tbl_name)
        column_schema = kwargs.get(
            PGARG_COL_SCHMA, PGDEF_ARGS_DICT[PGARG_COL_SCHMA])
        assert column_schema != ""

        create_script, create_values = pg_handler.pg_prepare_create_script(
            _tbl_name,
            column_schema)

        # execute table creation script
        self.cur.execute(create_script, create_values)
        LOGGER.debug(("(SQL) Create Script: {}").format(
            self.cur.query))

    def insert_one_record_dict(self, _tbl_name, _dict, **kwargs):
        """Handles single record insertion operation.

        Note: a record is a dictioanry and script uses named placeholder.

        Args:
            _tbl_name:  name of table to insert a record into
            _dict:      record to insert
            **kwargs:   optional keyword arguments
        """
        _dict = sanitize_dict(_dict)
        script = pg_handler.pg_prepare_insert_single_dict(
            _tbl_name,
            list(_dict.keys())
        )
        # execute insertion script
        self.cur.execute(script, _dict)

        LOGGER.debug(("(SQL) Insert A Record Script: {}").format(
            self.cur.query))

    def insert_one_record(self, _tbl_name, _columns, _values, **kwargs):
        """Handles single record insertion operation.

        Note: columns and values are two lists in same order and script
        uses general placeholder.

        Args:
            _tbl_name:  name of table to insert a record into
            _columns_:  column names
            _values:    values to insert in the same order as column names
            **kwargs:   optional keyword arguments
        """
        # strip whitespaces if applicable
        _values = sanitize_insert_values(_values)
        insert_script = pg_handler.pg_prepare_insert_single_list(
            _tbl_name,
            _columns,
            **kwargs
        )
        # execute insertion script
        self.cur.execute(insert_script, _values)

        LOGGER.debug(("(SQL) Insert A Record Script: {}").format(
            self.cur.query))

    def insert(self, _tbl_name, **kwargs):
        """Handles multiple record insertion operation

        Args:
            _tbl_name: name of table for insertion
            kwargs:    optional keyword arguments
        """
        # cleanup_before_insert = kwargs.get(
        #     PGARG_CLN_BFRE_INSRT, PGDEF_ARGS_DICT[PGARG_CLN_BFRE_INSRT])
        collection = kwargs.get(PGARG_COLL, PGDEF_ARGS_DICT[PGARG_COLL])

        assert type(collection) == list
        if len(collection) <= 0:
            LOGGER.debug(
                "No records to insert in table `{}`.".format(_tbl_name))
            return
        assert type(collection[0]) == dict

        LOGGER.debug(("START INSERTING A COLLECTION of {}-records."
                      " into a table: `{}`").
                     format(len(collection), _tbl_name))

        collection = sanitize_insert_collection(collection)
        for row in collection:
            # row is a dictionary
            self.insert_one_record_dict(_tbl_name, row, **kwargs)

        LOGGER.debug(("FINISH INSERTING A COLLECTION of {}-records."
                      " into a table: `{}`").
                     format(len(collection), _tbl_name))

    def select_count(self, _tbl_name, **kwargs):
        """Handles select and return count operation

        Args:
            _tbl_name:  name of table

        Returns:
            number of records in _tbl_name for the provided condition
        """
        where_conditions = kwargs.get(
            PGARG_WHR_CONDS, PGDEF_ARGS_DICT[PGARG_WHR_CONDS])

        select_script, select_values = pg_handler.pg_prepare_select_script(
            _tbl_name,
            where_conditions,
            [],
            count=True
        )
        self.cur.execute(select_script, select_values)
        LOGGER.debug(("(SQL) Select and Return Script: {}.").format(
            self.cur.query))
        return self.cur.fetchone()[0]

    def select_and_return(self, _tbl_name, **kwargs):
        """Handles record selection and returns records or count

        Args:
            _tbl_name: involved table
            kwargs:    optional keyword arguments

        Returns:
            list of records or count of records
        """
        select_result = None
        return_only_count = kwargs.get(
            PGARG_RTRN_ONLY_CNT, PGDEF_ARGS_DICT[PGARG_RTRN_ONLY_CNT])
        select_columns = kwargs.get(
            PGARG_SLCT_COLS, PGDEF_ARGS_DICT[PGARG_SLCT_COLS])
        where_conditions = kwargs.get(
            PGARG_WHR_CONDS, PGDEF_ARGS_DICT[PGARG_WHR_CONDS])

        select_script, select_values = pg_handler.pg_prepare_select_script(
            _tbl_name,
            where_conditions,
            select_columns
        )
        self.cur.execute(select_script, select_values)
        LOGGER.debug(("(SQL) Select and Return Script: {}.").format(
            self.cur.query))

        # return only the count of records
        if return_only_count:
            select_result = self.cur.rowcount  # len(all_records)
        # return whole records list
        else:
            select_result = []
            select_result = self.cur.fetchall()

        return select_result

    def select_and_print(self, _tbl_name, **kwargs):
        """Handles records selection from db table and print

        Args:
            _tbl_name: involved table
            kwargs:    optional keyword arguments
        """
        select_columns = kwargs.get(
            PGARG_SLCT_COLS, PGDEF_ARGS_DICT[PGARG_SLCT_COLS])
        where_conditions = kwargs.get(
            PGARG_WHR_CONDS, PGDEF_ARGS_DICT[PGARG_WHR_CONDS])

        select_script, select_values = pg_handler.pg_prepare_select_script(
            _tbl_name,
            where_conditions,
            select_columns
        )
        self.cur.execute(select_script, select_values)
        LOGGER.debug(("(SQL) Select and Print Script: {}.").format(
            self.cur.query))

        all_records = self.cur.fetchall()

        for i, record in enumerate(all_records):
            if (_tbl_name == "cu_featuremenu" and
                    FEATURE_NOT_FOUND in str(record)):
                LOGGER.warning("`{}` Code not found => {} => : "
                               .format(_tbl_name,
                                       FEATURE_NOT_FOUND) + str(record))
            else:
                LOGGER.debug("`{}` => {}.".format(
                    _tbl_name, str(record)))

    def delete(self, _tbl_name, **kwargs):
        """Handles table record deletion operation

        Args:
            _tbl_name: involved table
            kwargs:    optional keyword arguments
        """
        where_conditions = kwargs.get(
            PGARG_WHR_CONDS, PGDEF_ARGS_DICT[PGARG_WHR_CONDS])

        (delete_script,
         delete_values) = pg_handler.pg_prepare_table_cleanup_script(
            _tbl_name, where_conditions)

        self.cur.execute(delete_script, delete_values)
        LOGGER.debug(("(SQL) Delete Script: {}.").format(
            self.cur.query))

    def commit(self):
        """commit this transaction"""
        self.conn.commit()

    def rollback(self):
        """rollback this transaction"""
        self.conn.rollback()