Merge pull request #306 from powerapi-ng/feat/database-disconnect

feat(database): Add `disconnect` method

Author: gfieni

src/powerapi/database/base_db.py

@@ -26,86 +26,92 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-from typing import List, Type
-from powerapi.report import Report
+
 from powerapi.exception import PowerAPIExceptionWithMessage
+from powerapi.report import Report
 
 
 class DBError(PowerAPIExceptionWithMessage):
 
     """
-    Error raised when an error occuried when using a database
+    Error raised when an error occurred when using a database.
     """
     def __init__(self, msg: str):
         PowerAPIExceptionWithMessage.__init__(self, msg)
 
 
 class IterDB:
     """
-    IterDB class
-
-    This class allows to browse a database as an iterable
+    This class define the interface of a database results iterator.
     """
 
-    def __init__(self, db, report_type, stream_mode):
+    def __init__(self, db, report_type: type[Report], stream_mode: bool):
         """
+        :param db: Database instance
+        :param report_type: Report type to convert the database results into
+        :param stream_mode: Define if the iterator should stop when there is no data
         """
         self.db = db
-        self.stream_mode = stream_mode
         self.report_type = report_type
+        self.stream_mode = stream_mode
 
     def __iter__(self):
         """
+        Return an iterator for the database results.
         """
         raise NotImplementedError()
 
     def __next__(self) -> Report:
         """
+        Return and consume the next database result available.
         """
         raise NotImplementedError()
 
 
 class BaseDB:
     """
-    Abstract class which define every common function for database uses.
-
-    This class define every common function that need to be implemented
-    by each DB module. A database module correspond to a kind of BDD.
-    For example, Mongodb, influxdb, csv are different kind of BDD.
+    This class define the interface needed to fetch/save reports from/to a database.
     """
-    def __init__(self, report_type: Type[Report], exceptions: List[Type[Exception]] = None, asynchrone: bool = False):
-        self.exceptions = exceptions or []
-        self.asynchrone = asynchrone
+
+    def __init__(self, report_type: type[Report], exceptions: list[type[Exception]] = None, is_async: bool = False):
+        """
+        :param report_type: The type of report expected
+        :param exceptions: List of exception type raised by the database module
+        :param is_async: Whether the database use asyncio or not
+        """
         self.report_type = report_type
+        self.exceptions = exceptions or []
+        self.is_async = is_async
 
     def connect(self):
         """
-        Function that allow to load the database. Depending of the type,
-        different process can happen.
+        Connect to the database.
+        """
+        raise NotImplementedError()
 
-        .. note:: Need to be overrided
+    def disconnect(self):
+        """
+        Disconnect from the database.
         """
         raise NotImplementedError()
 
-    def iter(self, stream_mode: bool) -> IterDB:
+    def iter(self, stream_mode: bool = False) -> IterDB:
         """
-        Create the iterator for get the data
-        :param stream_mode: Define if we read in stream mode
+        Create and returns a database results iterator.
+        :param stream_mode: Define if we read continuously (streaming) or stop when no data is available
         """
         raise NotImplementedError()
 
     def save(self, report: Report):
         """
-        Allow to save a json input in the db
-
-        :param report: Report
+        Save a report to the database.
+        :param report: Report to be saved
         """
         raise NotImplementedError()
 
-    def save_many(self, reports: List[Report]):
+    def save_many(self, reports: list[Report]):
         """
-        Allow to save a batch of data
-
-        :param reports: Batch of Serialized Report
+        Save multiple reports to the database. (batch mode)
+        :param reports: List of Report to be saved
         """
         raise NotImplementedError()

src/powerapi/database/csvdb.py

@@ -245,17 +245,20 @@ def clean_files(self):
     # Override from BaseDB #
     ########################
 
-    def iter(self, stream_mode: bool) -> CsvIterDB:
+    def iter(self, stream_mode: bool = False) -> CsvIterDB:
         """
         Create the iterator for get the data
         """
         return CsvIterDB(self, self.filenames, self.report_type, stream_mode)
 
     def connect(self):
         """
-        Override from BaseDB.
+        Connect to the csv database.
+        """
 
-        Nothing to do with CSV, because it's just files operations.
+    def disconnect(self):
+        """
+        Disconnect from the csv database.
         """
 
     def save(self, report: Report):

src/powerapi/database/file_db.py

@@ -112,7 +112,12 @@ def connect(self):
         if not os.path.exists(self.filename):
             raise FileBadDBError(self.filename)
 
-    def iter(self, stream_mode: bool) -> FileIterDB:
+    def disconnect(self):
+        """
+        Disconnect from the file database.
+        """
+
+    def iter(self, stream_mode: bool = False) -> FileIterDB:
         """
         Create the iterator for get the data
         """

src/powerapi/database/influxdb2.py

@@ -94,12 +94,7 @@ def _ping_client(self):
 
     def connect(self):
         """
-            Override from BaseDB.
-
-            Create the connection to the influxdb database with the current
-            configuration (url:port/bucket_name/org), then check if the connection has
-            been created without failure.
-
+        Connect to the influxdb2 database.
         """
         # close connection if reload
         if self.client is not None:
@@ -123,6 +118,11 @@ def connect(self):
 
         self.buckets_api.create_bucket(bucket_name=self.bucket_name)
 
