ysql_select_for_update.py

# Dependencies:
# On CentOS you can install psycopg2 thus:
#
#  sudo yum install postgresql-libs
#  sudo yum install python-psycopg2

import psycopg2
from threading import Thread,Semaphore
import random
import time

# Test Params:
num_items=1000
percentage_red=20  # about 20% rows are red and remaining are green.

cluster_ip="localhost"
cluster_ip="172.151.29.251"

connect_string="host={} dbname=yugabyte user=yugabyte port=5433".format(cluster_ip)

def create_table():
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor()
  cur.execute("""DROP TABLE IF EXISTS items""");
  print("Dropped (if exists): items table")
  print("====================")
  cur.execute("""
        CREATE TABLE IF NOT EXISTS items (
           id     integer,
           color  text,
           PRIMARY KEY(id)
        )
    """)
  print("Created items table.")
  # cur.execute("""CREATE INDEX IF NOT EXISTS color_idx ON items(color)""")
  # print("Created color_idx on table")

  print("Sleeping 30 seconds for tablets to balance");
  time.sleep(30)
  cur.close()
  conn.close()

def init_rows():
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor();
  print("Populating rows")

  for idx in range(num_items):
    if (random.randint(1, 100) <= percentage_red):
      color = 'red'
    else:
      color = 'green'
    cur.execute("""INSERT INTO items(id, color) VALUES (%s, %s)""", (idx, color))
  print("====================")
  cur.close()
  conn.close()

def select_rows(num_rows_to_update, for_update):

  conn = psycopg2.connect(connect_string)
  cur = conn.cursor()
  b_delta = 0
  s_delta = 0
  c_delta = 0
  for idx in range(num_rows_to_update):
    t1 = time.time()
    cur.execute("BEGIN")
    t2 = time.time()
    b_delta = b_delta + (t2 - t1)

    t1 = time.time()
    if (for_update):
      cur.execute("SELECT id, color FROM items WHERE id = %s FOR UPDATE", (idx, ))
    else:
      cur.execute("SELECT id, color FROM items WHERE id = %s", (idx, ))
    t2 = time.time()
    s_delta = s_delta + (t2 - t1)

    t1 = time.time()
    cur.execute("COMMIT")
    t2 = time.time()
    c_delta = c_delta + (t2 - t1)


  print("Avg BEGIN Time: %s ms" % (((b_delta) * 1000) / num_rows_to_update))

  if (for_update):
    print("Avg SELECT .. FOR UPDATE Time: %s ms" % (((s_delta) * 1000) / num_rows_to_update))
  else:
    print("Avg SELECT Time: %s ms" % (((s_delta) * 1000) / num_rows_to_update))

  print("Avg COMMIT Time: %s ms" % (((c_delta) * 1000) / num_rows_to_update))

  cur.close()
  conn.close()

# Main
create_table()
init_rows()
select_rows(100, False) # Simple SELECT
select_rows(100, True)  # SELECT .. FOR UPDATE

Sample out against a RF=3 (multi-zone cluster):

$ python ~/notes/ysql_select_for_update.py
Dropped (if exists): items table
====================
Created items table.
Created name_idx on table
Sleeping 30 seconds for tablets to balance
Populating rows
====================
Avg BEGIN Time: 0.186944007874 ms
Avg SELECT Time: 1.05236291885 ms
Avg COMMIT Time: 0.187420845032 ms
Avg BEGIN Time: 0.226678848267 ms
Avg SELECT .. FOR UPDATE Time: 2.54352092743 ms 
Avg COMMIT Time: 0.870678424835 ms