test_prepare.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
importgc
importunittest

fromasyncpgimport_testbaseastb
fromasyncpgimportexceptions


classTestPrepare(tb.ConnectedTestCase):

asyncdeftest_prepare_01(self):
self.assertEqual(self.con._protocol.queries_count, 0)
st=awaitself.con.prepare('SELECT 1 = $1 AS test')
self.assertEqual(self.con._protocol.queries_count, 0)
self.assertEqual(st.get_query(), 'SELECT 1 = $1 AS test')

rec=awaitst.fetchrow(1)
self.assertEqual(self.con._protocol.queries_count, 1)
self.assertTrue(rec['test'])
self.assertEqual(len(rec), 1)

self.assertEqual(False, awaitst.fetchval(10))
self.assertEqual(self.con._protocol.queries_count, 2)

asyncdeftest_prepare_02(self):
withself.assertRaisesRegex(Exception, 'column "a" does not exist'):
awaitself.con.prepare('SELECT a')

asyncdeftest_prepare_03(self):
cases= [
 ('text', ("'NULL'", 'NULL'), [
'aaa',
None
 ]),

 ('decimal', ('0', 0), [
123,
123.5,
None
 ])
 ]

fortype, (none_name, none_val), valsincases:
st=awaitself.con.prepare('''
 SELECT CASE WHEN $1::{type} IS NULL THEN {default}
 ELSE $1::{type} END'''.format(
type=type, default=none_name))

forvalinvals:
withself.subTest(type=type, value=val):
res=awaitst.fetchval(val)
ifvalisNone:
self.assertEqual(res, none_val)
else:
self.assertEqual(res, val)

asyncdeftest_prepare_04(self):
s=awaitself.con.prepare('SELECT $1::smallint')
self.assertEqual(awaits.fetchval(10), 10)

s=awaitself.con.prepare('SELECT $1::smallint * 2')
self.assertEqual(awaits.fetchval(10), 20)

s=awaitself.con.prepare('SELECT generate_series(5,10)')
self.assertEqual(awaits.fetchval(), 5)
# Since the "execute" message was sent with a limit=1,
# we will receive a PortalSuspended message, instead of
# CommandComplete. Which means there will be no status
# message set.
self.assertIsNone(s.get_statusmsg())
# Repeat the same test for 'fetchrow()'.
self.assertEqual(awaits.fetchrow(), (5,))
self.assertIsNone(s.get_statusmsg())

asyncdeftest_prepare_05_unknownoid(self):
s=awaitself.con.prepare("SELECT 'test'")
self.assertEqual(awaits.fetchval(), 'test')

asyncdeftest_prepare_06_interrupted_close(self):
stmt=awaitself.con.prepare('''SELECT pg_sleep(10)''')
fut=self.loop.create_task(stmt.fetch())

awaitasyncio.sleep(0.2)

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

withself.assertRaises(asyncpg.QueryCanceledError):
awaitfut

# Test that it's OK to call close again
awaitself.con.close()

asyncdeftest_prepare_07_interrupted_terminate(self):
stmt=awaitself.con.prepare('''SELECT pg_sleep(10)''')
fut=self.loop.create_task(stmt.fetchval())

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

# Test that it's OK to call terminate again
self.con.terminate()

asyncdeftest_prepare_08_big_result(self):
stmt=awaitself.con.prepare('select generate_series(0,10000)')
result=awaitstmt.fetch()

self.assertEqual(len(result), 10001)
self.assertEqual(
 [r[0] forrinresult],
list(range(10001)))

asyncdeftest_prepare_09_raise_error(self):
# Stress test ReadBuffer.read_cstr()
msg='0'*1024*100
query="""
 DO language plpgsql $$
 BEGIN
 RAISE EXCEPTION '{}';
 END
 $$;""".format(msg)

stmt=awaitself.con.prepare(query)
withself.assertRaisesRegex(asyncpg.RaiseError, msg):
withtb.silence_asyncio_long_exec_warning():
awaitstmt.fetchval()

asyncdeftest_prepare_10_stmt_lru(self):
cache=self.con._stmt_cache

query='select {}'
cache_max=cache.get_max_size()
iter_max=cache_max*2+11

