def execute(self, context: 'Context'): hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, location=self.location ) self.log.info('Submitting manual transfer for %s', self.transfer_config_id) response = hook.start_manual_transfer_runs( transfer_config_id=self.transfer_config_id, requested_time_range=self.requested_time_range, requested_run_time=self.requested_run_time, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) transfer_config = _get_transfer_config_details(response.runs[0].name) BigQueryDataTransferConfigLink.persist( context=context, task_instance=self, region=transfer_config["region"], config_id=transfer_config["config_id"], project_id=transfer_config["project_id"], ) result = StartManualTransferRunsResponse.to_dict(response) run_id = get_object_id(result['runs'][0]) self.xcom_push(context, key="run_id", value=run_id) self.log.info('Transfer run %s submitted successfully.', run_id) return result
def setUp(self) -> None: with mock.patch( "airflow.providers.google.cloud.hooks.bigquery_dts.GoogleBaseHook.__init__", new=mock_base_gcp_hook_no_default_project_id, ): self.hook = BiqQueryDataTransferServiceHook() self.hook._get_credentials = mock.MagicMock(return_value=CREDENTIALS) # type: ignore
def execute(self, context: 'Context'): hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, location=self.location ) self.log.info("Creating DTS transfer config") response = hook.create_transfer_config( project_id=self.project_id, transfer_config=self.transfer_config, authorization_code=self.authorization_code, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) transfer_config = _get_transfer_config_details(response.name) BigQueryDataTransferConfigLink.persist( context=context, task_instance=self, region=transfer_config["region"], config_id=transfer_config["config_id"], project_id=transfer_config["project_id"], ) result = TransferConfig.to_dict(response) self.log.info("Created DTS transfer config %s", get_object_id(result)) self.xcom_push(context, key="transfer_config_id", value=get_object_id(result)) return result
def execute(self, context): hook = BiqQueryDataTransferServiceHook(gcp_conn_id=self.gcp_conn_id) hook.delete_transfer_config( transfer_config_id=self.transfer_config_id, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, )
def execute(self, context) -> None: hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain) hook.delete_transfer_config( transfer_config_id=self.transfer_config_id, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, )
def poke(self, context: 'Context') -> bool: hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_cloud_conn_id, impersonation_chain=self.impersonation_chain, ) run = hook.get_transfer_run( run_id=self.run_id, transfer_config_id=self.transfer_config_id, project_id=self.project_id, retry=self.retry, timeout=self.request_timeout, metadata=self.metadata, ) self.log.info("Status of %s run: %s", self.run_id, str(run.state)) return run.state in self.expected_statuses
def poke(self, context): hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_cloud_conn_id) run = hook.get_transfer_run( run_id=self.run_id, transfer_config_id=self.transfer_config_id, project_id=self.project_id, retry=self.retry, timeout=self.request_timeout, metadata=self.metadata, ) result = MessageToDict(run) state = result["state"] self.log.info("Status of %s run: %s", self.run_id, state) return state in self.expected_statuses
def execute(self, context): hook = BiqQueryDataTransferServiceHook(gcp_conn_id=self.gcp_conn_id) self.log.info("Creating DTS transfer config") response = hook.create_transfer_config( project_id=self.project_id, transfer_config=self.transfer_config, authorization_code=self.authorization_code, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) result = MessageToDict(response) self.log.info("Created DTS transfer config %s", get_object_id(result)) self.xcom_push(context, key="transfer_config_id", value=get_object_id(result)) return result
def poke(self, context: 'Context') -> bool: hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_cloud_conn_id, impersonation_chain=self.impersonation_chain, ) run = hook.get_transfer_run( run_id=self.run_id, transfer_config_id=self.transfer_config_id, project_id=self.project_id, retry=self.retry, timeout=self.request_timeout, metadata=self.metadata, ) self.log.info("Status of %s run: %s", self.run_id, str(run.state)) if run.state in (TransferState.FAILED, TransferState.CANCELLED): raise AirflowException(f"Transfer {self.run_id} did not succeed") return run.state in self.expected_statuses
def execute(self, context): hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain ) self.log.info('Submitting manual transfer for %s', self.transfer_config_id) response = hook.start_manual_transfer_runs( transfer_config_id=self.transfer_config_id, requested_time_range=self.requested_time_range, requested_run_time=self.requested_run_time, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) result = StartManualTransferRunsResponse.to_dict(response) run_id = get_object_id(result['runs'][0]) self.xcom_push(context, key="run_id", value=run_id) self.log.info('Transfer run %s submitted successfully.', run_id) return result