def test_basic(self):
query = 'select * from table where dt="{{ date }}"'
variable = {"date": "1970-01-01"}
self.assertEqual(
render_templated_query(query, variable),
'select * from table where dt="1970-01-01"',
)
def test_recursive_get_latest_partition_variable(self):
templated_query = render_templated_query(
'select * from table where dt="{{ latest_part }}"',
{"latest_part": '{{latest_partition("default.table", "dt")}}'},
self.DEFAULT_ENGINE_ID,
)
self.assertEqual(templated_query,
'select * from table where dt="2021-01-01"')
def test_render_templated_query(self):
templated_query = render_templated_query(
'select * from table where dt="{{ latest_partition("default.table", "dt") }}"',
{},
self.DEFAULT_ENGINE_ID,
)
self.assertEqual(templated_query,
'select * from table where dt="2021-01-01"')
def test_global_vars(self):
datetime_mock = mock.Mock(wraps=datetime)
datetime_mock.today.return_value = date(1970, 1, 1)
with mock.patch("lib.query_analysis.templating.datetime", new=datetime_mock):
query = 'select * from table where dt="{{ date }}"'
self.assertEqual(
render_templated_query(query, {"date": "{{ today }}"}),
'select * from table where dt="1970-01-01"',
)
def test_recursion(self):
query = 'select * from table where dt="{{ date }}"'
variable = {
"date": "{{ date2 }}",
"date2": "1970-{{ date3 }}-01",
"date3": "01",
}
self.assertEqual(
render_templated_query(query, variable),
'select * from table where dt="1970-01-01"',
)
def test_escape_comments(self):
query = """select * from
-- {{ end_date }}
/* {{ end_date }} */
sample_table limit 5;
-- '{{ end_date }}'
/*
{{ end_date}}
{{ end_date}}
*/
/*
{{ end_date}}*/
-- {{ end_date }}"""
self.assertEqual(render_templated_query(query, {}), query)
def test_escape_comments_non_greedy(self):
query_non_greedy = """select * from
/*
{{ end_date }}
*/
{{ test }}
/*
{{ end_date2 }}
*/
"""
self.assertEqual(
render_templated_query(query_non_greedy, {"test": "render"}),
"""select * from
/*
{{ end_date }}
*/
render
/*
{{ end_date2 }}
*/""",
)
def get_templated_query(query: str, variables: Dict[str, str], engine_id: int):
try:
return render_templated_query(query, variables, engine_id)
except QueryTemplatingError as e:
raise RequestException(e)
def run_datadoc_with_config(
self,
doc_id,
user_id=None,
# Notification related settings
notify_with=None,
notify_on=NotifyOn.ALL.value,
# Exporting related settings
exports=[],
*args,
**kwargs,
):
tasks_to_run = []
record_id = None
with DBSession() as session:
data_doc = datadoc_logic.get_data_doc_by_id(doc_id, session=session)
if not data_doc or data_doc.archived:
return
runner_id = user_id if user_id is not None else data_doc.owner_uid
query_cells = data_doc.get_query_cells()
# Preping chain jobs each unit is a [make_qe_task, run_query_task] combo
for index, query_cell in enumerate(query_cells):
engine_id = query_cell.meta["engine"]
query = render_templated_query(query_cell.context, data_doc.meta,
engine_id)
start_query_execution_kwargs = {
"cell_id": query_cell.id,
"query_execution_params": {
"query": query,
"engine_id": engine_id,
"uid": runner_id,
},
}
tasks_to_run.append(
_start_query_execution_task.si(
**start_query_execution_kwargs,
previous_query_status=QueryExecutionStatus.DONE.value,
) if index == 0 else _start_query_execution_task.s(
**start_query_execution_kwargs))
tasks_to_run.append(run_query_task.s())
# Create db entry record
record_id = create_task_run_record_for_celery_task(self,
session=session)
completion_params = {
"doc_id": doc_id,
"user_id": user_id,
"record_id": record_id,
"notify_with": notify_with,
"notify_on": notify_on,
"exports": exports,
}
chain(*tasks_to_run).apply_async(
link=on_datadoc_run_success.s(completion_params=completion_params, ),
link_error=on_datadoc_run_failure.s(
completion_params=completion_params),
)
