def _get_iterator(self, generator_asset, query_params=None):
raw_query = self._get_raw_query(generator_asset)
if raw_query is None:
logger.warning("No query defined for generator asset: %s" % generator_asset)
# There is no valid query path or temp query storage defined with the generator_asset
return None
if query_params is None:
query_params = {}
try:
substituted_query = Template(raw_query).substitute(query_params)
except KeyError:
raise BatchKwargsError(
"Unable to generate batch kwargs for asset '" + generator_asset + "': "
"missing template key",
{
"generator_asset": generator_asset,
"query_template": raw_query
}
)
return iter([
SqlAlchemyDatasourceQueryBatchKwargs(
query=substituted_query,
raw_query=raw_query,
query_params=query_params
)])
def build_batch_kwargs_from_partition_id(self,
generator_asset,
partition_id=None,
query_params=None):
"""Build batch kwargs from a partition id."""
raw_query = self._get_raw_query(generator_asset)
if "$partition_id" not in raw_query and "${partition_id}" not in raw_query:
raise BatchKwargsError(
"No partition_id parameter found in the requested query.", {})
try:
if query_params is None:
query_params = {}
query_params.update({'partition_id': partition_id})
substituted_query = Template(raw_query).substitute(query_params)
except KeyError:
raise BatchKwargsError(
"Unable to generate batch kwargs for asset '" +
generator_asset + "': "
"missing template key", {
"generator_asset": generator_asset,
"query_template": raw_query
})
return SqlAlchemyDatasourceQueryBatchKwargs(query=substituted_query,
raw_query=raw_query,
query_params=query_params)
def _get_iterator(self, data_asset_name, query_parameters=None):
raw_query = self._get_raw_query(data_asset_name=data_asset_name)
if raw_query is None:
logger.warning("No query defined for data asset: %s" %
data_asset_name)
# There is no valid query path or temp query storage defined with the data_asset_name
return None
if query_parameters is None:
iter_ = iter(
[SqlAlchemyDatasourceQueryBatchKwargs(query=raw_query)])
else:
iter_ = iter([
SqlAlchemyDatasourceQueryBatchKwargs(
query=raw_query, query_parameters=query_parameters)
])
return iter_
def _build_batch_kwargs(self, batch_parameters):
"""Build batch kwargs from a partition id."""
data_asset_name = batch_parameters.pop("data_asset_name")
raw_query = self._get_raw_query(data_asset_name=data_asset_name)
partition_id = batch_parameters.pop("partition_id", None)
batch_kwargs = self._datasource.process_batch_parameters(**batch_parameters)
batch_kwargs["query"] = raw_query
if partition_id:
if not batch_kwargs["query_parameters"]:
batch_kwargs["query_parameters"] = {}
batch_kwargs["query_parameters"]["partition_id"] = partition_id
return SqlAlchemyDatasourceQueryBatchKwargs(batch_kwargs)