Repository: great-expectations/great_expectations

Issue:
Allow passing a schema name to SqlAlchemyDataset
Currently, SqlAlchemyDataset accepts a table name but not a schema name. I'm currently working in Redshift. When I try to instantiate a `SqlAlchemyDataset`, I get a `KeyError` with `public.<tablename>`. After some inspection, I determined that the issue is popping up in this line:https://github.com/great-expectations/great_expectations/blob/d7a4ee03cfce8bea66c5f91573bb87e4385b1eb1/great_expectations/dataset/sqlalchemy_dataset.py#L222

Because I'm not passing a schema name, it uses the default schema, which is "public" for Redshift. When I run this line myself but specifying the schema name, the line runs successfully. Accessing non-default schemas seems like a common use case, so a schema argument would be helpful in many other cases.

Discussion:

Patch:
diff --git a/docs/source/custom_expectations.rst b/docs/source/custom_expectations.rst
index f7340871b..419814e8d 100644
--- a/docs/source/custom_expectations.rst
+++ b/docs/source/custom_expectations.rst
@@ -74,7 +74,7 @@ For SqlAlchemyDataset, the decorators work slightly differently. See the MetaSql
             mode_query = sa.select([
                 sa.column(column).label('value'),
                 sa.func.count(sa.column(column)).label('frequency')
-            ]).select_from(sa.table(self.table_name)).group_by(sa.column(column)).order_by(sa.desc(sa.column('frequency')))
+            ]).select_from(self._table).group_by(sa.column(column)).order_by(sa.desc(sa.column('frequency')))
 
             mode = self.engine.execute(mode_query).scalar()
             return {
diff --git a/docs/source/roadmap_changelog.rst b/docs/source/roadmap_changelog.rst
index 4563e1084..339696554 100644
--- a/docs/source/roadmap_changelog.rst
+++ b/docs/source/roadmap_changelog.rst
@@ -14,6 +14,7 @@ Planned Features
 
 v.0.4.4__develop
 ----------------
+* Add support for custom schema in SqlAlchemyDataset (#370)
 
 v.0.4.4
 ----------------
diff --git a/great_expectations/dataset/sqlalchemy_dataset.py b/great_expectations/dataset/sqlalchemy_dataset.py
index 8028e95fe..ec86b4416 100644
--- a/great_expectations/dataset/sqlalchemy_dataset.py
+++ b/great_expectations/dataset/sqlalchemy_dataset.py
@@ -67,7 +67,7 @@ class MetaSqlAlchemyDataset(Dataset):
                                       sa.column(column).is_(None) == False if None in ignore_values else True),
                               1)], else_=0)
                 ).label('unexpected_count')
-            ]).select_from(sa.table(self.table_name))
+            ]).select_from(self._table)
 
             count_results = dict(self.engine.execute(count_query).fetchone())
 
@@ -81,7 +81,7 @@ class MetaSqlAlchemyDataset(Dataset):
 
             # Retrieve unexpected  values
             unexpected_query_results = self.engine.execute(
-                sa.select([sa.column(column)]).select_from(sa.table(self.table_name)).where(
+                sa.select([sa.column(column)]).select_from(self._table).where(
                     sa.and_(sa.not_(expected_condition),
                             sa.or_(
                                 # SA normally evaluates `== None` as `IS NONE`. However `sa.in_()`
@@ -164,7 +164,7 @@ class MetaSqlAlchemyDataset(Dataset):
                     sa.func.sum(
                         sa.case([(sa.column(column) == None, 1)], else_=0)
                     ).label('null_count'),
-                ]).select_from(sa.table(self.table_name))
+                ]).select_from(self._table)
 
                 count_results = dict(self.engine.execute(count_query).fetchone())
 
@@ -204,11 +204,12 @@ class MetaSqlAlchemyDataset(Dataset):
 
 class SqlAlchemyDataset(MetaSqlAlchemyDataset):
 
-    def __init__(self, table_name=None, engine=None, connection_string=None, custom_sql=None, *args, **kwargs):
+    def __init__(self, table_name=None, engine=None, connection_string=None,
+                 custom_sql=None, schema=None, *args, **kwargs):
         if table_name is None:
             raise ValueError("No table_name provided.")
 
-        self.table_name = table_name
+        self._table = sa.Table(table_name, sa.MetaData(), schema=schema)
 
         if engine is None and connection_string is None:
             raise ValueError("Engine or connection_string must be provided.")
@@ -223,11 +224,17 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
                 # Currently we do no error handling if the engine doesn't work out of the box.
                 raise err
 
+        if schema is not None and custom_sql is not None:
+            # temporary table will be written to temp schema, so don't allow
+            # a user-defined schema
+            raise ValueError("Cannot specify both schema and custom_sql.")
+
+
         if custom_sql:
-            self.create_temporary_table(self.table_name, custom_sql)
+            self.create_temporary_table(table_name, custom_sql)
 
         insp = reflection.Inspector.from_engine(engine)
-        self.columns = insp.get_columns(self.table_name)
+        self.columns = insp.get_columns(table_name, schema=schema)
 
         # Only call super once connection is established and table_name and columns known to allow autoinspection
         super(SqlAlchemyDataset, self).__init__(*args, **kwargs)
@@ -282,7 +289,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
         if value is None:
             raise ValueError("value must be provided")
 
-        count_query = sa.select([sa.func.count()]).select_from(sa.table(self.table_name))
+        count_query = sa.select([sa.func.count()]).select_from(self._table)
         row_count = self.engine.execute(count_query).scalar()
 
         return {
@@ -310,7 +317,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
         except ValueError:
             raise ValueError("min_value and max_value must be integers")
 
-        count_query = sa.select([sa.func.count()]).select_from(sa.table(self.table_name))
+        count_query = sa.select([sa.func.count()]).select_from(self._table)
         row_count = self.engine.execute(count_query).scalar()
 
         if min_value != None and max_value != None:
@@ -521,7 +528,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
             raise ValueError("parse_strings_as_datetimes is not supported in SqlAlchemy")
 
         col_max = self.engine.execute(
-            sa.select([sa.func.max(sa.column(column))]).select_from(sa.table(self.table_name))
+            sa.select([sa.func.max(sa.column(column))]).select_from(self._table)
         ).scalar()
 
         # Handle possible missing values
@@ -568,7 +575,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
             raise ValueError("parse_strings_as_datetimes is not supported in SqlAlchemy")
 
         col_min = self.engine.execute(
-            sa.select([sa.func.min(sa.column(column))]).select_from(sa.table(self.table_name))
+            sa.select([sa.func.min(sa.column(column))]).select_from(self._table)
         ).scalar()
 
         # Handle possible missing values
@@ -609,7 +616,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
             raise ValueError("min_value and max_value cannot both be None")
 
         col_sum = self.engine.execute(
-            sa.select([sa.func.sum(sa.column(column))]).select_from(sa.table(self.table_name))
+            sa.select([sa.func.sum(sa.column(column))]).select_from(self._table)
         ).scalar()
 
         # Handle possible missing values
@@ -659,7 +666,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
             raise ValueError("column is not numeric")
 
         col_avg = self.engine.execute(
-            sa.select([sa.func.avg(sa.column(column))]).select_from(sa.table(self.table_name))
+            sa.select([sa.func.avg(sa.column(column))]).select_from(self._table)
         ).scalar()
 
         # Handle possible missing values
@@ -707,7 +714,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
                 sa.func.sum(
                     sa.case([(sa.column(column) == None, 1)], else_=0)
                 ).label('null_count')
-            ]).select_from(sa.table(self.table_name))
+            ]).select_from(self._table)
         )
 
         counts = dict(count_query.fetchone())
