Ignore custom data codec for internal introspection (#618)

Fixes: #617

Author: fantix

asyncpg/connection.py

@@ -342,13 +342,16 @@ async def _get_statement(
 *,
 named: bool=False,
 use_cache: bool=True,
+ignore_custom_codec=False,
 record_class=None
  ):
 ifrecord_classisNone:
 record_class=self._protocol.get_record_class()
 
 ifuse_cache:
-statement=self._stmt_cache.get((query, record_class))
+statement=self._stmt_cache.get(
+ (query, record_class, ignore_custom_codec)
+ )
 ifstatementisnotNone:
 returnstatement
 
@@ -371,6 +374,7 @@ async def _get_statement(
 query,
 timeout,
 record_class=record_class,
+ignore_custom_codec=ignore_custom_codec,
  )
 need_reprepare=False
 types_with_missing_codecs=statement._init_types()
@@ -415,7 +419,8 @@ async def _get_statement(
  )
 
 ifuse_cache:
-self._stmt_cache.put((query, record_class), statement)
+self._stmt_cache.put(
+ (query, record_class, ignore_custom_codec), statement)
 
 # If we've just created a new statement object, check if there
 # are any statements for GC.
@@ -426,7 +431,12 @@ async def _get_statement(
 
 asyncdef_introspect_types(self, typeoids, timeout):
 returnawaitself.__execute(
-self._intro_query, (list(typeoids),), 0, timeout)
+self._intro_query,
+ (list(typeoids),),
+0,
+timeout,
+ignore_custom_codec=True,
+ )
 
 asyncdef_introspect_type(self, typename, schema):
 if (
@@ -439,20 +449,22 @@ async def _introspect_type(self, typename, schema):
  [typeoid],
 limit=0,
 timeout=None,
+ignore_custom_codec=True,
  )
-ifrows:
-typeinfo=rows[0]
-else:
-typeinfo=None
 else:
-typeinfo=awaitself.fetchrow(
-introspection.TYPE_BY_NAME, typename, schema)
+rows=awaitself._execute(
+introspection.TYPE_BY_NAME,
+ [typename, schema],
+limit=1,
+timeout=None,
+ignore_custom_codec=True,
+ )
 
-ifnottypeinfo:
+ifnotrows:
 raiseValueError(
 'unknown type:{}.{}'.format(schema, typename))
 
-returntypeinfo
+returnrows[0]
 
 defcursor(
 self,
@@ -1325,7 +1337,9 @@ def _mark_stmts_as_closed(self):
 def_maybe_gc_stmt(self, stmt):
 if (
 stmt.refs==0
-andnotself._stmt_cache.has((stmt.query, stmt.record_class))
+andnotself._stmt_cache.has(
+ (stmt.query, stmt.record_class, stmt.ignore_custom_codec)
+ )
  ):
 # If low-level `stmt` isn't referenced from any high-level
 # `PreparedStatement` object and is not in the `_stmt_cache`:
@@ -1589,6 +1603,7 @@ async def _execute(
 timeout,
 *,
 return_status=False,
+ignore_custom_codec=False,
 record_class=None
  ):
 withself._stmt_exclusive_section:
@@ -1599,6 +1614,7 @@ async def _execute(
 timeout,
 return_status=return_status,
 record_class=record_class,
+ignore_custom_codec=ignore_custom_codec,
  )
 returnresult
 
@@ -1610,6 +1626,7 @@ async def __execute(
 timeout,
 *,
 return_status=False,
+ignore_custom_codec=False,
 record_class=None
  ):
 executor=lambdastmt, timeout: self._protocol.bind_execute(
@@ -1620,6 +1637,7 @@ async def __execute(
 executor,
 timeout,
 record_class=record_class,
+ignore_custom_codec=ignore_custom_codec,
  )
 
 asyncdef_executemany(self, query, args, timeout):
@@ -1637,20 +1655,23 @@ async def _do_execute(
 timeout,
 retry=True,
 *,
+ignore_custom_codec=False,
 record_class=None
  ):
 iftimeoutisNone:
 stmt=awaitself._get_statement(
 query,
 None,
 record_class=record_class,
+ignore_custom_codec=ignore_custom_codec,
  )
 else:
 before=time.monotonic()
 stmt=awaitself._get_statement(
 query,
 timeout,
 record_class=record_class,
+ignore_custom_codec=ignore_custom_codec,
  )
 after=time.monotonic()
 timeout-=after-before

asyncpg/protocol/codecs/base.pxd

@@ -166,5 +166,6 @@ cdef class DataCodecConfig:
 dict _derived_type_codecs
 dict _custom_type_codecs
 
- cdef inline Codec get_codec(self, uint32_t oid, ServerDataFormat format)
+ cdef inline Codec get_codec(self, uint32_t oid, ServerDataFormat format,
+ bint ignore_custom_codec=*)
  cdef inline Codec get_any_local_codec(self, uint32_t oid)

asyncpg/protocol/codecs/base.pyx

@@ -692,18 +692,20 @@ cdef class DataCodecConfig:
 
 return codec
 
- cdef inline Codec get_codec(self, uint32_t oid, ServerDataFormat format):
+ cdef inline Codec get_codec(self, uint32_t oid, ServerDataFormat format,
+ bint ignore_custom_codec=False):
  cdef Codec codec
 
- codec =self.get_any_local_codec(oid)
-if codec isnotNone:
-if codec.format != format:
-# The codec for this OID has been overridden by
-# set_{builtin}_type_codec with a different format.
-# We must respect that and not return a core codec.
-returnNone
-else:
-return codec
+ifnot ignore_custom_codec:
+ codec =self.get_any_local_codec(oid)
+if codec isnotNone:
+if codec.format != format:
+# The codec for this OID has been overridden by
+# set_{builtin}_type_codec with a different format.
+# We must respect that and not return a core codec.
+returnNone
+else:
+return codec
 
  codec = get_core_codec(oid, format)
 if codec isnotNone:

asyncpg/protocol/prepared_stmt.pxd

@@ -12,6 +12,7 @@ cdef class PreparedStatementState:
  readonly bint closed
  readonly int refs
  readonly type record_class
+ readonly bint ignore_custom_codec
 
 
 list row_desc

asyncpg/protocol/prepared_stmt.pyx

@@ -16,7 +16,8 @@ cdef class PreparedStatementState:
  str name,
  str query,
  BaseProtocol protocol,
- type record_class
+ type record_class,
+ bint ignore_custom_codec
  ):
 self.name = name
 self.query = query
@@ -28,6 +29,7 @@ cdef class PreparedStatementState:
 self.closed =False
 self.refs =0
 self.record_class = record_class
+self.ignore_custom_codec = ignore_custom_codec
 
 def_get_parameters(self):
  cdef Codec codec
@@ -205,7 +207,8 @@ cdef class PreparedStatementState:
  cols_mapping[col_name] = i
  cols_names.append(col_name)
  oid = row[3]
- codec =self.settings.get_data_codec(oid)
+ codec =self.settings.get_data_codec(
+ oid, ignore_custom_codec=self.ignore_custom_codec)
 if codec isNoneornot codec.has_decoder():
 raise exceptions.InternalClientError(
 'no decoder for OID{}'.format(oid))
@@ -230,7 +233,8 @@ cdef class PreparedStatementState:
 
 for i from0<= i <self.args_num:
  p_oid =self.parameters_desc[i]
- codec =self.settings.get_data_codec(p_oid)
+ codec =self.settings.get_data_codec(
+ p_oid, ignore_custom_codec=self.ignore_custom_codec)
 if codec isNoneornot codec.has_encoder():
 raise exceptions.InternalClientError(
 'no encoder for OID{}'.format(p_oid))

asyncpg/protocol/protocol.pyx

@@ -145,6 +145,7 @@ cdef class BaseProtocol(CoreProtocol):
  async def prepare(self, stmt_name, query, timeout,
 *,
  PreparedStatementState state=None,
+ignore_custom_codec=False,
  record_class):
 ifself.cancel_waiter isnotNone:
  await self.cancel_waiter
@@ -161,7 +162,7 @@ cdef class BaseProtocol(CoreProtocol):
 self.last_query = query
 if state isNone:
  state = PreparedStatementState(
- stmt_name, query, self, record_class)
+ stmt_name, query, self, record_class, ignore_custom_codec)
 self.statement = state
 exceptExceptionas ex:
  waiter.set_exception(ex)

asyncpg/protocol/settings.pxd

@@ -26,4 +26,5 @@ cdef class ConnectionSettings(pgproto.CodecContext):
  cpdef inline set_builtin_type_codec(
 self, typeoid, typename, typeschema, typekind, alias_to, format)
  cpdef inline Codec get_data_codec(
-self, uint32_t oid, ServerDataFormat format=*)
+self, uint32_t oid, ServerDataFormat format=*,
+ bint ignore_custom_codec=*)

