Create the workload class

We now create the Python class file that describes our workload.

This is going to be a regular .py file which we need to create from scratch.

Fortunately, we can use built-in function dbworkload util gen_stub to generate a skeleton which we can use as a base to get started.

Create the stub file

Execute the following command

dbworkload util gen_stub -i bank.sql

A new file, bank.py, will be created in your directory.

Review the workload class file

The file is pretty long to be pasted here all in one shot, so we will go through it in sections.

Library imports

The first few lines define the imports

import datetime as dt
import psycopg
import random
import time
from uuid import uuid4

These are the most common libraries used when using a workload class. More can be added as needed.

Class init

Here is the definition of class Bank.

class Bank:
    def __init__(self, args: dict):
        # args is a dict of string passed with the --args flag

        self.think_time: float = float(args.get("think_time", 5) / 1000)

        # you can arbitrarely add any variables you want
        self.my_var = 1

        # translation table for efficiently generating a string
        # -------------------------------------------------------
        # make translation table from 0..255 to A..Z, 0..9, a..z
        # the length must be 256
        self.tbl = bytes.maketrans(
            bytearray(range(256)),
            bytearray(
                [ord(b"a") + b % 26 for b in range(113)]
                + [ord(b"0") + b % 10 for b in range(30)]
                + [ord(b"A") + b % 26 for b in range(113)]
            ),
        )

The init adds some convenient examples on how to pass runtime arguments and set class variables. self.tbl is used to generate random string - we will review that in the next sections.

setup()

Next up is setup(). dbworkload invokes this function when it first starts up.

    # the setup() function is executed only once
    # when a new executing thread is started.
    # Also, the function is a vector to receive the excuting threads's unique id and the total thread count
    def setup(self, conn: psycopg.Connection, id: int, total_thread_count: int):
        with conn.cursor() as cur:
            print(
                f"My thread ID is {id}. The total count of threads is {total_thread_count}"
            )
            print(cur.execute(f"select version()").fetchone()[0])

loop()

Function loop() is what is repeatedly executed by dbworkload. In the return list you define the functions you want to execute.

    # the loop() function returns a list of functions
    # that dbworkload will execute, sequentially.
    # Once every func has been executed, loop() is re-evaluated.
    # This process continues until dbworkload exits.
    def loop(self):
        return [
            self.txn_0,
            self.txn_1,
            self.txn_2,
            self.txn_3,
        ]

Utility functions

Here the stub includes some commonly used functions, such as random_str to generate just that.

    #####################
    # Utility Functions #
    #####################
    def __think__(self, conn: psycopg.Connection):
        time.sleep(self.think_time)

    def random_str(self, size: int = 12):
        return (
            random.getrandbits(8 * size)
            .to_bytes(size, "big")
            .translate(self.tbl)
            .decode()
        )

Workload transactions

Next are the stub of the functions generated from the SQL statements in the bank.sql file.

This is the read operation

    def txn_0(self, conn: psycopg.Connection):
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT *
                FROM orders
                WHERE acc_no = %s
                  AND id = %s
                """,
                (
                    # add bind parameter,
                    # add bind parameter, 
                ), 
            ).fetchall()

And below are the transactions related to the order

    def txn_1(self, conn: psycopg.Connection):
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO orders (acc_no, status, amount)
                VALUES (%s, 'Pending', %s) RETURNING id
                """,
                (
                    # add bind parameter,
                    # add bind parameter, 
                ), 
            ).fetchall()

    def txn_2(self, conn: psycopg.Connection):
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT *
                FROM ref_data
                WHERE acc_no = %s
                """,
                (
                    # add bind parameter, 
                ), 
            ).fetchall()

    def txn_3(self, conn: psycopg.Connection):
        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE orders
                SET status = 'Complete'
                WHERE (acc_no,
                       id) = (%s,
                              %s)
                """,
                (
                    # add bind parameter,
                    # add bind parameter, 
                ), 
            )

Helpful tips

Finally, the stub finishes with some tips on how to generate random data

'''
# Quick random generators reminder

# random string of 25 chars
self.random_str(25),

# random int between 0 and 100k
random.randint(0, 100000),

# random float with 2 decimals 
round(random.random()*1000000, 2)

# now()
dt.datetime.utcnow()

# random timestamptz between certain dates,
# expressed as unix ts
dt.datetime.fromtimestamp(random.randint(1655032268, 1759232268))

# random UUID
uuid4()

# random bytes
size = 12
random.getrandbits(8 * size).to_bytes(size, "big")

'''

In the next session, we customize the file to fit our workload.

Psycopg 3 basic usage

It might be a good idea now to refresh how the psycopg driver works. Fortunately, there is a great intro doc in the official website 🚀