PostgreSQL drop/rename race condition

droprace.py

import logging
import os
import sys
import time
import threading
import psycopg2

"""
This script demonstrates the following race condition under PostgreSQL:
One connection, in transaction, drops table "mytable," then renames
"newtable" to "mytable."  A second connection queries "mytable" after the
first connection issues the DROP.  Since the first connection already has
an ACCESS EXCLUSIVE lock on the table, the second connection blocks until
the first transaction is committed.  Then under some versions of
PostgreSQL, the query fails with a message 'could not open relation with
OID nnn.'

This produces an error with PostgreSQL 8.4 (on ubuntu lucid) and
PostgreSQL 9.1 (on debian wheezy) but works with no error with
PostgreSQL 9.3 (on ubuntu trusty).

Even on the platforms where it generates an error, it runs without
error if thread a locks pg_class with ACCESS EXCLUSIVE before dropping
the table IFF thread b has not already queried the table.

The implication is that by the time it blocks waiting for access to the
table, the querying thread already refers to the table solely by oid.
Locking pg_class presumably prevents it from even discovering the oid,
but only if it hasn't already looked it up already.  In that case, it
apparently has cached the oid and locking pg_class subsequently has no
effect in this regard.  It seems that 9.3 maintains the indirection of
reference by relname until after the required lock on the table is
actually acquired.
"""

mutex = threading.Lock()

def execute(name, cursor, statement, params=None):
    logging.info('%s: %s' % (name, statement))
    if cursor:
        cursor.execute(statement, params)

def query_thread(dbname, preload):
    conn_b = psycopg2.connect(database=dbname)
    conn_b.autocommit = True
    cur_b = conn_b.cursor()

    with mutex:
        if 'yes' == preload:
            execute('b', cur_b, 'SELECT * from mytable')
            logging.info('b fetched: %r' % cur_b.fetchall())
        else:
            time.sleep(1)

    time.sleep(1)
    with mutex:
        try:
            execute('b', cur_b, 'SELECT * from mytable')
        except:
            logging.error('b caught exception', exc_info=True)
        else:
            logging.info('b fetched: %r' % cur_b.fetchall())
    conn_b.rollback()

    cur_b.close()
    conn_b.close()


def race(dbname, lock_pg_class=False, preload=False):
    conn_a = psycopg2.connect(database=dbname)
    cur_a = conn_a.cursor()

    with mutex:
        t = threading.Thread(target=query_thread, args=(dbname, preload))
        t.start()

        execute('a', cur_a, 'DROP TABLE IF EXISTS mytable')
        execute('a', cur_a, 'CREATE TABLE mytable (i integer)')
        for i in range(3):
            execute('a', cur_a, 'INSERT INTO mytable (i) VALUES (%s)', [i])
        conn_a.commit()

    time.sleep(1)
    with mutex:
        execute('a', cur_a, 'CREATE TABLE newtable AS TABLE mytable WITH NO DATA')
        for i in range(4, 7):
            execute('a', cur_a, 'INSERT INTO newtable (i) VALUES (%s)', [i])
        if 'yes' == lock_pg_class:
            execute('a', cur_a, 'LOCK pg_class IN ACCESS EXCLUSIVE MODE')
        execute('a', cur_a, 'DROP TABLE mytable')
        logging.info("now a has exclusive lock on the table")

    time.sleep(2)
    execute('a', cur_a, 'ALTER TABLE newtable RENAME TO mytable')
    logging.info("a COMMIT")
    conn_a.commit()

    t.join()

    cur_a.execute('DROP TABLE mytable')
    conn_a.commit()
    cur_a.close()
    conn_a.close()

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    if len(sys.argv) < 2:
        print """usage: python %s <dbname> [lock pg_class] [query first]

            lock pg_class: 'yes' or 'no'
                           Whether or not to lock pg_class before
                           executing DROP TABLE statement.

            query first: 'yes' or 'no'
                          Whether or not querying session should
                          query the table before the modifying
                          session locks the table.
            """ % os.path.basename(sys.argv[0])
    else:
        race(*sys.argv[1:])

droprace2.py

import os
import sys
import time
from textwrap import dedent
from subprocess import Popen, PIPE

#  This script is the same as droprace.py, but uses the psql client
#  in child processes instead of psycopg2/threading

def race(dbname, lock_pg_class=False, preload=False):
    sql_a = Popen(['psql', '-a', dbname], stdin=PIPE)
    sql_b = Popen(['psql', '-a', dbname], stdin=PIPE)

    sql_a.stdin.write(dedent("""
            -- A
            DROP TABLE IF EXISTS mytable;
            CREATE TABLE mytable (i integer);
            INSERT INTO mytable (i) VALUES (1), (2), (3);
            """))

    time.sleep(1)
    if 'yes' == preload:
        sql_b.stdin.write('-- B\nSELECT * from mytable;\n')
    time.sleep(1)
    sql_a.stdin.write(dedent("""
            -- A
            BEGIN;
            CREATE TABLE newtable AS TABLE mytable WITH NO DATA;
            INSERT INTO newtable (i) VALUES (4), (5), (6);
            """))
    if 'yes' == lock_pg_class:
        sql_a.stdin.write('LOCK pg_class IN ACCESS EXCLUSIVE MODE;\n')
    sql_a.stdin.write(dedent("""
            DROP TABLE mytable;
            -- now A has exclusive lock on the table
            --
            """))

    time.sleep(1)
    sql_b.stdin.write(dedent("""
            -- B (this will block until A commits)
            SELECT * from mytable;
            """))
    time.sleep(1)

    sql_a.stdin.write(dedent("""
            -- A
            ALTER TABLE newtable RENAME TO mytable;
            COMMIT;
            """))

    time.sleep(1)
    sql_a.stdin.write('DROP TABLE mytable;\n')
    sql_a.stdin.close()
    sql_b.stdin.close()
    sql_a.wait()
    sql_b.wait()

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print """usage: python %s <dbname> [lock pg_class] [query first]

            lock pg_class: 'yes' or 'no'
                           Whether or not to lock pg_class before
                           executing DROP TABLE statement.

            query first: 'yes' or 'no'
                          Whether or not querying session should
                          query the table before the modifying
                          session locks the table.
            """ % os.path.basename(sys.argv[0])
    else:
        race(*sys.argv[1:])