Add support for the text I/O fallback for arrays of unknown types.

Currently, asyncpg supports falling back to text I/O for types that do
not have a registered codec.  This commit extends this behaviour to
arrays of such types.

Additionally, add an explicit error for when the text I/O is attempted
for a range or a composite type, as these are not supported yet.

Fixes: #82.

Author: elprans

asyncpg/protocol/buffer.pxd

@@ -48,7 +48,7 @@ cdef class WriteBuffer:
     cdef write_bytes(self, bytes data)
     cdef write_bytestring(self, bytes string)
     cdef write_str(self, str string, str encoding)
-    cdef write_cstr(self, char *data, ssize_t len)
+    cdef write_cstr(self, const char *data, ssize_t len)
     cdef write_int16(self, int16_t i)
     cdef write_int32(self, int32_t i)
     cdef write_int64(self, int64_t i)

asyncpg/protocol/buffer.pyx

@@ -169,7 +169,7 @@ cdef class WriteBuffer:
     cdef write_str(self, str string, str encoding):
         self.write_bytestring(string.encode(encoding))
 
-    cdef write_cstr(self, char *data, ssize_t len):
+    cdef write_cstr(self, const char *data, ssize_t len):
         self._check_readonly()
         self._ensure_alloced(len)
 

asyncpg/protocol/codecs/array.pyx

@@ -140,6 +140,127 @@ cdef inline array_encode(ConnectionSettings settings, WriteBuffer buf,
     buf.write_buffer(elem_data)
 
 
+cdef _write_textarray_data(ConnectionSettings settings, object obj,
+                           int32_t ndims, int32_t dim, WriteBuffer array_data,
+                           encode_func_ex encoder, const void *encoder_arg,
+                           Py_UCS4 typdelim):
+    cdef:
+        ssize_t i = 0
+        int8_t delim = <int8_t>typdelim
+        WriteBuffer elem_data
+        Py_buffer pybuf
+        const char *elem_str
+        char ch
+        ssize_t elem_len
+        ssize_t quoted_elem_len
+        bint need_quoting
+
+    array_data.write_byte(b'{')
+
+    if dim < ndims - 1:
+        for item in obj:
+            if i > 0:
+                array_data.write_byte(delim)
+                array_data.write_byte(b' ')
+            _write_textarray_data(settings, item, ndims, dim + 1, array_data,
+                                  encoder, encoder_arg, typdelim)
+            i += 1
+    else:
+        for item in obj:
+            elem_data = WriteBuffer.new()
+
+            if i > 0:
+                array_data.write_byte(delim)
+                array_data.write_byte(b' ')
+
+            if item is None:
+                array_data.write_bytes(b'NULL')
+                i += 1
+                continue
+            else:
+                try:
+                    encoder(settings, elem_data, item, encoder_arg)
+                except TypeError as e:
+                    raise ValueError(
+                        'invalid array element: {}'.format(
+                            e.args[0])) from None
+
+            # element string length (first four bytes are the encoded length.)
+            elem_len = elem_data.len() - 4
+
+            if elem_len == 0:
+                # Empty string
+                array_data.write_bytes(b'""')
+            else:
+                cpython.PyObject_GetBuffer(
+                    elem_data, &pybuf, cpython.PyBUF_SIMPLE)
+
+                elem_str = <const char*>(pybuf.buf) + 4
+
+                try:
+                    if not apg_strcasecmp_char(elem_str, b'NULL'):
+                        array_data.write_bytes(b'"NULL"')
+                    else:
+                        quoted_elem_len = elem_len
+                        need_quoting = False
+
+                        for i in range(elem_len):
+                            ch = elem_str[i]
+                            if ch == b'"' or ch == b'\\':
+                                # Quotes and backslashes need escaping.
+                                quoted_elem_len += 1
+                                need_quoting = True
+                            elif (ch == b'{' or ch == b'}' or ch == delim or
+                                    apg_ascii_isspace(<uint32_t>ch)):
+                                need_quoting = True
+
+                        if need_quoting:
+                            array_data.write_byte(b'"')
+
+                            if quoted_elem_len == elem_len:
+                                array_data.write_cstr(elem_str, elem_len)
+                            else:
+                                # Escaping required.
+                                for i in range(elem_len):
+                                    ch = elem_str[i]
+                                    if ch == b'"' or ch == b'\\':
+                                        array_data.write_byte(b'\\')
+                                    array_data.write_byte(ch)
+
+                            array_data.write_byte(b'"')
+                        else:
+                            array_data.write_cstr(elem_str, elem_len)
+                finally:
+                    cpython.PyBuffer_Release(&pybuf)
+
+            i += 1
+
+    array_data.write_byte(b'}')
+
+
+cdef inline textarray_encode(ConnectionSettings settings, WriteBuffer buf,
+                             object obj, encode_func_ex encoder,
+                             const void *encoder_arg, Py_UCS4 typdelim):
+    cdef:
+        WriteBuffer array_data
+        int32_t dims[ARRAY_MAXDIM]
+        int32_t ndims = 1
+        int32_t i
+
+    if not _is_container(obj):
+        raise TypeError(
+            'a non-trivial iterable expected (got type {!r})'.format(
+                type(obj).__name__))
+
+    _get_array_shape(obj, dims, &ndims)
+
+    array_data = WriteBuffer.new()
+    _write_textarray_data(settings, obj, ndims, 0, array_data,
+                          encoder, encoder_arg, typdelim)
+    buf.write_int32(array_data.len())
+    buf.write_buffer(array_data)
+
+
 cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
                          decode_func_ex decoder, const void *decoder_arg):
     cdef:

