fix: BI-3279 use native UUID in asyncpg adapter (#578)

datalens-tech · Aug 15, 2024 · 9fab45c · 9fab45c
1 parent b1cb74d
commit 9fab45c
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/...nector_postgresql/dl_connector_postgresql/core/postgresql_base/async_adapters_postgres.py b/...nector_postgresql/dl_connector_postgresql/core/postgresql_base/async_adapters_postgres.py
@@ -18,6 +18,7 @@
     TypeVar,
 )
 from urllib.parse import quote_plus
+import uuid
 
 import asyncpg
 import asyncpg.exceptions
@@ -185,14 +186,23 @@ async def _get_connection(self, db_name_from_query: str) -> AsyncIterator[asyncp
                 # TODO ^ make up a new DL exc and wrap asyncpg.exceptions.ReadOnlySQLTransactionError
                 # There is some decimal-magic in asyncpg
                 # and this magic is incompatible with magic in sqlalchemy-psycopg2
-                # so lets disable it
+                # so let's disable it
                 await connection.set_type_codec(
                     "numeric",
                     encoder=str,
                     decoder=lambda x: x,
                     schema="pg_catalog",
                     format="text",
                 )
+                # asyncpg uses a custom wrapper for UUID, and we can't serialize it
+                # so use builtin UUID explicitly
+                await connection.set_type_codec(
+                    "uuid",
+                    encoder=str,
+                    decoder=lambda x: uuid.UUID(x),
+                    schema="pg_catalog",
+                    format="text",
+                )
                 # we set date-like values to params as strings
                 # but asyncpg expects python-objects
                 await connection.set_type_codec(

diff --git a/lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_dashsql.py b/lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_dashsql.py
@@ -98,6 +98,7 @@ async def test_result_extended(self, data_api_lowlevel_aiohttp_client: TestClien
             "num_timestamp",
             "num_timestamptz",
             "num_array",
+            "some_uuid",
             "some_nan",
         ]
         assert resp_data[0]["data"]["bi_types"] == [
@@ -122,9 +123,10 @@ async def test_result_extended(self, data_api_lowlevel_aiohttp_client: TestClien
             "genericdatetime",
             "genericdatetime",
             "array_int",
+            "uuid",
             "float",
         ]
-        assert resp_data[1]["data"][21] == "nan", "must not be a float('nan')"
+        assert resp_data[1]["data"][-1] == "nan", "must not be a float('nan')"
 
     @pytest.mark.asyncio
     async def test_result_with_params(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):

diff --git a/lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/config.py b/lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/config.py
@@ -72,6 +72,7 @@ class CoreSslConnectionSettings:
     ('2020-01-01T00:00:0' || number)::timestamp as num_timestamp,
     ('2020-01-01T00:00:0' || number)::timestamptz as num_timestamptz,
     ARRAY[number, 2, 3] as num_array,
+    gen_random_uuid() as some_uuid,
     'nan'::double precision + number as some_nan
 from base
 limit 10