mirror of
https://github.com/psycopg/psycopg2.git
synced 2024-11-27 03:13:43 +03:00
253 lines
7.7 KiB
Python
253 lines
7.7 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
|
|
import psycopg2
|
|
from psycopg2 import extensions
|
|
|
|
import select
|
|
import StringIO
|
|
|
|
import sys
|
|
if sys.version_info < (3,):
|
|
import tests
|
|
else:
|
|
import py3tests as tests
|
|
|
|
|
|
class AsyncTests(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.sync_conn = psycopg2.connect(tests.dsn)
|
|
self.conn = psycopg2.connect(tests.dsn, async=True)
|
|
|
|
state = psycopg2.extensions.POLL_WRITE
|
|
while state != psycopg2.extensions.POLL_OK:
|
|
if state == psycopg2.extensions.POLL_WRITE:
|
|
select.select([], [self.conn.fileno()], [])
|
|
elif state == psycopg2.extensions.POLL_READ:
|
|
select.select([self.conn.fileno()], [], [])
|
|
state = self.conn.poll()
|
|
|
|
curs = self.conn.cursor()
|
|
curs.execute('''
|
|
CREATE TEMPORARY TABLE table1 (
|
|
id int PRIMARY KEY
|
|
)''')
|
|
self.conn.commit()
|
|
|
|
def tearDown(self):
|
|
self.sync_conn.close()
|
|
self.conn.close()
|
|
|
|
def wait_for_query(self, cur):
|
|
state = cur.poll()
|
|
while state != psycopg2.extensions.POLL_OK:
|
|
if state == psycopg2.extensions.POLL_READ:
|
|
select.select([cur.fileno()], [], [])
|
|
elif state == psycopg2.extensions.POLL_WRITE:
|
|
select.select([], [cur.fileno()], [])
|
|
state = cur.poll()
|
|
|
|
def test_wrong_execution_type(self):
|
|
cur = self.conn.cursor()
|
|
sync_cur = self.sync_conn.cursor()
|
|
|
|
self.assertRaises(psycopg2.ProgrammingError, cur.execute,
|
|
"select 'a'", async=False)
|
|
self.assertRaises(psycopg2.ProgrammingError, sync_cur.execute,
|
|
"select 'a'", async=True)
|
|
|
|
# but this should work anyway
|
|
sync_cur.execute("select 'a'", async=False)
|
|
cur.execute("select 'a'", async=True)
|
|
|
|
def test_async_select(self):
|
|
cur = self.conn.cursor()
|
|
self.assertFalse(self.conn.executing())
|
|
cur.execute("select 'a'")
|
|
self.assertTrue(self.conn.executing())
|
|
|
|
self.wait_for_query(cur)
|
|
|
|
self.assertFalse(self.conn.executing())
|
|
self.assertEquals(cur.fetchone()[0], "a")
|
|
|
|
def test_async_callproc(self):
|
|
cur = self.conn.cursor()
|
|
try:
|
|
cur.callproc("pg_sleep", (0.1, ), True)
|
|
except psycopg2.ProgrammingError:
|
|
# PG <8.1 did not have pg_sleep
|
|
return
|
|
self.assertTrue(self.conn.executing())
|
|
|
|
self.wait_for_query(cur)
|
|
self.assertFalse(self.conn.executing())
|
|
self.assertEquals(cur.fetchall()[0][0], '')
|
|
|
|
def test_async_after_async(self):
|
|
cur = self.conn.cursor()
|
|
cur2 = self.conn.cursor()
|
|
|
|
cur.execute("insert into table1 values (1)")
|
|
|
|
# an async execute after an async one blocks and waits for completion
|
|
cur.execute("select * from table1")
|
|
self.wait_for_query(cur)
|
|
|
|
self.assertEquals(cur.fetchall()[0][0], 1)
|
|
|
|
cur.execute("delete from table1")
|
|
self.wait_for_query(cur)
|
|
|
|
cur.execute("select * from table1")
|
|
self.wait_for_query(cur)
|
|
|
|
self.assertEquals(cur.fetchone(), None)
|
|
|
|
def test_fetch_after_async(self):
|
|
cur = self.conn.cursor()
|
|
cur.execute("select 'a'")
|
|
|
|
# a fetch after an asynchronous query blocks and waits for completion
|
|
self.assertEquals(cur.fetchall()[0][0], "a")
|
|
|
|
def test_rollback_while_async(self):
|
|
cur = self.conn.cursor()
|
|
|
|
cur.execute("select 'a'")
|
|
|
|
# a rollback blocks and should leave the connection in a workable state
|
|
self.conn.rollback()
|
|
self.assertFalse(self.conn.executing())
|
|
|
|
# try a sync cursor first
|
|
sync_cur = self.sync_conn.cursor()
|
|
sync_cur.execute("select 'b'")
|
|
self.assertEquals(sync_cur.fetchone()[0], "b")
|
|
|
|
# now try the async cursor
|
|
cur.execute("select 'c'")
|
|
self.wait_for_query(cur)
|
|
self.assertEquals(cur.fetchmany()[0][0], "c")
|
|
|
|
def test_commit_while_async(self):
|
|
cur = self.conn.cursor()
|
|
|
|
cur.execute("insert into table1 values (1)")
|
|
|
|
# a commit blocks
|
|
self.conn.commit()
|
|
self.assertFalse(self.conn.executing())
|
|
|
|
cur.execute("select * from table1")
|
|
self.wait_for_query(cur)
|
|
self.assertEquals(cur.fetchall()[0][0], 1)
|
|
|
|
cur.execute("delete from table1")
|
|
self.conn.commit()
|
|
|
|
cur.execute("select * from table1")
|
|
self.wait_for_query(cur)
|
|
self.assertEquals(cur.fetchone(), None)
|
|
|
|
def test_set_parameters_while_async(self):
|
|
prev_encoding = self.conn.encoding
|
|
cur = self.conn.cursor()
|
|
|
|
cur.execute("select 'c'")
|
|
self.assertTrue(self.conn.executing())
|
|
|
|
# getting transaction status works
|
|
self.assertEquals(self.conn.get_transaction_status(),
|
|
extensions.TRANSACTION_STATUS_ACTIVE)
|
|
self.assertTrue(self.conn.executing())
|
|
|
|
# this issues a ROLLBACK internally
|
|
self.conn.set_client_encoding("LATIN1")
|
|
|
|
self.assertFalse(self.conn.executing())
|
|
self.assertEquals(self.conn.encoding, "LATIN1")
|
|
|
|
self.conn.set_client_encoding(prev_encoding)
|
|
|
|
def test_reset_while_async(self):
|
|
prev_encoding = self.conn.encoding
|
|
# pick something different than the current encoding
|
|
new_encoding = (prev_encoding == "LATIN1") and "UTF8" or "LATIN1"
|
|
|
|
self.conn.set_client_encoding(new_encoding)
|
|
|
|
cur = self.conn.cursor()
|
|
cur.execute("select 'c'")
|
|
self.assertTrue(self.conn.executing())
|
|
|
|
self.conn.reset()
|
|
self.assertFalse(self.conn.executing())
|
|
self.assertEquals(self.conn.encoding, prev_encoding)
|
|
|
|
def test_async_iter(self):
|
|
cur = self.conn.cursor()
|
|
|
|
cur.execute("insert into table1 values (1), (2), (3)")
|
|
self.wait_for_query(cur)
|
|
cur.execute("select id from table1 order by id")
|
|
|
|
# iteration just blocks
|
|
self.assertEquals(list(cur), [(1, ), (2, ), (3, )])
|
|
self.assertFalse(self.conn.executing())
|
|
|
|
def test_copy_while_async(self):
|
|
cur = self.conn.cursor()
|
|
cur.execute("select 'a'")
|
|
|
|
# copy just blocks
|
|
cur.copy_from(StringIO.StringIO("1\n3\n5\n\\.\n"), "table1")
|
|
|
|
cur.execute("select * from table1 order by id")
|
|
self.assertEquals(cur.fetchall(), [(1, ), (3, ), (5, )])
|
|
|
|
def test_async_executemany(self):
|
|
cur = self.conn.cursor()
|
|
self.assertRaises(
|
|
psycopg2.ProgrammingError,
|
|
cur.executemany, "insert into table1 values (%s)", [1, 2, 3])
|
|
|
|
def test_async_scroll(self):
|
|
cur = self.conn.cursor()
|
|
cur.execute("insert into table1 values (1), (2), (3)")
|
|
self.wait_for_query(cur)
|
|
cur.execute("select id from table1 order by id")
|
|
|
|
# scroll blocks, but should work
|
|
cur.scroll(1)
|
|
self.assertFalse(self.conn.executing())
|
|
self.assertEquals(cur.fetchall(), [(2, ), (3, )])
|
|
|
|
cur = self.conn.cursor()
|
|
cur.execute("select id from table1 order by id")
|
|
|
|
cur2 = self.conn.cursor()
|
|
self.assertRaises(psycopg2.ProgrammingError, cur2.scroll, 1)
|
|
|
|
self.assertRaises(psycopg2.ProgrammingError, cur.scroll, 4)
|
|
|
|
cur = self.conn.cursor()
|
|
cur.execute("select id from table1 order by id")
|
|
cur.scroll(2)
|
|
cur.scroll(-1)
|
|
self.assertEquals(cur.fetchall(), [(2, ), (3, )])
|
|
|
|
def test_async_dont_read_all(self):
|
|
cur = self.conn.cursor()
|
|
cur.execute("select 'a'; select 'b'")
|
|
|
|
# fetch the result
|
|
self.wait_for_query(cur)
|
|
|
|
# it should be the result of the second query
|
|
self.assertEquals(cur.fetchone()[0][0], "b")
|
|
|
|
def test_suite():
|
|
return unittest.TestLoader().loadTestsFromName(__name__)
|