fix: add get_ibis_table_schema (GoogleCloudPlatform#410)

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 [

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 = {}