asyncpg/protocol/settings.pyx

@@ -87,14 +87,18 @@ cdef class ConnectionSettings(pgproto.CodecContext):
  typekind, alias_to, _format)
 
  cpdef inline Codec get_data_codec(self, uint32_t oid,
- ServerDataFormat format=PG_FORMAT_ANY):
+ ServerDataFormat format=PG_FORMAT_ANY,
+ bint ignore_custom_codec=False):
 if format == PG_FORMAT_ANY:
- codec =self._data_codecs.get_codec(oid, PG_FORMAT_BINARY)
+ codec =self._data_codecs.get_codec(
+ oid, PG_FORMAT_BINARY, ignore_custom_codec)
 if codec isNone:
- codec =self._data_codecs.get_codec(oid, PG_FORMAT_TEXT)
+ codec =self._data_codecs.get_codec(
+ oid, PG_FORMAT_TEXT, ignore_custom_codec)
 return codec
 else:
-returnself._data_codecs.get_codec(oid, format)
+returnself._data_codecs.get_codec(
+ oid, format, ignore_custom_codec)
 
 def__getattr__(self, name):
 ifnot name.startswith('_'):

tests/test_introspection.py

@@ -43,6 +43,20 @@ def tearDownClass(cls):
 
 super().tearDownClass()
 
+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(
@@ -142,6 +156,7 @@ async def test_introspection_retries_after_cache_bust(self):
 # 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;