Passed
Pull Request — master (#78)
by
unknown
05:53
created

MaintenanceController.update_window()   A

Complexity

Conditions 1

Size

Total Lines 7
Code Lines 7

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 2
CRAP Score 1

Importance

Changes 0
Metric Value
cc 1
eloc 7
nop 2
dl 0
loc 7
ccs 2
cts 2
cp 1
crap 1
rs 10
c 0
b 0
f 0
1
"""MaintenanceController."""
2
3
# pylint: disable=invalid-name
4 1
from datetime import datetime
5 1
import os
6 1
import pytz
7 1
from typing import Optional
8
9 1
from bson.codec_options import CodecOptions
10 1
import pymongo
11 1
from pymongo.errors import AutoReconnect, DuplicateKeyError
12 1
from tenacity import retry_if_exception_type, stop_after_attempt, wait_random
13
14 1
from kytos.core import log
15 1
from kytos.core.db import Mongo
16 1
from kytos.core.retry import before_sleep, for_all_methods, retries
17 1
from napps.kytos.maintenance.models import (
18
    MaintenanceWindow,
19
    MaintenanceWindows,
20
    MaintenanceID,
21
    Status,
22
)
23
24
25 1
@for_all_methods(
26
    retries,
27
    stop=stop_after_attempt(
28
        int(os.environ.get("MONGO_AUTO_RETRY_STOP_AFTER_ATTEMPT", 3))
29
    ),
30
    wait=wait_random(
31
        min=int(os.environ.get("MONGO_AUTO_RETRY_WAIT_RANDOM_MIN", 0.1)),
32
        max=int(os.environ.get("MONGO_AUTO_RETRY_WAIT_RANDOM_MAX", 1)),
33
    ),
34
    before_sleep=before_sleep,
35
    retry=retry_if_exception_type((AutoReconnect,)),
36
)
37 1
class MaintenanceController:
38
    """MaintenanceController."""
39
40 1
    def __init__(self, get_mongo=lambda: Mongo()) -> None:
0 ignored issues
show
Comprehensibility Best Practice introduced by
The variable Mongo does not seem to be defined.
Loading history...
41
        """Constructor of MaintenanceController."""
42 1
        self.mongo = get_mongo()
43 1
        self.db_client = self.mongo.client
44 1
        self.db = self.db_client[self.mongo.db_name]
45 1
        self.windows = self.db['maintenance.windows'].with_options(
46
            codec_options=CodecOptions(
47
                tz_aware=True,
48
            )
49
        )
50
51 1
    def bootstrap_indexes(self) -> None:
52
        """Bootstrap all maintenance related indexes."""
53 1
        unique_index_tuples = [
54
            ("maintenance.windows", [("id", pymongo.ASCENDING)]),
55
        ]
56 1
        for collection, keys in unique_index_tuples:
57 1
            if self.mongo.bootstrap_index(collection, keys, unique=True):
58 1
                log.info(
59
                    f"Created DB unique index {keys}, collection: {collection})"
60
                )
61
62 1
    def insert_window(self, window: MaintenanceWindow):
63 1
        now = datetime.now(pytz.utc)
64 1
        try:
65 1
            self.windows.insert_one({
66
                        **window.dict(exclude={'inserted_at', 'updated_at'}),
67
                        'inserted_at': now,
68
                        'updated_at': now,
69
            })
70
        except DuplicateKeyError as err:
71
            raise ValueError(f'Window with id: {window.id} already exists') from err
72
73 1
    def update_window(self, window: MaintenanceWindow):
74 1
        self.windows.update_one(
75
            {'id': window.id},
76
            [{
77
                '$set': {
78
                    **window.dict(exclude={'inserted_at', 'updated_at'}),
79
                    'updated_at': '$$NOW',
80
                },
81
            }],
82
        )
83
84 1
    def get_window(self, mw_id: MaintenanceID) -> Optional[MaintenanceWindow]:
85 1
        window = self.windows.find_one(
86
            {'id': mw_id},
87
            {'_id': False},
88
        )
89 1
        if window is None:
90 1
            return None
91
        else:
92 1
            return MaintenanceWindow.construct(**window)
93
94 1
    def start_window(self, mw_id: MaintenanceID) -> MaintenanceWindow:
95
        window = self.windows.find_one_and_update(
96
            {'id': mw_id},
97
            [{
98
                '$set': {
99
                    'status': Status.RUNNING,
100
                    'last_modified': '$$NOW',
101
                },
102
            }],
103
            {'_id': False},
104
            return_document=pymongo.ReturnDocument.AFTER,
105
        )
106
        return MaintenanceWindow.construct(**window)
107
108 1
    def end_window(self, mw_id: MaintenanceID) -> MaintenanceWindow:
109
        window = self.windows.find_one_and_update(
110
            {'id': mw_id},
111
            [{
112
                '$set': {
113
                    'status': Status.FINISHED,
114
                    'last_modified': '$$NOW',
115
                },
116
            }],
117
            {'_id': False},
118
            return_document=pymongo.ReturnDocument.AFTER,
119
        )
120
        return MaintenanceWindow.construct(**window)
121
122 1
    def check_overlap(self, window):
123
        # If two time periods are overlapping,
124
        # then the start of one time period must occur in the other time period
125
        windows = self.windows.find(
126
            {
127
                '$or': [
128
                    {'$and': [
129
                        {'start': {'$lte': window.start}},
130
                        {'end': {'$gt': window.start}},
131
                    ]},
132
                    {'$and': [
133
                        {'start': {'$gte': window.start}},
134
                        {'start': {'$lt': window.end}},
135
                    ]}
136
                ]
137
            },
138
            {'_id': False}
139
        )
140
        return MaintenanceWindows.construct(
141
            __root__ = [MaintenanceWindow.construct(**window) for window in windows]
142
        )
143
144 1
    def get_windows(self) -> MaintenanceWindows:
145 1
        windows = self.windows.find(projection={'_id': False})
146 1
        return MaintenanceWindows.construct(
147
            __root__ = [MaintenanceWindow.construct(**window) for window in windows]
148
        )
149
150 1
    def remove_window(self, mw_id: MaintenanceID):
151
        self.windows.delete_one({'id': mw_id})
152
153 1
    def prepare_start(self):
154
        now = datetime.now(pytz.utc)
155
        self.windows.update_many(
156
            {'$and': [
157
                {'status': {'$eq': Status.PENDING}},
158
                {'start': {'$lte': now}},
159
            ]},
160
            {
161
                '$set': {
162
                    'status': Status.RUNNING,
163
                    'last_modified': now,
164
                },
165
            }
166
        )
167
        self.windows.update_many(
168
            {'$and': [
169
                {'status': {'$eq': Status.RUNNING}},
170
                {'end': {'$lte': now}},
171
            ]},
172
            {
173
                '$set': {
174
                    'status': Status.FINISHED,
175
                    'last_modified': now,
176
                },
177
            }
178
        )
179