def test_column_fallback(sa):
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)
assert set(dataset.get_table_columns()) == {"name", "age", "pet"}
fallback_dataset = SqlAlchemyDataset("test_sql_data", engine=engine)
# override columns attribute to test fallback
fallback_dataset.columns = fallback_dataset.column_reflection_fallback()
assert set(fallback_dataset.get_table_columns()) == {"name", "age", "pet"}
# 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)
def test_sqlalchemydataset_with_custom_sql():
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)
custom_sql = "SELECT name, pet FROM test_sql_data WHERE age > 25"
custom_sql_dataset = SqlAlchemyDataset('test_sql_data',
engine=engine,
custom_sql=custom_sql)
custom_sql_dataset._initialize_expectations()
custom_sql_dataset.set_default_expectation_argument(
"result_format", {"result_format": "COMPLETE"})
result = custom_sql_dataset.expect_column_values_to_be_in_set(
"pet", ["fish", "cat", "python"])
assert result['success'] == True
result = custom_sql_dataset.expect_column_to_exist("age")
assert result['success'] == False