diff --git a/pype/ftrack/lib/custom_db_connector.py b/pype/ftrack/lib/custom_db_connector.py index 505ac96610..f4cc00444b 100644 --- a/pype/ftrack/lib/custom_db_connector.py +++ b/pype/ftrack/lib/custom_db_connector.py @@ -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 " @@ -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 " @@ -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(