# First, we have no cached statements.
self.assertEqual(len(cache), 0)

stmts= []
foriinrange(iter_max):
s=awaitself.con._prepare(query.format(i), use_cache=True)
self.assertEqual(awaits.fetchval(), i)
stmts.append(s)

# At this point our cache should be full.
self.assertEqual(len(cache), cache_max)
self.assertTrue(all(nots.closedforsincache.iter_statements()))

# Since there are references to the statements (`stmts` list),
# no statements are scheduled to be closed.
self.assertEqual(len(self.con._stmts_to_close), 0)

# Removing refs to statements and preparing a new statement
# will cause connection to cleanup any stale statements.
stmts.clear()
gc.collect()

# Now we have a bunch of statements that have no refs to them
# scheduled to be closed.
self.assertEqual(len(self.con._stmts_to_close), iter_max-cache_max)
self.assertTrue(all(s.closedforsinself.con._stmts_to_close))
self.assertTrue(all(nots.closedforsincache.iter_statements()))

zero=awaitself.con.prepare(query.format(0))
# Hence, all stale statements should be closed now.
self.assertEqual(len(self.con._stmts_to_close), 0)

# The number of cached statements will stay the same though.
self.assertEqual(len(cache), cache_max)
self.assertTrue(all(nots.closedforsincache.iter_statements()))

# After closing all statements will be closed.
awaitself.con.close()
self.assertEqual(len(self.con._stmts_to_close), 0)
self.assertEqual(len(cache), 0)

# An attempt to perform an operation on a closed statement
# will trigger an error.
withself.assertRaisesRegex(asyncpg.InterfaceError, 'is closed'):
awaitzero.fetchval()

asyncdeftest_prepare_11_stmt_gc(self):
# Test that prepared statements should stay in the cache after
# they are GCed.

cache=self.con._stmt_cache

# First, we have no cached statements.
self.assertEqual(len(cache), 0)
self.assertEqual(len(self.con._stmts_to_close), 0)

# The prepared statement that we'll create will be GCed
# right await. However, its state should be still in
# in the statements LRU cache.
awaitself.con._prepare('select 1', use_cache=True)
gc.collect()

self.assertEqual(len(cache), 1)
self.assertEqual(len(self.con._stmts_to_close), 0)

asyncdeftest_prepare_12_stmt_gc(self):
# Test that prepared statements are closed when there is no space
# for them in the LRU cache and there are no references to them.

cache=self.con._stmt_cache
cache_max=cache.get_max_size()

# First, we have no cached statements.
self.assertEqual(len(cache), 0)
self.assertEqual(len(self.con._stmts_to_close), 0)

stmt=awaitself.con._prepare('select 100000000', use_cache=True)
self.assertEqual(len(cache), 1)
self.assertEqual(len(self.con._stmts_to_close), 0)

foriinrange(cache_max):
awaitself.con._prepare('select {}'.format(i), use_cache=True)

self.assertEqual(len(cache), cache_max)
self.assertEqual(len(self.con._stmts_to_close), 0)

delstmt
gc.collect()

self.assertEqual(len(cache), cache_max)
self.assertEqual(len(self.con._stmts_to_close), 1)

asyncdeftest_prepare_13_connect(self):
v=awaitself.con.fetchval(
'SELECT $1::smallint AS foo', 10, column='foo')
self.assertEqual(v, 10)

r=awaitself.con.fetchrow('SELECT $1::smallint * 2 AS test', 10)
self.assertEqual(r['test'], 20)

rows=awaitself.con.fetch('SELECT generate_series(0,$1::int)', 3)
self.assertEqual([r[0] forrinrows], [0, 1, 2, 3])

asyncdeftest_prepare_14_explain(self):
# Test simple EXPLAIN.
stmt=awaitself.con.prepare('SELECT typname FROM pg_type')
plan=awaitstmt.explain()
self.assertEqual(plan[0]['Plan']['Relation Name'], 'pg_type')

