test_introspection.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
importjson

fromasyncpgimport_testbaseastb
fromasyncpgimportconnectionasapg_con


MAX_RUNTIME=0.25


classSlowIntrospectionConnection(apg_con.Connection):
"""Connection class to test introspection races."""
introspect_count=0

asyncdef_introspect_types(self, *args, **kwargs):
self.introspect_count+=1
awaitasyncio.sleep(0.4)
returnawaitsuper()._introspect_types(*args, **kwargs)


classTestIntrospection(tb.ConnectedTestCase):
@classmethod
defsetUpClass(cls):
super().setUpClass()
cls.adminconn=cls.loop.run_until_complete(cls.connect())
cls.loop.run_until_complete(
cls.adminconn.execute('CREATE DATABASE asyncpg_intro_test'))

@classmethod
deftearDownClass(cls):
cls.loop.run_until_complete(
cls.adminconn.execute('DROP DATABASE asyncpg_intro_test'))

cls.loop.run_until_complete(cls.adminconn.close())
cls.adminconn=None

super().tearDownClass()

@classmethod
defget_server_settings(cls):
settings=super().get_server_settings()
settings.pop('jit', None)
returnsettings

defsetUp(self):
super().setUp()
self.loop.run_until_complete(self._add_custom_codec(self.con))

asyncdef_add_custom_codec(self, conn):
# mess up with the codec - builtin introspection shouldn't be affected
awaitconn.set_type_codec(
"oid",
schema="pg_catalog",
encoder=lambdavalue: None,
decoder=lambdavalue: None,
format="text",
 )

@tb.with_connection_options(database='asyncpg_intro_test')
asyncdeftest_introspection_on_large_db(self):
awaitself.con.execute(
'CREATE TABLE base ({})'.format(
','.join('c{:02} varchar'.format(n) forninrange(50))
 )
 )
forninrange(1000):
awaitself.con.execute(
'CREATE TABLE child_{:04} () inherits (base)'.format(n)
 )

withself.assertRunUnder(MAX_RUNTIME):
awaitself.con.fetchval('SELECT $1::int[]', [1, 2])

@tb.with_connection_options(statement_cache_size=0)
asyncdeftest_introspection_no_stmt_cache_01(self):
old_uid=apg_con._uid

self.assertEqual(self.con._stmt_cache.get_max_size(), 0)
awaitself.con.fetchval('SELECT $1::int[]', [1, 2])

awaitself.con.execute('''
 CREATE EXTENSION IF NOT EXISTS hstore
 ''')

try:
awaitself.con.set_builtin_type_codec(
'hstore', codec_name='pg_contrib.hstore')
finally:
awaitself.con.execute('''
 DROP EXTENSION hstore
 ''')

self.assertEqual(apg_con._uid, old_uid)

@tb.with_connection_options(max_cacheable_statement_size=1)
asyncdeftest_introspection_no_stmt_cache_02(self):
# max_cacheable_statement_size will disable caching both for
# the user query and for the introspection query.
old_uid=apg_con._uid

awaitself.con.fetchval('SELECT $1::int[]', [1, 2])

awaitself.con.execute('''
 CREATE EXTENSION IF NOT EXISTS hstore
 ''')

try:
awaitself.con.set_builtin_type_codec(
'hstore', codec_name='pg_contrib.hstore')
finally:
awaitself.con.execute('''
 DROP EXTENSION hstore
 ''')

self.assertEqual(apg_con._uid, old_uid)

@tb.with_connection_options(max_cacheable_statement_size=10000)
asyncdeftest_introspection_no_stmt_cache_03(self):
# max_cacheable_statement_size will disable caching for
# the user query but not for the introspection query.
old_uid=apg_con._uid

awaitself.con.fetchval(
"SELECT $1::int[], '{foo}'".format(foo='a'*10000), [1, 2])

self.assertGreater(apg_con._uid, old_uid)

asyncdeftest_introspection_sticks_for_ps(self):
# Test that the introspected codec pipeline for a prepared
# statement is not affected by a subsequent codec cache bust.

ps=awaitself.con._prepare('SELECT $1::json[]', use_cache=True)

try:
# Setting a custom codec blows the codec cache for derived types.
awaitself.con.set_type_codec(
'json', encoder=lambdav: v, decoder=json.loads,
schema='pg_catalog', format='text'
 )

# The originally prepared statement should still be OK and
# use the previously selected codec.
self.assertEqual(awaitps.fetchval(['{"foo": 1}']), ['{"foo": 1}'])

# The new query uses the custom codec.
v=awaitself.con.fetchval('SELECT $1::json[]', ['{"foo": 1}'])
self.assertEqual(v, [{'foo': 1}])

finally:
awaitself.con.reset_type_codec(
'json', schema='pg_catalog')

asyncdeftest_introspection_retries_after_cache_bust(self):
# Test that codec cache bust racing with the introspection
# query would cause introspection to retry.
slow_intro_conn=awaitself.connect(
connection_class=SlowIntrospectionConnection)
awaitself._add_custom_codec(slow_intro_conn)
try:
awaitself.con.execute('''
 CREATE DOMAIN intro_1_t AS int;
 CREATE DOMAIN intro_2_t AS int;
 ''')

awaitslow_intro_conn.fetchval('''
 SELECT $1::intro_1_t
 ''', 10)
# slow_intro_conn cache is now populated with intro_1_t

asyncdefwait_and_drop():
awaitasyncio.sleep(0.1)
awaitslow_intro_conn.reload_schema_state()

# Now, in parallel, run another query that
# references both intro_1_t and intro_2_t.
awaitasyncio.gather(
slow_intro_conn.fetchval('''
 SELECT $1::intro_1_t, $2::intro_2_t
 ''', 10, 20),
wait_and_drop()
 )

# Initial query + two tries for the second query.
self.assertEqual(slow_intro_conn.introspect_count, 3)

finally:
awaitself.con.execute('''
 DROP DOMAIN intro_1_t;
 DROP DOMAIN intro_2_t;
 ''')
awaitslow_intro_conn.close()

@tb.with_connection_options(database='asyncpg_intro_test')
asyncdeftest_introspection_loads_basetypes_of_domains(self):
# Test that basetypes of domains are loaded to the
# client encode/decode cache
awaitself.con.execute('''
 DROP TABLE IF EXISTS test;
 DROP DOMAIN IF EXISTS num_array;
 CREATE DOMAIN num_array numeric[];
 CREATE TABLE test (
 num num_array
 );
 ''')

try:
# if domain basetypes are not loaded, this insert will fail
awaitself.con.execute(
'INSERT INTO test (num) VALUES ($1)', ([1, 2],))
finally:
awaitself.con.execute('''
 DROP TABLE IF EXISTS test;
 DROP DOMAIN IF EXISTS num_array;
 ''')