From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 24 Aug 2021 15:20:09 +0000 (+0200)
Subject: Drop TODO point now working
X-Git-Tag: 3.0.beta1~49
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=21e43d39a170667ddcdf65fd21e0b5fb425f2af0;p=thirdparty%2Fpsycopg.git

Drop TODO point now working

Add tests to prove it works.
---

diff --git a/psycopg/psycopg/copy.py b/psycopg/psycopg/copy.py
index 1d52451df..0e889dacb 100644
--- a/psycopg/psycopg/copy.py
+++ b/psycopg/psycopg/copy.py
@@ -91,12 +91,6 @@ class BaseCopy(Generic[ConnectionType]):
 
         The types must be specified as a sequence of oid or PostgreSQL type
         names (e.g. ``int4``, ``timestamptz[]``).
-
-        .. admonition:: TODO
-
-            Only builtin names are supprted for the moment. In order to specify
-            custom data types you must use their oid.
-
         """
         registry = self.cursor.adapters.types
         oids = [
diff --git a/tests/fix_db.py b/tests/fix_db.py
index d826b2a5f..6d301dc5f 100644
--- a/tests/fix_db.py
+++ b/tests/fix_db.py
@@ -138,3 +138,14 @@ def check_connection_version(got, function):
     want = [m.args[0] for m in function.pytestmark if m.name == "pg"]
     if want:
         return check_server_version(got, want[0])
+
+
+@pytest.fixture
+def hstore(svcconn):
+    from psycopg import Error
+
+    try:
+        with svcconn.transaction():
+            svcconn.execute("create extension if not exists hstore")
+    except Error as e:
+        pytest.skip(str(e))
diff --git a/tests/test_copy.py b/tests/test_copy.py
index 0f794584e..f0d15dcf8 100644
--- a/tests/test_copy.py
+++ b/tests/test_copy.py
@@ -12,6 +12,8 @@ from psycopg import sql
 from psycopg import errors as e
 from psycopg.pq import Format
 from psycopg.adapt import PyFormat as PgFormat
+from psycopg.types import TypeInfo
+from psycopg.types.hstore import register_adapters as register_hstore
 from psycopg.types.numeric import Int4
 
 from .utils import gc_collect
@@ -116,6 +118,23 @@ def test_rows(conn, format):
     assert conn.pgconn.transaction_status == conn.TransactionStatus.INTRANS
 
 
+def test_set_custom_type(conn, hstore):
+    command = """copy (select '"a"=>"1", "b"=>"2"'::hstore) to stdout"""
+    cur = conn.cursor()
+
+    with cur.copy(command) as copy:
+        rows = list(copy.rows())
+
+    assert rows == [('"a"=>"1", "b"=>"2"',)]
+
+    register_hstore(TypeInfo.fetch(conn, "hstore"), cur)
+    with cur.copy(command) as copy:
+        copy.set_types(["hstore"])
+        rows = list(copy.rows())
+
+    assert rows == [({"a": "1", "b": "2"},)]
+
+
 @pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY])
 def test_copy_out_allchars(conn, format):
     cur = conn.cursor()
diff --git a/tests/test_copy_async.py b/tests/test_copy_async.py
index 7c755ba0e..2efb7a1a3 100644
--- a/tests/test_copy_async.py
+++ b/tests/test_copy_async.py
@@ -11,7 +11,9 @@ from psycopg import pq
 from psycopg import sql
 from psycopg import errors as e
 from psycopg.pq import Format
+from psycopg.types import TypeInfo
 from psycopg.adapt import PyFormat as PgFormat
+from psycopg.types.hstore import register_adapters as register_hstore
 
 from .utils import gc_collect
 from .test_copy import sample_text, sample_binary, sample_binary_rows  # noqa
@@ -99,6 +101,23 @@ async def test_rows(aconn, format):
     assert aconn.pgconn.transaction_status == aconn.TransactionStatus.INTRANS
 
 
+async def test_set_custom_type(aconn, hstore):
+    command = """copy (select '"a"=>"1", "b"=>"2"'::hstore) to stdout"""
+    cur = aconn.cursor()
+
+    async with cur.copy(command) as copy:
+        rows = [row async for row in copy.rows()]
+
+    assert rows == [('"a"=>"1", "b"=>"2"',)]
+
+    register_hstore(await TypeInfo.fetch_async(aconn, "hstore"), cur)
+    async with cur.copy(command) as copy:
+        copy.set_types(["hstore"])
+        rows = [row async for row in copy.rows()]
+
+    assert rows == [({"a": "1", "b": "2"},)]
+
+
 @pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY])
 async def test_copy_out_allchars(aconn, format):
     cur = aconn.cursor()
diff --git a/tests/types/test_hstore.py b/tests/types/test_hstore.py
index e52aae7e7..f81a49291 100644
--- a/tests/types/test_hstore.py
+++ b/tests/types/test_hstore.py
@@ -97,12 +97,3 @@ def test_roundtrip_array(hstore, conn):
     register_adapters(TypeInfo.fetch(conn, "hstore"), conn)
     samp1 = conn.execute("select %s", (samp,)).fetchone()[0]
     assert samp1 == samp
-
-
-@pytest.fixture
-def hstore(svcconn):
-    try:
-        with svcconn.transaction():
-            svcconn.execute("create extension if not exists hstore")
-    except psycopg.Error as e:
-        pytest.skip(str(e))