def _run_image(self): self.log.info('Starting docker service from image %s', self.image) self.service = self.cli.create_service( types.TaskTemplate( container_spec=types.ContainerSpec(image=self.image, command=self.get_command(), env=self.environment, user=self.user), restart_policy=types.RestartPolicy(condition='none'), resources=types.Resources(mem_limit=self.mem_limit)), name='airflow-%s' % get_random_string(), labels={'name': 'airflow__%s__%s' % (self.dag_id, self.task_id)}) self.log.info('Service started: %s', str(self.service)) status = None # wait for the service to start the task while not self.cli.tasks(filters={'service': self.service['ID']}): continue while True: status = self.cli.tasks( filters={'service': self.service['ID']})[0]['Status']['State'] if status in ['failed', 'complete']: self.log.info('Service status before exiting: %s', status) break if self.auto_remove: self.cli.remove_service(self.service['ID']) if status == 'failed': raise AirflowException('Service failed: ' + repr(self.service))
def _run_service(self) -> None: self.log.info('Starting docker service from image %s', self.image) if not self.cli: raise Exception("The 'cli' should be initialized before!") self.service = self.cli.create_service( types.TaskTemplate( container_spec=types.ContainerSpec( image=self.image, command=self.get_command(), env=self.environment, user=self.user, tty=self.tty, ), restart_policy=types.RestartPolicy(condition='none'), resources=types.Resources(mem_limit=self.mem_limit), ), name='airflow-%s' % get_random_string(), labels={'name': 'airflow__%s__%s' % (self.dag_id, self.task_id)}, ) self.log.info('Service started: %s', str(self.service)) # wait for the service to start the task while not self.cli.tasks(filters={'service': self.service['ID']}): continue if self.enable_logging: self._stream_logs_to_output() while True: if self._has_service_terminated(): self.log.info('Service status before exiting: %s', self._service_status()) break if self.auto_remove: if not self.service: raise Exception("The 'service' should be initialized before!") self.cli.remove_service(self.service['ID']) if self._service_status() == 'failed': raise AirflowException('Service failed: ' + repr(self.service))
def test_connection() -> APIResponse: """ To test a connection, this method first creates an in-memory dummy conn_id & exports that to an env var, as some hook classes tries to find out the conn from their __init__ method & errors out if not found. It also deletes the conn id env variable after the test. """ body = request.json dummy_conn_id = get_random_string() conn_env_var = f'{CONN_ENV_PREFIX}{dummy_conn_id.upper()}' try: data = connection_schema.load(body) data['conn_id'] = dummy_conn_id conn = Connection(**data) os.environ[conn_env_var] = conn.get_uri() status, message = conn.test_connection() return connection_test_schema.dump({"status": status, "message": message}) except ValidationError as err: raise BadRequest(detail=str(err.messages)) finally: if conn_env_var in os.environ: del os.environ[conn_env_var]
def _get_service_name(self): return '%s__af_%s' % (self.task_id, get_random_string())