# Test "EXPLAIN ANALYZE".
stmt=awaitself.con.prepare(
'SELECT typname, typlen FROM pg_type WHERE typlen > $1')
plan=awaitstmt.explain(2, analyze=True)
self.assertEqual(plan[0]['Plan']['Relation Name'], 'pg_type')
self.assertIn('Actual Total Time', plan[0]['Plan'])

# Test that 'EXPLAIN ANALYZE' is executed in a transaction
# that gets rollbacked.
tr=self.con.transaction()
awaittr.start()
try:
awaitself.con.execute('CREATE TABLE mytab (a int)')
stmt=awaitself.con.prepare(
'INSERT INTO mytab (a) VALUES (1), (2)')
plan=awaitstmt.explain(analyze=True)
self.assertEqual(plan[0]['Plan']['Operation'], 'Insert')

# Check that no data was inserted
res=awaitself.con.fetch('SELECT * FROM mytab')
self.assertEqual(res, [])
finally:
awaittr.rollback()

asyncdeftest_prepare_15_stmt_gc_cache_disabled(self):
# Test that even if the statements cache is off, we're still
# cleaning up GCed statements.

cache=self.con._stmt_cache

self.assertEqual(len(cache), 0)
self.assertEqual(len(self.con._stmts_to_close), 0)

# Disable cache
cache.set_max_size(0)

stmt=awaitself.con._prepare('select 100000000', use_cache=True)
self.assertEqual(len(cache), 0)
self.assertEqual(len(self.con._stmts_to_close), 0)

delstmt
gc.collect()

# After GC, _stmts_to_close should contain stmt's state
self.assertEqual(len(cache), 0)
self.assertEqual(len(self.con._stmts_to_close), 1)

# Next "prepare" call will trigger a cleanup
stmt=awaitself.con._prepare('select 1', use_cache=True)
self.assertEqual(len(cache), 0)
self.assertEqual(len(self.con._stmts_to_close), 0)

delstmt

asyncdeftest_prepare_16_command_result(self):
asyncdefstatus(query):
stmt=awaitself.con.prepare(query)
awaitstmt.fetch()
returnstmt.get_statusmsg()

try:
self.assertEqual(
awaitstatus('CREATE TABLE mytab (a int)'),
'CREATE TABLE')

self.assertEqual(
awaitstatus('INSERT INTO mytab (a) VALUES (1), (2)'),
'INSERT 0 2')

self.assertEqual(
awaitstatus('SELECT a FROM mytab'),
'SELECT 2')

self.assertEqual(
awaitstatus('UPDATE mytab SET a = 3 WHERE a = 1'),
'UPDATE 1')
finally:
self.assertEqual(
awaitstatus('DROP TABLE mytab'),
'DROP TABLE')

asyncdeftest_prepare_17_stmt_closed_lru(self):
st=awaitself.con.prepare('SELECT 1')
st._state.mark_closed()
withself.assertRaisesRegex(asyncpg.InterfaceError, 'is closed'):
awaitst.fetch()

st=awaitself.con.prepare('SELECT 1')
self.assertEqual(awaitst.fetchval(), 1)

asyncdeftest_prepare_18_empty_result(self):
# test EmptyQueryResponse protocol message
st=awaitself.con.prepare('')
self.assertEqual(awaitst.fetch(), [])
self.assertIsNone(awaitst.fetchval())
self.assertIsNone(awaitst.fetchrow())

self.assertEqual(awaitself.con.fetch(''), [])
self.assertIsNone(awaitself.con.fetchval(''))
self.assertIsNone(awaitself.con.fetchrow(''))

asyncdeftest_prepare_19_concurrent_calls(self):
st=self.loop.create_task(self.con.fetchval(
'SELECT ROW(pg_sleep(0.1), 1)'))

# Wait for some time to make sure the first query is fully
# prepared (!) and is now awaiting the results (!!).
awaitasyncio.sleep(0.01)

withself.assertRaisesRegex(asyncpg.InterfaceError,
'another operation'):
awaitself.con.execute('SELECT 2')

self.assertEqual(awaitst, (None, 1))

asyncdeftest_prepare_20_concurrent_calls(self):
expected= ((None, 1),)

formethname, valin [('fetch', [expected]),
 ('fetchval', expected[0]),
 ('fetchrow', expected)]:

withself.subTest(meth=methname):

