fix: accommodate Decimal NaN != Decimal NaN (#573)

* test: added test case addressing #571 * fix: accommodate Decimal NaN != Decimal NaN
Autodesk · Sep 27, 2024 · 3602266 · 3602266
1 parent 74e2598
commit 3602266
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 12 deletions.
diff --git a/pgbelt/util/postgres.py b/pgbelt/util/postgres.py
@@ -1,5 +1,6 @@
 from logging import Logger
 
+from decimal import Decimal
 from asyncpg import Pool
 from asyncpg import Record
 from asyncpg.exceptions import UndefinedObjectError
@@ -171,12 +172,33 @@ async def compare_data(
         # Check each row for exact match
         for src_row, dst_row in zip(src_rows, dst_rows):
             if src_row != dst_row:
-                raise AssertionError(
-                    "Row match failure between source and destination.\n"
-                    f"Table: {full_table_name}\n"
-                    f"Source Row: {src_row}\n"
-                    f"Dest Row: {dst_row}"
-                )
+
+                # Addresses #571, AsyncPG is decoding numeric NaN as Python Decimal('NaN').
+                # Decimal('NaN') != Decimal('NaN'), breaks comparison. Convert those NaNs to None.
+                src_row_d = {
+                    key: (
+                        value
+                        if not (isinstance(value, Decimal) and value.is_nan())
+                        else None
+                    )
+                    for key, value in row.items()
+                }
+                dst_row_d = {
+                    key: (
+                        value
+                        if not (isinstance(value, Decimal) and value.is_nan())
+                        else None
+                    )
+                    for key, value in row.items()
+                }
+
+                if src_row_d != dst_row_d:
+                    raise AssertionError(
+                        "Row match failure between source and destination.\n"
+                        f"Table: {full_table_name}\n"
+                        f"Source Row: {src_row}\n"
+                        f"Dest Row: {dst_row}"
+                    )
 
     # Just a paranoia check. If this throws, then it's possible pgbelt didn't migrate any data.
     # This was found in issue #420, and previous commands threw errors before this issue could arise.

diff --git a/tests/integration/files/test_schema_data.sql b/tests/integration/files/test_schema_data.sql
@@ -19,6 +19,7 @@ CREATE TABLE public."UsersCapital" (
     hash_firstname text NOT NULL,
     hash_lastname text NOT NULL,
     gender character varying(6) NOT NULL,
+    numericnan numeric(19,4), -- Testing for #571 Numeric NaN
     CONSTRAINT users_gender_check CHECK (((gender)::text = ANY (ARRAY[('male'::character varying)::text, ('female'::character varying)::text])))
 );
 
@@ -106,12 +107,12 @@ INSERT INTO public.fruits (id, name)
 -- Data for Name: UsersCapital; Type: TABLE DATA; Schema: public; Owner: owner
 --
 
-INSERT INTO public."UsersCapital" (id, hash_firstname, hash_lastname, gender)
-    VALUES (1, 'garbagefirst', 'garbagelast', 'male'),
-    (2, 'garbagefirst1', 'garbagelast1', 'female'),
-    (3, 'sdgarbagefirst', 'dgsadsrbagelast', 'male'),
-    (4, 'dsdssdgarbagefirst', 'dgsaggggdjjjsrbagelast', 'female'),
-    (5, 'dsdssdgarbagefirt', 'dgsagggdjjjsrbagelast', 'female');
+INSERT INTO public."UsersCapital" (id, hash_firstname, hash_lastname, gender, numericnan)
+    VALUES (1, 'garbagefirst', 'garbagelast', 'male', 1),
+    (2, 'garbagefirst1', 'garbagelast1', 'female', 0),
+    (3, 'sdgarbagefirst', 'dgsadsrbagelast', 'male', 'NaN'),
+    (4, 'dsdssdgarbagefirst', 'dgsaggggdjjjsrbagelast', 'female', 1),
+    (5, 'dsdssdgarbagefirt', 'dgsagggdjjjsrbagelast', 'female', 0);
 
 
 --