ody_migr_benchmarks.py

#!/usr/bin/env python
"""Script to profile execution times for multiple bulk insertion methods

Findings from this script align to the conclusions made in the following
tutorial: https://trvrm.github.io/bulk-psycopg2-inserts.html
Conclusion:

Using insertion script with select + unnest option is the most efficient way
of bulk insertion. Using unnest to load multiple rows simultaneously has the
following advantages:

- It is significantly faster than a regular insert loop, especially when
  inserting thousands of rows. Also aster with psycopg2.extra.
  execute_values and inserting list of values with a normal single insertion
  script with cursor.mogrify.
- The benefits of using unnest() increase at least up to 50,000 rows
- It still allows us to write straightforward parameterised and safe SQL
  with no (injection prone) string concatenations - we use psycopg2.sql
  module to generate safe sql scripts.

But using `unnest` requires explicit type conversion incase of NULL values.
A little less efficient (but still order of magnitude faster than interative
approach) but more pythonic option is to use `execute_values`. `execute_values`
require a template of placeholders to be specified beforehand which is why
we should explicitly specify default values for the missing columns during
the source to destionation table mappings.

psycopg2 API resources: http://initd.org/psycopg/docs/extras.html#fast-exec
"""


import os
import ody_migr_db_handler as pg_handler

import pandas
import psycopg2
import psycopg2.extras
from psycopg2.extras import execute_values, execute_batch
# psycopg2's SQL composer against sql injection
from psycopg2 import sql
import time
import contextlib
import json


# decorator to compute execution time of the function being tested
@contextlib.contextmanager
def timer(name="duration"):
    '''Utility function for timing execution'''
    start = time.time()
    yield
    duration = time.time() - start
    print("{}: {:.5f} second(s)".format(name, duration))


# import database environment variables
DATABASE_HOST = os.environ.get("DATABASE_HOST")
DATABASE_PORT = os.environ.get("DATABASE_PORT")
DATABASE_NAME = os.environ.get("DATABASE_NAME")
DATABASE_USER = os.environ.get("DATABASE_USER")
DATABASE_PASSWORD = os.environ.get("DATABASE_PASSWORD")

TEST_TABLE_NAME_STR = "psycopg2_insert_time_test"

# database table setup script
SETUP_SQL = sql.SQL("""
    DROP TABLE IF EXISTS {};

    CREATE TABLE {}(
        id serial primary key,
        first_name char(50) NOT NULL,
        last_name char(50) NOT NULL,
        email char(50),
        gender char(50) NOT NULL,
        ip_address char(50),
        ssn char(50) NOT NULL,
        company_name char(50),
        race char(50) NOT NULL,
        department char(50)
    );

    GRANT ALL ON {} TO odyssey;
""").format(
    sql.SQL(TEST_TABLE_NAME_STR),
    sql.SQL(TEST_TABLE_NAME_STR),
    sql.SQL(TEST_TABLE_NAME_STR)
)

# database cleanup script
CLEANUP_SQL = sql.SQL("""
    DROP TABLE IF EXISTS {};
""").format(
    sql.SQL(TEST_TABLE_NAME_STR)
)


COL_FNAME = "first_name"
COL_LNAME = "last_name"
COL_EMAIL = "email"
COL_GENDER = "gender"
COL_IP = "ip_address"
COL_SSN = "ssn"
COL_CMPNYNAME = "company_name"
COL_RACE = "race"
COL_DPMT = "department"

fixture_columns = [
    COL_FNAME,
    COL_LNAME,
    COL_EMAIL,
    COL_GENDER,
    COL_IP,
    COL_SSN,
    COL_CMPNYNAME,
    COL_RACE,
    COL_DPMT,
]


