def register_service(cls, service_def):
valid, errors = cls.validate_service(service_def)
if not valid:
raise cls.NotValid(errors)
service = ServiceModel(**service_def)
service.save()
return service.as_dict()
def __enter__(self):
self.mock.start()
ServiceModel.create_table(1, 1, True)
TargetGroupModel.create_table(1, 1, True)
BackendModel.create_table(1, 1, True)
ResolverModel.create_table(1, 1, True)
self._server_time_getter = NTP.get_server_time
NTP.get_server_time = datetime.now
return self
def register_target_group(cls, target_group_definition):
valid, errors = cls.validate_target_group(target_group_definition)
if not valid:
raise cls.NotValid(errors)
service_name = target_group_definition.get('service_name')
try:
ServiceModel.get(service_name)
except ServiceModel.DoesNotExist:
raise Service.DoesNotExist(service_name)
target_group = TargetGroupModel(**target_group_definition)
target_group.save()
return target_group.as_dict()
def query_services(cls):
services = {s.name: s.as_dict() for s in ServiceModel.scan()}
for target_group in TargetGroupModel.scan():
service = services[target_group.service_name]
groups = service['target_groups'] = service.get(
'target_groups', [])
groups.append(target_group.as_dict())
service['target_groups'] = sorted(
groups, key=itemgetter('target_group_name'))
with BackendModel.batch_write() as batch:
for backend in BackendModel.scan():
service = services[backend.service_name]
if backend.status.upper() == "DRAINING":
now = NTP.get_server_time()
diff = (now - backend.updated_at.replace(tzinfo=None)
).total_seconds()
if diff > service["connection_draining"]:
backend.status = "DRAINED"
backend.updated_at = now
batch.save(backend)
group = next(
filter(
lambda tg: tg["target_group_name"] == backend.
target_group_name, service['target_groups']))
backends = group['backends'] = group.get('backends', [])
backends.append(backend.as_dict())
group['backends'] = sorted(backends, key=itemgetter('host'))
return list(services.values())
def register_backend(cls, backend_definition):
valid, errors = cls.validate_backend(backend_definition)
if not valid:
raise cls.NotValid(errors)
service_name = backend_definition.get('service_name')
target_group_name = backend_definition.get('target_group_name')
try:
ServiceModel.get(service_name)
except ServiceModel.DoesNotExist:
raise Service.DoesNotExist(service_name)
try:
TargetGroupModel.get(service_name, target_group_name)
except TargetGroupModel.DoesNotExist:
raise Service.DoesNotExist(target_group_name)
backend_definition["updated_at"] = NTP.get_server_time()
backend = BackendModel(**backend_definition)
backend.save()
return backend.as_dict()
def deregister_service(cls, name):
try:
service = ServiceModel.get(name)
except ServiceModel.DoesNotExist:
raise Service.DoesNotExist(name)
service_description = cls.describe_service(name)
for target_group in service_description['target_groups']:
target_group_name = target_group['target_group_name']
for backend in target_group['backends']:
cls.deregister_backend(name, target_group_name,
backend['host'])
cls.deregister_target_group(name, target_group_name)
service.delete()
return True
def update_service(cls, service_name, update_def):
try:
service = ServiceModel.get(service_name)
except ServiceModel.DoesNotExist:
raise Service.DoesNotExist(service_name)
for k, v in update_def.items():
if k in service.attribute_values.keys():
service.attribute_values[k] = v
service_dict = service.as_dict()
valid, errors = cls.validate_service(service_dict)
if not valid:
raise cls.NotValid(errors)
service.save()
return service.as_dict()
def describe_service(cls, service_name):
try:
service = ServiceModel.get(service_name)
except ServiceModel.DoesNotExist:
raise Service.DoesNotExist(service_name)
target_groups = TargetGroupModel.query(service_name)
target_groups = sorted([tg.as_dict() for tg in target_groups],
key=itemgetter('target_group_name'))
for target_group in target_groups:
target_group_name = target_group.get('target_group_name')
target_group["backends"] = sorted(list(
map(
lambda backend: backend.as_dict(),
BackendModel.scan(
service_name__eq=service_name,
target_group_name__eq=target_group_name))),
key=itemgetter('host'))
data = service.as_dict()
data['target_groups'] = target_groups
return data