meth=getattr(self.con, methname)

vf=self.loop.create_task(
meth('SELECT ROW(pg_sleep(0.1), 1)'))

awaitasyncio.sleep(0.01)

withself.assertRaisesRegex(asyncpg.InterfaceError,
'another operation'):
awaitmeth('SELECT 2')

self.assertEqual(awaitvf, val)

asyncdeftest_prepare_21_errors(self):
stmt=awaitself.con.prepare('SELECT 10 / $1::int')

withself.assertRaises(asyncpg.DivisionByZeroError):
awaitstmt.fetchval(0)

self.assertEqual(awaitstmt.fetchval(5), 2)

asyncdeftest_prepare_22_empty(self):
# Support for empty target list was added in PostgreSQL 9.4
ifself.server_version< (9, 4):
raiseunittest.SkipTest(
'PostgreSQL servers < 9.4 do not support empty target list.')

result=awaitself.con.fetchrow('SELECT')
self.assertEqual(result, ())
self.assertEqual(repr(result), '<Record>')

asyncdeftest_prepare_statement_invalid(self):
awaitself.con.execute('CREATE TABLE tab1(a int, b int)')

try:
awaitself.con.execute('INSERT INTO tab1 VALUES (1, 2)')

stmt=awaitself.con.prepare('SELECT * FROM tab1')

awaitself.con.execute(
'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')

withself.assertRaisesRegex(asyncpg.InvalidCachedStatementError,
'cached statement plan is invalid'):
awaitstmt.fetchrow()

finally:
awaitself.con.execute('DROP TABLE tab1')

@tb.with_connection_options(statement_cache_size=0)
asyncdeftest_prepare_23_no_stmt_cache_seq(self):
self.assertEqual(self.con._stmt_cache.get_max_size(), 0)

asyncdefcheck_simple():
# Run a simple query a few times.
self.assertEqual(awaitself.con.fetchval('SELECT 1'), 1)
self.assertEqual(awaitself.con.fetchval('SELECT 2'), 2)
self.assertEqual(awaitself.con.fetchval('SELECT 1'), 1)

awaitcheck_simple()

# Run a query that timeouts.
withself.assertRaises(asyncio.TimeoutError):
awaitself.con.fetchrow('select pg_sleep(10)', timeout=0.02)

# Check that we can run new queries after a timeout.
awaitcheck_simple()

# Try a cursor/timeout combination. Cursors should always use
# named prepared statements.
asyncwithself.con.transaction():
withself.assertRaises(asyncio.TimeoutError):
asyncfor_inself.con.cursor( # NOQA
'select pg_sleep(10)', timeout=0.1):
pass

# Check that we can run queries after a failed cursor
# operation.
awaitcheck_simple()

@tb.with_connection_options(max_cached_statement_lifetime=142)
asyncdeftest_prepare_24_max_lifetime(self):
cache=self.con._stmt_cache

self.assertEqual(cache.get_max_lifetime(), 142)
cache.set_max_lifetime(1)

s=awaitself.con._prepare('SELECT 1', use_cache=True)
state=s._state

s=awaitself.con._prepare('SELECT 1', use_cache=True)
self.assertIs(s._state, state)

s=awaitself.con._prepare('SELECT 1', use_cache=True)
self.assertIs(s._state, state)

awaitasyncio.sleep(1)

s=awaitself.con._prepare('SELECT 1', use_cache=True)
self.assertIsNot(s._state, state)

@tb.with_connection_options(max_cached_statement_lifetime=0.5)
asyncdeftest_prepare_25_max_lifetime_reset(self):
cache=self.con._stmt_cache

s=awaitself.con._prepare('SELECT 1', use_cache=True)
state=s._state

# Disable max_lifetime
cache.set_max_lifetime(0)

awaitasyncio.sleep(1)

# The statement should still be cached (as we disabled the timeout).
s=awaitself.con._prepare('SELECT 1', use_cache=True)
self.assertIs(s._state, state)

@tb.with_connection_options(max_cached_statement_lifetime=0.5)
asyncdeftest_prepare_26_max_lifetime_max_size(self):
cache=self.con._stmt_cache

