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('''