diff --git a/great_expectations/data_context/base.py b/great_expectations/data_context/base.py index 30ad372fd546..9dce04b3aa03 100644 --- a/great_expectations/data_context/base.py +++ b/great_expectations/data_context/base.py @@ -14,5 +14,5 @@ def connect(self, options): def list_datasets(self): return NotImplementedError - def get_data_set(self, dataset_name): + def get_dataset(self, dataset_name): return NotImplementedError diff --git a/great_expectations/data_context/sqlalchemy_context.py b/great_expectations/data_context/sqlalchemy_context.py index 0d0276b075b7..91e41a561616 100644 --- a/great_expectations/data_context/sqlalchemy_context.py +++ b/great_expectations/data_context/sqlalchemy_context.py @@ -24,5 +24,5 @@ def list_datasets(self): tables = [str(table) for table in self.meta.sorted_tables] return tables - def get_dataset(self, dataset_name, custom_sql=None): - return SqlAlchemyDataset(table_name=dataset_name, engine=self.engine, custom_sql=custom_sql) + def get_dataset(self, dataset_name, custom_sql=None, schema=None): + return SqlAlchemyDataset(table_name=dataset_name, engine=self.engine, custom_sql=custom_sql, schema=schema) diff --git a/tests/test_data_contexts/test_data_contexts.py b/tests/test_data_contexts/test_data_contexts.py index 0be70fb86e9b..8d763f4cf3b0 100644 --- a/tests/test_data_contexts/test_data_contexts.py +++ b/tests/test_data_contexts/test_data_contexts.py @@ -17,7 +17,7 @@ def test_db_connection_string(tmpdir_factory): path = tmpdir_factory.mktemp("db_context").join("test.db") engine = sa.create_engine('sqlite:///' + str(path)) df1.to_sql('table_1', con=engine, index=True) - df2.to_sql('table_2', con=engine, index=True) + df2.to_sql('table_2', con=engine, index=True, schema='main') # Return a connection string to this newly-created db return 'sqlite:///' + str(path) @@ -45,8 +45,10 @@ def test_sqlalchemy_data_context(test_db_connection_string): 'SqlAlchemy', test_db_connection_string, echo=False) assert context.list_datasets() == ['table_1', 'table_2'] - dataset = context.get_dataset('table_1') - assert isinstance(dataset, SqlAlchemyDataset) + dataset1 = context.get_dataset('table_1') + dataset2 = context.get_dataset('table_2', schema='main') + assert isinstance(dataset1, SqlAlchemyDataset) + assert isinstance(dataset2, SqlAlchemyDataset) def test_pandas_data_context(test_folder_connection_path):