From 57896f430967f9726ccdf8c82b1c3b9833f0062c Mon Sep 17 00:00:00 2001
From: Neha Nene <nehanene@google.com>
Date: Thu, 28 Apr 2022 01:30:53 -0500
Subject: [PATCH] fix: supports NULL datetime/timestamps, fixes bug with
 validation_status in PR 455 (#460)

---
 data_validation/combiner.py                | 12 ++++++++++++
 tests/system/data_sources/test_mysql.py    |  2 +-
 tests/system/data_sources/test_postgres.py |  2 +-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/data_validation/combiner.py b/data_validation/combiner.py
index 41532a8e8..bbc29ada5 100644
--- a/data_validation/combiner.py
+++ b/data_validation/combiner.py
@@ -116,6 +116,18 @@ def _calculate_difference(field_differences, datatype, validation, is_value_comp
             .else_(consts.VALIDATION_STATUS_FAIL)
             .end()
         )
+    # String data types i.e "None" can be returned for NULL timestamp/datetime aggs
+    elif isinstance(datatype, ibis.expr.datatypes.String):
+        difference = pct_difference = ibis.null().cast("float64")
+        validation_status = (
+            ibis.case()
+            .when(
+                target_value.isnull() & source_value.isnull(),
+                consts.VALIDATION_STATUS_SUCCESS,
+            )
+            .else_(consts.VALIDATION_STATUS_FAIL)
+            .end()
+        )
     else:
         difference = (target_value - source_value).cast("float64")
 
diff --git a/tests/system/data_sources/test_mysql.py b/tests/system/data_sources/test_mysql.py
index 1973e3f5d..a0aeba9ac 100644
--- a/tests/system/data_sources/test_mysql.py
+++ b/tests/system/data_sources/test_mysql.py
@@ -81,7 +81,7 @@ def test_schema_validation():
         df = data_validator.execute()
 
         for validation in df.to_dict(orient="records"):
-            assert validation["status"] == consts.VALIDATION_STATUS_SUCCESS
+            assert validation["validation_status"] == consts.VALIDATION_STATUS_SUCCESS
     except exceptions.DataClientConnectionFailure:
         # Local Testing will not work for MySQL
         pass
diff --git a/tests/system/data_sources/test_postgres.py b/tests/system/data_sources/test_postgres.py
index 26c410349..7e2753bea 100644
--- a/tests/system/data_sources/test_postgres.py
+++ b/tests/system/data_sources/test_postgres.py
@@ -124,4 +124,4 @@ def test_schema_validation(cloud_sql):
     df = data_validator.execute()
 
     for validation in df.to_dict(orient="records"):
-        assert validation["status"] == consts.VALIDATION_STATUS_SUCCESS
+        assert validation["validation_status"] == consts.VALIDATION_STATUS_SUCCESS