def test_get_all(self, mock_pager, mock_server, mock_tableau_auth): # pylint: disable=unused-argument
with TableauHook(
tableau_conn_id='tableau_test_password') as tableau_hook:
jobs = tableau_hook.get_all(resource_name='jobs')
assert jobs == mock_pager.return_value
mock_pager.assert_called_once_with(mock_server.return_value.jobs.get)
def _get_workbook_by_name(self, tableau_hook: TableauHook) -> WorkbookItem:
for workbook in tableau_hook.get_all(resource_name='workbooks'):
if workbook.name == self.workbook_name:
self.log.info('Found matching workbook with id %s', workbook.id)
return workbook
raise AirflowException(f'Workbook {self.workbook_name} not found!')
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.salesforce.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_get_conn_auth_via_token_and_site_in_init(self, mock_server, mock_tableau_auth):
with TableauHook(site_id='test', tableau_conn_id='tableau_test_token') as tableau_hook:
mock_server.assert_called_once_with(tableau_hook.conn.host, use_server_version=True)
mock_tableau_auth.assert_called_once_with(
token_name=tableau_hook.conn.extra_dejson['token_name'],
personal_access_token=tableau_hook.conn.extra_dejson['personal_access_token'],
site_id=tableau_hook.site_id
)
mock_server.return_value.auth.sign_in_with_personal_access_token.assert_called_once_with(
mock_tableau_auth.return_value
)
mock_server.return_value.auth.sign_out.assert_called_once_with()
def test_get_conn_auth_via_password_and_site_in_connection(self, mock_server, mock_tableau_auth):
with TableauHook(tableau_conn_id='tableau_test_password') as tableau_hook:
mock_server.assert_called_once_with(tableau_hook.conn.host, use_server_version=True)
mock_tableau_auth.assert_called_once_with(
username=tableau_hook.conn.login,
password=tableau_hook.conn.password,
site_id=tableau_hook.conn.extra_dejson['site_id']
)
mock_server.return_value.auth.sign_in.assert_called_once_with(
mock_tableau_auth.return_value
)
mock_server.return_value.auth.sign_out.assert_called_once_with()
def poke(self, context: dict) -> bool:
"""
Pokes until the job has successfully finished.
:param context: The task context during execution.
:type context: dict
:return: True if it succeeded and False if not.
:rtype: bool
"""
with TableauHook(self.site_id, self.tableau_conn_id) as tableau_hook:
finish_code = TableauJobFinishCode(
int(tableau_hook.server.jobs.get_by_id(self.job_id).finish_code)
)
self.log.info('Current finishCode is %s (%s)', finish_code.name, finish_code.value)
if finish_code in [TableauJobFinishCode.ERROR, TableauJobFinishCode.CANCELED]:
raise TableauJobFailedException('The Tableau Refresh Workbook Job failed!')
return finish_code == TableauJobFinishCode.SUCCESS