# Copyright (c) 2019 Mateusz Pawlik
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

# ------------------------------------------------------------------------------
# DESCRIPTION
# ------------------------------------------------------------------------------
# This script issues multiple concurrent transaction to a database. The main
# purpose is to play with PostgreSQL's commands to find and cancel blocking
# transactions.
# 
# Most of the code is takes from the cookie jar code.
# 
# To keep the code short, many comments are deleted.
# ------------------------------------------------------------------------------

import psycopg2
import psycopg2.errorcodes
import random
from multiprocessing import Pool
import argparse
import sys
import os

# Initialize the table the we will query.
def db_init(num_rows, max_value):
    try:
        db = psycopg2.connect(default_conn_string)
        cur = db.cursor()
        cur.execute("DROP TABLE IF EXISTS blocking_data;")
        cur.execute("CREATE TABLE blocking_data (\
                        data_id integer,\
                        data_value float\
                    );")
        db.commit()
        # Generate random integer in PostgreSQL.
        # www.postgresqltutorial.com/postgresql-random-range/
        cur.execute("INSERT INTO blocking_data (\
                        SELECT generate_series(0,%s), floor(random() * (%s + 1))::int);",\
                        (num_rows - 1, max_value))
        db.commit()
        cur.close()
        db.close()
    except Exception as e:
        print(type(e))
        print(e.args)
        sys.exit("Error occurred while initializing the database.")

# Executes a select statement.
# We use `FOR UPDATE` keyword to explicitely lock rows and cause queries to
# wait for each other.
# 
# NOTE: This leads to multiple deadlocks.
def do_something(max_value, debug):
    proc = os.getpid()
    result = (proc, -1)

    try:
        db = psycopg2.connect(default_conn_string)
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        return (proc, -1)

    try:
        cur = db.cursor()
        
        # Select random rows by range of `data_value` values.
        start = random.randrange(max_value + 1)
        stop = random.randrange(start, max_value + 1)
        cur.execute(
                "SELECT * FROM blocking_data WHERE data_value >= %s AND data_value <= %s FOR UPDATE;",\
                (start, stop)
        )
        
        # If the query executes, return number of rows as part of the result.
        result = (proc, len(cur.fetchall()))
        db.commit()
    # Catch `deadlock_detected` and `query_canceled` errors and return their
    # names as part of the result.
    except psycopg2.OperationalError as e:
        if debug:
            print(type(e))
            print(e.args)
        if e.pgcode == psycopg2.errorcodes.DEADLOCK_DETECTED:
            result = (proc, 'deadlock detected')
        elif e.pgcode == psycopg2.errorcodes.QUERY_CANCELED:
            result = (proc, 'query cancelled')
        else:
            result = (proc, -1)
    # Catch other errors.
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        db.rollback()
        result = (proc, -1)

    try:
        cur.close()
        db.close()
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        return (proc, -1)

    return result

# ------------------------------------------------------------------------------
# EXECUTE SIMULATOR
# ------------------------------------------------------------------------------
if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Advanced Databases - Blocking queries simulator. \
                     Description can be found in the source code.')
    parser.add_argument(
        '--num-rows', '-r',
        type=int,
        dest='num_rows',
        help='Number of rows (default: 10).',
        default=10)
    parser.add_argument(
        '--max-value', '-m',
        type=int,
        dest='max_value',
        help='Max value in a row (default: 5).',
        default=5)
    parser.add_argument(
        '--num-workers', '-w',
        type=int,
        dest='num_workers',
        help='Number of workers (default: 2).',
        default=2)
    parser.add_argument(
        '--num-queries', '-q',
        type=int,
        dest='num_queries',
        help='Number of issued queries (default: 10).',
        default=10)
    parser.add_argument(
        '--connection-string',
        type=str,
        dest='conn_string',
        action='store',
        help="Database connection string (default: empty string)",
        default="")
    parser.add_argument(
        '-d', '--debug',
        dest='debug',
        action='store_true',
        help="Debug mode - prints database errors if any.")
    args = parser.parse_args()
    
    default_conn_string = args.conn_string
    num_rows = args.num_rows
    max_value = args.max_value
    num_workers = args.num_workers
    num_queries = args.num_queries
    
    debug=args.debug

    db_init(num_rows, max_value)
    
    # NOTE: Tested configuration `python3 blocking.py -r 1000000 -m 20 -q 20 -w 10 -d`
    #       results in blocking transactions but also deadlocks.
    #       Deadlocks seem to be handled internally by PostgreSQL.
    # 
    #       In `psql`, you can use the following commands:
    #       `SHOW deadlock_timeout;` Shows time to wait on a lock before checking to see if there is a deadlock condition.
    #       `SET deadlock_timeout = '5s';` Sets the time to 5 seconds - only for the current process.
    
    # NOTE: Set `deadlock_timeout` for al the workers such that we can test cancelling blocking transactions.
    
    # IDEA: If transaction waits too long, cancell the blocking using a student's function.
    
    pool = Pool(num_workers)
    multiple_results = [
        pool.apply_async(do_something, (max_value, debug))
        for _ in range(num_queries)
    ]
    pool.close()
    pool.join()
    # Sum up the number of retries.
    # all_return_values = []
    for r in multiple_results:
        res = r.get()
        print(res)
# ------------------------------------------------------------------------------