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 _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:
finish_code = TableauJobFinishCode.PENDING
negative_codes = (TableauJobFinishCode.ERROR,
TableauJobFinishCode.CANCELED)
while not finish_code == TableauJobFinishCode.SUCCESS:
return_code = int(
tableau_hook.server.jobs.get_by_id(job_id).finish_code)
finish_code = TableauJobFinishCode(return_code)
if finish_code in negative_codes:
raise TableauJobFailedException(
'The Tableau Refresh Workbook Job failed!')
self.log.info('Workbook %s has been successfully refreshed.',
self.workbook_name)
return 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:
if not tableau_hook.wait_for_state(
job_id=job_id,
check_interval=self.check_interval,
target_state=TableauJobFinishCode.SUCCESS,
):
raise TableauJobFailedException(
'The Tableau Refresh Workbook Job failed!')
self.log.info('Workbook %s has been successfully refreshed.',
self.workbook_name)
return job_id
def test_get_all(self, mock_pager, mock_server, mock_tableau_auth): # pylint: disable=unused-argument
"""
Test get all
"""
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_resource_id(self, tableau_hook: TableauHook) -> str:
if self.match_with == 'id':
return self.find
for resource in tableau_hook.get_all(resource_name=self.resource):
if getattr(resource, self.match_with) == self.find:
resource_id = resource.id
self.log.info('Found matching with id %s', resource_id)
return resource_id
raise AirflowException(f'{self.resource} with {self.match_with} {self.find} not found!')
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
def test_get_conn_auth_via_password_and_site_in_connection(
self, mock_server, mock_tableau_auth):
"""
Test get conn auth via password
"""
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 test_get_conn_auth_via_token_and_site_in_init(self, mock_server,
mock_tableau_auth):
"""
Test get conn auth via token
"""
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 execute(self, context: dict) -> str:
"""
Executes the Tableau API resource and pushes the job id or downloaded file URI to xcom.
:param context: The task context during execution.
:type context: dict
:return: the id of the job that executes the extract refresh or downloaded file URI.
:rtype: str
"""
available_resources = RESOURCES_METHODS.keys()
if self.resource not in available_resources:
error_message = f'Resource not found! Available Resources: {available_resources}'
raise AirflowException(error_message)
available_methods = RESOURCES_METHODS[self.resource]
if self.method not in available_methods:
error_message = f'Method not found! Available methods for {self.resource}: {available_methods}'
raise AirflowException(error_message)
with TableauHook(self.site_id, self.tableau_conn_id) as tableau_hook:
resource = getattr(tableau_hook.server, self.resource)
method = getattr(resource, self.method)
resource_id = self._get_resource_id(tableau_hook)
response = method(resource_id)
job_id = response.id
if self.method == 'refresh':
if self.blocking_refresh:
if not tableau_hook.wait_for_state(
job_id=job_id,
check_interval=self.check_interval,
target_state=TableauJobFinishCode.SUCCESS,
):
raise TableauJobFailedException(
f'The Tableau Refresh {self.resource} Job failed!')
return job_id