ensure-indexes.py 4.45 KB
Newer Older
1
2
3
4
# -*- coding: utf-8 -*-
# ----------------------------------------------------------------------
# Ensure MongoDB indexes
# ----------------------------------------------------------------------
5
# Copyright (C) 2007-2019 The NOC Project
6
7
8
9
10
11
12
13
# See LICENSE for details
# ----------------------------------------------------------------------

# Python modules
from __future__ import print_function
# NOC modules
from noc.core.management.base import BaseCommand
from noc.models import get_model, iter_model_id, is_document
Dmitry Volodin's avatar
Dmitry Volodin committed
14
from noc.core.datastream.loader import loader as ds_loader
15
16
from noc.core.gridvcs.base import GridVCS
from noc.core.gridvcs.utils import REPOS
Dmitry Volodin's avatar
Dmitry Volodin committed
17
from noc.config import config
18
19
20
21


class Command(BaseCommand):
    def handle(self, host=None, port=None, *args, **options):
22
        from noc.lib.nosql import get_db
23

24
        db = get_db()
25
        collections = set(db.list_collection_names())
26
27
28
29
30
31
        for model_id in iter_model_id():
            model = get_model(model_id)
            if not model:
                self.die("Invalid model: %s" % model_id)
            if not is_document(model):
                continue
32
33
34
35
36
37
38
39
            # Rename collections when necessary
            legacy_collections = model._meta.get("legacy_collections", [])
            for old_name in legacy_collections:
                if old_name in collections:
                    new_name = model._meta["collection"]
                    self.print("[%s] Renaming %s to %s" % (model_id, old_name, new_name))
                    db[old_name].rename(new_name)
                    break
40
41
42
43
            # Ensure only documents with auto_create_index == False
            if model._meta.get('auto_create_index', True):
                continue
            # Index model
44
            self.index_model(model_id, model)
45
46
        # Index datastreams
        self.index_datastreams()
47
48
        # Index GridVCS
        self.index_gridvcs()
49
50
        # Index mongo cache
        self.index_cache()
51
52
53
        # @todo: Detect changes
        self.print("OK")

54
55
56
57
58
59
60
61
62
63
64
    def index_model(self, model_id, model):
        """
        Create necessary indexes for model
        :param model_id: model id
        :param model: model class
        :return:
        """
        self.print("[%s] Checking indexes" % model_id)
        coll = model._get_collection()
        # Get existing unique indexes
        idx_info = coll.index_information()
65
66
67
        x_name = {}  # fields -> name
        x_unique = {}  # fields -> bool(is_unique)
        left_unique = set()
68
        for xn in idx_info:
69
70
71
72
73
74
75
76
            fields = ",".join(str(k[0]) for k in idx_info[xn]["key"])
            x_name[fields] = xn
            is_unique = idx_info[xn].get("unique", False)
            x_unique[fields] = is_unique
            if is_unique:
                left_unique.add(fields)
        if x_name:
            # Get declared indexes
77
78
79
            xspecs = model._meta["index_specs"]
            for xi in xspecs:
                fields = ",".join(str(k[0]) for k in xi["fields"])
80
81
82
83
84
85
86
87
88
89
90
                if fields in x_name:
                    # Check for uniqueness match
                    du = xi.get("unique", False)
                    if du != x_unique[fields]:
                        # Uniqueness mismatch
                        self.print("[%s] Dropping mismatched index %s" % (
                            model_id, x_name[fields]))
                        coll.drop_index(x_name[fields])
                    elif du and x_unique[fields]:
                        # Remove unique index from left
                        left_unique.remove(fields)
91
            # Delete state unique indexes
92
93
94
95
            for fields in left_unique:
                self.print("[%s] Dropping stale unique index %s" % (
                    model_id, x_name[fields]))
                coll.drop_index(x_name[fields])
96
97
        # Apply indexes
        model.ensure_indexes()
Dmitry Volodin's avatar
PEP8    
Dmitry Volodin committed
98

Andrey Vertiprahov's avatar
Andrey Vertiprahov committed
99
    def index_datastreams(self):
Andrey Vertiprahov's avatar
Andrey Vertiprahov committed
100
        for name in ds_loader:
Dmitry Volodin's avatar
Dmitry Volodin committed
101
102
            if not getattr(config.datastream, "enable_%s" % name, False):
                continue
Dmitry Volodin's avatar
Dmitry Volodin committed
103
            ds = ds_loader[name]
Dmitry Volodin's avatar
Dmitry Volodin committed
104
105
            self.print("[%s] Indexing datastream" % ds.name)
            ds.ensure_collection()
106

107
108
109
110
    def index_gridvcs(self):
        for repo in REPOS:
            GridVCS(repo).ensure_collection()

111
112
113
114
115
116
117
    def index_cache(self):
        from noc.core.cache.base import cache
        if not hasattr(cache, "ensure_indexes"):
            return
        self.print("[%s] Indexing cache" % cache.__class__.__name__)
        cache.ensure_indexes()

118
119
120

if __name__ == "__main__":
    Command().run()