registry.py
  1  """
  2  Registry for dynamically binding middleware implementations.
  3  
  4  This module provides a registry system that allows middleware implementations
  5  to be bound at runtime, enabling pluggable behavior for configuration resolution
  6  and other middleware functions.
  7  """
  8  
  9  import logging
 10  from typing import Optional, Type, Dict, Any, List, Callable
 11  
 12  log = logging.getLogger(__name__)
 13  
 14  LOG_IDENTIFIER = "[MiddlewareRegistry]"
 15  
 16  
 17  class MiddlewareRegistry:
 18      """
 19      Registry for middleware implementations that can be overridden at runtime.
 20  
 21      This registry allows different implementations of middleware to be bound
 22      dynamically, enabling extensibility and customization of system behavior.
 23      """
 24  
 25      _config_resolver: Optional[Type] = None
 26      _resource_sharing_service: Optional[Type] = None
 27      _initialization_callbacks: List[callable] = []
 28      _post_migration_hooks: List[Callable[[str], None]] = []
 29  
 30      @classmethod
 31      def bind_config_resolver(cls, resolver_class: Type):
 32          """
 33          Bind a custom config resolver implementation.
 34  
 35          Args:
 36              resolver_class: Class that implements the ConfigResolver interface
 37          """
 38          cls._config_resolver = resolver_class
 39          log.info(
 40              "%s Bound custom config resolver: %s",
 41              LOG_IDENTIFIER,
 42              resolver_class.__name__,
 43          )
 44  
 45      @classmethod
 46      def get_config_resolver(cls) -> Type:
 47          """
 48          Get the current config resolver implementation.
 49  
 50          Returns:
 51              The bound config resolver class, or the default ConfigResolver if none bound.
 52          """
 53          if cls._config_resolver:
 54              return cls._config_resolver
 55  
 56          from .config_resolver import ConfigResolver
 57  
 58          return ConfigResolver
 59  
 60      @classmethod
 61      def bind_resource_sharing_service(cls, service_class: Type):
 62          """
 63          Bind a custom resource sharing service implementation.
 64  
 65          Args:
 66              service_class: Class that implements the ResourceSharingService interface
 67          """
 68          cls._resource_sharing_service = service_class
 69          log.info(
 70              "%s Bound custom resource sharing service: %s",
 71              LOG_IDENTIFIER,
 72              service_class.__name__,
 73          )
 74  
 75      @classmethod
 76      def get_resource_sharing_service(cls) -> Type:
 77          """
 78          Get the current resource sharing service implementation.
 79  
 80          Returns:
 81              The bound resource sharing service class, or the default DefaultResourceSharingService if none bound.
 82          """
 83          if cls._resource_sharing_service:
 84              return cls._resource_sharing_service
 85  
 86          from ..services.default_resource_sharing_service import DefaultResourceSharingService
 87  
 88          return DefaultResourceSharingService
 89  
 90      @classmethod
 91      def register_initialization_callback(cls, callback: callable):
 92          """
 93          Register a callback to be called during system initialization.
 94  
 95          Args:
 96              callback: Function to call during initialization
 97          """
 98          cls._initialization_callbacks.append(callback)
 99          log.debug(
100              "%s Registered initialization callback: %s",
101              LOG_IDENTIFIER,
102              callback.__name__,
103          )
104  
105      @classmethod
106      def register_post_migration_hook(cls, hook: Callable[[str], None]):
107          """
108          Register a hook to be called after community database migrations complete.
109  
110          This allows enterprise/plugin packages to run their own migrations without
111          the community code being aware of them. Hooks receive the database URL
112          as a parameter.
113  
114          Args:
115              hook: Callable that takes database_url (str) and runs migrations
116          """
117          cls._post_migration_hooks.append(hook)
118          log.info(
119              "%s Registered post-migration hook: %s",
120              LOG_IDENTIFIER,
121              getattr(hook, '__name__', repr(hook)),
122          )
123  
124      @classmethod
125      def run_post_migration_hooks(cls, database_url: str):
126          """
127          Execute all registered post-migration hooks.
128  
129          Called by community code after its migrations complete. Enterprise/plugin
130          packages can register hooks to run their own migrations.
131  
132          Args:
133              database_url: Database URL to pass to hooks
134          """
135          if not cls._post_migration_hooks:
136              log.debug("%s No post-migration hooks registered", LOG_IDENTIFIER)
137              return
138  
139          log.info(
140              "%s Running %d post-migration hook(s)...",
141              LOG_IDENTIFIER,
142              len(cls._post_migration_hooks),
143          )
144  
145          for hook in cls._post_migration_hooks:
146              hook_name = getattr(hook, '__name__', repr(hook))
147              try:
148                  log.info("%s Executing post-migration hook: %s", LOG_IDENTIFIER, hook_name)
149                  hook(database_url)
150                  log.info("%s Post-migration hook completed: %s", LOG_IDENTIFIER, hook_name)
151              except Exception as e:
152                  log.error(
153                      "%s Error executing post-migration hook %s: %s",
154                      LOG_IDENTIFIER,
155                      hook_name,
156                      e,
157                  )
158                  log.exception("%s Full traceback:", LOG_IDENTIFIER)
159                  raise
160  
161      @classmethod
162      def initialize_middleware(cls):
163          """
164          Initialize all registered middleware components.
165  
166          This should be called during system startup to initialize any
167          bound middleware implementations.
168          """
169          log.info("%s Initializing middleware components...", LOG_IDENTIFIER)
170  
171          for callback in cls._initialization_callbacks:
172              try:
173                  callback()
174                  log.debug(
175                      "%s Executed initialization callback: %s",
176                      LOG_IDENTIFIER,
177                      callback.__name__,
178                  )
179              except Exception as e:
180                  log.error(
181                      "%s Error executing initialization callback %s: %s",
182                      LOG_IDENTIFIER,
183                      callback.__name__,
184                      e,
185                  )
186  
187          log.info("%s Middleware initialization complete.", LOG_IDENTIFIER)
188  
189      @classmethod
190      def reset_bindings(cls):
191          """
192          Reset all bindings to defaults.
193  
194          This is useful for testing or when switching between different
195          middleware configurations.
196          """
197          cls._config_resolver = None
198          cls._resource_sharing_service = None
199          cls._initialization_callbacks = []
200          cls._post_migration_hooks = []
201          log.info("%s Reset all middleware bindings", LOG_IDENTIFIER)
202  
203      @classmethod
204      def get_registry_status(cls) -> Dict[str, Any]:
205          """
206          Get the current status of the middleware registry.
207  
208          Returns:
209              Dict containing information about bound middleware implementations.
210          """
211          return {
212              "config_resolver": (
213                  cls._config_resolver.__name__ if cls._config_resolver else "default"
214              ),
215              "resource_sharing_service": (
216                  cls._resource_sharing_service.__name__ if cls._resource_sharing_service else "default"
217              ),
218              "initialization_callbacks": len(cls._initialization_callbacks),
219              "post_migration_hooks": len(cls._post_migration_hooks),
220              "has_custom_bindings": cls._config_resolver is not None or cls._resource_sharing_service is not None,
221          }