asyncpg/protocol/codecs/base.pxd

@@ -81,6 +81,9 @@ cdef class Codec:
     cdef encode_array(self, ConnectionSettings settings, WriteBuffer buf,
                       object obj)
 
+    cdef encode_array_text(self, ConnectionSettings settings, WriteBuffer buf,
+                           object obj)
+
     cdef encode_range(self, ConnectionSettings settings, WriteBuffer buf,
                       object obj)
 
@@ -137,6 +140,7 @@ cdef class Codec:
     cdef Codec new_composite_codec(uint32_t oid,
                                    str name,
                                    str schema,
+                                   CodecFormat format,
                                    list element_codecs,
                                    tuple element_type_oids,
                                    object element_names)

asyncpg/protocol/codecs/base.pyx

@@ -50,15 +50,24 @@ cdef class Codec:
             self.encoder = <codec_encode_func>&self.encode_scalar
             self.decoder = <codec_decode_func>&self.decode_scalar
         elif type == CODEC_ARRAY:
-            self.encoder = <codec_encode_func>&self.encode_array
             if format == PG_FORMAT_BINARY:
+                self.encoder = <codec_encode_func>&self.encode_array
                 self.decoder = <codec_decode_func>&self.decode_array
             else:
+                self.encoder = <codec_encode_func>&self.encode_array_text
                 self.decoder = <codec_decode_func>&self.decode_array_text
         elif type == CODEC_RANGE:
+            if format != PG_FORMAT_BINARY:
+                raise RuntimeError(
+                    'cannot encode type "{}"."{}": text encoding of '
+                    '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_COMPOSITE:
+            if format != PG_FORMAT_BINARY:
+                raise RuntimeError(
+                    'cannot encode type "{}"."{}": text encoding of '
+                    'composite types is not supported'.format(schema, name))
             self.encoder = <codec_encode_func>&self.encode_composite
             self.decoder = <codec_decode_func>&self.decode_composite
         elif type == CODEC_PY:
@@ -91,6 +100,13 @@ cdef class Codec:
                      codec_encode_func_ex,
                      <void*>(<cpython.PyObject>self.element_codec))
 
+    cdef encode_array_text(self, ConnectionSettings settings, WriteBuffer buf,
+                           object obj):
+        return textarray_encode(settings, buf, obj,
+                                codec_encode_func_ex,
+                                <void*>(<cpython.PyObject>self.element_codec),
+                                self.element_delimiter)
+
     cdef encode_range(self, ConnectionSettings settings, WriteBuffer buf,
                       object obj):
         range_encode(settings, buf, obj, self.element_codec.oid,
@@ -269,22 +285,22 @@ cdef class Codec:
                                Codec element_codec):
         cdef Codec codec
         codec = Codec(oid)
-        codec.init(name, schema, 'range', CODEC_RANGE, PG_FORMAT_BINARY,
-                   NULL, NULL, None, None, element_codec, None, None, None,
-                   0)
+        codec.init(name, schema, 'range', CODEC_RANGE, element_codec.format,
+                   NULL, NULL, None, None, element_codec, None, None, None, 0)
         return codec
 
     @staticmethod
     cdef Codec new_composite_codec(uint32_t oid,
                                    str name,
                                    str schema,
+                                   CodecFormat format,
                                    list element_codecs,
                                    tuple element_type_oids,
                                    object element_names):
         cdef Codec codec
         codec = Codec(oid)
         codec.init(name, schema, 'composite', CODEC_COMPOSITE,
-                   PG_FORMAT_BINARY, NULL, NULL, None, None, None,
+                   format, NULL, NULL, None, None, None,
                    element_type_oids, element_names, element_codecs, 0)
         return codec
 
@@ -368,11 +384,12 @@ cdef class DataCodecConfig:
                     elem_format = PG_FORMAT_BINARY
                 else:
                     elem_format = PG_FORMAT_TEXT
+
                 elem_codec = self.get_codec(array_element_oid, elem_format)
                 if elem_codec is None:
-                    raise RuntimeError(
-                        'no codec for array element type {}'.format(
-                            array_element_oid))
+                    elem_format = PG_FORMAT_TEXT
+                    elem_codec = self.declare_fallback_codec(
+                        array_element_oid, name, schema)
 
                 elem_delim = <Py_UCS4>ti['elemdelim'][0]
 
@@ -410,9 +427,8 @@ cdef class DataCodecConfig:
 
                 self._type_codecs_cache[oid, format] = \
                     Codec.new_composite_codec(
-                        oid, name, schema, comp_elem_codecs,
-                        comp_type_attrs,
-                        element_names)
+                        oid, name, schema, format, comp_elem_codecs,
+                        comp_type_attrs, element_names)
 
             elif ti['kind'] == b'd':
                 # Domain type
