mongo.py 2.31 KB
Newer Older
Dmitry Volodin's avatar
Dmitry Volodin committed
1
# -*- coding: utf-8 -*-
Dmitry Volodin's avatar
Dmitry Volodin committed
2
3
4
# ----------------------------------------------------------------------
# Mongo backend
# ----------------------------------------------------------------------
5
# Copyright (C) 2007-2019 The NOC Project
Dmitry Volodin's avatar
Dmitry Volodin committed
6
7
# See LICENSE for details
# ----------------------------------------------------------------------
Dmitry Volodin's avatar
Dmitry Volodin committed
8

Dmitry Volodin's avatar
Dmitry Volodin committed
9
10
# Python modules
from __future__ import absolute_import
Dmitry Volodin's avatar
Dmitry Volodin committed
11
import datetime
Dmitry Volodin's avatar
Dmitry Volodin committed
12
# Third-party modules
Dmitry Volodin's avatar
Dmitry Volodin committed
13
import bson
Dmitry Volodin's avatar
Dmitry Volodin committed
14
15
16
from six.moves.cPickle import loads, dumps, HIGHEST_PROTOCOL
# NOC modules
from .base import BaseCache
Dmitry Volodin's avatar
Dmitry Volodin committed
17
from noc.lib.nosql import get_db
18
from noc.config import config
Dmitry Volodin's avatar
Dmitry Volodin committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34


class MongoCache(BaseCache):
    collection_name = "noc.caches.mongo"
    KEY_FIELD = "_id"
    VALUE_FIELD = "v"
    EXPIRES_FIELD = "x"

    @classmethod
    def get_collection(cls):
        if not hasattr(cls, "_collection"):
            cls._collection = get_db()[cls.collection_name]
        return cls._collection

    @classmethod
    def ensure_indexes(cls):
35
36
37
38
39
40
41
        """
        Create all necessary indexes. Called by ensure-index
        :return:
        """
        cls.get_collection().create_index(
            cls.EXPIRES_FIELD,
            expireAfterSeconds=0
Dmitry Volodin's avatar
Dmitry Volodin committed
42
43
44
45
46
47
48
49
50
51
        )

    def get(self, key, default=None, version=None):
        k = self.make_key(key, version)
        now = datetime.datetime.now()
        d = self.get_collection().find_one({
            self.KEY_FIELD: k
        })
        if d and d[self.EXPIRES_FIELD] > now:
            # Found, not expired
Dmitry Volodin's avatar
Dmitry Volodin committed
52
            return loads(d[self.VALUE_FIELD])
Dmitry Volodin's avatar
Dmitry Volodin committed
53
54
55
56
57
58
59
60
61
62
63
64
        else:
            return default

    def set(self, key, value, ttl=None, version=None):
        """
        Set key
        :param key:
        :param value:
        :param ttl:
        :return:
        """
        k = self.make_key(key, version)
65
        ttl = ttl or config.memcached.default_ttl
Dmitry Volodin's avatar
Dmitry Volodin committed
66
        expires = datetime.datetime.now() + datetime.timedelta(seconds=ttl)
Dmitry Volodin's avatar
Dmitry Volodin committed
67
        self.get_collection().update_one({
Dmitry Volodin's avatar
Dmitry Volodin committed
68
69
70
            self.KEY_FIELD: k
        }, {
            "$set": {
Dmitry Volodin's avatar
Dmitry Volodin committed
71
72
                self.VALUE_FIELD: bson.Binary(
                    dumps(value, HIGHEST_PROTOCOL)),
Dmitry Volodin's avatar
Dmitry Volodin committed
73
74
75
76
77
78
                self.EXPIRES_FIELD: expires
            }
        }, upsert=True)

    def delete(self, key, version=None):
        k = self.make_key(key, version)
Dmitry Volodin's avatar
Dmitry Volodin committed
79
        self.get_collection().delete_one({
Dmitry Volodin's avatar
Dmitry Volodin committed
80
81
            self.KEY_FIELD: k
        })