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