trytond: Add final init hooks [CUSTOM]

coopengo_modules/debug/__init__.py

@@ -42,13 +42,12 @@ def register():
         module='debug', type_='wizard')
 
     try:
-        Pool.register_post_init_hooks(
-            tryton_syntax_analysis,
-            set_method_names_for_profiling,
-            name_one2many_gets,
-            activate_auto_profile,
-            enable_debug_views,
+        Pool.register_post_init_hooks(tryton_syntax_analysis, module='debug')
+        Pool.register_post_init_hooks(set_method_names_for_profiling,
             module='debug')
+        Pool.register_post_init_hooks(name_one2many_gets, module='debug')
+        Pool.register_post_init_hooks(activate_auto_profile, module='debug')
+        Pool.register_post_init_hooks(enable_debug_views, module='debug')
     except AttributeError:
         logger.warning('Post init hooks disabled')
 

trytond/trytond/modules/__init__.py

@@ -316,6 +316,8 @@ def create_indexes(concurrently):
             # Ensure cache is clear for other instances
             Cache.clear_all()
             Cache.refresh_pool(transaction)
+
+        pool.setup_complete(update)
     logger.info('all modules loaded')
 
 

trytond/trytond/pool.py

@@ -61,7 +61,9 @@ class Pool(object):
     _pool_instances = WeakSet()
     test = False
     _init_hooks = {}
+    _final_init_hooks = {}
     _post_init_calls = {}
+    _final_init_calls = {}
     _registered_migration_hooks = {}
     _final_migrations = {}
     _registered_notifications = {}
@@ -127,20 +129,53 @@ def register_mixin(mixin, classinfo, module):
         Pool.classes_mixin[module].append((classinfo, mixin))
 
     @staticmethod
-    def register_post_init_hooks(*hooks, **kwargs):
-        if kwargs['module'] not in Pool._init_hooks:
-            Pool._init_hooks[kwargs['module']] = []
-        Pool._init_hooks[kwargs['module']] += hooks
+    def register_post_init_hooks(hook, *, module):
+        '''
+        Add the "hook" to be called at the end of the setup (after
+        __post_setup__ calls) if <module> is installed.
+
+        This can be used to patch standard functions, or add global behaviours
+        via added inheritance.
+        '''
+        if module not in Pool._init_hooks:
+            Pool._init_hooks[module] = []
+        Pool._init_hooks[module].append(hook)
+
+    @staticmethod
+    def register_final_init_hooks(hook, *, module):
+        '''
+        Add the "hook" to be called once the pool is ready, if <module> is
+        installed.
+
+        This can be used to trigger additional commands that require a ready
+        pool before executing
+        '''
+        if module not in Pool._final_init_hooks:
+            Pool._final_init_hooks[module] = []
+        Pool._final_init_hooks[module].append(hook)
 
     @staticmethod
-    def register_final_migration(*hooks, module=None):
+    def register_final_migration(hook, *, module):
+        '''
+        Add the "hook" to be called at the end of the upgrade process, if
+        <module> is installed and upgraded.
+
+        This can be used to run modular, multi-module migrations
+        '''
         assert module is not None
         if module not in Pool._registered_migration_hooks:
             Pool._registered_migration_hooks[module] = []
-        Pool._registered_migration_hooks[module] += hooks
+        Pool._registered_migration_hooks[module].append(hook)
 
     @staticmethod
-    def register_notification_callbacks(keyword, callback, *, module=None):
+    def register_notification_callbacks(keyword, callback, *, module):
+        '''
+        Register the <callback> to be triggered if the <keyword> is received
+        via a postgres channel.
+
+        This can be used to update some globals when some pre-defined actions
+        are taken in the application
+        '''
         if module not in Pool._registered_notifications:
             Pool._registered_notifications[module] = {}
         Pool._registered_notifications[module][keyword] = callback
@@ -154,6 +189,7 @@ def start(cls):
             for classes in Pool.classes.values():
                 classes.clear()
             cls._init_hooks = {}
+            cls._final_init_hooks = {}
             cls._registered_notifications = {}
             register_classes(with_test=cls.test)
             cls._started = True
@@ -190,6 +226,7 @@ def init(self, update=None, lang=None, activatedeps=False, indexes=None):
             self._pool = defaultdict(dict)
             self._modules = []
             self._post_init_calls[self.database_name] = []
+            self._final_init_calls[self.database_name] = []
             self._final_migrations[self.database_name] = []
             self._notification_callbacks[self.database_name] = {}
             try:
@@ -283,6 +320,8 @@ def fill(self, module, modules):
                 classes[type_].append(cls)
         self._post_init_calls[self.database_name] += self._init_hooks.get(
             module, [])
+        self._final_init_calls[self.database_name] += \
+            self._final_init_hooks.get(module, [])
         self._final_migrations[self.database_name] += \
             self._registered_migration_hooks.get(module, [])
         self._notification_callbacks[self.database_name].update(
@@ -302,6 +341,10 @@ def setup(self, classes=None):
             for cls in lst:
                 cls.__post_setup__()
 
+    def setup_complete(self, update):
+        for hook in self._final_init_calls[self.database_name]:
+            hook(self, update)
+
     def setup_mixin(self, type=None, name=None):
         logger.info('setup mixin for "%s"', self.database_name)
         if type is not None: