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