Skip to content

Commit c4cc89f

Browse files
committed
Disable custom data codec for internal introspection
Fixes: #617
1 parent 2bac166 commit c4cc89f

File tree

8 files changed

+81
-24
lines changed

8 files changed

+81
-24
lines changed

asyncpg/connection.py

Lines changed: 34 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -340,7 +340,8 @@ async def _get_statement(
340340
*,
341341
named: bool=False,
342342
use_cache: bool=True,
343-
record_class=None
343+
record_class=None,
344+
disable_custom_codec=False,
344345
):
345346
if record_class is None:
346347
record_class = self._protocol.get_record_class()
@@ -371,7 +372,9 @@ async def _get_statement(
371372
record_class=record_class,
372373
)
373374
need_reprepare = False
374-
types_with_missing_codecs = statement._init_types()
375+
types_with_missing_codecs = statement._init_types(
376+
disable_custom_codec=disable_custom_codec
377+
)
375378
tries = 0
376379
while types_with_missing_codecs:
377380
settings = self._protocol.get_settings()
@@ -387,7 +390,9 @@ async def _get_statement(
387390
# which has blown away the anonymous statement we've prepared
388391
# for the query, so we need to re-prepare it.
389392
need_reprepare = not intro_stmt.name and not statement.name
390-
types_with_missing_codecs = statement._init_types()
393+
types_with_missing_codecs = statement._init_types(
394+
disable_custom_codec=disable_custom_codec
395+
)
391396
tries += 1
392397
if tries > 5:
393398
# In the vast majority of cases there will be only
@@ -401,7 +406,7 @@ async def _get_statement(
401406

402407
# Now that types have been resolved, populate the codec pipeline
403408
# for the statement.
404-
statement._init_codecs()
409+
statement._init_codecs(disable_custom_codec)
405410

406411
if need_reprepare:
407412
await self._protocol.prepare(
@@ -424,7 +429,12 @@ async def _get_statement(
424429

425430
async def _introspect_types(self, typeoids, timeout):
426431
return await self.__execute(
427-
self._intro_query, (list(typeoids),), 0, timeout)
432+
self._intro_query,
433+
(list(typeoids),),
434+
0,
435+
timeout,
436+
disable_custom_codec=True,
437+
)
428438

429439
def cursor(
430440
self,
@@ -571,7 +581,8 @@ async def fetchrow(
571581
query,
572582
*args,
573583
timeout=None,
574-
record_class=None
584+
record_class=None,
585+
_disable_custom_codec=False,
575586
):
576587
"""Run a query and return the first row.
577588
@@ -601,6 +612,7 @@ async def fetchrow(
601612
1,
602613
timeout,
603614
record_class=record_class,
615+
disable_custom_codec=_disable_custom_codec,
604616
)
605617
if not data:
606618
return None
@@ -1110,7 +1122,8 @@ async def set_type_codec(self, typename, *,
11101122
self._check_open()
11111123

11121124
typeinfo = await self.fetchrow(
1113-
introspection.TYPE_BY_NAME, typename, schema)
1125+
introspection.TYPE_BY_NAME, typename, schema,
1126+
_disable_custom_codec=True)
11141127
if not typeinfo:
11151128
raise ValueError('unknown type: {}.{}'.format(schema, typename))
11161129

@@ -1141,7 +1154,8 @@ async def reset_type_codec(self, typename, *, schema='public'):
11411154
"""
11421155

11431156
typeinfo = await self.fetchrow(
1144-
introspection.TYPE_BY_NAME, typename, schema)
1157+
introspection.TYPE_BY_NAME, typename, schema,
1158+
_disable_custom_codec=True)
11451159
if not typeinfo:
11461160
raise ValueError('unknown type: {}.{}'.format(schema, typename))
11471161

@@ -1191,7 +1205,8 @@ async def set_builtin_type_codec(self, typename, *,
11911205
self._check_open()
11921206

11931207
typeinfo = await self.fetchrow(
1194-
introspection.TYPE_BY_NAME, typename, schema)
1208+
introspection.TYPE_BY_NAME, typename, schema,
1209+
_disable_custom_codec=True)
11951210
if not typeinfo:
11961211
raise exceptions.InterfaceError(
11971212
'unknown type: {}.{}'.format(schema, typename))
@@ -1578,7 +1593,8 @@ async def _execute(
15781593
timeout,
15791594
*,
15801595
return_status=False,
1581-
record_class=None
1596+
record_class=None,
1597+
disable_custom_codec=False,
15821598
):
15831599
with self._stmt_exclusive_section:
15841600
result, _ = await self.__execute(
@@ -1588,6 +1604,7 @@ async def _execute(
15881604
timeout,
15891605
return_status=return_status,
15901606
record_class=record_class,
1607+
disable_custom_codec=disable_custom_codec,
15911608
)
15921609
return result
15931610

@@ -1599,7 +1616,8 @@ async def __execute(
15991616
timeout,
16001617
*,
16011618
return_status=False,
1602-
record_class=None
1619+
record_class=None,
1620+
disable_custom_codec=False,
16031621
):
16041622
executor = lambda stmt, timeout: self._protocol.bind_execute(
16051623
stmt, args, '', limit, return_status, timeout)
@@ -1609,6 +1627,7 @@ async def __execute(
16091627
executor,
16101628
timeout,
16111629
record_class=record_class,
1630+
disable_custom_codec=disable_custom_codec,
16121631
)
16131632

16141633
async def _executemany(self, query, args, timeout):
@@ -1626,20 +1645,23 @@ async def _do_execute(
16261645
timeout,
16271646
retry=True,
16281647
*,
1629-
record_class=None
1648+
record_class=None,
1649+
disable_custom_codec=False,
16301650
):
16311651
if timeout is None:
16321652
stmt = await self._get_statement(
16331653
query,
16341654
None,
16351655
record_class=record_class,
1656+
disable_custom_codec=disable_custom_codec,
16361657
)
16371658
else:
16381659
before = time.monotonic()
16391660
stmt = await self._get_statement(
16401661
query,
16411662
timeout,
16421663
record_class=record_class,
1664+
disable_custom_codec=disable_custom_codec,
16431665
)
16441666
after = time.monotonic()
16451667
timeout -= after - before

asyncpg/protocol/codecs/base.pxd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,7 @@ cdef class DataCodecConfig:
165165
cdef:
166166
dict _derived_type_codecs
167167
dict _custom_type_codecs
168+
bint _custom_type_codecs_enabled
168169

169170
cdef inline Codec get_codec(self, uint32_t oid, ServerDataFormat format)
170171
cdef inline Codec get_any_local_codec(self, uint32_t oid)

asyncpg/protocol/codecs/base.pyx

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,8 @@ cdef class DataCodecConfig:
427427
self._derived_type_codecs = {}
428428
# Codec instances set up by the user for the connection.
429429
self._custom_type_codecs = {}
430+
# Switch to disable custom type codecs for built-in introspection
431+
self._custom_type_codecs_enabled = True
430432

431433
def add_types(self, types):
432434
cdef:
@@ -695,15 +697,16 @@ cdef class DataCodecConfig:
695697
cdef inline Codec get_codec(self, uint32_t oid, ServerDataFormat format):
696698
cdef Codec codec
697699

698-
codec = self.get_any_local_codec(oid)
699-
if codec is not None:
700-
if codec.format != format:
701-
# The codec for this OID has been overridden by
702-
# set_{builtin}_type_codec with a different format.
703-
# We must respect that and not return a core codec.
704-
return None
705-
else:
706-
return codec
700+
if self._custom_type_codecs_enabled:
701+
codec = self.get_any_local_codec(oid)
702+
if codec is not None:
703+
if codec.format != format:
704+
# The codec for this OID has been overridden by
705+
# set_{builtin}_type_codec with a different format.
706+
# We must respect that and not return a core codec.
707+
return None
708+
else:
709+
return codec
707710

708711
codec = get_core_codec(oid, format)
709712
if codec is not None:

asyncpg/protocol/prepared_stmt.pxd

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ cdef class PreparedStatementState:
2929
tuple rows_codecs
3030

3131
cdef _encode_bind_msg(self, args)
32-
cpdef _init_codecs(self)
32+
cpdef _init_codecs(self, bint disable_custom_codec)
3333
cdef _ensure_rows_decoder(self)
3434
cdef _ensure_args_encoder(self)
3535
cdef _set_row_desc(self, object desc)

asyncpg/protocol/prepared_stmt.pyx

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,14 @@ cdef class PreparedStatementState:
6767

6868
return tuple(result)
6969

70-
def _init_types(self):
70+
def _init_types(self, disable_custom_codec=False):
7171
cdef:
7272
Codec codec
7373
set missing = set()
7474

75+
if disable_custom_codec:
76+
self.settings.use_custom_codec(False)
77+
7578
if self.parameters_desc:
7679
for p_oid in self.parameters_desc:
7780
codec = self.settings.get_data_codec(<uint32_t>p_oid)
@@ -84,12 +87,22 @@ cdef class PreparedStatementState:
8487
if codec is None or not codec.has_decoder():
8588
missing.add(rdesc[3])
8689

90+
if disable_custom_codec:
91+
self.settings.use_custom_codec(True)
92+
8793
return missing
8894

89-
cpdef _init_codecs(self):
95+
cpdef _init_codecs(self, bint disable_custom_codec):
96+
97+
if disable_custom_codec:
98+
self.settings.use_custom_codec(False)
99+
90100
self._ensure_args_encoder()
91101
self._ensure_rows_decoder()
92102

103+
if disable_custom_codec:
104+
self.settings.use_custom_codec(True)
105+
93106
def attach(self):
94107
self.refs += 1
95108

asyncpg/protocol/settings.pxd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ cdef class ConnectionSettings(pgproto.CodecContext):
1515

1616
cdef add_setting(self, str name, str val)
1717
cdef is_encoding_utf8(self)
18+
cdef use_custom_codec(self, bint enabled)
1819
cpdef get_text_codec(self)
1920
cpdef inline register_data_types(self, types)
2021
cpdef inline add_python_codec(

asyncpg/protocol/settings.pyx

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ cdef class ConnectionSettings(pgproto.CodecContext):
2929
cdef is_encoding_utf8(self):
3030
return self._is_utf8
3131

32+
cdef use_custom_codec(self, bint enabled):
33+
self._data_codecs._custom_type_codecs_enabled = enabled
34+
3235
cpdef get_text_codec(self):
3336
return self._codec
3437

tests/test_introspection.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,19 @@ def tearDownClass(cls):
4343

4444
super().tearDownClass()
4545

46+
def setUp(self):
47+
super().setUp()
48+
self.loop.run_until_complete(self._add_custom_codec(self.con))
49+
50+
async def _add_custom_codec(self, conn):
51+
await conn.set_type_codec(
52+
"char",
53+
schema="pg_catalog",
54+
encoder=lambda value: value,
55+
decoder=lambda value: value,
56+
format="text",
57+
)
58+
4659
@tb.with_connection_options(database='asyncpg_intro_test')
4760
async def test_introspection_on_large_db(self):
4861
await self.con.execute(
@@ -142,6 +155,7 @@ async def test_introspection_retries_after_cache_bust(self):
142155
# query would cause introspection to retry.
143156
slow_intro_conn = await self.connect(
144157
connection_class=SlowIntrospectionConnection)
158+
await self._add_custom_codec(slow_intro_conn)
145159
try:
146160
await self.con.execute('''
147161
CREATE DOMAIN intro_1_t AS int;

0 commit comments

Comments
 (0)