def test_add_query(basic_sqlalchemy_datasource):
generator = QueryBatchKwargsGenerator(
datasource=basic_sqlalchemy_datasource)
generator.add_query("my_asset",
"select * from my_table where val > $condition")
batch_kwargs = generator.yield_batch_kwargs(
"my_asset", query_parameters={"condition": 5})
assert isinstance(batch_kwargs, SqlAlchemyDatasourceQueryBatchKwargs)
assert batch_kwargs.query == "select * from my_table where val > $condition"
assert batch_kwargs.query_parameters == {"condition": 5}
def test_basic_operation(basic_sqlalchemy_datasource):
# We should be able to include defined queries as part of configuration
generator = QueryBatchKwargsGenerator(
datasource=basic_sqlalchemy_datasource,
queries={
"my_asset": "SELECT * FROM my_table WHERE value = $condition",
"my_simple_asset": "SELECT c1, c2 FROM my_table"
})
# Returned assets should be typed and processed by template language
batch_kwargs = generator.yield_batch_kwargs(
"my_asset", query_parameters={'condition': "foo"})
assert isinstance(batch_kwargs, SqlAlchemyDatasourceQueryBatchKwargs)
assert batch_kwargs.query == "SELECT * FROM my_table WHERE value = $condition"
assert batch_kwargs.query_parameters == {'condition': "foo"}
# Without a template, everything should still work
batch_kwargs = generator.yield_batch_kwargs("my_simple_asset")
assert isinstance(batch_kwargs, SqlAlchemyDatasourceQueryBatchKwargs)
assert batch_kwargs.query == "SELECT c1, c2 FROM my_table"