summaryrefslogtreecommitdiff
path: root/src/db.py
blob: c4244e111ac83db3da85d6fea2e8e57307b543c6 (plain)
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
# A database storing dictionaries, keyed on a timestamp.
# value = A dict which will be stored as a JSON object encoded in
# UTF-8. Note that dict keys of type integer or float will become
# strings while values will keep their type.

# Note that there's a (slim) chance that you'd stomp on the previous
# value if you're too quick with generating the timestamps, ie
# invoking time.time() several times quickly enough.

import time
import couch
import index


class DictDB():
    def __init__(self, database, hostname, username, password):
        """
        Check if the database exists, otherwise we will create it together
        with the indexes specified in index.py.
        """

        self.server = couch.client.Server(
            f"http://{username}:{password}@{hostname}:5984/")

        try:
            self.couchdb = self.server.database(database)
        except couch.exceptions.NotFound:
            print("Creating database and indexes.")
            self.couchdb = self.server.create(database)
            self.server.create('_users')

            for i in index.indexes:
                self.couchdb.index(i)

        self._ts = time.time()

    def unique_key(self):
        """
        Create a unique key based on the current time. We will use this as
        the ID for any new documents we store in CouchDB.
        """

        ts = time.time()
        while ts == self._ts:
            ts = time.time()
        self._ts = round(ts * 1000)

        return self._ts

    def add(self, data, batch_write=False):
        """
        Store a document in CouchDB.
        """

        key = str(self.unique_key())
        if type(data) is list:
            for item in data:
                item['_id'] = key

                self.couchdb.save(item)
        else:
            data['_id'] = key
            self.couchdb.save(data)

        return key

    def get(self, key):
        """
        Get a document based on its ID, return an empty dict if not found.
        """

        try:
            doc = self.couchdb.get(key)
        except couch.exceptions.NotFound:
            doc = {}

        return doc

    def slice(self, key_from=None, key_to=None):
        pass

    def search(self, **kwargs):
        """
        Execute a Mango query, ideally we should have an index matching the,
        the query otherwise things will be slow.
        """

        data = list()
        selector = dict()

        if kwargs:
            selector = {"selector": {}}

            for key in kwargs:
                if kwargs[key].isnumeric():
                    kwargs[key] = int(kwargs[key])
                selector['selector'][key] = {'$eq': kwargs[key]}

        for doc in self.couchdb.find(selector):
            data.append(doc)

        return data