Merge "Add callbacks for set_global_engine"

keystone/common/sql/core.py

@@ -32,10 +32,12 @@ from keystone import exception
 from keystone.openstack.common import jsonutils
 
 
+LOG = logging.getLogger(__name__)
 CONF = config.CONF
 
-# maintain a single engine reference for sqlite in-memory
+# maintain a single engine reference for sqlalchemy engine
 GLOBAL_ENGINE = None
+GLOBAL_ENGINE_CALLBACKS = set()
 
 
 ModelBase = declarative.declarative_base()
@@ -95,9 +97,49 @@ ModelBase.__init__ = initialize_decorator(ModelBase.__init__)
 
 
 def set_global_engine(engine):
+    """Set the global engine.
+
+    This sets the current global engine, which is returned by
+    Base.get_engine(allow_global_engine=True).
+
+    When the global engine is changed, all of the callbacks registered via
+    register_global_engine_callback since the last time set_global_engine was
+    changed are called. The callback functions are invoked with no arguments.
+
+    """
+
     global GLOBAL_ENGINE
+    global GLOBAL_ENGINE_CALLBACKS
+
+    if engine is GLOBAL_ENGINE:
+        # It's the same engine so nothing to do.
+        return
+
     GLOBAL_ENGINE = engine
 
+    cbs = GLOBAL_ENGINE_CALLBACKS
+    GLOBAL_ENGINE_CALLBACKS = set()
+    for cb in cbs:
+        try:
+            cb()
+        except Exception:
+            LOG.exception(_("Global engine callback raised."))
+            # Just logging the exception so can process other callbacks.
+
+
+def register_global_engine_callback(cb_fn):
+    """Register a function to be called when the global engine is set.
+
+    Note that the callback will be called only once or not at all, so to get
+    called each time the global engine is changed the function must be
+    re-registered.
+
+    """
+
+    global GLOBAL_ENGINE_CALLBACKS
+
+    GLOBAL_ENGINE_CALLBACKS.add(cb_fn)
+
 
 # Special Fields
 class JsonBlob(sql_types.TypeDecorator):

tests/test_sql_core.py

@@ -17,6 +17,130 @@ from keystone.common import sql
 from keystone import test
 
 
+class CallbackMonitor:
+    def __init__(self, expect_called=True, raise_=False):
+        self.expect_called = expect_called
+        self.called = False
+        self._complete = False
+        self._raise = raise_
+
+    def call_this(self):
+        if self._complete:
+            return
+
+        if not self.expect_called:
+            raise Exception("Did not expect callback.")
+
+        if self.called:
+            raise Exception("Callback already called.")
+
+        self.called = True
+
+        if self._raise:
+            raise Exception("When called, raises.")
+
+    def check(self):
+        if self.expect_called:
+            if not self.called:
+                raise Exception("Expected function to be called.")
+        self._complete = True
+
+
+class TestGlobalEngine(test.TestCase):
+
+    def tearDown(self):
+        sql.set_global_engine(None)
+        super(TestGlobalEngine, self).tearDown()
+
+    def test_notify_on_set(self):
+        # If call sql.set_global_engine(), notify callbacks get called.
+
+        cb_mon = CallbackMonitor()
+
+        sql.register_global_engine_callback(cb_mon.call_this)
+        fake_engine = object()
+        sql.set_global_engine(fake_engine)
+
+        cb_mon.check()
+
+    def test_multi_notify(self):
+        # You can also set multiple notify callbacks and they each get called.
+
+        cb_mon1 = CallbackMonitor()
+        cb_mon2 = CallbackMonitor()
+
+        sql.register_global_engine_callback(cb_mon1.call_this)
+        sql.register_global_engine_callback(cb_mon2.call_this)
+
+        fake_engine = object()
+        sql.set_global_engine(fake_engine)
+
+        cb_mon1.check()
+        cb_mon2.check()
+
+    def test_notify_once(self):
+        # After a callback is called, it's not called again if set global
+        # engine again.
+
+        cb_mon = CallbackMonitor()
+
+        sql.register_global_engine_callback(cb_mon.call_this)
+        fake_engine = object()
+        sql.set_global_engine(fake_engine)
+
+        fake_engine = object()
+        # Note that cb_mon.call_this would raise if it's called again.
+        sql.set_global_engine(fake_engine)
+
+        cb_mon.check()
+
+    def test_set_same_engine(self):
+        # If you set the global engine to the same engine, callbacks don't get
+        # called.
+
+        fake_engine = object()
+
+        sql.set_global_engine(fake_engine)
+
+        cb_mon = CallbackMonitor(expect_called=False)
+        sql.register_global_engine_callback(cb_mon.call_this)
+
+        # Note that cb_mon.call_this would raise if it's called.
+        sql.set_global_engine(fake_engine)
+
+        cb_mon.check()
+
+    def test_notify_register_same(self):
+        # If you register the same callback twice, only gets called once.
+        cb_mon = CallbackMonitor()
+
+        sql.register_global_engine_callback(cb_mon.call_this)
+        sql.register_global_engine_callback(cb_mon.call_this)
+
+        fake_engine = object()
+        # Note that cb_mon.call_this would raise if it's called twice.
+        sql.set_global_engine(fake_engine)
+
+        cb_mon.check()
+
+    def test_callback_throws(self):
+        # If a callback function raises,
+        # a) the caller doesn't know about it,
+        # b) other callbacks are still called
+
+        cb_mon1 = CallbackMonitor(raise_=True)
+        cb_mon2 = CallbackMonitor()
+
+        sql.register_global_engine_callback(cb_mon1.call_this)
+        sql.register_global_engine_callback(cb_mon2.call_this)
+
+        fake_engine = object()
+        sql.set_global_engine(fake_engine)
+
+        cb_mon1.check()
+        cb_mon2.check()
+
+
 class TestBase(test.TestCase):
 
     def tearDown(self):