From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 11 Jun 2022 14:56:13 +0000 (+0200)
Subject: feat(copy): add 'writer' parameter to Cursor.copy()
X-Git-Tag: 3.1~44^2~12
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=bd4496ce0a78704259462e8d3fad51de4cdcc4c8;p=thirdparty%2Fpsycopg.git

feat(copy): add 'writer' parameter to Cursor.copy()
---

diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py
index 422af8a12..bab46c2d9 100644
--- a/psycopg/psycopg/cursor.py
+++ b/psycopg/psycopg/cursor.py
@@ -15,7 +15,7 @@ from . import pq
 from . import adapt
 from . import errors as e
 from .abc import ConnectionType, Query, Params, PQGen
-from .copy import Copy
+from .copy import Copy, Writer as CopyWriter
 from .rows import Row, RowMaker, RowFactory
 from ._column import Column
 from ._cmodule import _psycopg
@@ -873,7 +873,13 @@ class Cursor(BaseCursor["Connection[Any]", Row]):
         self._scroll(value, mode)
 
     @contextmanager
-    def copy(self, statement: Query, params: Optional[Params] = None) -> Iterator[Copy]:
+    def copy(
+        self,
+        statement: Query,
+        params: Optional[Params] = None,
+        *,
+        writer: Optional[CopyWriter[Any]] = None,
+    ) -> Iterator[Copy]:
         """
         Initiate a :sql:`COPY` operation and return an object to manage it.
 
@@ -883,7 +889,7 @@ class Cursor(BaseCursor["Connection[Any]", Row]):
             with self._conn.lock:
                 self._conn.wait(self._start_copy_gen(statement, params))
 
-            with Copy(self) as copy:
+            with Copy(self, writer=writer) as copy:
                 yield copy
         except e.Error as ex:
             raise ex.with_traceback(None)