From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 22 May 2022 00:12:01 +0000 (+0200)
Subject: fix(crdb): configure dumpers to deal correctly with Enum types
X-Git-Tag: 3.1~49^2~56
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=5d2085f3c4264d3611d18b0e5d592d967feb1b1d;p=thirdparty%2Fpsycopg.git

fix(crdb): configure dumpers to deal correctly with Enum types
---

diff --git a/psycopg/psycopg/crdb.py b/psycopg/psycopg/crdb.py
index e509d78ec..acbd99a17 100644
--- a/psycopg/psycopg/crdb.py
+++ b/psycopg/psycopg/crdb.py
@@ -5,13 +5,15 @@ Types configuration specific for CockroachDB.
 # Copyright (C) 2022 The Psycopg Team
 
 import re
+from enum import Enum
 from typing import Any, Optional, Union, TYPE_CHECKING
 
 from . import errors as e
 from .abc import AdaptContext
-from .postgres import adapters as pg_adapters
-from ._adapters_map import AdaptersMap
+from .postgres import adapters as pg_adapters, TEXT_OID
 from .conninfo import ConnectionInfo
+from ._adapters_map import AdaptersMap
+from .types.enum import EnumDumper, EnumBinaryDumper
 
 adapters = AdaptersMap(pg_adapters)
 
@@ -55,6 +57,14 @@ class CrdbConnectionInfo(ConnectionInfo):
         return int(m.group(1)) * 10000 + int(m.group(2)) * 100 + int(m.group(3))
 
 
+class CrdbEnumDumper(EnumDumper):
+    oid = TEXT_OID
+
+
+class CrdbEnumBinaryDumper(EnumBinaryDumper):
+    oid = TEXT_OID
+
+
 def register_crdb_adapters(context: AdaptContext) -> None:
     from .types import string
 
@@ -63,6 +73,8 @@ def register_crdb_adapters(context: AdaptContext) -> None:
     # Dump strings with text oid instead of unknown.
     # Unlike PostgreSQL, CRDB seems able to cast text to most types.
     adapters.register_dumper(str, string.StrDumper)
+    adapters.register_dumper(Enum, CrdbEnumBinaryDumper)
+    adapters.register_dumper(Enum, CrdbEnumDumper)
 
 
 register_crdb_adapters(adapters)
diff --git a/tests/types/test_enum.py b/tests/types/test_enum.py
index 0284dcf53..e5b15d132 100644
--- a/tests/types/test_enum.py
+++ b/tests/types/test_enum.py
@@ -7,6 +7,8 @@ from psycopg.adapt import PyFormat
 from psycopg.types import TypeInfo
 from psycopg.types.enum import EnumInfo, register_enum
 
+from ..fix_crdb import crdb_encoding
+
 
 class PureTestEnum(Enum):
     FOO = auto()
@@ -34,7 +36,7 @@ class IntTestEnum(int, Enum):
 
 
 enum_cases = [PureTestEnum, StrTestEnum, IntTestEnum]
-encodings = ["utf8", "latin1"]
+encodings = ["utf8", crdb_encoding("latin1")]
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -49,7 +51,7 @@ def ensure_enum(enum, conn):
     conn.execute(
         sql.SQL(
             """
-            drop type if exists {name} cascade;
+            drop type if exists {name};
             create type {name} as enum ({labels});
             """
         ).format(name=sql.Identifier(name), labels=sql.SQL(",").join(labels))
@@ -116,6 +118,7 @@ def test_enum_loader_nonascii(conn, encoding, fmt_in, fmt_out):
         assert cur.fetchone()[0] == enum[label]
 
 
+@pytest.mark.crdb("skip", reason="encoding")
 @pytest.mark.parametrize("enum", enum_cases)
 @pytest.mark.parametrize("fmt_in", PyFormat)
 @pytest.mark.parametrize("fmt_out", pq.Format)
@@ -158,6 +161,7 @@ def test_enum_dumper_nonascii(conn, encoding, fmt_in, fmt_out):
         assert cur.fetchone()[0] == item
 
 
+@pytest.mark.crdb("skip", reason="encoding")
 @pytest.mark.parametrize("enum", enum_cases)
 @pytest.mark.parametrize("fmt_in", PyFormat)
 @pytest.mark.parametrize("fmt_out", pq.Format)