def execute(self, context: dict) -> str: """ Executes the Tableau Extract Refresh and pushes the job id to xcom. :param context: The task context during execution. :type context: dict :return: the id of the job that executes the extract refresh :rtype: str """ with TableauHook(self.site_id, self.tableau_conn_id) as tableau_hook: workbook = self._get_workbook_by_name(tableau_hook) job_id = self._refresh_workbook(tableau_hook, workbook.id) if self.blocking: from airflow.providers.tableau.sensors.tableau_job_status import TableauJobStatusSensor TableauJobStatusSensor( job_id=job_id, site_id=self.site_id, tableau_conn_id=self.tableau_conn_id, task_id='wait_until_succeeded', dag=None, ).execute(context={}) self.log.info('Workbook %s has been successfully refreshed.', self.workbook_name) return job_id
def test_poke_failed(self, finish_code, mock_tableau_hook): """ Test poke failed """ mock_tableau_hook.return_value.__enter__ = Mock(return_value=mock_tableau_hook) mock_get = mock_tableau_hook.server.jobs.get_by_id mock_get.return_value.finish_code = finish_code sensor = TableauJobStatusSensor(**self.kwargs) with pytest.raises(TableauJobFailedException): sensor.poke({}) mock_get.assert_called_once_with(sensor.job_id)
def test_poke(self, mock_tableau_hook): """ Test poke """ mock_tableau_hook.return_value.__enter__ = Mock(return_value=mock_tableau_hook) mock_get = mock_tableau_hook.server.jobs.get_by_id mock_get.return_value.finish_code = '0' sensor = TableauJobStatusSensor(**self.kwargs) job_finished = sensor.poke(context={}) assert job_finished mock_get.assert_called_once_with(sensor.job_id)
with DAG( dag_id='example_tableau_refresh_workbook', dagrun_timeout=timedelta(hours=2), schedule_interval=None, start_date=days_ago(2), tags=['example'], ) as dag: # Refreshes a workbook and waits until it succeeds. task_refresh_workbook_blocking = TableauRefreshWorkbookOperator( site_id='my_site', workbook_name='MyWorkbook', blocking=True, task_id='refresh_tableau_workbook_blocking', ) # Refreshes a workbook and does not wait until it succeeds. task_refresh_workbook_non_blocking = TableauRefreshWorkbookOperator( site_id='my_site', workbook_name='MyWorkbook', blocking=False, task_id='refresh_tableau_workbook_non_blocking', ) # The following task queries the status of the workbook refresh job until it succeeds. task_check_job_status = TableauJobStatusSensor( site_id='my_site', job_id= "{{ ti.xcom_pull(task_ids='refresh_tableau_workbook_non_blocking') }}", task_id='check_tableau_job_status', ) task_refresh_workbook_non_blocking >> task_check_job_status
start_date=datetime(2021, 1, 1), default_args={'site_id': 'my_site'}, tags=['example'], ) as dag: # Refreshes a workbook and waits until it succeeds. task_refresh_workbook_blocking = TableauOperator( resource='workbooks', method='refresh', find='MyWorkbook', match_with='name', blocking_refresh=True, task_id='refresh_tableau_workbook_blocking', ) # Refreshes a workbook and does not wait until it succeeds. task_refresh_workbook_non_blocking = TableauOperator( resource='workbooks', method='refresh', find='MyWorkbook', match_with='name', blocking_refresh=False, task_id='refresh_tableau_workbook_non_blocking', ) # The following task queries the status of the workbook refresh job until it succeeds. task_check_job_status = TableauJobStatusSensor( job_id=task_refresh_workbook_non_blocking.output, task_id='check_tableau_job_status', ) # Task dependency created via XComArgs: # task_refresh_workbook_non_blocking >> task_check_job_status