test_execute.py

# Copyright (C) 2016-present the asyncpg authors and contributors
# <see AUTHORS file>
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


importasyncio
importasyncpg

fromasyncpgimport_testbaseastb
fromasyncpgimportexceptions


classTestExecuteScript(tb.ConnectedTestCase):

asyncdeftest_execute_script_1(self):
self.assertEqual(self.con._protocol.queries_count, 0)
status=awaitself.con.execute('''
 SELECT 1;

 SELECT true FROM pg_type WHERE false = true;

 SELECT generate_series(0, 9);
 ''')
self.assertEqual(self.con._protocol.queries_count, 1)
self.assertEqual(status, 'SELECT 10')

asyncdeftest_execute_script_2(self):
status=awaitself.con.execute('''
 CREATE TABLE mytab (a int);
 ''')
self.assertEqual(status, 'CREATE TABLE')

try:
status=awaitself.con.execute('''
 INSERT INTO mytab (a) VALUES ($1), ($2)
 ''', 10, 20)
self.assertEqual(status, 'INSERT 0 2')
finally:
awaitself.con.execute('DROP TABLE mytab')

asyncdeftest_execute_script_3(self):
withself.assertRaisesRegex(asyncpg.PostgresSyntaxError,
'cannot insert multiple commands'):

awaitself.con.execute('''
 CREATE TABLE mytab (a int);
 INSERT INTO mytab (a) VALUES ($1), ($2);
 ''', 10, 20)

asyncdeftest_execute_script_check_transactionality(self):
withself.assertRaises(asyncpg.PostgresError):
awaitself.con.execute('''
 CREATE TABLE mytab (a int);
 SELECT * FROM mytab WHERE 1 / 0 = 1;
 ''')

withself.assertRaisesRegex(asyncpg.PostgresError,
'"mytab" does not exist'):

awaitself.con.prepare('''
 SELECT * FROM mytab
 ''')

asyncdeftest_execute_exceptions_1(self):
withself.assertRaisesRegex(asyncpg.PostgresError,
'relation "__dne__" does not exist'):

awaitself.con.execute('select * from __dne__')

asyncdeftest_execute_script_interrupted_close(self):
fut=self.loop.create_task(
self.con.execute('''SELECT pg_sleep(10)'''))

awaitasyncio.sleep(0.2)

self.assertFalse(self.con.is_closed())
awaitself.con.close()
self.assertTrue(self.con.is_closed())

withself.assertRaises(asyncpg.QueryCanceledError):
awaitfut

asyncdeftest_execute_script_interrupted_terminate(self):
fut=self.loop.create_task(
self.con.execute('''SELECT pg_sleep(10)'''))

awaitasyncio.sleep(0.2)

self.assertFalse(self.con.is_closed())
self.con.terminate()
self.assertTrue(self.con.is_closed())

withself.assertRaisesRegex(asyncpg.ConnectionDoesNotExistError,
'closed in the middle'):
awaitfut

self.con.terminate()


classTestExecuteMany(tb.ConnectedTestCase):
defsetUp(self):
super().setUp()
self.loop.run_until_complete(self.con.execute(
'CREATE TABLE exmany (a text, b int PRIMARY KEY)'))

deftearDown(self):
self.loop.run_until_complete(self.con.execute('DROP TABLE exmany'))
super().tearDown()

asyncdeftest_executemany_basic(self):
result=awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])

self.assertIsNone(result)

result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')

self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])

# Empty set
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', ())

result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')

self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])

asyncdeftest_executemany_returning(self):
result=awaitself.con.fetchmany('''
 INSERT INTO exmany VALUES($1, $2) RETURNING a, b
 ''', [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])
self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])
result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')
self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])

# Empty set
awaitself.con.fetchmany('''
 INSERT INTO exmany VALUES($1, $2) RETURNING a, b
 ''', ())
result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')
self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])

# Without "RETURNING"
result=awaitself.con.fetchmany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', [('e', 5), ('f', 6)])
self.assertEqual(result, [])
result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')
self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4), ('e', 5), ('f', 6)
 ])

asyncdeftest_executemany_bad_input(self):
withself.assertRaisesRegex(
exceptions.DataError,
r"invalid input in executemany\(\) argument sequence element #1: "
r"expected a sequence",
 ):
awaitself.con.executemany('''
 INSERT INTO exmany (b) VALUES($1)
 ''', [(0,), {1: 0}])

withself.assertRaisesRegex(
exceptions.DataError,
r"invalid input for query argument \$1 in element #1 of "
r"executemany\(\) sequence: 'bad'",
 ):
awaitself.con.executemany('''
 INSERT INTO exmany (b) VALUES($1)
 ''', [(0,), ("bad",)])

asyncdeftest_executemany_error_in_input_gen(self):
bad_data= ([1/0] forvinrange(10))

