Add support for WHERE clause in copy_to methods

Author: redgoldlace

asyncpg/connection.py

@@ -827,7 +827,7 @@ async def copy_to_table(self, table_name, *, source,
                             delimiter=None, null=None, header=None,
                             quote=None, escape=None, force_quote=None,
                             force_not_null=None, force_null=None,
-                            encoding=None):
+                            encoding=None, where=None):
         """Copy data to the specified table.
 
         :param str table_name:
@@ -846,6 +846,15 @@ async def copy_to_table(self, table_name, *, source,
         :param str schema_name:
             An optional schema name to qualify the table.
 
+        :param str where:
+            An optional condition used to filter rows when copying.
+
+            .. note::
+
+                Usage of this parameter requires support for the
+                ``COPY FROM ... WHERE`` syntax, introduced in
+                PostgreSQL version 12.
+
         :param float timeout:
             Optional timeout value in seconds.
 
@@ -873,6 +882,9 @@ async def copy_to_table(self, table_name, *, source,
             https://www.postgresql.org/docs/current/static/sql-copy.html
 
         .. versionadded:: 0.11.0
+
+        .. versionadded:: 0.27.0
+            Added ``where`` parameter.
         """
         tabname = utils._quote_ident(table_name)
         if schema_name:
@@ -884,21 +896,22 @@ async def copy_to_table(self, table_name, *, source,
         else:
             cols = ''
 
+        cond = self._format_copy_where(where)
         opts = self._format_copy_opts(
             format=format, oids=oids, freeze=freeze, delimiter=delimiter,
             null=null, header=header, quote=quote, escape=escape,
             force_not_null=force_not_null, force_null=force_null,
             encoding=encoding
         )
 
-        copy_stmt = 'COPY {tab}{cols} FROM STDIN {opts}'.format(
-            tab=tabname, cols=cols, opts=opts)
+        copy_stmt = 'COPY {tab}{cols} FROM STDIN {opts} {cond}'.format(
+            tab=tabname, cols=cols, opts=opts, cond=cond)
 
         return await self._copy_in(copy_stmt, source, timeout)
 
     async def copy_records_to_table(self, table_name, *, records,
                                     columns=None, schema_name=None,
-                                    timeout=None):
+                                    timeout=None, where=None):
         """Copy a list of records to the specified table using binary COPY.
 
         :param str table_name:
@@ -915,6 +928,16 @@ async def copy_records_to_table(self, table_name, *, records,
         :param str schema_name:
             An optional schema name to qualify the table.
 
+        :param str where:
+            An optional condition used to filter rows when copying.
+
+            .. note::
+
+                Usage of this parameter requires support for the
+                ``COPY FROM ... WHERE`` syntax, introduced in
+                PostgreSQL version 12.
+
+
         :param float timeout:
             Optional timeout value in seconds.
 
@@ -959,6 +982,9 @@ async def copy_records_to_table(self, table_name, *, records,
 
         .. versionchanged:: 0.24.0
             The ``records`` argument may be an asynchronous iterable.
+
+        .. versionadded:: 0.27.0
+            Added ``where`` parameter.
         """
         tabname = utils._quote_ident(table_name)
         if schema_name:
@@ -976,14 +1002,27 @@ async def copy_records_to_table(self, table_name, *, records,
 
         intro_ps = await self._prepare(intro_query, use_cache=True)
 
+        cond = self._format_copy_where(where)
         opts = '(FORMAT binary)'
 
-        copy_stmt = 'COPY {tab}{cols} FROM STDIN {opts}'.format(
-            tab=tabname, cols=cols, opts=opts)
+        copy_stmt = 'COPY {tab}{cols} FROM STDIN {opts} {cond}'.format(
+            tab=tabname, cols=cols, opts=opts, cond=cond)
 
         return await self._protocol.copy_in(
             copy_stmt, None, None, records, intro_ps._state, timeout)
 
+    def _format_copy_where(self, where):
+        if where and not self._server_caps.sql_copy_from_where:
+            raise exceptions.UnsupportedServerFeatureError(
+                'the `where` parameter requires PostgreSQL 12 or later')
+
+        if where:
+            where_clause = 'WHERE ' + where
+        else:
+            where_clause = ''
+
+        return where_clause
+
     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,
@@ -2308,7 +2347,7 @@ class _ConnectionProxy:
 ServerCapabilities = collections.namedtuple(
     'ServerCapabilities',
     ['advisory_locks', 'notifications', 'plpgsql', 'sql_reset',
-     'sql_close_all'])
+     'sql_close_all', 'sql_copy_from_where'])
 ServerCapabilities.__doc__ = 'PostgreSQL server capabilities.'
 
 
@@ -2320,27 +2359,31 @@ def _detect_server_capabilities(server_version, connection_settings):
         plpgsql = False
         sql_reset = True
         sql_close_all = False
+        sql_copy_from_where = False
     elif hasattr(connection_settings, 'crdb_version'):
         # CockroachDB detected.
         advisory_locks = False
         notifications = False
         plpgsql = False
         sql_reset = False
         sql_close_all = False
+        sql_copy_from_where = False
     elif hasattr(connection_settings, 'crate_version'):
         # CrateDB detected.
         advisory_locks = False
         notifications = False
         plpgsql = False
         sql_reset = False
         sql_close_all = False
+        sql_copy_from_where = False
     else:
         # Standard PostgreSQL server assumed.
         advisory_locks = True
         notifications = True
         plpgsql = True
         sql_reset = True
         sql_close_all = True
+        sql_copy_from_where = server_version.major >= 12
 
     return ServerCapabilities(
         advisory_locks=advisory_locks,

asyncpg/exceptions/_base.py

@@ -13,7 +13,7 @@
 __all__ = ('PostgresError', 'FatalPostgresError', 'UnknownPostgresError',
            'InterfaceError', 'InterfaceWarning', 'PostgresLogMessage',
            'InternalClientError', 'OutdatedSchemaCacheError', 'ProtocolError',
-           'UnsupportedClientFeatureError')
+           'UnsupportedClientFeatureError', 'UnsupportedServerFeatureError')
 
 
 def _is_asyncpg_class(cls):
@@ -228,6 +228,10 @@ class UnsupportedClientFeatureError(InterfaceError):
     """Requested feature is unsupported by asyncpg."""
 
 
+class UnsupportedServerFeatureError(InterfaceError):
+    """Requested feature is unsupported by PostgreSQL server."""
+
+
 class InterfaceWarning(InterfaceMessage, UserWarning):
     """A warning caused by an improper use of asyncpg API."""
 

asyncpg/pool.py

@@ -736,7 +736,8 @@ async def copy_to_table(
         force_quote=None,
         force_not_null=None,
         force_null=None,
-        encoding=None
+        encoding=None,
+        where=None
     ):
         """Copy data to the specified table.
 
@@ -765,7 +766,8 @@ async def copy_to_table(
                 force_quote=force_quote,
                 force_not_null=force_not_null,
                 force_null=force_null,
-                encoding=encoding
+                encoding=encoding,
+                where=where
             )
 
     async def copy_records_to_table(
@@ -775,7 +777,8 @@ async def copy_records_to_table(
         records,
         columns=None,
         schema_name=None,
-        timeout=None
+        timeout=None,
+        where=None
     ):
         """Copy a list of records to the specified table using binary COPY.
 
@@ -792,7 +795,8 @@ async def copy_records_to_table(
                 records=records,
                 columns=columns,
                 schema_name=schema_name,
-                timeout=timeout
+                timeout=timeout,
+                where=where
             )
 
     def acquire(self, *, timeout=None):