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), )