From ed4629594eeaf7ca3a3a914b4c07e1b9c977f07c Mon Sep 17 00:00:00 2001
From: AJ Welch <95496513+ajwelch4@users.noreply.github.com>
Date: Wed, 16 Mar 2022 16:49:34 -0400
Subject: [PATCH] fix: support schema validation for more clients (#355) (#380)

* fix: support schema validation for more clients (#355)

* fix: use consts.VALIDATION_STATUS_SUCCESS
---
 tests/system/data_sources/test_sql_server.py | 26 ++++++++++++++++++++
 third_party/ibis/ibis_DB2/client.py          |  3 +++
 third_party/ibis/ibis_mssql/client.py        |  3 +++
 third_party/ibis/ibis_oracle/client.py       |  2 ++
 third_party/ibis/ibis_snowflake/client.py    |  3 +++
 5 files changed, 37 insertions(+)

diff --git a/tests/system/data_sources/test_sql_server.py b/tests/system/data_sources/test_sql_server.py
index 757213651..51aab155d 100644
--- a/tests/system/data_sources/test_sql_server.py
+++ b/tests/system/data_sources/test_sql_server.py
@@ -87,3 +87,29 @@ def test_sql_server_count(cloud_sql):
     data_validator = data_validation.DataValidation(config_count_valid, verbose=False,)
     df = data_validator.execute()
     assert df["source_agg_value"][0] == df["target_agg_value"][0]
+
+
+def test_schema_validation():
+    conn = {
+        "source_type": "MSSQL",
+        "host": SQL_SERVER_HOST,
+        "user": SQL_SERVER_USER,
+        "password": SQL_SERVER_PASSWORD,
+        "port": 1433,
+        "database": "guestbook",
+    }
+
+    config = {
+        consts.CONFIG_SOURCE_CONN: conn,
+        consts.CONFIG_TARGET_CONN: conn,
+        consts.CONFIG_TYPE: "Schema",
+        consts.CONFIG_SCHEMA_NAME: "dbo",
+        consts.CONFIG_TABLE_NAME: "entries",
+        consts.CONFIG_FORMAT: "table",
+    }
+
+    validator = data_validation.DataValidation(config, verbose=True)
+    df = validator.execute()
+
+    for validation in df.to_dict(orient="records"):
+        assert validation["status"] == consts.VALIDATION_STATUS_SUCCESS
diff --git a/third_party/ibis/ibis_DB2/client.py b/third_party/ibis/ibis_DB2/client.py
index 05df4aaa7..cd2d12e64 100644
--- a/third_party/ibis/ibis_DB2/client.py
+++ b/third_party/ibis/ibis_DB2/client.py
@@ -196,3 +196,6 @@ def list_tables(self, like=None, database=None, schema=None):
         else:
             parent = super(DB2Client, self)
             return parent.list_tables(like=like, schema=schema)
+
+    def get_schema(self, name, schema=None):
+        return self.table(name, schema=schema).schema()
\ No newline at end of file
diff --git a/third_party/ibis/ibis_mssql/client.py b/third_party/ibis/ibis_mssql/client.py
index 14c3c8941..cec859414 100644
--- a/third_party/ibis/ibis_mssql/client.py
+++ b/third_party/ibis/ibis_mssql/client.py
@@ -264,3 +264,6 @@ def _get_schema_using_query(self, limited_query):
                 type_map[row[1]] for row in cur.proxy._cursor_description()
             ]
         return sch.Schema(names, ibis_types)
+
+    def get_schema(self, name, schema=None):
+        return self.table(name, schema=schema).schema()
\ No newline at end of file
diff --git a/third_party/ibis/ibis_oracle/client.py b/third_party/ibis/ibis_oracle/client.py
index 2ed185463..081b3ac76 100644
--- a/third_party/ibis/ibis_oracle/client.py
+++ b/third_party/ibis/ibis_oracle/client.py
@@ -189,3 +189,5 @@ def list_tables(self, like=None, database=None, schema=None):
             parent = super(OracleClient, self)
             return parent.list_tables(like=like, schema=schema)
 
+    def get_schema(self, name, schema=None):
+        return self.table(name, schema=schema).schema()
diff --git a/third_party/ibis/ibis_snowflake/client.py b/third_party/ibis/ibis_snowflake/client.py
index b338bd616..eb6d0fe83 100644
--- a/third_party/ibis/ibis_snowflake/client.py
+++ b/third_party/ibis/ibis_snowflake/client.py
@@ -160,3 +160,6 @@ def executeSQL(self, sql):
         finally:
             connection.close()
         return response
+
+    def get_schema(self, name, schema=None):
+        return self.table(name, schema=schema).schema()
\ No newline at end of file