added check active table to custom db connector

This commit is contained in:
iLLiCiTiT 2019-10-23 15:01:51 +02:00
parent 4290af23e7
commit 19f810ff57

View file

@ -20,6 +20,9 @@ import requests
import pymongo
from pymongo.client_session import ClientSession
class NotActiveTable(Exception):
pass
def auto_reconnect(func):
"""Handling auto reconnect in 3 retry times"""
@functools.wraps(func)
@ -37,12 +40,23 @@ def auto_reconnect(func):
return decorated
def check_active_table(func):
"""Handling auto reconnect in 3 retry times"""
@functools.wraps(func)
def decorated(obj, *args, **kwargs):
if not obj.active_table:
raise NotActiveTable("Active table is not set. (This is bug)")
return func(obj, *args, **kwargs)
return decorated
class DbConnector:
log = logging.getLogger(__name__)
timeout = 1000
def __init__(self, mongo_url, database_name, table_name):
def __init__(self, mongo_url, database_name, table_name=None):
self._mongo_client = None
self._sentry_client = None
self._sentry_logging_handler = None
@ -115,6 +129,7 @@ class DbConnector:
def collections(self):
return self._database.collection_names()
@check_active_table
@auto_reconnect
def insert_one(self, item, session=None):
assert isinstance(item, dict), "item must be of type <dict>"
@ -123,6 +138,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def insert_many(self, items, ordered=True, session=None):
# check if all items are valid
@ -136,6 +152,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def find(self, filter, projection=None, sort=None, session=None):
return self._database[self.active_table].find(
@ -145,6 +162,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def find_one(self, filter, projection=None, sort=None, session=None):
assert isinstance(filter, dict), "filter must be <dict>"
@ -156,6 +174,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def replace_one(self, filter, replacement, session=None):
return self._database[self.active_table].replace_one(
@ -163,6 +182,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def update_one(self, filter, update, session=None):
return self._database[self.active_table].update_one(
@ -170,6 +190,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def update_many(self, filter, update, session=None):
return self._database[self.active_table].update_many(
@ -177,12 +198,14 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def distinct(self, *args, **kwargs):
return self._database[self.active_table].distinct(
*args, **kwargs
)
@check_active_table
@auto_reconnect
def drop_collection(self, name_or_collection, session=None):
return self._database[self.active_table].drop(
@ -190,6 +213,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def delete_one(filter, collation=None, session=None):
return self._database[self.active_table].delete_one(
@ -198,6 +222,7 @@ class DbConnector:
session=session
)
@check_active_table
@auto_reconnect
def delete_many(filter, collation=None, session=None):
return self._database[self.active_table].delete_many(