s=awaitself.con._prepare('SELECT 1', use_cache=True)
state=s._state

# Disable max_lifetime
cache.set_max_size(0)

s=awaitself.con._prepare('SELECT 1', use_cache=True)
self.assertIsNot(s._state, state)

# Check that nothing crashes after the initial timeout
awaitasyncio.sleep(1)

@tb.with_connection_options(max_cacheable_statement_size=50)
asyncdeftest_prepare_27_max_cacheable_statement_size(self):
cache=self.con._stmt_cache

awaitself.con._prepare('SELECT 1', use_cache=True)
self.assertEqual(len(cache), 1)

# Test that long and explicitly created prepared statements
# are not cached.
awaitself.con._prepare("SELECT \'"+"a"*50+"\'", use_cache=True)
self.assertEqual(len(cache), 1)

# Test that implicitly created long prepared statements
# are not cached.
awaitself.con.fetchval("SELECT \'"+"a"*50+"\'")
self.assertEqual(len(cache), 1)

# Test that short prepared statements can still be cached.
awaitself.con._prepare('SELECT 2', use_cache=True)
self.assertEqual(len(cache), 2)

asyncdeftest_prepare_28_max_args(self):
N=32768
args=','.join('${}'.format(i) foriinrange(1, N+1))
query='SELECT ARRAY[{}]'.format(args)

withself.assertRaisesRegex(
exceptions.InterfaceError,
'the number of query arguments cannot exceed 32767'):
awaitself.con.fetchval(query, *range(1, N+1))

asyncdeftest_prepare_29_duplicates(self):
# In addition to test_record.py, let's have a full functional
# test for records with duplicate keys.
r=awaitself.con.fetchrow('SELECT 1 as a, 2 as b, 3 as a')
self.assertEqual(list(r.items()), [('a', 1), ('b', 2), ('a', 3)])
self.assertEqual(list(r.keys()), ['a', 'b', 'a'])
self.assertEqual(list(r.values()), [1, 2, 3])
self.assertEqual(r['a'], 3)
self.assertEqual(r['b'], 2)
self.assertEqual(r[0], 1)
self.assertEqual(r[1], 2)
self.assertEqual(r[2], 3)

asyncdeftest_prepare_30_invalid_arg_count(self):
withself.assertRaisesRegex(
exceptions.InterfaceError,
'the server expects 1 argument for this query, 0 were passed'):
awaitself.con.fetchval('SELECT $1::int')

withself.assertRaisesRegex(
exceptions.InterfaceError,
'the server expects 0 arguments for this query, 1 was passed'):
awaitself.con.fetchval('SELECT 1', 1)

asyncdeftest_prepare_31_pgbouncer_note(self):
try:
awaitself.con.execute("""
 DO $$ BEGIN
 RAISE EXCEPTION
 'duplicate statement' USING ERRCODE = '42P05';
 END; $$ LANGUAGE plpgsql;
 """)
exceptasyncpg.DuplicatePreparedStatementErrorase:
self.assertTrue('pgbouncer'ine.hint)
else:
self.fail('DuplicatePreparedStatementError not raised')

try:
awaitself.con.execute("""
 DO $$ BEGIN
 RAISE EXCEPTION
 'invalid statement' USING ERRCODE = '26000';
 END; $$ LANGUAGE plpgsql;
 """)
exceptasyncpg.InvalidSQLStatementNameErrorase:
self.assertTrue('pgbouncer'ine.hint)
else:
self.fail('InvalidSQLStatementNameError not raised')

asyncdeftest_prepare_does_not_use_cache(self):
cache=self.con._stmt_cache

# prepare with disabled cache
awaitself.con.prepare('select 1')
self.assertEqual(len(cache), 0)

asyncdeftest_prepare_explicitly_named(self):
ps=awaitself.con.prepare('select 1', name='foobar')
self.assertEqual(ps.get_name(), 'foobar')
self.assertEqual(awaitself.con.fetchval('EXECUTE foobar'), 1)

withself.assertRaisesRegex(
exceptions.DuplicatePreparedStatementError,
'prepared statement "foobar" already exists',
 ):
awaitself.con.prepare('select 1', name='foobar')