def pg_prepare_insert_script_cursor_mogrify(_tbl_name,
                                            _columns,
                                            _cursor,
                                            _values_list):
    """Prepare insertion script

    A little arduous but reasonably safe sql with placeholders and
    cursor.mogrify formatted columns and values for bulk insertion.

    Example script:
        INSERT INTO psycopg2_insert_time_test ("first_name", "last_name",
        "email", "gender", "ip_address", "ssn", "company_name", "race",
        "department")
        VALUES
        ('254.70.235.138','Product Management','Noir','Marlyn','Livefish',
        'Female','mnoir0@bizjournals.com','Guamanian','651-05-8691'),
        ('151.11.81.28','Sales','Doxey','Berty','Skaboo','Male',
        'bdoxey1@amazon.com','Malaysian','658-15-9923')
        ...
    """
    col_names = sql.SQL(", ").join(map(sql.Identifier, _columns))
    placeholders = ",".join(["%s"] * len(fixture_columns))

    args_str = ','.join(_cursor.mogrify("({})".format(
        placeholders), x).decode('utf-8') for x in _values_list)

    return sql.SQL("INSERT INTO {} ({}) VALUES {}").format(
        sql.SQL(_tbl_name),
        col_names,
        sql.SQL(args_str)
    )


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

    Uses named placeholders to insert list of records (dictonary)

    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` attribute of cursor.executemany() must be a list of named dict
    """
    return pg_handler.pg_prepare_insert_single_dict(_tbl_name, _columns)


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

    Uses named placeholders to insert list of records (list)

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


    `values` attribute of cursor.executemany() must be a list of multiple
    records - list of list (each record values must have same order
    as _columns)
    """
    return pg_handler.pg_prepare_insert_single_list(_tbl_name, _columns)


def get_data():
    # load the fixture data for testing purpose from disk
    json_file = ("smb://192.168.169.33/devops/"
                 "psycopg2_bulk/MOCK_DATA_FINAL.json")
    fixture_data = {}
    fixture_data_count = 0
    with open(json_file) as jf:
        fixture_data = json.load(jf)
        fixture_data_count = len(fixture_data)
    assert fixture_data_count > 0
    return fixture_data, fixture_data_count


# connection
def connect():
    connection = psycopg2.connect(
        host=DATABASE_HOST,
        database=DATABASE_NAME,
        user=DATABASE_USER,
        password=DATABASE_PASSWORD
    )
    return connection


# execute a script in a transaction
def execute(sql, params={}):
    with connect() as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql, params)


class BulkInsertionTester():
    def __init__(self, count):
        # create a fresh table
        execute(SETUP_SQL)
        assert count > 0, "experimental row count must be greater than zero!"

        fixture_data, fixture_data_count = get_data()
        # prepare data for experiment
        self.data = fixture_data
        if fixture_data_count < count:
            self.count = fixture_data_count
        else:
            self.count = count
            self.data = self.data[:count]

    def slow_multiple_conns_list(self):
        '''Creates a new connection for each insertion'''
        for row in self.data:
            cols = list(row.keys())
            vals = list(row.values())
            # get insertion script to insert values in list/tuple form
            insert_script = pg_handler.pg_prepare_insert_single_list(
                TEST_TABLE_NAME_STR, cols)
            execute(insert_script, vals)

    def slow_multiple_conns_dict(self):
        '''
            Creates a new connection for each insertion
            Insertion script incorporate dictionary values
        '''
        for row in self.data:
            # row is a dictionary of columns and values
            single_insert = pg_handler.pg_prepare_insert_single_dict(
                TEST_TABLE_NAME_STR, fixture_columns)
            execute(single_insert, row)

    def insert_iterative_single_conn_list(self):
        '''
            One connection multiple queries (iterative).
            Insertion script incorporate list/tuple values
        '''
        with connect() as connection:
            with connection.cursor() as cursor:
                for row in self.data:
                    cols = list(row.keys())
                    vals = list(row.values())
                    insert_script = pg_handler.pg_prepare_insert_single_list(
                        TEST_TABLE_NAME_STR, cols)
                    cursor.execute(insert_script, vals)

    def insert_iterative_single_conn_dict(self):
        '''
            One connection multiple queries (iterative).
            Insertion script incorporate dictionary values
        '''
        with connect() as connection:
            with connection.cursor() as cursor:
                for row in self.data:
                    # row is a dictionary
                    single_insert = pg_handler.pg_prepare_insert_single_dict(
                        TEST_TABLE_NAME_STR, fixture_columns)
                    cursor.execute(single_insert, row)

    def insert_executemany_dict(self):
        '''
            One connection multiple queries.
            Uses cursor.executemany() - list of values (dictionary)

        '''
        with connect() as connection:
            with connection.cursor() as cursor:
                executemany_script = pg_prepare_insert_script_executemany_dict(
                    TEST_TABLE_NAME_STR, fixture_columns)
                cursor.executemany(executemany_script, self.data)

    def insert_executemany_list(self):
        '''
            One connection multiple queries.
            Uses cursor.executemany() - list of values (list)
        '''
        values = []
        for row in self.data:
            # still linear, done to match the order of cols and values
            row_vals = [row[col] for col in fixture_columns]
            values.append(row_vals)

        with connect() as connection:
            with connection.cursor() as cursor:
                executemany_script = pg_prepare_insert_script_executemany_list(
                    TEST_TABLE_NAME_STR, fixture_columns)

                cursor.executemany(executemany_script, values)

    def fast_insert_execute_values_dict(self):
        '''
            One connection, one query - multiple values insertions.
            Uses psycopg2.extras.execute_values
        '''
        (insert_query,
         template) = pg_handler.pg_prepare_insert_script_execute_values_dict(
            TEST_TABLE_NAME_STR, fixture_columns)

        with connect() as connection:
            with connection.cursor() as cursor:
                execute_values(
                    cursor,
                    insert_query.as_string(connection),
                    self.data,
                    template=template.as_string(connection),
                    page_size=self.count
                )

    def fast_insert_execute_values_list(self):
        '''
            One connection, one query - multiple values insertions.
            Uses psycopg2.extras.execute_values
        '''
        values = []
        for row in self.data:
            # still linear, done to match the order of cols and values
            row_vals = [row[col] for col in fixture_columns]
            values.append(row_vals)

        (insert_query,
         template) = pg_handler.pg_prepare_insert_script_execute_values_list(
            TEST_TABLE_NAME_STR, fixture_columns)

        with connect() as connection:
            with connection.cursor() as cursor:
                execute_values(cursor,
                               insert_query.as_string(connection),
                               values,
                               template=template.as_string(connection),
                               page_size=self.count
                               )

    def insert_execute_batch_dict(self):
        '''
            One connection, one query - multiple values insertions.
            Uses psycopg2.extras.execute_batch() - list of dict
        '''
        # query similar to executemany()
        insert_query = pg_prepare_insert_script_executemany_dict(
            TEST_TABLE_NAME_STR, fixture_columns)

        with connect() as connection:
            with connection.cursor() as cursor:
                execute_batch(cursor,
                              insert_query,
                              self.data,
                              page_size=self.count
                              )

    def insert_execute_batch_list(self):
        '''
             One connection, one query - multiple values insertions.
            Uses psycopg2.extras.execute_batch() - list of list

        '''
        values = []
        for row in self.data:
            # still linear, done to match the order of cols and values
            row_vals = [row[col] for col in fixture_columns]
            values.append(row_vals)

        with connect() as connection:
            with connection.cursor() as cursor:
                # query similar to executemany()
                insert_query = pg_prepare_insert_script_executemany_list(
                    TEST_TABLE_NAME_STR, fixture_columns)
                execute_batch(cursor,
                              insert_query,
                              values,
                              page_size=self.count
                              )

    def fast_insert_many_mogrify(self):
        '''
        One connection, one query - multiple values insertions.
        Uses cursor.mogrify to validate scripts

        '''
        values = []
        for row in self.data:
            # still linear, done to match the order of cols and values
            row_vals = [row[col] for col in fixture_columns]
            values.append(row_vals)

        with connect() as connection:
            with connection.cursor() as cursor:
                insert_query = pg_prepare_insert_script_cursor_mogrify(
                    TEST_TABLE_NAME_STR, fixture_columns, cursor, values)
                cursor.execute(insert_query)

    def fast_insert_unnest(self):
        """
        One connection, one query - multiple values insertions.
        Insertion script incorporate unnest + dictionary keys.

        `values` can also be passed as:
            first_name = [r['first_name'] for r in self.data]
            last_name = [r['last_name'] for r in self.data]
            email = [r['email'] for r in self.data]
            gender = [r['gender'] for r in self.data]
            ip_address = [r['ip_address'] for r in self.data]
            ssn = [r['ssn'] for r in self.data]
            company_name = [r['company_name'] for r in self.data]
            race = [r['race'] for r in self.data]
            department = [r['department'] for r in self.data]
            values = locals()
        """
        insert_sql = pg_handler.pg_prepare_insert_script_unnest(
            TEST_TABLE_NAME_STR, fixture_columns)

        values = {}
        for col in fixture_columns:
            values[col] = [r[col] for r in self.data]

        with connect() as connection:
            with connection.cursor() as cursor:
                cursor.execute(insert_sql, values)

    def run(self):
        print("")
        with timer('[slow] 1 transaction, 1 insertion '
                   '(a record is list)'):
            self.slow_multiple_conns_list()

        print("")

        with timer('[slow] 1 transaction, 1 insertion '
                   '(a record is a named dict)'):
            self.slow_multiple_conns_dict()

        print("")

        with timer('[normal] Iterative insertion - 1 transaction, '
                   'many insertions (a record is list)'):
            self.insert_iterative_single_conn_list()

        print("")

        with timer('[normal] Iterative insertion - 1 transaction, '
                   'many insertions (a record is named dict)'):
            self.insert_iterative_single_conn_dict()

        print("")

        with timer('[normal] cursor.executemany() (records are list of list)'):
            self.insert_executemany_list()

        print("")

        with timer('[normal] cursor.executemany() '
                   '(records are list of named dict)'):
            self.insert_executemany_dict()

        print("")

        with timer('[fast] psycopg2.extra'
                   '.execute_batch() (values are list of list)'):
            self.insert_execute_batch_list()

        print("")

        with timer('[fast] psycopg2.extra'
                   '.execute_batch() (values are list of named dict)'):
            self.insert_execute_batch_dict()

        print("")

        with timer('[fast] cursor.mogrify and'
                   ' cursor.execute() (many records results to'
                   ' very large script)'):
            self.fast_insert_many_mogrify()

        print("")

        with timer('[fast] psycopg2.extra'
                   '.execute_values() (values are list of list)'):
            self.fast_insert_execute_values_list()

        print("")

        with timer('[fast] psycopg2.extra'
                   '.execute_values() (values are list of named dict)'):
            self.fast_insert_execute_values_dict()

        print("")

        with timer('[fast] unnest and cursor.execute()'):
            self.fast_insert_unnest()

        print("")
        print("****************************************************")
        print("")


# insertion rate table
def fast_insertion_rate(count):
    tester = BulkInsertionTester(count)
    start = time.time()
    tester.fast_insert_unnest()
    duration = time.time() - start
    return count / duration


def normal_insertion_rate(count):
    tester = BulkInsertionTester(count)
    start = time.time()
    tester.insert_iterative_single_conn_dict()
    duration = time.time() - start
    return count / duration


if __name__ == "__main__":
    experiment_row_counts = [
        1, 50, 100, 200, 500, 1000, 2000,
        5000, 10000, 20000, 35000, 50000
    ]

    for rec_count in experiment_row_counts:
        print("Stats for {} row insertions!".format(rec_count))
        tester = BulkInsertionTester(rec_count)
        print("Data count loaded from fixture: {}".format(tester.count))
        tester.run()

    rates = [
        {
            "count": count,
            'rate of insertion (fast)': fast_insertion_rate(count),
            'rate of insertion (normal)': normal_insertion_rate(count)

        }
        for count in experiment_row_counts
    ]

    # display result table (following two lines require `pandas`)
    frame = pandas.DataFrame(rates).set_index('count')
    print(frame)

    # plot (requires matplotlib)
    # frame.plot(logx=True)

    # CLEANUP: drop table
    execute(CLEANUP_SQL)