ody_migr_db_handler.py

#!/usr/bin/env python
""" Module for safe sql generation and provides psycopg2 context manager.

This module contains class PGSession for psycopg2 connection session
handling and database utility functions for sql query generation.
psycopg2.sql module is being used to generate all, safe sql statements.
"""


import os
import psycopg2
import logging
# psycopg2's SQL composer against sql injection
from psycopg2 import sql
from ody_migr_utils import get_valid_json

LOGGER = logging.getLogger(__name__)

# import database environment variables
PG_DB_HOST = os.environ.get("DATABASE_HOST")
PG_DB_PORT = os.environ.get("DATABASE_PORT")
PG_DB_NAME = os.environ.get("DATABASE_NAME")
PG_DB_USER = os.environ.get("DATABASE_USER")
PG_DB_PWD = os.environ.get("DATABASE_PASSWORD")


class PGSession:
    """Context manager for PG database connection

    TODO: investigate connection pooling:
    http://initd.org/psycopg/docs/pool.html
    """

    def __init__(self):
        """initialize database manager"""
        self.conn = None

    def __enter__(self):
        """open database connection"""
        try:
            self.conn = psycopg2.connect(("host={} dbname={} "
                                          "user={} password={}")
                                         .format(
                                             PG_DB_HOST,
                                             PG_DB_NAME,
                                             PG_DB_USER,
                                             PG_DB_PWD
            ))
            return self.conn

        except psycopg2.OperationalError as err:
            raise SystemExit("Unable to connect!{}".format(err))

    def __exit__(self, exception_type, exception_value, _traceback):
        """close database connection"""
        if exception_type:
            LOGGER.error(exception_value)
        if self.conn is not None:
            self.conn.close()


# DATABASE UTILITIES

def _unnest_col(col_identifier):
    """unnest lambda operation"""
    return sql.SQL("unnest( {} )").format(col_identifier)


def pg_prepare_insert_script_unnest(_tbl_name, _columns):
    """Prepare insertion script - efficient unnest operation

    Uses SELECT unnest property to efficiently read column wise values one
    at a time from the {col1: list of values, col2: list of values,
    col3: list of values} formatted dictionary for efficient bulk insertion.

    Example script:
        INSERT INTO psycopg2_insert_time_test
            ("first_name", "last_name", "email", "gender", "ip_address",
            "ssn", "company_name", "race", "department")
        SELECT unnest( %(first_name)s ), unnest( %(last_name)s ),
            unnest( %(email)s ),
            unnest( %(gender)s ),
            unnest( %(ip_address)s ),
            unnest( %(ssn)s ),
            unnest( %(company_name)s ),
            unnest( %(race)s ),
            unnest( %(department)s )

    `values` in cursor.execute() would be a dictionary of following
    structure:
        {
            ...
            email: list of email values,
            gender: list of gender values,
            ...
        }
        All the values (list) must have same number of elements for each
        key (columns).
        For example, first element of each of the key in the dictionary is
        first record to be inserted in the database table, and so on.
    """
    return sql.SQL("INSERT INTO {} ({}) SELECT {}").format(
        sql.SQL(_tbl_name),
        sql.SQL(', ').join(map(sql.Identifier, _columns)),
        sql.SQL(", ").join(
            map(_unnest_col, map(sql.Placeholder, _columns))))


def pg_prepare_insert_single_dict(_tbl_name, _columns, **kwargs):
    """Prepare insertion script - dictionary placeholders

    Same for a single insertion and multiple insertion using executemany()

    Example script:
    INSERT INTO psycopg2_insert_time_test ("first_name", "last_name", "email",
    "gender", "ip_address", "ssn", "company_name", "race", "department")
    VALUES (%(first_name)s, %(last_name)s, %(email)s, %(gender)s,
    %(ip_address)s, %(ssn)s, %(company_name)s, %(race)s, %(department)s)


    `values` for cursor.execute() must be a single record (dict of values)
    `values` for cursor.executemany must be a list/tuple of many records
    (list/tuple of named dict)

    Args:
        _tbl_name:      name of the table to insert
        _columns:  list/tuple of column names
        **kwargs:       returning_col: (optional) specify column name
                        you want to return after the insert execution script
    Returns::
        record insertion script with named placeholders
    """
    returning_col = kwargs.get("returning_col", "")

    insert_script = sql.SQL("INSERT INTO {} ({}) VALUES ({})").format(
        sql.SQL(_tbl_name),
        sql.SQL(', ').join(map(sql.Identifier, _columns)),
        sql.SQL(', ').join(map(sql.Placeholder, _columns)))

    if returning_col == "":
        return insert_script
    else:
        script_returning = sql.SQL("RETURNING {}").format(
            sql.Identifier(returning_col))
        final_script = sql.SQL(' ').join([insert_script, script_returning])
        return final_script