+    def disconnect(self):
+        """
+        Disconnect from the influxdb2 database.
+        """
+
     def get_db_by_name(self, db_name: str):
         """
             Get the database (bucket) with the given name

src/powerapi/database/mongodb.py

@@ -145,7 +145,12 @@ def connect(self):
 
         self.collection = self.mongo_client[self.db_name][self.collection_name]
 
-    def iter(self, stream_mode: bool) -> MongoIterDB:
+    def disconnect(self):
+        """
+        Disconnect from the mongodb database.
+        """
+
+    def iter(self, stream_mode: bool = False) -> MongoIterDB:
         """
         Create the iterator for get the data
         """

src/powerapi/database/opentsdb.py

@@ -91,6 +91,11 @@ def connect(self):
         if not self.client.is_connected() and not self.client.is_alive():
             raise CantConnectToOpenTSDBException('connexion error')
 
+    def disconnect(self):
+        """
+        Disconnect from the OpenTSDB database.
+        """
+
     def save(self, report: PowerReport):
         """
         Override from BaseDB

src/powerapi/database/prometheus_db.py

@@ -73,6 +73,11 @@ def connect(self):
         """
         start_http_server(port=self.port, addr=self.address)
 
+    def disconnect(self):
+        """
+        Disconnect from the Prometheus database.
+        """
+
 
 class PrometheusDB(BasePrometheusDB):
     """

src/powerapi/database/socket_db.py

@@ -31,13 +31,36 @@
 from typing import Type, List
 import json
 
-
 from powerapi.utils import JsonStream
 from powerapi.report import Report
 from .base_db import IterDB, BaseDB, DBError
 
-BUFFER_SIZE = 4096
-SOCKET_TIMEOUT = 0.5
+
+class IterSocketDB(IterDB):
+    """
+    iterator connected to a socket that receive report from a sensor
+    """
+
+    def __init__(self, report_type, stream_mode, queue):
+        """
+        """
+        IterDB.__init__(self, None, report_type, stream_mode)
+
+        self.queue = queue
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            json_str = await asyncio.wait_for(self.queue.get(), 2)
+            # json = self.queue.get_nowait()
+            # self.queue.get()
+            report = self.report_type.from_json(json.loads(json_str))
+            return report
+        # except Empty:
+        except asyncio.TimeoutError:
+            return None
 
 
 class SocketDB(BaseDB):
@@ -46,23 +69,31 @@ class SocketDB(BaseDB):
     """
 
     def __init__(self, report_type: Type[Report], port: int):
-        BaseDB.__init__(self, report_type, asynchrone=True)
+        BaseDB.__init__(self, report_type, is_async=True)
         self.queue = None
         self.port = port
         self.server = None
 
     async def connect(self):
+        """
+        Connect to the socket database.
+        """
         self.queue = asyncio.Queue()
         self.server = await asyncio.start_server(self._gen_server_callback(), host='127.0.0.1', port=self.port)
 
+    async def disconnect(self):
+        """
+        Disconnect from the socket database.
+        """
+
     async def stop(self):
         """
         stop server connection
         """
         self.server.close()
         await self.server.wait_closed()
 
-    def iter(self, stream_mode):
+    def iter(self, stream_mode: bool = False) -> IterSocketDB:
         return IterSocketDB(self.report_type, stream_mode, self.queue)
 
     def _gen_server_callback(self):
@@ -91,30 +122,3 @@ def save(self, report: Report):
 
     def save_many(self, reports: List[Report]):
         raise DBError('Socket db don\'t support save_many method')
-
-
-class IterSocketDB(IterDB):
-    """
-    iterator connected to a socket that receive report from a sensor
-    """
-
-    def __init__(self, report_type, stream_mode, queue):
-        """
-        """
-        IterDB.__init__(self, None, report_type, stream_mode)
-
-        self.queue = queue
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        try:
-            json_str = await asyncio.wait_for(self.queue.get(), 2)
-            # json = self.queue.get_nowait()
-            # self.queue.get()
-            report = self.report_type.from_json(json.loads(json_str))
-            return report
-        # except Empty:
-        except asyncio.TimeoutError:
-            return None

src/powerapi/database/virtiofs_db.py

@@ -79,6 +79,11 @@ def connect(self):
         if not os.path.exists(self.root_directory_name):
             raise DirectoryDoesNotExistForVirtioFS(self.root_directory_name)
 
+    def disconnect(self):
+        """
+        Disconnect from the virtiosfs database.
+        """
+
     def save(self, report: Report):
         directory_name = self._generate_vm_directory_name(report.target)
         if directory_name is None:

src/powerapi/puller/handlers.py

@@ -72,7 +72,7 @@ def _connect(self):
 
     def _pull_database(self):
         try:
-            if self.state.asynchrone:
+            if self.state.database.is_async:
                 report = self.loop.run_until_complete(anext(self.state.database_it))
                 if report is None:
                     raise StopIteration()
@@ -95,7 +95,7 @@ def run(self):
 
         :param None msg: None.
         """
-        if self.state.asynchrone:
+        if self.state.database.is_async:
             self.loop = asyncio.new_event_loop()
             asyncio.set_event_loop(self.loop)
             self.state.loop = self.loop
@@ -200,7 +200,7 @@ def pull_db(self):
 
     def _database_connection(self):
         try:
-            if not self.state.asynchrone:
+            if not self.state.database.is_async:
                 self.state.database.connect()
                 self.state.database_it = self.state.database.iter(stream_mode=self.state.stream_mode)