Passed
Pull Request — master (#78)
by
unknown
02:40
created

MaintenanceScheduler.new_scheduler()   A

Complexity

Conditions 1

Size

Total Lines 10
Code Lines 7

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 2
CRAP Score 1.3644

Importance

Changes 0
Metric Value
cc 1
eloc 7
nop 2
dl 0
loc 10
ccs 2
cts 7
cp 0.2857
crap 1.3644
rs 10
c 0
b 0
f 0
1 1
import pytz
2 1
from dataclasses import dataclass
3
4 1
from apscheduler.jobstores.base import JobLookupError
5 1
from apscheduler.schedulers.background import BackgroundScheduler
6 1
from apscheduler.schedulers.base import BaseScheduler
7
8
9 1
from .deployer import MaintenanceDeployer
10 1
from ..controllers import MaintenanceController
11 1
from ..models import (
12
    MaintenanceID,
13
    MaintenanceWindow,
14
    MaintenanceWindows,
15
    OverlapError,
16
    Status,
17
)
18
19 1
from kytos.core import log
20
21 1
@dataclass
22 1
class MaintenanceStart:
23
    """
24
    Callable used for starting maintenance windows
25
    """
26 1
    maintenance_scheduler: 'MaintenanceScheduler'
27 1
    mw_id: MaintenanceID
28
29 1
    def __call__(self):
30 1
        self.maintenance_scheduler.start_maintenance(self.mw_id)
31
32
33 1
@dataclass
34 1
class MaintenanceEnd:
35
    """
36
    Callable used for ending maintenance windows
37
    """
38 1
    maintenance_scheduler: 'MaintenanceScheduler'
39 1
    mw_id: MaintenanceID
40
41 1
    def __call__(self):
42 1
        self.maintenance_scheduler.end_maintenance(self.mw_id)
43
44 1
@dataclass
45 1
class MaintenanceScheduler:
46
    """Class for scheduling maintenance windows."""
47 1
    deployer: MaintenanceDeployer
48 1
    db_controller: MaintenanceController
49 1
    scheduler: BaseScheduler
50
51 1
    @classmethod
52 1
    def new_scheduler(cls, deployer: MaintenanceDeployer):
53
        """
54
        Creates a new scheduler from the given MaintenanceDeployer
55
        """
56
        scheduler = BackgroundScheduler(timezone=pytz.utc)
57
        db_controller = MaintenanceController()
58
        db_controller.bootstrap_indexes()
59
        instance = cls(deployer, db_controller, scheduler)
60
        return instance
61
62 1
    def start(self):
63
        """
64
        Begin running the scheduler.
65
        """
66 1
        self.db_controller.prepare_start()
67
68
        # Populate the scheduler with all pending tasks
69 1
        windows = self.db_controller.get_windows()
70 1
        for window in windows:
71 1
            if window.status == Status.RUNNING:
72 1
                self.deployer.start_mw(window)
73 1
            self._schedule(window)
74
75
        # Start the scheduler
76 1
        self.scheduler.start()
77
78 1
    def shutdown(self):
79
        """
80
        Stop running the scheduler.
81
        """
82 1
        windows = self.db_controller.get_windows()
83
84
        # Depopulate the scheduler
85 1
        for window in windows:
86 1
            self._unschedule(window)
87
88 1
        self.scheduler.remove_all_jobs()
89 1
        self.scheduler.shutdown()
90
91 1
    def start_maintenance(self, mw_id: MaintenanceID):
92
        """Begins executing the maintenance window
93
        """
94
        # Get Maintenance from DB and Update
95 1
        window = self.db_controller.start_window(mw_id)
96
97
        # Activate Running
98 1
        self.deployer.start_mw(window)
99
100
        # Schedule next task
101 1
        self._schedule(window)
102
103 1
    def end_maintenance(self, mw_id: MaintenanceID):
104
        """Ends execution of the maintenance window
105
        """
106
        # Get Maintenance from DB
107 1
        window = self.db_controller.end_window(mw_id)
108
109
        # Set to Ending
110 1
        self.deployer.end_mw(window)
111
112 1
    def end_maintenance_early(self, mw_id: MaintenanceID):
113
        """Ends execution of the maintenance window early
114
        """
115
        # Get Maintenance from DB
116
        window = self.db_controller.end_window(mw_id)
117
118
        # Unschedule tasks
119
        self._unschedule(window)
120
121 1
    def add(self, window: MaintenanceWindow, force=False):
122
        """Add jobs to start and end a maintenance window."""
123
124
        if force is False:
125
            overlapping_windows = self.db_controller.check_overlap(window)
126
            if overlapping_windows:
127
                raise OverlapError(window, overlapping_windows)
128
129
        # Add window to DB
130
        self.db_controller.insert_window(window)
131
132
        # Schedule next task
133
        self._schedule(window)
134
135 1
    def update(self, window: MaintenanceWindow):
136
        """Update an existing Maintenance Window."""
137
138
        # Update window
139 1
        self.db_controller.update_window(window)
140
141
        # Reschedule any pending tasks
142 1
        self._reschedule(window)
143
144 1
    def remove(self, mw_id: MaintenanceID):
145
        """Remove jobs that start and end a maintenance window."""
146
        # Get Maintenance from DB
147
        window = self.db_controller.get_window(mw_id)
148
149
        # Remove from schedule
150
        self._unschedule(window)
151
152
        # Remove from DB
153
        self.db_controller.remove_window(mw_id)
154
155 1
    def _schedule(self, window: MaintenanceWindow):
156 1
        log.info(f'Scheduling "{window.id}"')
157 1
        if window.status == Status.PENDING:
158 1
            self.scheduler.add_job(
159
                MaintenanceStart(self, window.id),
160
                'date',
161
                id=f'{window.id}-start',
162
                run_date=window.start
163
            )
164 1
            log.info(f'Scheduled "{window.id}" start at {window.start}')
165 1
        if window.status == Status.RUNNING:
166 1
            self.scheduler.add_job(
167
                MaintenanceEnd(self, window.id),
168
                'date',
169
                id=f'{window.id}-end',
170
                run_date=window.end
171
            )
172 1
            log.info(f'Scheduled "{window.id}" end at {window.end}')
173
174 1
    def _reschedule(self, window: MaintenanceWindow):
175 1
        log.info(f'Rescheduling "{window.id}"')
176 1
        try:
177 1
            self.scheduler.remove_job(
178
                f'{window.id}-start',
179
            )
180 1
            self.scheduler.add_job(
181
                MaintenanceStart(self, window.id),
182
                'date',
183
                id=f'{window.id}-start',
184
                run_date=window.start
185
            )
186 1
            log.info(f'Rescheduled "{window.id}" start to {window.start}')
187
        except JobLookupError:
188
            log.info(f'Could not reschedule "{window.id}" start, no start job')
189 1
        try:
190 1
            self.scheduler.remove_job(
191
                f'{window.id}-end',
192
            )
193 1
            self.scheduler.add_job(
194
                MaintenanceEnd(self, window.id),
195
                'date',
196
                id=f'{window.id}-end',
197
                run_date=window.end
198
            )
199 1
            log.info(f'Rescheduled "{window.id}" end to {window.end}')
200
        except JobLookupError:
201
            log.info(f'Could not reschedule "{window.id}" end, no end job')
202
203 1
    def _unschedule(self, window: MaintenanceWindow):
204
        """Remove maintenance events from scheduler.
205
        Does not update DB, due to being
206
        primarily for shutdown startup cases.
207
        """
208 1
        started = False
209 1
        ended = False
210 1
        try:
211 1
            self.scheduler.remove_job(f'{window.id}-start')
212 1
        except JobLookupError:
213 1
            started = True
214 1
            log.info(f'Job to start "{window.id}" already removed.')
215 1
        try:
216 1
            self.scheduler.remove_job(f'{window.id}-end')
217 1
        except JobLookupError:
218 1
            ended = True
219 1
            log.info(f'Job to end "{window.id}" already removed.')
220 1
        if started and not ended:
221 1
            self.deployer.end_mw(window)
222
223 1
    def get_maintenance(self, mw_id: MaintenanceID) -> MaintenanceWindow:
224
        """Get a single maintenance by id"""
225
        return self.db_controller.get_window(mw_id)
226
227 1
    def list_maintenances(self) -> MaintenanceWindows:
228
        """Returns a list of all maintenances"""
229
        return self.db_controller.get_windows()
230