Implement support for multirange types

Author: elprans

asyncpg/introspection.py

@@ -39,7 +39,9 @@
  END) AS basetype,
  t.typelem AS elemtype,
  elem_t.typdelim AS elemdelim,
- range_t.rngsubtype AS range_subtype,
+ COALESCE(
+ range_t.rngsubtype,
+ multirange_t.rngsubtype) AS range_subtype,
  (CASE WHEN t.typtype = 'c' THEN
  (SELECT
  array_agg(ia.atttypid ORDER BY ia.attnum)
@@ -78,6 +80,9 @@
  LEFT JOIN pg_range range_t ON (
  t.oid = range_t.rngtypid
  )
+ LEFT JOIN pg_range multirange_t ON (
+ t.oid = multirange_t.rngmultitypid
+ )
  )
 '''
 

asyncpg/protocol/codecs/array.pyx

@@ -858,19 +858,7 @@ cdef arraytext_decode(ConnectionSettings settings, FRBuffer *buf):
  return array_decode(settings, buf, <decode_func_ex>&text_decode_ex, NULL)
 
 
-cdef anyarray_decode(ConnectionSettings settings, FRBuffer *buf):
- # Instances of anyarray (or any other polymorphic pseudotype) are
- # never supposed to be returned from actual queries.
- raise exceptions.ProtocolError(
- 'unexpected instance of \'anyarray\' type')
-
-
 cdef init_array_codecs():
- register_core_codec(ANYARRAYOID,
- NULL,
- <decode_func>&anyarray_decode,
- PG_FORMAT_BINARY)
-
  # oid[] and text[] are registered as core codecs
  # to make type introspection query work
  #

asyncpg/protocol/codecs/base.pxd

@@ -23,12 +23,13 @@ ctypedef object (*codec_decode_func)(Codec codec,
 
 
 cdef enum CodecType:
- CODEC_UNDEFINED = 0
- CODEC_C = 1
- CODEC_PY = 2
- CODEC_ARRAY = 3
- CODEC_COMPOSITE = 4
- CODEC_RANGE = 5
+ CODEC_UNDEFINED = 0
+ CODEC_C = 1
+ CODEC_PY = 2
+ CODEC_ARRAY = 3
+ CODEC_COMPOSITE = 4
+ CODEC_RANGE = 5
+ CODEC_MULTIRANGE = 6
 
 
 cdef enum ServerDataFormat:
@@ -95,6 +96,9 @@ cdef class Codec:
  cdef encode_range(self, ConnectionSettings settings, WriteBuffer buf,
  object obj)
 
+ cdef encode_multirange(self, ConnectionSettings settings, WriteBuffer buf,
+ object obj)
+
  cdef encode_composite(self, ConnectionSettings settings, WriteBuffer buf,
  object obj)
 
@@ -109,6 +113,8 @@ cdef class Codec:
 
  cdef decode_range(self, ConnectionSettings settings, FRBuffer *buf)
 
+ cdef decode_multirange(self, ConnectionSettings settings, FRBuffer *buf)
+
  cdef decode_composite(self, ConnectionSettings settings, FRBuffer *buf)
 
  cdef decode_in_python(self, ConnectionSettings settings, FRBuffer *buf)
@@ -139,6 +145,12 @@ cdef class Codec:
  str schema,
  Codec element_codec)
 
+ @staticmethod
+ cdef Codec new_multirange_codec(uint32_t oid,
+ str name,
+ str schema,
+ Codec element_codec)
+
  @staticmethod
  cdef Codec new_composite_codec(uint32_t oid,
  str name,

asyncpg/protocol/codecs/base.pyx

@@ -71,6 +71,13 @@ cdef class Codec:
  'range types is not supported'.format(schema, name))
  self.encoder = <codec_encode_func>&self.encode_range
  self.decoder = <codec_decode_func>&self.decode_range
+ elif type == CODEC_MULTIRANGE:
+ if format != PG_FORMAT_BINARY:
+ raise exceptions.UnsupportedClientFeatureError(
+ 'cannot decode type "{}"."{}": text encoding of '
+ 'range types is not supported'.format(schema, name))
+ self.encoder = <codec_encode_func>&self.encode_multirange
+ self.decoder = <codec_decode_func>&self.decode_multirange
  elif type == CODEC_COMPOSITE:
  if format != PG_FORMAT_BINARY:
  raise exceptions.UnsupportedClientFeatureError(
@@ -122,6 +129,12 @@ cdef class Codec:
  codec_encode_func_ex,
  <void*>(<cpython.PyObject>self.element_codec))
 
+ cdef encode_multirange(self, ConnectionSettings settings, WriteBuffer buf,
+ object obj):
+ multirange_encode(settings, buf, obj, self.element_codec.oid,
+ codec_encode_func_ex,
+ <void*>(<cpython.PyObject>self.element_codec))
+
  cdef encode_composite(self, ConnectionSettings settings, WriteBuffer buf,
  object obj):
  cdef:
@@ -209,6 +222,10 @@ cdef class Codec:
  return range_decode(settings, buf, codec_decode_func_ex,
  <void*>(<cpython.PyObject>self.element_codec))
 
+ cdef decode_multirange(self, ConnectionSettings settings, FRBuffer *buf):
+ return multirange_decode(settings, buf, codec_decode_func_ex,
+ <void*>(<cpython.PyObject>self.element_codec))
+
  cdef decode_composite(self, ConnectionSettings settings,
  FRBuffer *buf):
  cdef:
@@ -294,7 +311,11 @@ cdef class Codec:
  if self.c_encoder is not NULL or self.py_encoder is not None:
  return True
 
- elif self.type == CODEC_ARRAY or self.type == CODEC_RANGE:
+ elif (
+ self.type == CODEC_ARRAY
+ or self.type == CODEC_RANGE
+ or self.type == CODEC_MULTIRANGE
+ ):
  return self.element_codec.has_encoder()
 
  elif self.type == CODEC_COMPOSITE:
@@ -312,7 +333,11 @@ cdef class Codec:
  if self.c_decoder is not NULL or self.py_decoder is not None:
  return True
 
- elif self.type == CODEC_ARRAY or self.type == CODEC_RANGE:
+ elif (
+ self.type == CODEC_ARRAY
+ or self.type == CODEC_RANGE
+ or self.type == CODEC_MULTIRANGE
+ ):
  return self.element_codec.has_decoder()
 
  elif self.type == CODEC_COMPOSITE:
@@ -358,6 +383,18 @@ cdef class Codec:
  None, None, None, 0)
  return codec
 
+ @staticmethod
+ cdef Codec new_multirange_codec(uint32_t oid,
+ str name,
+ str schema,
+ Codec element_codec):
+ cdef Codec codec
+ codec = Codec(oid)
+ codec.init(name, schema, 'multirange', CODEC_MULTIRANGE,
+ element_codec.format, PG_XFORMAT_OBJECT, NULL, NULL,
+ None, None, element_codec, None, None, None, 0)
+ return codec
+
  @staticmethod
  cdef Codec new_composite_codec(uint32_t oid,
  str name,
@@ -536,6 +573,21 @@ cdef class DataCodecConfig:
  self._derived_type_codecs[oid, elem_codec.format] = \
  Codec.new_range_codec(oid, name, schema, elem_codec)
 
+ elif ti['kind'] == b'm':
+ # Multirange type
+
+ if not range_subtype_oid:
+ raise exceptions.InternalClientError(
+ f'type record missing base type for multirange {oid}')
+
+ elem_codec = self.get_codec(range_subtype_oid, PG_FORMAT_ANY)
+ if elem_codec is None:
+ elem_codec = self.declare_fallback_codec(
+ range_subtype_oid, ti['range_subtype_name'], schema)
+
+ self._derived_type_codecs[oid, elem_codec.format] = \
+ Codec.new_multirange_codec(oid, name, schema, elem_codec)
+
  elif ti['kind'] == b'e':
  # Enum types are essentially text
  self._set_builtin_type_codec(oid, name, schema, 'scalar',

asyncpg/protocol/codecs/pgproto.pyx

@@ -273,8 +273,9 @@ cdef init_pseudo_codecs():
  FDW_HANDLEROID, TSM_HANDLEROID, INTERNALOID, OPAQUEOID,
  ANYELEMENTOID, ANYNONARRAYOID, ANYCOMPATIBLEOID,
  ANYCOMPATIBLEARRAYOID, ANYCOMPATIBLENONARRAYOID,
- ANYCOMPATIBLERANGEOID, PG_DDL_COMMANDOID, INDEX_AM_HANDLEROID,
- TABLE_AM_HANDLEROID,
+ ANYCOMPATIBLERANGEOID, ANYCOMPATIBLEMULTIRANGEOID,
+ ANYRANGEOID, ANYMULTIRANGEOID, ANYARRAYOID,
+ PG_DDL_COMMANDOID, INDEX_AM_HANDLEROID, TABLE_AM_HANDLEROID,
  ]
 
  register_core_codec(ANYENUMOID,
@@ -330,6 +331,19 @@ cdef init_pseudo_codecs():
  <decode_func>pgproto.bytea_decode,
  PG_FORMAT_BINARY)
 
+ # These two are internal to BRIN index support and are unlikely
+ # to be sent, but since I/O functions for these exist, add decoders
+ # nonetheless.
+ register_core_codec(PG_BRIN_BLOOM_SUMMARYOID,
+ NULL,
+ <decode_func>pgproto.bytea_decode,
+ PG_FORMAT_BINARY)
+
+ register_core_codec(PG_BRIN_MINMAX_MULTI_SUMMARYOID,
+ NULL,
+ <decode_func>pgproto.bytea_decode,
+ PG_FORMAT_BINARY)
+
 
 cdef init_text_codecs():
  textoids = [

asyncpg/protocol/codecs/range.pyx

@@ -7,6 +7,8 @@
 
 from asyncpg import types as apg_types
 
+from collections.abc import Sequence as SequenceABC
+
 # defined in postgresql/src/include/utils/rangetypes.h
 DEF RANGE_EMPTY = 0x01 # range is empty
 DEF RANGE_LB_INC = 0x02 # lower bound is inclusive
@@ -139,11 +141,55 @@ cdef range_decode(ConnectionSettings settings, FRBuffer *buf,
  empty=(flags & RANGE_EMPTY) != 0)
 
 
-cdef init_range_codecs():
- register_core_codec(ANYRANGEOID,
- NULL,
- <decode_func>pgproto.text_decode,
- PG_FORMAT_TEXT)
+cdef multirange_encode(ConnectionSettings settings, WriteBuffer buf,
+ object obj, uint32_t elem_oid,
+ encode_func_ex encoder, const void *encoder_arg):
+ cdef:
+ WriteBuffer elem_data
+
+ if not isinstance(obj, SequenceABC):
+ raise TypeError(
+ 'expected a sequence (got type {!r})'.format(type(obj).__name__)
+ )
+
+ elem_data = WriteBuffer.new()
+
+ for elem in obj:
+ range_encode(settings, elem_data, elem, elem_oid, encoder, encoder_arg)
+
+ # Datum length
+ buf.write_int32(4 + elem_data.len())
+ # Number of elements in multirange
+ buf.write_int32(len(obj))
+ buf.write_buffer(elem_data)
+
 
+cdef multirange_decode(ConnectionSettings settings, FRBuffer *buf,
+ decode_func_ex decoder, const void *decoder_arg):
+ cdef:
+ int32_t nelems = hton.unpack_int32(frb_read(buf, 4))
+ FRBuffer elem_buf
+ int32_t elem_len
+ int i
+ list result
+
+ if nelems == 0:
+ return []
+
+ if nelems < 0:
+ raise exceptions.ProtocolError(
+ 'unexpected multirange size value: {}'.format(nelems))
+
+ result = cpython.PyList_New(nelems)
+ for i in range(nelems):
+ elem_len = hton.unpack_int32(frb_read(buf, 4))
+ if elem_len == -1:
+ raise exceptions.ProtocolError(
+ 'unexpected NULL element in multirange value')
+ else:
+ frb_slice_from(&elem_buf, buf, elem_len)
+ elem = range_decode(settings, &elem_buf, decoder, decoder_arg)
+ cpython.Py_INCREF(elem)
+ cpython.PyList_SET_ITEM(result, i, elem)
 
-init_range_codecs()
+ return result

asyncpg/protocol/pgtypes.pxi

@@ -101,6 +101,10 @@ DEF JSONPATHOID = 4072
 DEF REGNAMESPACEOID = 4089
 DEF REGROLEOID = 4096
 DEF REGCOLLATIONOID = 4191
+DEF ANYMULTIRANGEOID = 4537
+DEF ANYCOMPATIBLEMULTIRANGEOID = 4538
+DEF PG_BRIN_BLOOM_SUMMARYOID = 4600
+DEF PG_BRIN_MINMAX_MULTI_SUMMARYOID = 4601
 DEF PG_MCV_LISTOID = 5017
 DEF PG_SNAPSHOTOID = 5038
 DEF XID8OID = 5069
@@ -116,11 +120,13 @@ BUILTIN_TYPE_OID_MAP = {
  ACLITEMOID: 'aclitem',
  ANYARRAYOID: 'anyarray',
  ANYCOMPATIBLEARRAYOID: 'anycompatiblearray',
+ ANYCOMPATIBLEMULTIRANGEOID: 'anycompatiblemultirange',
  ANYCOMPATIBLENONARRAYOID: 'anycompatiblenonarray',
  ANYCOMPATIBLEOID: 'anycompatible',
  ANYCOMPATIBLERANGEOID: 'anycompatiblerange',
  ANYELEMENTOID: 'anyelement',
  ANYENUMOID: 'anyenum',
+ ANYMULTIRANGEOID: 'anymultirange',
  ANYNONARRAYOID: 'anynonarray',
  ANYOID: 'any',
  ANYRANGEOID: 'anyrange',
@@ -161,6 +167,8 @@ BUILTIN_TYPE_OID_MAP = {
  OIDOID: 'oid',
  OPAQUEOID: 'opaque',
  PATHOID: 'path',
+ PG_BRIN_BLOOM_SUMMARYOID: 'pg_brin_bloom_summary',
+ PG_BRIN_MINMAX_MULTI_SUMMARYOID: 'pg_brin_minmax_multi_summary',
  PG_DDL_COMMANDOID: 'pg_ddl_command',
  PG_DEPENDENCIESOID: 'pg_dependencies',
  PG_LSNOID: 'pg_lsn',

tests/test_codecs.py

@@ -1042,6 +1042,59 @@ async def test_range_types(self):
  dic = {obj_a: 1, obj_b: 2}
  self.assertEqual(len(dic), count)
 
+ async def test_multirange_types(self):
+ """Test encoding/decoding of multirange types."""
+
+ if self.server_version < (14, 0):
+ self.skipTest("this server does not support multirange types")
+
+ cases = [
+ ('int4multirange', [
+ [
+ [],
+ []
+ ],
+ [
+ [()],
+ []
+ ],
+ [
+ [asyncpg.Range(empty=True)],
+ []
+ ],
+ [
+ [asyncpg.Range(0, 9, lower_inc=False, upper_inc=True)],
+ [asyncpg.Range(1, 10)]
+ ],
+ [
+ [(1, 9), (9, 11)],
+ [asyncpg.Range(1, 12)]
+ ],
+ [
+ [(1, 9), (20, 30)],
+ [asyncpg.Range(1, 10), asyncpg.Range(20, 31)]
+ ],
+ [
+ [(None, 2)],
+ [asyncpg.Range(None, 3)],
+ ]
+ ])
+ ]
+
+ for (typname, sample_data) in cases:
+ st = await self.con.prepare(
+ "SELECT $1::" + typname
+ )
+
+ for sample, expected in sample_data:
+ with self.subTest(sample=sample, typname=typname):
+ result = await st.fetchval(sample)
+ self.assertEqual(result, expected)
+
+ with self.assertRaisesRegex(
+ asyncpg.DataError, 'expected a sequence'):
+ await self.con.fetch("SELECT $1::int4multirange", 1)
+
  async def test_extra_codec_alias(self):
  """Test encoding/decoding of a builtin non-pg_catalog codec."""
  await self.con.execute('''