osm_nbi/dbmongo.py

   1 #import pymongo
   2 import logging
   3 from pymongo import MongoClient, errors
   4 from dbbase import DbException, DbBase
   5 from http import HTTPStatus
   6 from time import time, sleep
   7
   8 __author__ = "Alfonso Tierno <alfonso.tiernosepulveda@telefonica.com>"
   9
  10 # TODO consider use this decorator for database access retries
  11 # @retry_mongocall
  12 # def retry_mongocall(call):
  13 #     def _retry_mongocall(*args, **kwargs):
  14 #         retry = 1
  15 #         while True:
  16 #             try:
  17 #                 return call(*args, **kwargs)
  18 #             except pymongo.AutoReconnect as e:
  19 #                 if retry == 4:
  20 #                     raise DbException(str(e))
  21 #                 sleep(retry)
  22 #     return _retry_mongocall
  23
  24
  25 class DbMongo(DbBase):
  26     conn_initial_timout = 120
  27     conn_timout = 10
  28
  29     def __init__(self, logger_name='db'):
  30         self.logger = logging.getLogger(logger_name)
  31
  32     def db_connect(self, config):
  33         try:
  34             if "logger_name" in config:
  35                 self.logger = logging.getLogger(config["logger_name"])
  36             self.client = MongoClient(config["host"], config["port"])
  37             self.db = self.client[config["name"]]
  38             if "loglevel" in config:
  39                 self.logger.setLevel(getattr(logging, config['loglevel']))
  40             # get data to try a connection
  41             now = time()
  42             while True:
  43                 try:
  44                     self.db.users.find_one({"username": "admin"})
  45                     return
  46                 except errors.ConnectionFailure as e:
  47                     if time() - now >= self.conn_initial_timout:
  48                         raise
  49                     self.logger.info("Waiting to database up {}".format(e))
  50                     sleep(2)
  51         except errors.PyMongoError as e:
  52             raise DbException(str(e))
  53
  54     def db_disconnect(self):
  55         pass  # TODO
  56
  57     @staticmethod
  58     def _format_filter(filter):
  59         try:
  60             db_filter = {}
  61             for query_k, query_v in filter.items():
  62                 dot_index = query_k.rfind(".")
  63                 if dot_index > 1 and query_k[dot_index+1:] in ("eq", "ne", "gt", "gte", "lt", "lte", "cont",
  64                                                                "ncont", "neq"):
  65                     operator = "$" + query_k[dot_index+1:]
  66                     if operator == "$neq":
  67                         operator = "$ne"
  68                     k = query_k[:dot_index]
  69                 else:
  70                     operator = "$eq"
  71                     k = query_k
  72
  73                 v = query_v
  74                 if isinstance(v, list):
  75                     if operator in ("$eq", "$cont"):
  76                         operator = "$in"
  77                         v = query_v
  78                     elif operator in ("$ne", "$ncont"):
  79                         operator = "$nin"
  80                         v = query_v
  81                     else:
  82                         v = query_v.join(",")
  83
  84                 if operator in ("$eq", "$cont"):
  85                     # v cannot be a comma separated list, because operator would have been changed to $in
  86                     db_filter[k] = v
  87                 elif operator == "$ncount":
  88                     # v cannot be a comma separated list, because operator would have been changed to $nin
  89                     db_filter[k] = {"$ne": v}
  90                 else:
  91                     # maybe db_filter[k] exist. e.g. in the query string for values between 5 and 8: "a.gt=5&a.lt=8"
  92                     if k not in db_filter:
  93                         db_filter[k] = {}
  94                     db_filter[k][operator] = v
  95
  96             return db_filter
  97         except Exception as e:
  98             raise DbException("Invalid query string filter at {}:{}. Error: {}".format(query_k, v, e),
  99                               http_code=HTTPStatus.BAD_REQUEST)
 100
 101     def get_list(self, table, filter={}):
 102         try:
 103             l = []
 104             collection = self.db[table]
 105             rows = collection.find(self._format_filter(filter))
 106             for row in rows:
 107                 l.append(row)
 108             return l
 109         except DbException:
 110             raise
 111         except Exception as e:  # TODO refine
 112             raise DbException(str(e))
 113
 114     def get_one(self, table, filter={}, fail_on_empty=True, fail_on_more=True):
 115         try:
 116             if filter:
 117                 filter = self._format_filter(filter)
 118             collection = self.db[table]
 119             if not (fail_on_empty and fail_on_more):
 120                 return collection.find_one(filter)
 121             rows = collection.find(filter)
 122             if rows.count() == 0:
 123                 if fail_on_empty:
 124                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 125                                       HTTPStatus.NOT_FOUND)
 126                 return None
 127             elif rows.count() > 1:
 128                 if fail_on_more:
 129                     raise DbException("Found more than one {} with filter='{}'".format(table[:-1], filter),
 130                                       HTTPStatus.CONFLICT)
 131             return rows[0]
 132         except Exception as e:  # TODO refine
 133             raise DbException(str(e))
 134
 135     def del_list(self, table, filter={}):
 136         try:
 137             collection = self.db[table]
 138             rows = collection.delete_many(self._format_filter(filter))
 139             return {"deleted": rows.deleted_count}
 140         except DbException:
 141             raise
 142         except Exception as e:  # TODO refine
 143             raise DbException(str(e))
 144
 145     def del_one(self, table, filter={}, fail_on_empty=True):
 146         try:
 147             collection = self.db[table]
 148             rows = collection.delete_one(self._format_filter(filter))
 149             if rows.deleted_count == 0:
 150                 if fail_on_empty:
 151                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 152                                       HTTPStatus.NOT_FOUND)
 153                 return None
 154             return {"deleted": rows.deleted_count}
 155         except Exception as e:  # TODO refine
 156             raise DbException(str(e))
 157
 158     def create(self, table, indata):
 159         try:
 160             collection = self.db[table]
 161             data = collection.insert_one(indata)
 162             return data.inserted_id
 163         except Exception as e:  # TODO refine
 164             raise DbException(str(e))
 165
 166     def set_one(self, table, filter, update_dict, fail_on_empty=True):
 167         try:
 168             collection = self.db[table]
 169             rows = collection.update_one(self._format_filter(filter), {"$set": update_dict})
 170             if rows.updated_count == 0:
 171                 if fail_on_empty:
 172                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 173                                       HTTPStatus.NOT_FOUND)
 174                 return None
 175             return {"deleted": rows.deleted_count}
 176         except Exception as e:  # TODO refine
 177             raise DbException(str(e))
 178
 179     def replace(self, table, id, indata, fail_on_empty=True):
 180         try:
 181             collection = self.db[table]
 182             rows = collection.replace_one({"_id": id}, indata)
 183             if rows.modified_count == 0:
 184                 if fail_on_empty:
 185                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 186                                       HTTPStatus.NOT_FOUND)
 187                 return None
 188             return {"replace": rows.modified_count}
 189         except Exception as e:  # TODO refine
 190             raise DbException(str(e))