Add support for COPY OUT

This commit adds two new Connection methods: copy_from_table() and
copy_from_query() that allow copying the contents of a table or
the results of a query using PostgreSQL COPY protocol.

With help from @1st1.

Issue MagicStack#21.

Author: elprans

Diff for: asyncpg/compat.py

@@ -6,9 +6,13 @@
 
 
 import functools
+import os
 import sys
 
 
+PY_36 = sys.version_info >= (3, 6)
+
+
 if sys.version_info < (3, 5, 2):
     def aiter_compat(func):
         @functools.wraps(func)
@@ -18,3 +22,26 @@ async def wrapper(self):
 else:
     def aiter_compat(func):
         return func
+
+
+if PY_36:
+    fspath = os.fspath
+else:
+    def fspath(path):
+        fsp = getattr(path, '__fspath__', None)
+        if fsp is not None and callable(fsp):
+            path = fsp()
+            if not isinstance(path, (str, bytes)):
+                raise TypeError(
+                    'expected {}() to return str or bytes, not {}'.format(
+                        fsp.__qualname__, type(path).__name__
+                    ))
+            return path
+        elif isinstance(path, (str, bytes)):
+            return path
+        else:
+            raise TypeError(
+                'expected str, bytes or path-like object, not {}'.format(
+                    type(path).__name__
+                )
+            )

Diff for: asyncpg/connection.py

@@ -10,6 +10,7 @@
 import struct
 import time
 
+from . import compat
 from . import connect_utils
 from . import cursor
 from . import exceptions
@@ -18,6 +19,7 @@
 from . import protocol
 from . import serverversion
 from . import transaction
+from . import utils
 
 
 class ConnectionMeta(type):
@@ -345,6 +347,178 @@ async def fetchrow(self, query, *args, timeout=None):
             return None
         return data[0]
 
