mirror of https://github.com/jumpserver/jumpserver
139 lines
4.3 KiB
Python
139 lines
4.3 KiB
Python
import multiprocessing
|
|
from django.core.management.base import BaseCommand, CommandError
|
|
from django.db.models import TextChoices
|
|
from .utils import ServicesUtil
|
|
from .hands import *
|
|
|
|
|
|
class Services(TextChoices):
|
|
gunicorn = 'gunicorn', 'gunicorn'
|
|
celery_ansible = 'celery_ansible', 'celery_ansible'
|
|
celery_default = 'celery_default', 'celery_default'
|
|
beat = 'beat', 'beat'
|
|
flower = 'flower', 'flower'
|
|
ws = 'ws', 'ws'
|
|
web = 'web', 'web'
|
|
celery = 'celery', 'celery'
|
|
task = 'task', 'task'
|
|
all = 'all', 'all'
|
|
|
|
@classmethod
|
|
def get_service_object_class(cls, name):
|
|
from . import services
|
|
services_map = {
|
|
cls.gunicorn.value: services.GunicornService,
|
|
cls.flower: services.FlowerService,
|
|
cls.celery_default: services.CeleryDefaultService,
|
|
cls.celery_ansible: services.CeleryAnsibleService,
|
|
cls.beat: services.BeatService
|
|
}
|
|
return services_map.get(name)
|
|
|
|
@classmethod
|
|
def web_services(cls):
|
|
return [cls.gunicorn, cls.flower]
|
|
|
|
@classmethod
|
|
def celery_services(cls):
|
|
return [cls.celery_ansible, cls.celery_default]
|
|
|
|
@classmethod
|
|
def task_services(cls):
|
|
return cls.celery_services() + [cls.beat]
|
|
|
|
@classmethod
|
|
def all_services(cls):
|
|
return cls.web_services() + cls.task_services()
|
|
|
|
@classmethod
|
|
def export_services_values(cls):
|
|
return [cls.all.value, cls.web.value, cls.task.value] + [s.value for s in cls.all_services()]
|
|
|
|
@classmethod
|
|
def get_service_objects(cls, service_names, **kwargs):
|
|
services = set()
|
|
for name in service_names:
|
|
method_name = f'{name}_services'
|
|
if hasattr(cls, method_name):
|
|
_services = getattr(cls, method_name)()
|
|
elif hasattr(cls, name):
|
|
_services = [getattr(cls, name)]
|
|
else:
|
|
continue
|
|
services.update(set(_services))
|
|
|
|
service_objects = []
|
|
for s in services:
|
|
service_class = cls.get_service_object_class(s.value)
|
|
if not service_class:
|
|
continue
|
|
kwargs.update({
|
|
'name': s.value
|
|
})
|
|
service_object = service_class(**kwargs)
|
|
service_objects.append(service_object)
|
|
return service_objects
|
|
|
|
|
|
class Action(TextChoices):
|
|
start = 'start', 'start'
|
|
status = 'status', 'status'
|
|
stop = 'stop', 'stop'
|
|
restart = 'restart', 'restart'
|
|
|
|
|
|
class BaseActionCommand(BaseCommand):
|
|
help = 'Service Base Command'
|
|
|
|
action = None
|
|
util = None
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def add_arguments(self, parser):
|
|
cores = 10
|
|
if (multiprocessing.cpu_count() * 2 + 1) < cores:
|
|
cores = multiprocessing.cpu_count() * 2 + 1
|
|
|
|
parser.add_argument(
|
|
'services', nargs='+', choices=Services.export_services_values(), help='Service',
|
|
)
|
|
parser.add_argument('-d', '--daemon', nargs="?", const=True)
|
|
parser.add_argument('-w', '--worker', type=int, nargs="?", default=cores)
|
|
parser.add_argument('-f', '--force', nargs="?", const=True)
|
|
|
|
def initial_util(self, *args, **options):
|
|
service_names = options.get('services')
|
|
service_kwargs = {
|
|
'worker_gunicorn': options.get('worker')
|
|
}
|
|
services = Services.get_service_objects(service_names=service_names, **service_kwargs)
|
|
|
|
kwargs = {
|
|
'services': services,
|
|
'run_daemon': options.get('daemon', False),
|
|
'stop_daemon': self.action == Action.stop.value and Services.all.value in service_names,
|
|
'force_stop': options.get('force') or False,
|
|
}
|
|
self.util = ServicesUtil(**kwargs)
|
|
|
|
def handle(self, *args, **options):
|
|
self.initial_util(*args, **options)
|
|
assert self.action in Action.values, f'The action {self.action} is not in the optional list'
|
|
_handle = getattr(self, f'_handle_{self.action}', lambda: None)
|
|
_handle()
|
|
|
|
def _handle_start(self):
|
|
self.util.start_and_watch()
|
|
os._exit(0)
|
|
|
|
def _handle_stop(self):
|
|
self.util.stop()
|
|
|
|
def _handle_restart(self):
|
|
self.util.restart()
|
|
|
|
def _handle_status(self):
|
|
self.util.show_status()
|