generate_postgresql_data.py

import psycopg2
from faker import Faker
from random import randint
import logging
from time import sleep
import sys
from math import floor
from argparse import ArgumentParser, BooleanOptionalAction

parser = ArgumentParser()
parser.add_argument("--clear", action=BooleanOptionalAction)

LOGGER = logging.getLogger(__name__)
LOGGER.setLevel(logging.DEBUG)
handler = logging.StreamHandler()
fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
formatter = logging.Formatter(fmt)
handler.setFormatter(formatter)
LOGGER.addHandler(handler)

USERNAME = "testdb"
PASSWORD = "testdb"
HOSTNAME = "localhost"
PORT = 5432
DATABASE = "testdb"

TEST_TABLE = "test_table"
TEST_TRIGGER = TEST_TABLE + "_update_updated_at"


def setup():
    f"""Create a test table {TEST_TABLE}

    Create a simple test table with 2 data columns & 3 timestamp columns.

    Define a trigger to ensure the updated_at column is set on each
    update to `CURRENT_TIMESTAMP`.
    """

    create_test_table_stmt = f"""
    CREATE TABLE IF NOT EXISTS {TEST_TABLE} (
        id BIGINT GENERATED BY DEFAULT AS IDENTITY
      , col1 TEXT
      , col2 INT
      , created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
      , updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
      , deleted_at TIMESTAMPTZ
    )
    """

    create_procedure_update_updated_at_stmt = """
    CREATE OR REPLACE FUNCTION update_updated_at()
    RETURNS TRIGGER AS $$
    BEGIN
        NEW.updated_at = CURRENT_TIMESTAMP;
        RETURN NEW;
    END;
    $$ language 'plpgsql';
    """

    create_trigger_test_table_updated_at_stmt = f"""
    CREATE TRIGGER {TEST_TRIGGER}
    BEFORE UPDATE ON {TEST_TABLE}
    FOR EACH ROW EXECUTE PROCEDURE update_updated_at();
    """

    check_table_exists_stmt = """
    SELECT COUNT(*) > 0 table_exists
    FROM information_schema.tables
    WHERE table_name = %s
    """

    try:
        cursor = DB.cursor()
        cursor.execute(check_table_exists_stmt, (TEST_TABLE,))
        exists = cursor.fetchone()
        if not exists[0]:
            cursor.execute(create_test_table_stmt)
            cursor.execute(create_procedure_update_updated_at_stmt)
            cursor.execute(create_trigger_test_table_updated_at_stmt)
            logging.info("Created database table %s", TEST_TABLE)
        else:
            logging.info(
                "Database table %s already exists. Skipping table creation",
                TEST_TABLE,
            )
        DB.commit()
    except Exception as e:
        DB.rollback()
        raise e


def generate_activity_test_table():
    """
    Adds, Updates & Deletes records to simulate database activity.

    Delete upto 5% of existing records (set deleted_at to now)
    Update upto 10% of existing records
    Add between 1 - 10 new records
    """

    count_rows_stmt = f"SELECT COUNT(*) nrows FROM {TEST_TABLE}"
    del_row_stmt = f"""
    UPDATE {TEST_TABLE}
    SET deleted_at = CURRENT_TIMESTAMP
    WHERE id = %s
    """

    upd_col1_stmt = f"""
    UPDATE {TEST_TABLE}
    SET col1 = %s
    WHERE id = %s
    """

    sel_active_rand_rows_stmt = f"""
    SELECT id
    FROM {TEST_TABLE}
    WHERE deleted_at IS NULL
    ORDER BY RANDOM()
    LIMIT %s
    """

    ins_stmt = f"""
    INSERT INTO {TEST_TABLE} (col1, col2)
    VALUES (%s, %s)
    """

    pct_rows_to_upd = randint(0, 10) / 100
    pct_rows_to_del = randint(0, 5) / 100

    try:
        cursor = DB.cursor()
        cursor.execute(count_rows_stmt)
        num_rows = cursor.fetchone()[0]
        LOGGER.info("There are %i rows in %s", num_rows, TEST_TABLE)
        DB.commit()
        num_to_del = floor(num_rows * pct_rows_to_del)
        num_to_upd = floor((num_rows - num_to_del) * pct_rows_to_upd)
        num_to_ins = randint(1, 10)
        LOGGER.info(
            'Simulating activity on %s: {"deletes": %i, "updates": %i, "inserts": %i}',
            TEST_TABLE,
            num_to_del,
            num_to_upd,
            num_to_ins,
        )
        if num_to_del > 0:
            cursor.execute(sel_active_rand_rows_stmt, (num_to_del,))
            rows_to_del = cursor.fetchall()
            for row in rows_to_del:
                cursor.execute(del_row_stmt, (row[0],))
            DB.commit()
        if num_to_upd > 0:
            cursor.execute(sel_active_rand_rows_stmt, (num_to_upd,))
            rows_to_upd = cursor.fetchall()
            for row in rows_to_upd:
                cursor.execute(upd_col1_stmt, (fake.name(), row[0]))
            DB.commit()

        for _ in range(num_to_ins):
            record = (fake.name(), randint(1, 100))
            cursor.execute(ins_stmt, record)
        DB.commit()

    except Exception as e:
        DB.rollback()
        raise e


def main():
    f"""Set up the database for test and simulate activity.

    Create a test table called `{TEST_TABLE}, and generate
    inserts, updates and soft deletes for 10 minutes
    or until user presses CTRL+C"""

    setup()
    iterations = 20
    sleep_seconds = 30
    for _ in range(iterations):
        try:
            generate_activity_test_table()
            sleep(sleep_seconds)
        except KeyboardInterrupt:
            print("Program terminated by user.")
            sys.exit(130)
    print(f"Program terminated after {iterations * sleep_seconds} seconds")


def teardown():
    cursor = DB.cursor()
    cursor.execute(f"DROP TABLE IF EXISTS {TEST_TABLE} CASCADE")


if __name__ == "__main__":
    args = parser.parse_args()
    DB = psycopg2.connect(
        host=HOSTNAME, port=PORT, user=USERNAME, password=PASSWORD, dbname=DATABASE
    )
    fake = Faker()
    if args.clear:
        teardown()
    else:
        main()