@@ -424,8 +440,9 @@ cdef class DataCodecConfig:
 
                 elem_codec = self.get_codec(base_type, format)
                 if elem_codec is None:
-                    raise RuntimeError(
-                        'no codec for domain base type {}'.format(base_type))
+                    format = PG_FORMAT_TEXT
+                    elem_codec = self.declare_fallback_codec(
+                        base_type, name, schema)
 
                 self._type_codecs_cache[oid, format] = elem_codec
 
@@ -441,34 +458,18 @@ cdef class DataCodecConfig:
                     elem_format = PG_FORMAT_BINARY
                 else:
                     elem_format = PG_FORMAT_TEXT
+
                 elem_codec = self.get_codec(range_subtype_oid, elem_format)
                 if elem_codec is None:
-                    raise RuntimeError(
-                        'no codec for range element type {}'.format(
-                            range_subtype_oid))
+                    elem_format = PG_FORMAT_TEXT
+                    elem_codec = self.declare_fallback_codec(
+                        range_subtype_oid, name, schema)
 
                 self._type_codecs_cache[oid, elem_format] = \
                     Codec.new_range_codec(oid, name, schema, elem_codec)
 
             else:
-                if oid <= MAXBUILTINOID:
-                    # This is a non-BKI type, for which ayncpg has no
-                    # defined codec.  This should only happen for newly
-                    # added builtin types, for which this version of
-                    # asyncpg is lacking support.
-                    #
-                    raise NotImplementedError(
-                        'unhandled standard data type {!r} (OID {})'.format(
-                            name, oid))
-                else:
-                    # This is a non-BKI type, and as such, has no
-                    # stable OID, so no possibility of a builtin codec.
-                    # In this case, fallback to text format.  Applications
-                    # can avoid this by specifying a codec for this type
-                    # using Connection.set_type_codec().
-                    #
-                    self.set_builtin_type_codec(oid, name, schema, 'scalar',
-                                                UNKNOWNOID)
+                self.declare_fallback_codec(oid, name, schema)
 
     def add_python_codec(self, typeoid, typename, typeschema, typekind,
                          encoder, decoder, binary):
@@ -478,13 +479,20 @@ cdef class DataCodecConfig:
             Codec.new_python_codec(typeoid, typename, typeschema, typekind,
                                    encoder, decoder, format)
 
+        self.clear_type_cache()
+
     def set_builtin_type_codec(self, typeoid, typename, typeschema, typekind,
-                               alias_to):
+                               alias_to, format=PG_FORMAT_ANY):
         cdef:
             Codec codec
             Codec target_codec
 
-        for format in (PG_FORMAT_BINARY, PG_FORMAT_TEXT):
+        if format == PG_FORMAT_ANY:
+            formats = (PG_FORMAT_BINARY, PG_FORMAT_TEXT)
+        else:
+            formats = (format,)
+
+        for format in formats:
             if self.get_codec(typeoid, format) is not None:
                 raise ValueError('cannot override codec for type {}'.format(
                     typeoid))
@@ -509,9 +517,41 @@ cdef class DataCodecConfig:
                 (typeoid, PG_FORMAT_TEXT) not in self._local_type_codecs):
             raise ValueError('unknown alias target: {}'.format(alias_to))
 
+        self.clear_type_cache()
+
     def clear_type_cache(self):
         self._type_codecs_cache.clear()
 
+    def declare_fallback_codec(self, uint32_t oid, str name, str schema):
+        cdef Codec codec
+
+        codec = self.get_codec(oid, PG_FORMAT_TEXT)
+        if codec is not None:
+            return codec
+
+        if oid <= MAXBUILTINOID:
+            # This is a BKI type, for which ayncpg has no
+            # defined codec.  This should only happen for newly
+            # added builtin types, for which this version of
+            # asyncpg is lacking support.
+            #
+            raise NotImplementedError(
+                'unhandled standard data type {!r} (OID {})'.format(
+                    name, oid))
+        else:
+            # This is a non-BKI type, and as such, has no
+            # stable OID, so no possibility of a builtin codec.
+            # In this case, fallback to text format.  Applications
+            # can avoid this by specifying a codec for this type
+            # using Connection.set_type_codec().
+            #
+            self.set_builtin_type_codec(oid, name, schema, 'scalar',
+                                        TEXTOID, PG_FORMAT_TEXT)
+
+            codec = self.get_codec(oid, PG_FORMAT_TEXT)
+
+        return codec
+
     cdef inline Codec get_codec(self, uint32_t oid, CodecFormat format):
         cdef Codec codec
 

asyncpg/protocol/codecs/text.pyx

@@ -63,5 +63,9 @@ cdef init_text_codecs():
                             <decode_func>&text_decode,
                             PG_FORMAT_BINARY)
 
+        register_core_codec(oid,
+                            <encode_func>&text_encode,
+                            <decode_func>&text_decode,
+                            PG_FORMAT_TEXT)
 
 init_text_codecs()