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())
