From 3602266ed1421dfe4efb5dabe8c9303e653f05d1 Mon Sep 17 00:00:00 2001 From: Varjitt Jeeva Date: Fri, 27 Sep 2024 15:03:37 -0400 Subject: [PATCH] fix: accommodate Decimal NaN != Decimal NaN (#573) * test: added test case addressing #571 * fix: accommodate Decimal NaN != Decimal NaN --- pgbelt/util/postgres.py | 34 ++++++++++++++++---- tests/integration/files/test_schema_data.sql | 13 ++++---- 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/pgbelt/util/postgres.py b/pgbelt/util/postgres.py index 223cb8c..c808d1a 100644 --- 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 index a975f7c..3dfb9af 100644 --- 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); --