pybind/mgr: add convenient context mgrs for accessing module db

Add two new functions returning context managers for accessing the
sqlite db from the module base class. exclusive_db_access takes
the db lock and yields the sqlite db connection. exclusive_db_cursor
does the same as the previous function but returns a cursor instead
of the connection.

Signed-off-by: John Mulligan <jmulligan@redhat.com>

diff --git a/src/pybind/mgr/mgr_module.py b/src/pybind/mgr/mgr_module.py
index 36f591332d1c585ea476b254c6c48da204f32f56..23f790cf3c238725e4c7f693522a85eb7baeb618 100644 (file)
--- a/src/pybind/mgr/mgr_module.py
+++ b/src/pybind/mgr/mgr_module.py
@@ -1,7 +1,22 @@
 import ceph_module  # noqa
 
-from typing import cast, Tuple, Any, Dict, Generic, Optional, Callable, List, \
-    Mapping, NamedTuple, Sequence, Union, Set, TYPE_CHECKING
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Set,
+    TYPE_CHECKING,
+    Tuple,
+    Union,
+    cast,
+)
 if TYPE_CHECKING:
     import sys
     if sys.version_info >= (3, 8):
@@ -17,6 +32,7 @@ import json
 import subprocess
 import threading
 from collections import defaultdict
+from contextlib import contextmanager
 from enum import IntEnum, Enum
 import os
 import rados
@@ -1346,6 +1362,23 @@ class MgrModule(ceph_module.BaseMgrModule, MgrModuleLoggingMixin):
             raise MgrDBNotReady()
         return self._db
 
+    @contextmanager
+    def exclusive_db_access(self) -> Iterator[sqlite3.Connection]:
+        """Context manager that grants exclusive access to the manager module sqlite3
+        db connection, while establishing a new db transaction.
+        """
+        with self._db_lock, self.db:
+            yield self.db
+
+    @contextmanager
+    def exclusive_db_cursor(self) -> Iterator[sqlite3.Cursor]:
+        """Context manager that yields a db cursor after getting exclusive
+        access to the manager module sqlite3 connection and a new db
+        transaction.
+        """
+        with self.exclusive_db_access() as db:
+            yield db.cursor()
+
     @property
     def release_name(self) -> str:
         """