#!/usr/bin/env python
# -*- coding: utf-8 -*-
# vim:set ts=4 sts=4 sw=4 expandtab:

"""

The UPDATE statement below should always return only 1 row, due to the LIMIT 1
clause in the sub-select and unique values of the joined id column.

Most of the time it correctly returns only one row. But when run repeatedly,
after on average about 2 minutes, it will return two rows with different
id values. This should never happen.

This problem has been repeatedly reproduced on the following platforms:

PG 9.3.1  on Debian Wheezy, Opteron 2347, Python 2.7.3, psycopg2 2.5.1
PG 9.0.13 on Debian Wheezy, Opteron 2347, Python 2.7.3, psycopg2 2.5.1
PG 9.1.10 on Ubuntu 13.04, Core 2 E8400, Python 2.7.4, psycopg2 2.4.5
PG 9.0.14 on OS X 10.8.5, Core i5, Python 2.7.2, psycopg2 2.5.1

The Debian packages are from the PGDG APT repository. The Ubuntu package is the
distribution default. The OS X package is from MacPorts.

The Ubuntu and OS X machines were running the default PG config.

The issue could not be reproduced when autovacuum is disabled!

Run the script with a PG_BUG_CONNSTR environment variable containing a
postgresql connection string to an empty database:

$ PG_BUG_CONNSTR="dbname=pgbug" python pgbug.py

Make sure the psycopg2 module is available.

"""

import os
import psycopg2


def test_bug():
    conn = psycopg2.connect(os.environ['PG_BUG_CONNSTR'])
    conn.autocommit = True

    cur = conn.cursor()
    cur.execute("DROP TABLE IF EXISTS t1")
    cur.execute("CREATE TABLE t1 (id INTEGER)")

    for i in range(1, 1000000):
        update(cur, i)


def update(cur, i):
    cur.execute("DELETE FROM t1")
    cur.execute("INSERT INTO t1 VALUES (0), (1)")
    cur.execute("""
        UPDATE t1 SET id = t1.id
        FROM (SELECT id FROM t1 LIMIT 1 FOR UPDATE) AS subset
        WHERE t1.id = subset.id
        RETURNING t1.id
    """)

    rows = cur.fetchall()
    print(rows, i)
    assert(len(rows) == 1)


if __name__ == '__main__':
    test_bug()