osm_common/dbmongo.py

   1
   2 import logging
   3 from pymongo import MongoClient, errors
   4 from osm_common.dbbase import DbException, DbBase
   5 from http import HTTPStatus
   6 from time import time, sleep
   7 from copy import deepcopy
   8
   9 __author__ = "Alfonso Tierno <alfonso.tiernosepulveda@telefonica.com>"
  10
  11 # TODO consider use this decorator for database access retries
  12 # @retry_mongocall
  13 # def retry_mongocall(call):
  14 #     def _retry_mongocall(*args, **kwargs):
  15 #         retry = 1
  16 #         while True:
  17 #             try:
  18 #                 return call(*args, **kwargs)
  19 #             except pymongo.AutoReconnect as e:
  20 #                 if retry == 4:
  21 #                     raise DbException(str(e))
  22 #                 sleep(retry)
  23 #     return _retry_mongocall
  24
  25
  26 def deep_update(to_update, update_with):
  27     """
  28     Update 'to_update' dict with the content 'update_with' dict recursively
  29     :param to_update: must be a dictionary to be modified
  30     :param update_with: must be a dictionary. It is not changed
  31     :return: to_update
  32     """
  33     for key in update_with:
  34         if key in to_update:
  35             if isinstance(to_update[key], dict) and isinstance(update_with[key], dict):
  36                 deep_update(to_update[key], update_with[key])
  37                 continue
  38         to_update[key] = deepcopy(update_with[key])
  39     return to_update
  40
  41
  42 class DbMongo(DbBase):
  43     conn_initial_timout = 120
  44     conn_timout = 10
  45
  46     def __init__(self, logger_name='db'):
  47         self.logger = logging.getLogger(logger_name)
  48
  49     def db_connect(self, config):
  50         try:
  51             if "logger_name" in config:
  52                 self.logger = logging.getLogger(config["logger_name"])
  53             self.client = MongoClient(config["host"], config["port"])
  54             self.db = self.client[config["name"]]
  55             if "loglevel" in config:
  56                 self.logger.setLevel(getattr(logging, config['loglevel']))
  57             # get data to try a connection
  58             now = time()
  59             while True:
  60                 try:
  61                     self.db.users.find_one({"username": "admin"})
  62                     return
  63                 except errors.ConnectionFailure as e:
  64                     if time() - now >= self.conn_initial_timout:
  65                         raise
  66                     self.logger.info("Waiting to database up {}".format(e))
  67                     sleep(2)
  68         except errors.PyMongoError as e:
  69             raise DbException(str(e))
  70
  71     def db_disconnect(self):
  72         pass  # TODO
  73
  74     @staticmethod
  75     def _format_filter(q_filter):
  76         """
  77         Translate query string filter into mongo database filter
  78         :param q_filter: Query string content. Follows SOL005 section 4.3.2 guidelines, with the follow extensions and
  79         differences:
  80             It accept ".nq" (not equal) in addition to ".neq".
  81             For arrays you can specify index (concrete index must match), nothing (any index may match) or 'ANYINDEX'
  82             (two or more matches applies for the same array element). Examples:
  83                 with database register: {A: [{B: 1, C: 2}, {B: 6, C: 9}]}
  84                 query 'A.B=6' matches because array A contains one element with B equal to 6
  85                 query 'A.0.B=6' does no match because index 0 of array A contains B with value 1, but not 6
  86                 query 'A.B=6&A.C=2' matches because one element of array matches B=6 and other matchesC=2
  87                 query 'A.ANYINDEX.B=6&A.ANYINDEX.C=2' does not match because it is needed the same element of the
  88                     array matching both
  89
  90         Examples of translations from SOL005 to  >> mongo  # comment
  91             A=B; A.eq=B         >> A: B             # must contain key A and equal to B or be a list that contains B
  92             A.cont=B            >> A: B
  93             A=B&A=C; A=B,C      >> A: {$in: [B, C]} # must contain key A and equal to B or C or be a list that contains
  94                 # B or C
  95             A.cont=B&A.cont=C; A.cont=B,C  >> A: {$in: [B, C]}
  96             A.ncont=B           >> A: {$nin: B}     # must not contain key A or if present not equal to B or if a list,
  97                 # it must not not contain B
  98             A.ncont=B,C; A.ncont=B&A.ncont=C    >> A: {$nin: [B,C]}     # must not contain key A or if present not equal
  99                 # neither B nor C; or if a list, it must not contain neither B nor C
 100             A.ne=B&A.ne=C; A.ne=B,C             >> A: {$nin: [B, C]}
 101             A.gt=B              >> A: {$gt: B}      # must contain key A and greater than B
 102             A.ne=B; A.neq=B         >> A: {$ne: B}          # must not contain key A or if present not equal to B, or if
 103                 # an array not contain B
 104             A.ANYINDEX.B=C          >> A: {$elemMatch: {B=C}
 105         :return: database mongo filter
 106         """
 107         try:
 108             db_filter = {}
 109             for query_k, query_v in q_filter.items():
 110                 dot_index = query_k.rfind(".")
 111                 if dot_index > 1 and query_k[dot_index+1:] in ("eq", "ne", "gt", "gte", "lt", "lte", "cont",
 112                                                                "ncont", "neq"):
 113                     operator = "$" + query_k[dot_index + 1:]
 114                     if operator == "$neq":
 115                         operator = "$ne"
 116                     k = query_k[:dot_index]
 117                 else:
 118                     operator = "$eq"
 119                     k = query_k
 120
 121                 v = query_v
 122                 if isinstance(v, list):
 123                     if operator in ("$eq", "$cont"):
 124                         operator = "$in"
 125                         v = query_v
 126                     elif operator in ("$ne", "$ncont"):
 127                         operator = "$nin"
 128                         v = query_v
 129                     else:
 130                         v = query_v.join(",")
 131
 132                 if operator in ("$eq", "$cont"):
 133                     # v cannot be a comma separated list, because operator would have been changed to $in
 134                     db_v = v
 135                 elif operator == "$ncount":
 136                     # v cannot be a comma separated list, because operator would have been changed to $nin
 137                     db_v = {"$ne": v}
 138                 else:
 139                     db_v = {operator: v}
 140
 141                 # process the ANYINDEX word at k.
 142                 kleft, _, kright = k.rpartition(".ANYINDEX.")
 143                 while kleft:
 144                     k = kleft
 145                     db_v = {"$elemMatch": {kright: db_v}}
 146                     kleft, _, kright = k.rpartition(".ANYINDEX.")
 147
 148                 # insert in db_filter
 149                 # maybe db_filter[k] exist. e.g. in the query string for values between 5 and 8: "a.gt=5&a.lt=8"
 150                 deep_update(db_filter, {k: db_v})
 151
 152             return db_filter
 153         except Exception as e:
 154             raise DbException("Invalid query string filter at {}:{}. Error: {}".format(query_k, v, e),
 155                               http_code=HTTPStatus.BAD_REQUEST)
 156
 157     def get_list(self, table, filter={}):
 158         try:
 159             result = []
 160             collection = self.db[table]
 161             db_filter = self._format_filter(filter)
 162             rows = collection.find(db_filter)
 163             for row in rows:
 164                 result.append(row)
 165             return result
 166         except DbException:
 167             raise
 168         except Exception as e:  # TODO refine
 169             raise DbException(str(e))
 170
 171     def get_one(self, table, filter={}, fail_on_empty=True, fail_on_more=True):
 172         try:
 173             if filter:
 174                 filter = self._format_filter(filter)
 175             collection = self.db[table]
 176             if not (fail_on_empty and fail_on_more):
 177                 return collection.find_one(filter)
 178             rows = collection.find(filter)
 179             if rows.count() == 0:
 180                 if fail_on_empty:
 181                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 182                                       HTTPStatus.NOT_FOUND)
 183                 return None
 184             elif rows.count() > 1:
 185                 if fail_on_more:
 186                     raise DbException("Found more than one {} with filter='{}'".format(table[:-1], filter),
 187                                       HTTPStatus.CONFLICT)
 188             return rows[0]
 189         except Exception as e:  # TODO refine
 190             raise DbException(str(e))
 191
 192     def del_list(self, table, filter={}):
 193         try:
 194             collection = self.db[table]
 195             rows = collection.delete_many(self._format_filter(filter))
 196             return {"deleted": rows.deleted_count}
 197         except DbException:
 198             raise
 199         except Exception as e:  # TODO refine
 200             raise DbException(str(e))
 201
 202     def del_one(self, table, filter={}, fail_on_empty=True):
 203         try:
 204             collection = self.db[table]
 205             rows = collection.delete_one(self._format_filter(filter))
 206             if rows.deleted_count == 0:
 207                 if fail_on_empty:
 208                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 209                                       HTTPStatus.NOT_FOUND)
 210                 return None
 211             return {"deleted": rows.deleted_count}
 212         except Exception as e:  # TODO refine
 213             raise DbException(str(e))
 214
 215     def create(self, table, indata):
 216         try:
 217             collection = self.db[table]
 218             data = collection.insert_one(indata)
 219             return data.inserted_id
 220         except Exception as e:  # TODO refine
 221             raise DbException(str(e))
 222
 223     def set_one(self, table, filter, update_dict, fail_on_empty=True):
 224         try:
 225             collection = self.db[table]
 226             rows = collection.update_one(self._format_filter(filter), {"$set": update_dict})
 227             if rows.matched_count == 0:
 228                 if fail_on_empty:
 229                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], filter),
 230                                       HTTPStatus.NOT_FOUND)
 231                 return None
 232             return {"modified": rows.modified_count}
 233         except Exception as e:  # TODO refine
 234             raise DbException(str(e))
 235
 236     def replace(self, table, id, indata, fail_on_empty=True):
 237         try:
 238             _filter = {"_id": id}
 239             collection = self.db[table]
 240             rows = collection.replace_one(_filter, indata)
 241             if rows.matched_count == 0:
 242                 if fail_on_empty:
 243                     raise DbException("Not found any {} with filter='{}'".format(table[:-1], _filter),
 244                                       HTTPStatus.NOT_FOUND)
 245                 return None
 246             return {"replaced": rows.modified_count}
 247         except Exception as e:  # TODO refine
 248             raise DbException(str(e))