def pg_prepare_insert_single_list(_tbl_name, _columns, **kwargs):
    """Prepare database table insertion script - general placeholders

    Same for a single insertion and multiple insertion using executemany()

    Example script:
    INSERT INTO "psycopg2_insert_time_test" ("ssn", "race", "first_name",
    "last_name", "email", "company_name", "gender", "department", "ip_address")
     VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)

    values` for cursor.execute() must be a single record (list/tuple of values)
    `values` for cursor.executemany must be a list/tuple of many records
    (list/tuple of list/tuple)

    Args:
        _tbl_name:      name of the table to insert
        _column_names:  list/tuple of column names
        **kwargs:       returning_col: (optional) specify column name
                        you want to return after the insert execution script
    Returns::
        record insertion script with general placeholders
    """
    returning_col = kwargs.get("returning_col", "")

    insert_script = sql.SQL("INSERT INTO {} ({}) VALUES ({})").format(
        sql.Identifier(_tbl_name),
        sql.SQL(', ').join(map(sql.Identifier, _columns)),
        sql.SQL(', ').join(sql.Placeholder() * len(_columns))
    )

    if returning_col == "":
        return insert_script
    else:
        script_returning = sql.SQL("RETURNING {}").format(
            sql.Identifier(returning_col))
        final_script = sql.SQL(' ').join([insert_script, script_returning])
        return final_script


def pg_prepare_insert_script_execute_values_dict(_tbl_name, _columns):
    """Prepare insertion script

    Uses single placeholder for list of multiple records as required by
    psycopg2.extra.execute_values()

    Example script:
    INSERT INTO psycopg2_insert_time_test ("first_name", "last_name",
    "email", "gender", "ip_address", "ssn", "company_name", "race",
    "department") VALUES %s

    Example template snippet:
    ( %(first_name)s, %(last_name)s, %(email)s, %(gender)s, %(ip_address)s,
    %(ssn)s, %(company_name)s, %(race)s, %(department)s )

    `argslist` for psycopg2.extra.execute_values is a sequence of dictionary

    Args:
        _tbl_name:      name of the table to insert
        _column_names:  list/tuple of column names

    Returns::
        record insertion script a single placeholders and template placeholder
        for each row
    """
    template = sql.SQL('( {} )').format(
        sql.SQL(', ').join(map(sql.Placeholder, _columns)))
    script = sql.SQL(
        "INSERT INTO {} ({}) VALUES {}").format(
        sql.SQL(_tbl_name),
        sql.SQL(", ").join(map(sql.Identifier, _columns)),
        sql.Placeholder())
    return script, template


def pg_prepare_insert_script_execute_values_list(_tbl_name, _columns):
    """Prepare insertion script

    Uses single placeholder for list of multiple records as required by
    psycopg2.extra.execute_values()

    Example script:
    INSERT INTO psycopg2_insert_time_test ("first_name", "last_name",
    "email", "gender", "ip_address", "ssn", "company_name", "race",
    "department") VALUES %s

    Example template snippet:
    ( %s, %s, %s, %s, %s, %s, %s, %s, %s )

    `argslist` for psycopg2.extra.execute_values is a sequence of sequence

    Args:
        _tbl_name:      name of the table to insert
        _column_names:  list/tuple of column names

    Returns::
        record insertion script a single placeholders and template placeholder
        for each row
    """
    template = sql.SQL('( {} )').format(
        sql.SQL(', ').join(sql.Placeholder() * len(_columns)))

    script = sql.SQL(
        "INSERT INTO {} ({}) VALUES {}").format(
        sql.SQL(_tbl_name),
        sql.SQL(", ").join(map(sql.Identifier, _columns)),
        sql.Placeholder())
    return script, template