withself.assertRaises(ZeroDivisionError):
asyncwithself.con.transaction():
awaitself.con.executemany('''
 INSERT INTO exmany (b)VALUES($1)
 ''', bad_data)

good_data= ([v] forvinrange(10))
asyncwithself.con.transaction():
awaitself.con.executemany('''
 INSERT INTO exmany (b)VALUES($1)
 ''', good_data)

asyncdeftest_executemany_server_failure(self):
withself.assertRaises(exceptions.UniqueViolationError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', [
 ('a', 1), ('b', 2), ('c', 2), ('d', 4)
 ])
result=awaitself.con.fetch('SELECT * FROM exmany')
self.assertEqual(result, [])

asyncdeftest_executemany_server_failure_after_writes(self):
withself.assertRaises(exceptions.UniqueViolationError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', [('a'*32768, x) forxinrange(10)] + [
 ('b', 12), ('c', 12), ('d', 14)
 ])
result=awaitself.con.fetch('SELECT b FROM exmany')
self.assertEqual(result, [])

asyncdeftest_executemany_server_failure_during_writes(self):
# failure at the beginning, server error detected in the middle
pos=0

defgen():
nonlocalpos
whilepos<128:
pos+=1
ifpos<3:
yield ('a', 0)
else:
yield'a'*32768, pos

withself.assertRaises(exceptions.UniqueViolationError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', gen())
result=awaitself.con.fetch('SELECT b FROM exmany')
self.assertEqual(result, [])
self.assertLess(pos, 128, 'should stop early')

asyncdeftest_executemany_client_failure_after_writes(self):
withself.assertRaises(ZeroDivisionError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', (('a'*32768, y+y/y) foryinrange(10, -1, -1)))
result=awaitself.con.fetch('SELECT b FROM exmany')
self.assertEqual(result, [])

asyncdeftest_executemany_timeout(self):
withself.assertRaises(asyncio.TimeoutError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES(pg_sleep(0.1) || $1, $2)
 ''', [('a'*32768, x) forxinrange(128)], timeout=0.5)
result=awaitself.con.fetch('SELECT * FROM exmany')
self.assertEqual(result, [])

asyncdeftest_executemany_timeout_flow_control(self):
event=asyncio.Event()

asyncdeflocker():
test_func=getattr(self, self._testMethodName).__func__
opts=getattr(test_func, '__connect_options__', {})
conn=awaitself.connect(**opts)
try:
tx=conn.transaction()
awaittx.start()
awaitconn.execute("UPDATE exmany SET a = '1' WHERE b = 10")
event.set()
awaitasyncio.sleep(1)
awaittx.rollback()
finally:
event.set()
awaitconn.close()

awaitself.con.executemany('''
 INSERT INTO exmany VALUES(NULL, $1)
 ''', [(x,) forxinrange(128)])
fut=asyncio.ensure_future(locker())
awaitevent.wait()
withself.assertRaises(asyncio.TimeoutError):
awaitself.con.executemany('''
 UPDATE exmany SET a = $1 WHERE b = $2
 ''', [('a'*32768, x) forxinrange(128)], timeout=0.5)
awaitfut
result=awaitself.con.fetch(
'SELECT * FROM exmany WHERE a IS NOT NULL')
self.assertEqual(result, [])

asyncdeftest_executemany_client_failure_in_transaction(self):
tx=self.con.transaction()
awaittx.start()
withself.assertRaises(ZeroDivisionError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, $2)
 ''', (('a'*32768, y+y/y) foryinrange(10, -1, -1)))
result=awaitself.con.fetch('SELECT b FROM exmany')
# only 2 batches executed (2 x 4)
self.assertEqual(
 [x[0] forxinresult], [y+1foryinrange(10, 2, -1)])
awaittx.rollback()
result=awaitself.con.fetch('SELECT b FROM exmany')
self.assertEqual(result, [])

asyncdeftest_executemany_client_server_failure_conflict(self):
self.con._transport.set_write_buffer_limits(65536*64, 16384*64)
withself.assertRaises(exceptions.UniqueViolationError):
awaitself.con.executemany('''
 INSERT INTO exmany VALUES($1, 0)
 ''', (('a'*32768,) foryinrange(4, -1, -1) ify/y))
result=awaitself.con.fetch('SELECT b FROM exmany')
self.assertEqual(result, [])

asyncdeftest_executemany_prepare(self):
stmt=awaitself.con.prepare('''
 INSERT INTO exmany VALUES($1, $2)
 ''')
result=awaitstmt.executemany([
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])
self.assertIsNone(result)
result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')
self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])
# Empty set
awaitstmt.executemany(())
result=awaitself.con.fetch('''
 SELECT * FROM exmany
 ''')
self.assertEqual(result, [
 ('a', 1), ('b', 2), ('c', 3), ('d', 4)
 ])