@@ -724,7 +731,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
         element_values = self.engine.execute(
             sa.select([sa.column(column)]).order_by(sa.column(column)).where(
                 sa.column(column) != None
-            ).offset(nonnull_count // 2 - 1).limit(2).select_from(sa.table(self.table_name))
+            ).offset(nonnull_count // 2 - 1).limit(2).select_from(self._table)
         )
 
         column_values = list(element_values.fetchall())
@@ -761,7 +768,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
                                           result_format=None, include_config=False, catch_exceptions=None, meta=None):
         # Duplicates are found by filtering a group by query
         dup_query = sa.select([sa.column(column)]).\
-            select_from(sa.table(self.table_name)).\
+            select_from(self._table).\
             group_by(sa.column(column)).\
             having(sa.func.count(sa.column(column)) > 1)
 
@@ -777,7 +784,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
             raise ValueError("min_value and max_value cannot both be None")
 
         unique_value_count = self.engine.execute(
-            sa.select([sa.func.count(sa.func.distinct(sa.column(column)))]).select_from(sa.table(self.table_name))
+            sa.select([sa.func.count(sa.func.distinct(sa.column(column)))]).select_from(self._table)
         ).scalar()
 
         # Handle possible missing values
@@ -814,7 +821,7 @@ class SqlAlchemyDataset(MetaSqlAlchemyDataset):
                     sa.case([(sa.column(column) == None, 1)], else_=0)
                 ).label('null_count'),
                 sa.func.count(sa.func.distinct(sa.column(column))).label('unique_value_count')
-            ]).select_from(sa.table(self.table_name))
+            ]).select_from(self._table)
         )
 
         counts = count_query.fetchone()


Test patch:
diff --git a/tests/sqlalchemy_dataset/test_sqlalchemydataset.py b/tests/sqlalchemy_dataset/test_sqlalchemydataset.py
index 4c32c47b3..59c2c3eff 100644
--- a/tests/sqlalchemy_dataset/test_sqlalchemydataset.py
+++ b/tests/sqlalchemy_dataset/test_sqlalchemydataset.py
@@ -18,7 +18,7 @@ def custom_dataset():
             mode_query = sa.select([
                 sa.column(column).label('value'),
                 sa.func.count(sa.column(column)).label('frequency')
-            ]).select_from(sa.table(self.table_name)).group_by(sa.column(column)).order_by(
+            ]).select_from(self._table).group_by(sa.column(column)).order_by(
                 sa.desc(sa.column('frequency')))
 
             mode = self.engine.execute(mode_query).scalar()
@@ -85,6 +85,21 @@ def test_broken_decorator_errors(custom_dataset):
         assert "Column aggregate expectation failed to return required information: observed_value" in str(err)
 
 
+def test_missing_engine_error():
+    with pytest.raises(ValueError) as err:
+        SqlAlchemyDataset('test_engine', schema='example')
+        assert "Engine or connection_string must be provided." in str(err)
+
+
+def test_schema_custom_sql_error():
+    engine = sa.create_engine('sqlite://')
+
+    with pytest.raises(ValueError) as err:
+        SqlAlchemyDataset('test_schema_custom', schema='example', engine=engine,
+                          custom_sql='SELECT * FROM example.fake')
+        assert "Cannot specify both schema and custom_sql." in str(err)
+
+
 def test_sqlalchemydataset_with_custom_sql():
     engine = sa.create_engine('sqlite://')