From 9ad529a73a3d53c74b69a5f5fc7e005d0e389207 Mon Sep 17 00:00:00 2001
From: Helen Cristina <helencristina@google.com>
Date: Tue, 29 Aug 2023 13:53:38 -0300
Subject: [PATCH] fix: Add missing SnowflakeDialect mapping for BINARY data
 type (#959)

* fix: Add missing SnowflakeDialect mapping for BINARY data type

* Add unit test for schema validation with specific Snowflake data types
---
 noxfile.py                                   |  2 +-
 tests/system/data_sources/test_snowflake.py  | 26 ++++++++++++++++++++
 third_party/ibis/ibis_snowflake/datatypes.py |  6 ++++-
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/noxfile.py b/noxfile.py
index 8a5a6d6fd..698ced0e9 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -261,7 +261,7 @@ def integration_hive(session):
 @nox.session(python=PYTHON_VERSIONS, venv_backend="venv")
 def integration_snowflake(session):
     """Run Snowflake integration tests.
-    Ensure Hive validation is running as expected.
+    Ensure Snowflake validation is running as expected.
     """
     _setup_session_requirements(
         session, extra_packages=["snowflake-sqlalchemy", "snowflake-connector-python"]
diff --git a/tests/system/data_sources/test_snowflake.py b/tests/system/data_sources/test_snowflake.py
index 82bb26b1f..cf86076e8 100644
--- a/tests/system/data_sources/test_snowflake.py
+++ b/tests/system/data_sources/test_snowflake.py
@@ -145,6 +145,32 @@ def test_schema_validation_core_types():
     assert len(df) == 0
 
 
+@mock.patch(
+    "data_validation.state_manager.StateManager.get_connection_config",
+    new=mock_get_connection_config,
+)
+def test_schema_validation_specific_types():
+    """Snowflake to Snowflake test_specific_data_types schema validation"""
+    parser = cli_tools.configure_arg_parser()
+    args = parser.parse_args(
+        [
+            "validate",
+            "schema",
+            "-sc=mock-conn",
+            "-tc=mock-conn",
+            "-tbls=PSO_DATA_VALIDATOR.PUBLIC.TEST_SPECIFIC_DATA_TYPES",
+            "--filter-status=fail",
+        ]
+    )
+    config_managers = main.build_config_managers_from_args(args)
+    assert len(config_managers) == 1
+    config_manager = config_managers[0]
+    validator = data_validation.DataValidation(config_manager.config, verbose=False)
+    df = validator.execute()
+    # With filter on failures the data frame should be empty
+    assert len(df) == 0
+
+
 @mock.patch(
     "data_validation.state_manager.StateManager.get_connection_config",
     new=mock_get_connection_config,
diff --git a/third_party/ibis/ibis_snowflake/datatypes.py b/third_party/ibis/ibis_snowflake/datatypes.py
index 80770075b..6e47f7168 100644
--- a/third_party/ibis/ibis_snowflake/datatypes.py
+++ b/third_party/ibis/ibis_snowflake/datatypes.py
@@ -18,7 +18,7 @@
 from ibis.backends.snowflake import Backend as SnowflakeBackend
 from ibis.backends.snowflake.datatypes import parse
 from snowflake.connector.constants import FIELD_ID_TO_NAME
-from snowflake.sqlalchemy import NUMBER
+from snowflake.sqlalchemy import NUMBER, BINARY
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect
 
 
@@ -30,6 +30,10 @@ def sa_sf_numeric(_, satype, nullable=True):
         nullable=nullable,
     )
 
+@dt.dtype.register(SnowflakeDialect, BINARY)
+def sa_sf_binary(_, satype, nullable=True):
+    return dt.Binary(nullable=nullable)
+
 
 def _metadata(self, query: str) -> Iterable[Tuple[str, dt.DataType]]:
     with self.begin() as con, con.connection.cursor() as cur: