Repository: great-expectations/great_expectations

Issue:
SqlAlchemyDatasets can't be used with custom sql for databases with multiple schemas
Instantiating a `SqlAlchemyDataset` with a database which supports multiple schemas fails on the following line:

https://github.com/great-expectations/great_expectations/blob/d7a4ee03cfce8bea66c5f91573bb87e4385b1eb1/great_expectations/dataset/sqlalchemy_dataset.py#L222

The issue is that the inspector prepends the default schema when one is not provided. The custom sql tables are written to a temporary schema, so the default schema is incorrect. I looked into this a bit, and as far as I can tell, the inspector just won't work unless you look up and pass the name of the temporary schema.

I have a fix for this issue that I'm going to submit, and I'm just filing an issue for documentation purposes. My proposed solution is to get the column names from a simple query rather than the inspector.

Discussion:

Patch:
diff --git a/docs/source/roadmap_changelog.rst b/docs/source/roadmap_changelog.rst
index 7656d1f47..5e62bc5a5 100644
--- a/docs/source/roadmap_changelog.rst
+++ b/docs/source/roadmap_changelog.rst
@@ -14,6 +14,8 @@ Planned Features
 
 v.0.4.4__develop
 ----------------
+* Enable custom_sql datasets for databases with multiple schemas, by
+  adding a fallback for column reflection (#387)
 * Remove `IF NOT EXISTS` check for custom sql temporary tables, for
   Redshift compatibility (#372)
 * Allow users to pass args/kwargs for engine creation in
diff --git a/great_expectations/dataset/sqlalchemy_dataset.py b/great_expectations/dataset/sqlalchemy_dataset.py
index 7a4f440cd..8d98a7660 100644
--- a/great_expectations/dataset/sqlalchemy_dataset.py
+++ b/great_expectations/dataset/sqlalchemy_dataset.py
@@ -255,8 +255,14 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
         if custom_sql:
             self.create_temporary_table(table_name, custom_sql)
 
-        insp = reflection.Inspector.from_engine(engine)
-        self.columns = insp.get_columns(table_name, schema=schema)
+
+        try:
+            insp = reflection.Inspector.from_engine(engine)
+            self.columns = insp.get_columns(table_name, schema=schema)
+        except KeyError:
+            # we will get a KeyError for temporary tables, since
+            # reflection will not find the temporary schema
+            self.columns = self.column_reflection_fallback()
 
         # Only call super once connection is established and table_name and columns known to allow autoinspection
         super(SqlAlchemyDataset, self).__init__(*args, **kwargs)
@@ -272,17 +278,13 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
             table_name=table_name, custom_sql=custom_sql)
         self.engine.execute(stmt)
 
-    def _is_numeric_column(self, column):
-        for col in self.columns:
-            if (col['name'] == column and
-                        isinstance(col['type'],
-                                   (sa.types.Integer, sa.types.BigInteger, sa.types.Float,
-                                    sa.types.Numeric, sa.types.SmallInteger, sa.types.Boolean)
-                                   )
-                    ):
-                return True
+    def column_reflection_fallback(self):
+        """If we can't reflect the table, use a query to at least get column names."""
+        sql = sa.select([sa.text("*")]).select_from(self._table)
+        col_names = self.engine.execute(sql).keys()
+        col_dict = [{'name': col_name} for col_name in col_names]
+        return col_dict
 
-        return False
 
     ###
     ###
@@ -689,9 +691,6 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
         if max_value is not None and not isinstance(max_value, (Number)):
             raise ValueError("max_value must be a number")
 
-        if not self._is_numeric_column(column):
-            raise ValueError("column is not numeric")
-
         col_avg = self.engine.execute(
             sa.select([sa.func.avg(sa.column(column))]).select_from(
                 self._table)


Test patch:
diff --git a/tests/column_aggregate_expectations/expect_column_mean_to_be_between.json b/tests/column_aggregate_expectations/expect_column_mean_to_be_between.json
index 01b00b791..159966849 100644
--- a/tests/column_aggregate_expectations/expect_column_mean_to_be_between.json
+++ b/tests/column_aggregate_expectations/expect_column_mean_to_be_between.json
@@ -112,6 +112,7 @@
       "tests": [{
           "title": "type mismatch: null observed_value",
           "exact_match_out": false,
+	  "suppress_test_for": ["SQLAlchemy"],
           "in": {
             "column": "s",
             "min_value": 0,
@@ -204,4 +205,4 @@
         }
       }]
   }]
-}
\ No newline at end of file
+}
diff --git a/tests/sqlalchemy_dataset/test_sqlalchemydataset.py b/tests/sqlalchemy_dataset/test_sqlalchemydataset.py
index c418d709f..51fcbb60e 100644
--- a/tests/sqlalchemy_dataset/test_sqlalchemydataset.py
+++ b/tests/sqlalchemy_dataset/test_sqlalchemydataset.py
@@ -129,3 +129,30 @@ def test_sqlalchemydataset_with_custom_sql():
 
     result = custom_sql_dataset.expect_column_to_exist("age")
     assert result['success'] == False
+
+
+def test_column_fallback():
+    engine = sa.create_engine('sqlite://')
+
+    data = pd.DataFrame({
+        "name": ["Frank", "Steve", "Jane", "Frank", "Michael"],
+        "age": [16, 21, 38, 22, 10],
+        "pet": ["fish", "python", "cat", "python", "frog"]
+    })
+
+    data.to_sql(name='test_sql_data', con=engine, index=False)
+    dataset = SqlAlchemyDataset('test_sql_data', engine=engine)
+    fallback_dataset = SqlAlchemyDataset('test_sql_data', engine=engine)
+    # override columns attribute to test fallback
+    fallback_dataset.columns = fallback_dataset.column_reflection_fallback()
+
+    # check that the results are the same for a few expectations
+    assert (dataset.expect_column_to_exist('age') == 
+            fallback_dataset.expect_column_to_exist('age'))
+
+    assert (dataset.expect_column_mean_to_be_between('age', min_value=10) == 
+            fallback_dataset.expect_column_mean_to_be_between('age', min_value=10))
+
+    # Test a failing expectation
+    assert (dataset.expect_table_row_count_to_equal(value=3) == 
+            fallback_dataset.expect_table_row_count_to_equal(value=3))
diff --git a/tests/test_utils.py b/tests/test_utils.py
index c9dcc8f0f..74c916a22 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -128,7 +128,7 @@ def evaluate_json_test(dataset, expectation_type, test):
     """
     This method will evaluate the result of a test build using the Great Expectations json test format.
 
-    NOTE: Tests can be suppressed for certain data types if the test contains the Key 'suppress_tests_for' with a list
+    NOTE: Tests can be suppressed for certain data types if the test contains the Key 'suppress_test_for' with a list
         of Dataset types to suppress, such as ['SQLAlchemy', 'Pandas'].
 
     :param dataset: (Dataset) A great expectations Dataset