+    async def copy_from_table(self, table_name, *, output,
+                              columns=None, schema_name=None, timeout=None,
+                              format=None, oids=None, delimiter=None,
+                              null=None, header=None, quote=None,
+                              escape=None, force_quote=None, encoding=None):
+        """Copy table contents to a file or file-like object.
+
+        :param str table_name:
+            The name of the table to copy data from.
+
+        :param output:
+            A :term:`path-like object <python:path-like object>`,
+            or a :term:`file-like object <python:file-like object>`, or
+            a :term:`coroutine function <python:coroutine function>`
+            that takes a ``bytes`` instance as a sole argument.
+
+        :param list columns:
+            An optional list of column names to copy.
+
+        :param str schema_name:
+            An optional schema name to qualify the table.
+
+        :param float timeout:
+            Optional timeout value in seconds.
+
+        The remaining kewyword arguments are ``COPY`` statement options,
+        see `COPY statement documentation`_ for details.
+
+        :return: The status string of the COPY command.
+
+        .. versionadded:: 0.11.0
+
+        .. _`COPY statement documentation`: https://www.postgresql.org/docs/\
+                                            current/static/sql-copy.html
+
+        """
+        tabname = utils._quote_ident(table_name)
+        if schema_name:
+            tabname = utils._quote_ident(schema_name) + '.' + tabname
+
+        if columns:
+            cols = '({})'.format(
+                ', '.join(utils._quote_ident(c) for c in columns))
+        else:
+            cols = ''
+
+        opts = self._format_copy_opts(
+            format=format, oids=oids, delimiter=delimiter,
+            null=null, header=header, quote=quote, escape=escape,
+            force_quote=force_quote, encoding=encoding
+        )
+
+        copy_stmt = 'COPY {tab}{cols} TO STDOUT {opts}'.format(
+            tab=tabname, cols=cols, opts=opts)
+
+        return await self._copy_out(copy_stmt, output, timeout)
+
+    async def copy_from_query(self, query, *args, output,
+                              timeout=None, format=None, oids=None,
+                              delimiter=None, null=None, header=None,
+                              quote=None, escape=None, force_quote=None,
+                              encoding=None):
+        """Copy the results of a query to a file or file-like object.
+
+        :param str query:
+            The query to copy the results of.
+
+        :param *args:
+            Query arguments.
+
+        :param output:
+            A :term:`path-like object <python:path-like object>`,
+            or a :term:`file-like object <python:file-like object>`, or
+            a :term:`coroutine function <python:coroutine function>`
+            that takes a ``bytes`` instance as a sole argument.
+
+        :param float timeout:
+            Optional timeout value in seconds.
+
+        The remaining kewyword arguments are ``COPY`` statement options,
+        see `COPY statement documentation`_ for details.
+
+        :return: The status string of the COPY command.
+
+        .. versionadded:: 0.11.0
+
+        .. _`COPY statement documentation`: https://www.postgresql.org/docs/\
+                                            current/static/sql-copy.html
+
+        """
+        opts = self._format_copy_opts(
+            format=format, oids=oids, delimiter=delimiter,
+            null=null, header=header, quote=quote, escape=escape,
+            force_quote=force_quote, encoding=encoding
+        )
+
+        if args:
+            query = await utils._mogrify(self, query, args)
+
+        copy_stmt = 'COPY ({query}) TO STDOUT {opts}'.format(
+            query=query, opts=opts)
+
+        return await self._copy_out(copy_stmt, output, timeout)
+
+    def _format_copy_opts(self, *, format=None, oids=None, freeze=None,
+                          delimiter=None, null=None, header=None, quote=None,
+                          escape=None, force_quote=None, force_not_null=None,
+                          force_null=None, encoding=None):
+        kwargs = dict(locals())
+        kwargs.pop('self')
+        opts = []
+
+        if force_quote is not None and isinstance(force_quote, bool):
+            kwargs.pop('force_quote')
+            if force_quote:
+                opts.append('FORCE_QUOTE *')
+
+        for k, v in kwargs.items():
+            if v is not None:
+                if k in ('force_not_null', 'force_null', 'force_quote'):
+                    v = '(' + ', '.join(utils._quote_ident(c) for c in v) + ')'
+                elif k in ('oids', 'freeze', 'header'):
+                    v = str(v)
+                else:
+                    v = utils._quote_literal(v)
+
+                opts.append('{} {}'.format(k.upper(), v))
+
+        if opts:
+            return '(' + ', '.join(opts) + ')'
+        else:
+            return ''
+
+    async def _copy_out(self, copy_stmt, output, timeout):
+        try:
+            path = compat.fspath(output)
+        except TypeError:
+            # output is not a path-like object
+            path = None
+
+        writer = None
+        opened_by_us = False
+        run_in_executor = self._loop.run_in_executor
+
+        if path is not None:
+            # a path
+            f = await run_in_executor(None, open, path, 'wb')
+            opened_by_us = True
+        elif hasattr(output, 'write'):
+            # file-like
+            f = output
+        elif callable(output):
+            # assuming calling output returns an awaitable.
+            writer = output
+        else:
+            raise TypeError(
+                'output is expected to be a file-like object, '
+                'a path-like object or a coroutine function, '
+                'not {}'.format(type(output).__name__)
+            )
+
+        if writer is None:
+            async def _writer(data):
+                await run_in_executor(None, f.write, data)
+            writer = _writer
+
+        try:
+            return await self._protocol.copy_out(copy_stmt, writer, timeout)
+        finally:
+            if opened_by_us:
+                f.close()
+
     async def set_type_codec(self, typename, *,
                              schema='public', encoder, decoder, binary=False):
         """Set an encoder/decoder pair for the specified data type.

Diff for: asyncpg/protocol/buffer.pxd

@@ -99,14 +99,17 @@ cdef class ReadBuffer:
     cdef _switch_to_next_buf(self)
     cdef inline read_byte(self)
     cdef inline char* _try_read_bytes(self, ssize_t nbytes)
-    cdef inline read(self, ssize_t nbytes)
+    cdef inline _read(self, char *buf, ssize_t nbytes)
+    cdef read(self, ssize_t nbytes)
     cdef inline read_bytes(self, ssize_t n)
     cdef inline read_int32(self)
     cdef inline read_int16(self)
     cdef inline read_cstr(self)
     cdef int32_t has_message(self) except -1
+    cdef inline int32_t has_message_type(self, char mtype) except -1
     cdef inline char* try_consume_message(self, ssize_t* len)
     cdef Memory consume_message(self)
+    cdef bytearray consume_messages(self, char mtype)
     cdef discard_message(self)
     cdef inline _discard_message(self)
     cdef inline char get_message_type(self)

Diff for: asyncpg/protocol/buffer.pyx

@@ -296,6 +296,11 @@ cdef class ReadBuffer:
                     'debug: second buffer of ReadBuffer is empty')
 
     cdef inline char* _try_read_bytes(self, ssize_t nbytes):
+        # Try to read *nbytes* from the first buffer.
+        #
+        # Returns pointer to data if there is at least *nbytes*
+        # in the buffer, NULL otherwise.
+        #
         # Important: caller must call _ensure_first_buf() prior
         # to calling try_read_bytes, and must not overread
 
@@ -321,9 +326,32 @@ cdef class ReadBuffer:
         else:
             return NULL
 
-    cdef inline read(self, ssize_t nbytes):
+    cdef inline _read(self, char *buf, ssize_t nbytes):
         cdef:
-            object result
+            ssize_t nread
+            char *buf0
+
+        while True:
+            buf0 = cpython.PyBytes_AS_STRING(self._buf0)
+
+            if self._pos0 + nbytes > self._len0:
+                nread = self._len0 - self._pos0
+                memcpy(buf, buf0 + self._pos0, <size_t>nread)
+                self._pos0 = self._len0
+                self._length -= nread
+                nbytes -= nread
+                buf += nread
+                self._ensure_first_buf()
+
+            else:
+                memcpy(buf, buf0 + self._pos0, <size_t>nbytes)
+                self._pos0 += nbytes
+                self._length -= nbytes
+                break
+
+    cdef read(self, ssize_t nbytes):
+        cdef:
+            bytearray result
             ssize_t nread
             char *buf
 
@@ -341,24 +369,11 @@ cdef class ReadBuffer:
             if self._current_message_len_unread < 0:
                 raise BufferError('buffer overread')
 
-        result = bytearray()
-        while True:
-            if self._pos0 + nbytes > self._len0:
-                result.extend(self._buf0[self._pos0:])
-                nread = self._len0 - self._pos0
-                self._pos0 = self._len0
-                self._length -= nread
-                nbytes -= nread
-                self._ensure_first_buf()
+        result = PyByteArray_FromStringAndSize(NULL, nbytes)
+        buf = PyByteArray_AsString(result)
+        self._read(buf, nbytes)
 
-            else:
-                result.extend(self._buf0[self._pos0:self._pos0 + nbytes])
-                self._pos0 += nbytes
-                self._length -= nbytes
-                return Memory.new(
-                    PyByteArray_AsString(result),
-                    result,
-                    len(result))
+        return Memory.new(buf, result, nbytes)
 
     cdef inline read_byte(self):
         cdef char* first_byte
@@ -506,6 +521,9 @@ cdef class ReadBuffer:
         self._current_message_ready = 1
         return 1
 
+    cdef inline int32_t has_message_type(self, char mtype) except -1:
+        return self.has_message() and self.get_message_type() == mtype
+
     cdef inline char* try_consume_message(self, ssize_t* len):
         cdef ssize_t buf_len
 
@@ -530,6 +548,35 @@ cdef class ReadBuffer:
         self._discard_message()
         return mem
 
+    cdef bytearray consume_messages(self, char mtype):
+        """Consume consecutive messages of the same type."""
+        cdef:
+            char *buf
+            ssize_t nbytes
+            ssize_t total_bytes = 0
+            bytearray result
+
+        if not self.has_message_type(mtype):
+            return None
+
+        # consume_messages is a volume-oriented method, so
+        # we assume that the remainder of the buffer will contain
+        # messages of the requested type.
+        result = PyByteArray_FromStringAndSize(NULL, self._length)
+        buf = PyByteArray_AsString(result)
+
+        while self.has_message_type(mtype):
+            nbytes = self._current_message_len_unread
+            self._read(buf, nbytes)
+            buf += nbytes
+            total_bytes += nbytes
+            self._discard_message()
+
+        # Clamp the result to an actual size read.
+        PyByteArray_Resize(result, total_bytes)
+
+        return result
+
     cdef discard_message(self):
         if self._current_message_type == 0:
             # Already discarded