From 39ffc5ceeae24c2fc053685806f90b2e84e1d5f5 Mon Sep 17 00:00:00 2001
From: Fuxin Hao <haofxpro@gmail.com>
Date: Wed, 30 Mar 2022 11:34:10 +0800
Subject: [PATCH] fix: add get_ibis_table_schema (#410)

---
 data_validation/clients.py           | 17 +++++++++++++++++
 data_validation/schema_validation.py | 14 +++++++++-----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/data_validation/clients.py b/data_validation/clients.py
index 1a03bc26e..a89ec2782 100644
--- a/data_validation/clients.py
+++ b/data_validation/clients.py
@@ -138,6 +138,23 @@ def get_ibis_table(client, schema_name, table_name, database_name=None):
         return client.table(table_name, database=schema_name)
 
 
+def get_ibis_table_schema(client, schema_name, table_name):
+    """Return Ibis Table Schema for Supplied Client.
+
+    client (IbisClient): Client to use for table
+    schema_name (str): Schema name of table object
+    table_name (str): Table name of table object
+    database_name (str): Database name (generally default is used)
+    """
+    if type(client) in [
+        MySQLClient,
+        PostgreSQLClient
+    ]:
+        return client.schema(schema_name).table(table_name).schema()
+    else:
+        return client.get_schema(table_name, schema_name)
+
+
 def list_schemas(client):
     """Return a list of schemas in the DB."""
     if type(client) in [
diff --git a/data_validation/schema_validation.py b/data_validation/schema_validation.py
index 7596ab65a..ac5d668dd 100644
--- a/data_validation/schema_validation.py
+++ b/data_validation/schema_validation.py
@@ -15,7 +15,7 @@
 import datetime
 import pandas
 
-from data_validation import metadata, consts
+from data_validation import metadata, consts, clients
 
 
 class SchemaValidation(object):
@@ -33,11 +33,15 @@ def __init__(self, config_manager, run_metadata=None, verbose=False):
 
     def execute(self):
         """ Performs a validation between source and a target schema"""
-        ibis_source_schema = self.config_manager.source_client.get_schema(
-            self.config_manager.source_table, self.config_manager.source_schema
+        ibis_source_schema = clients.get_ibis_table_schema(
+            self.config_manager.source_client,
+            self.config_manager.source_schema,
+            self.config_manager.source_table
         )
-        ibis_target_schema = self.config_manager.target_client.get_schema(
-            self.config_manager.target_table, self.config_manager.target_schema
+        ibis_target_schema = clients.get_ibis_table_schema(
+            self.config_manager.target_client,
+            self.config_manager.target_schema,
+            self.config_manager.target_table
         )
 
         source_fields = {}