MaxKB/apps/common/management/commands/services/command.py at 8c640e556293a1af0f963b6775f0862f3be2d608 · 1Panel-dev/MaxKB · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import math
import os

from django.core.management.base import BaseCommand
from django.db.models import TextChoices

from .utils import ServicesUtil


class Services(TextChoices):
    gunicorn = 'gunicorn', 'gunicorn'
    celery_default = 'celery_default', 'celery_default'
    local_model = 'local_model', 'local_model'
    scheduler = 'scheduler', 'scheduler'
    web = 'web', 'web'
    celery = 'celery', 'celery'
    celery_model = 'celery_model', 'celery_model'
    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.celery_default: services.CeleryDefaultService,
            cls.local_model: services.GunicornLocalModelService,
            cls.scheduler: services.SchedulerService,
        }
        return services_map.get(name)

    @classmethod
    def web_services(cls):
        return [cls.gunicorn, cls.local_model]

    @classmethod
    def celery_services(cls):
        return [cls.celery_default, cls.celery_model]

    @classmethod
    def task_services(cls):
        return cls.celery_services()

    @classmethod
    def scheduler_services(cls):
        return [cls.scheduler]

    @classmethod
    def all_services(cls):
        return cls.web_services() + cls.task_services() + cls.scheduler_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):
        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=2 if os.cpu_count() > 6 else math.floor(os.cpu_count() / 2))
        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()