def test_build_batch_kwargs_for_query_path(basic_sqlalchemy_datasource,
data_context):
base_directory = data_context.root_directory
basic_sqlalchemy_datasource._data_context = data_context
generator_name = "my_generator"
query_str = "SELECT * FROM my_table"
default_sql_files_directory = os.path.join(
base_directory,
"datasources",
basic_sqlalchemy_datasource.name,
"generators",
generator_name,
)
os.makedirs(default_sql_files_directory)
with open(os.path.join(default_sql_files_directory, "example.sql"),
"w") as outfile:
outfile.write(query_str)
generator = QueryBatchKwargsGenerator(
name=generator_name, datasource=basic_sqlalchemy_datasource)
batch_kwargs = generator.build_batch_kwargs(name="example")
assert isinstance(batch_kwargs, SqlAlchemyDatasourceQueryBatchKwargs)
assert batch_kwargs.query == query_str
assert batch_kwargs.query_parameters == None
def test_partition_id(basic_sqlalchemy_datasource):
generator = QueryBatchKwargsGenerator(
datasource=basic_sqlalchemy_datasource,
queries={
"my_asset": "SELECT * FROM my_table WHERE value = $partition_id",
})
batch_kwargs = generator.build_batch_kwargs("my_asset", partition_id="foo")
assert isinstance(batch_kwargs, SqlAlchemyDatasourceQueryBatchKwargs)
assert batch_kwargs.query == "SELECT * FROM my_table WHERE value = $partition_id"
assert batch_kwargs.query_parameters == {"partition_id": "foo"}
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_get_available_data_asset_names_for_query_path(empty_data_context):
# create queries path
context_path = empty_data_context.root_directory
query_path = os.path.join(
context_path, "datasources/mydatasource/generators/mygenerator")
os.makedirs(query_path, exist_ok=True)
shutil.copy(file_relative_path(__file__, "../../test_fixtures/dummy.sql"),
query_path)
data_source = Datasource(name="mydatasource",
data_context=empty_data_context)
generator = QueryBatchKwargsGenerator(name="mygenerator",
datasource=data_source)
sql_list = generator.get_available_data_asset_names()
assert ("dummy", "query") in sql_list["names"]
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"