def pg_prepare_select_script(_tbl_name,
                             _where_conditions,
                             select_columns,
                             **kwargs):
    """Prepare table selection script

    Args:
        _tbl_name:          name of the table to be created
        _where_conditions:  dict of where clauses
        select_columns:     list of columns name to select

    Returns:
        table selection script, and values for placeholders, if any
    """
    get_count = kwargs.get("count", False)
    assert type(_where_conditions) == dict

    if select_columns == []:
        if get_count:
            select_sql = sql.SQL(
                "SELECT count(*) FROM {}").format(sql.Identifier(_tbl_name))
        else:
            select_sql = sql.SQL(
                "SELECT * FROM {}").format(sql.Identifier(_tbl_name))
    else:
        assert type(select_columns) == list
        select_columns_str = sql.SQL(', ').join(
            [sql.Identifier(i) for i in select_columns])
        select_sql = sql.SQL(
            "SELECT {} FROM {}").format(select_columns_str,
                                        sql.Identifier(_tbl_name))

    where_values = [] if len(_where_conditions) > 0 else None
    where_conditions_sequence = []

    for column_cond, value_cond in _where_conditions.items():
        wh = sql.SQL("{}={}").format(
            sql.Identifier(column_cond), sql.Placeholder())
        where_conditions_sequence.append(wh)
        where_values.append(value_cond)

    where_sql = sql.SQL(' AND ').join(where_conditions_sequence)

    if where_sql != sql.Composed([]):
        select_sql = sql.SQL(' WHERE ').join([select_sql, where_sql])

    return select_sql, where_values


def pg_check_table_exists(_tbl_name):
    """Returns a sql script to check if a table exist in database

    Args:
        _tbl_name:  name of table to check existence of
    """
    query = """
    SELECT EXISTS(
        SELECT 1
        FROM   pg_catalog.pg_class c
        JOIN   pg_catalog.pg_namespace n ON n.oid=c.relnamespace
        WHERE  n.nspname='{}'
        AND    c.relname='{}'
        AND    c.relkind='r'
    )
    """.format('public', _tbl_name)
    return query, None


def pg_prepare_create_script(_tbl_name, _col_name_type_pair_str):
    """Prepare database table creation script

    Args:
        _tbl_name: name of the table to be created
        _col_name_type_pair_str: name and type of the columns

    Returns:
        table creation script, and values for placeholders, if any
    """
    script = sql.SQL("CREATE TABLE {} ({})").format(
        sql.Identifier(_tbl_name), _col_name_type_pair_str)
    return script, None


def pg_prepare_table_cleanup_script(_tbl_name, _where_conditions):
    """Prepare delete record from table script.

    Args:
        _tbl_name:      table name to delete the records from
        _where_conditions: dictionary of where conditions

    Returns:
        table cleanup script, and values for placeholders, if any
    """
    assert type(_where_conditions) == dict

    delete_script = sql.SQL("DELETE FROM {}").format(
        sql.Identifier(_tbl_name))

    values = [] if len(_where_conditions) > 0 else None
    where_conditions_sequence = []

    for column_cond, value_cond in _where_conditions.items():
        wh = sql.SQL("{}={}").format(
            sql.Identifier(column_cond), sql.Placeholder())
        where_conditions_sequence.append(wh)
        values.append(value_cond)

    where_sql = sql.SQL(' AND ').join(where_conditions_sequence)

    if where_sql != sql.Composed([]):
        delete_script = sql.SQL(' WHERE ').join([delete_script, where_sql])

    return delete_script, values


def pg_prepare_diagnostic_error(_diag):
    """Generate a detailed diagnostic report on database failure

    Args:
        _diag: Diagnostics object exposed by psycopg2.DatabaseError

    Returns:
        detailed error message (compact json string)
    """
    diag_dict = {
        "Severity": _diag.severity,
        "Column name": _diag.column_name,
        "Constraint name": _diag.constraint_name,
        "Context": _diag.context,
        "Datatype name": _diag.datatype_name,
        "Internal position": _diag.internal_position,
        "Internal query": _diag.internal_query,
        "Message detail": _diag.message_detail,
        "Message hint": _diag.message_hint,
        "Message primary": _diag.message_primary,
        "Schema name": _diag.schema_name,
        "Source file": _diag.source_file,
        "Source function": _diag.source_function,
        "Source line": _diag.source_line,
        "Sqlstate": _diag.sqlstate,
        "Statement position": _diag.statement_position,
        "Table name": _diag.table_name
    }
    return get_valid_json(diag_dict)