From 9fab45c1b3977978cd4b6c59f270c190d6eeee0d Mon Sep 17 00:00:00 2001
From: Nick Proskurin <42863572+MCPN@users.noreply.github.com>
Date: Thu, 15 Aug 2024 12:10:17 +0200
Subject: [PATCH] fix: BI-3279 use native UUID in asyncpg adapter (#578)

---
 .../core/postgresql_base/async_adapters_postgres.py  | 12 +++++++++++-
 .../db/api/test_dashsql.py                           |  4 +++-
 .../dl_connector_postgresql_tests/db/config.py       |  1 +
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/async_adapters_postgres.py b/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/async_adapters_postgres.py
index af38af413..3c9b06535 100644
--- a/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/async_adapters_postgres.py
+++ b/lib/dl_connector_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,7 +186,7 @@ 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,
@@ -193,6 +194,15 @@ async def _get_connection(self, db_name_from_query: str) -> AsyncIterator[asyncp
                     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
index 40350fcc7..a65dfebfb 100644
--- 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
index 9cc1b7737..93dc9bddb 100644
--- 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