Session manager Implementation (#7)

This pull request enables swarms to persist their state between runs and pick up where they left off if a stop is required.

* NEW: Swarms can now write their state to file as a ".pkl" file
* NEW: CLI updated to implement new behaviour. Swarm sessions can be resumed by providing the --resume-session and --swarm-name arguments.
* NEW: Sessions are stored in '/swarms/`
* NEW: Implemented CLI commands for managing swarms
iot-swarm sessions has commands to initialise, list, or remove sessions
* NEW: Tests updated for new behaviour.
* Swarm.load_swarm(<swarm-id>) can also be used to load a swarm from file

Author: lewis-chambers

pyproject.toml

@@ -4,18 +4,19 @@ requires = ["setuptools >= 61.0", "autosemver"]
 
 [project]
 dependencies = [
-    "platformdirs",
-    "boto3",
     "autosemver",
-    "config",
-    "click",
-    "docutils<0.17",
     "awscli",
     "awscrt",
     "awsiotsdk",
-    "oracledb",
     "backoff",
+    "boto3",
+    "click",
+    "config",
+    "dill",
+    "docutils<0.17",
+    "oracledb",
     "pandas",
+    "platformdirs",
 ]
 name = "iot-swarm"
 dynamic = ["version"]

src/iotswarm/db.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 from math import nan
 import sqlite3
+from typing import List
 
 logger = logging.getLogger(__name__)
 
@@ -39,15 +40,18 @@ def __repr__(self):
             logger_arg = f"inherit_logger={self._instance_logger.parent}"
         return f"{self.__class__.__name__}({logger_arg})"
 
+    def __eq__(self, obj):
+        return self._instance_logger == obj._instance_logger
+
     @abc.abstractmethod
-    def query_latest_from_site(self):
+    def query_latest_from_site(self) -> List:
         pass
 
 
 class MockDB(BaseDatabase):
 
     @staticmethod
-    def query_latest_from_site():
+    def query_latest_from_site() -> List:
         return []
 
 
@@ -63,6 +67,14 @@ class CosmosDB(BaseDatabase):
     site_id_query: CosmosQuery
     """SQL query for retrieving list of site IDs"""
 
+    def __eq__(self, obj):
+        return (
+            type(self.connection) == type(obj.connection)
+            and self.site_data_query == obj.site_data_query
+            and self.site_id_query == obj.site_id_query
+            and BaseDatabase.__eq__(self, obj)
+        )
+
     @staticmethod
     def _validate_table(table: CosmosTable) -> None:
         """Validates that the query is legal"""
@@ -99,6 +111,9 @@ def _validate_max_sites(max_sites: int) -> int:
 
         return max_sites
 
+    def query_latest_from_site(self):
+        pass
+
 
 class Oracle(CosmosDB):
     """Class for handling oracledb logic and retrieving values from DB."""
@@ -225,8 +240,16 @@ class LoopingCsvDB(BaseDatabase):
     connection: pd.DataFrame
     """Connection to the pd object holding data."""
 
-    cache: dict
-    """Cache object containing current index of each site queried."""
+    db_file: str | Path
+    """Path to the database file."""
+
+    def __eq__(self, obj):
+
+        return (
+            type(self.connection) == type(obj.connection)
+            and self.db_file == obj.db_file
+            and BaseDatabase.__eq__(self, obj)
+        )
 
     @staticmethod
     def _get_connection(*args) -> pd.DataFrame:
@@ -241,27 +264,30 @@ def __init__(self, csv_file: str | Path):
         """
 
         BaseDatabase.__init__(self)
+
+        if not isinstance(csv_file, Path):
+            csv_file = Path(csv_file)
+
+        self.db_file = csv_file
         self.connection = self._get_connection(csv_file)
-        self.cache = dict()
 
-    def query_latest_from_site(self, site_id: str) -> dict:
+    def query_latest_from_site(self, site_id: str, index: int) -> dict:
         """Queries the datbase for a `SITE_ID` incrementing by 1 each time called
         for a specific site. If the end is reached, it loops back to the start.
 
         Args:
             site_id: ID of the site to query for.
+            index: An offset index to query.
         Returns:
             A dict of the data row.
         """
 
         data = self.connection.query("SITE_ID == @site_id").replace({nan: None})
 
-        if site_id not in self.cache or self.cache[site_id] >= len(data):
-            self.cache[site_id] = 1
-        else:
-            self.cache[site_id] += 1
+        # Automatically loops back to start
+        db_index = index % len(data)
 
-        return data.iloc[self.cache[site_id] - 1].to_dict()
+        return data.iloc[db_index].to_dict()
 
     def query_site_ids(self, max_sites: int | None = None) -> list:
         """query_site_ids returns a list of site IDs from the database
@@ -316,32 +342,49 @@ def __init__(self, db_file: str | Path):
 
         self.cursor = self.connection.cursor()
 
-    def query_latest_from_site(self, site_id: str, table: CosmosTable) -> dict:
+    def __eq__(self, obj) -> bool:
+        return CosmosDB.__eq__(self, obj) and super(LoopingCsvDB, self).__eq__(obj)
+
+    def __getstate__(self) -> object:
+
+        state = self.__dict__.copy()
+
+        del state["connection"]
+        del state["cursor"]
+
+        return state
+
+    def __setstate__(self, state) -> object:
+
+        self.__dict__.update(state)
+
+        self.connection = self._get_connection(self.db_file)
+        self.cursor = self.connection.cursor()
+
+    def query_latest_from_site(
+        self, site_id: str, table: CosmosTable, index: int
+    ) -> dict:
         """Queries the datbase for a `SITE_ID` incrementing by 1 each time called
         for a specific site. If the end is reached, it loops back to the start.
 
         Args:
             site_id: ID of the site to query for.
             table: A valid table from the database
+            index: Offset of index.
         Returns:
             A dict of the data row.
         """
         query = self._fill_query(self.site_data_query, table)
 
-        if site_id not in self.cache:
-            self.cache[site_id] = 0
-        else:
-            self.cache[site_id] += 1
-
         data = self._query_latest_from_site(
-            query, {"site_id": site_id, "offset": self.cache[site_id]}
+            query, {"site_id": site_id, "offset": index}
         )
 
         if data is None:
-            self.cache[site_id] = 0
+            index = 0
 
         data = self._query_latest_from_site(
-            query, {"site_id": site_id, "offset": self.cache[site_id]}
+            query, {"site_id": site_id, "offset": index}
         )
 
         return data

src/iotswarm/devices.py

@@ -57,6 +57,9 @@ class BaseDevice:
     mqtt_suffix: str
     """Suffix added to mqtt message."""
 
+    swarm: object | None = None
+    """The session applied"""
+
     @property
     def mqtt_topic(self) -> str:
         "Builds the mqtt topic."
@@ -75,6 +78,30 @@ def mqtt_topic(self, value):
         self._mqtt_topic = value
         self.mqtt_base_topic = value
 
+    def __eq__(self, obj) -> bool:
+
+        base_equality = (
+            self.device_type == obj.device_type
+            and self.cycle == obj.cycle
+            and self.max_cycles == obj.max_cycles
+            and self.sleep_time == obj.sleep_time
+            and self.device_id == obj.device_id
+            and self.delay_start == obj.delay_start
+            and self._instance_logger == obj._instance_logger
+            and self.data_source == obj.data_source
+            and self.connection == obj.connection
+        )
+
+        table_equality = True
+        if hasattr(self, "table") and not self.table == obj.table:
+            table_equality = False
+
+        mqtt_equality = True
+        if hasattr(self, "mqtt_topic") and not self.mqtt_topic == obj.mqtt_topic:
+            mqtt_equality = False
+
+        return base_equality and table_equality and mqtt_equality
+
     def __init__(
         self,
         device_id: str,
@@ -235,12 +262,19 @@ async def _add_delay(self):
         self._instance_logger.debug(f"Delaying first cycle for: {delay}s.")
         await asyncio.sleep(delay)
 
-    def _send_payload(self, payload: dict):
+    def _send_payload(self, payload: dict) -> bool:
+        """Forwards the payload submission request to the connection
+
+        Args:
+            payload: The data to send.
+        Returns:
+            bool: True if sent sucessfully, else false.
+        """
 
         if isinstance(self.connection, IotCoreMQTTConnection):
-            self.connection.send_message(payload, topic=self.mqtt_topic)
+            return self.connection.send_message(payload, topic=self.mqtt_topic)
         else:
-            self.connection.send_message(payload)
+            return self.connection.send_message(payload)
 
     async def run(self):
         """The main invocation of the method. Expects a Oracle object to do work on
@@ -251,26 +285,32 @@ async def run(self):
         """
 
         while True:
+            if self.max_cycles > 0 and self.cycle >= self.max_cycles:
+                break
 
             if self.delay_start and self.cycle == 0:
                 await self._add_delay()
 
             payload = await self._get_payload()
-            payload = self._format_payload(payload)
 
-            if payload:
+            if payload is not None:
+                payload = self._format_payload(payload)
+
                 self._instance_logger.debug("Requesting payload submission.")
-                self._send_payload(payload)
-                self._instance_logger.info(
-                    f'Message sent{f" to topic: {self.mqtt_topic}" if self.mqtt_topic else ""}'
-                )
+
+                send_status = self._send_payload(payload)
+
+                if send_status == True:
+                    self._instance_logger.info(
+                        f'Message sent{f" to topic: {self.mqtt_topic}" if self.mqtt_topic else ""}'
+                    )
+                    self.cycle += 1
+
+                    if self.swarm is not None:
+                        self.swarm.write_self(replace=True)
             else:
                 self._instance_logger.warning(f"No data found.")
 
-            self.cycle += 1
-            if self.max_cycles > 0 and self.cycle >= self.max_cycles:
-                break
-
             await asyncio.sleep(self.sleep_time)
 
     async def _get_payload(self):
@@ -280,16 +320,21 @@ async def _get_payload(self):
                 self.device_id, self.table
             )
         elif isinstance(self.data_source, LoopingSQLite3):
-            return self.data_source.query_latest_from_site(self.device_id, self.table)
+            return self.data_source.query_latest_from_site(
+                self.device_id, self.table, self.cycle
+            )
         elif isinstance(self.data_source, LoopingCsvDB):
-            return self.data_source.query_latest_from_site(self.device_id)
+            return self.data_source.query_latest_from_site(self.device_id, self.cycle)
         elif isinstance(self.data_source, BaseDatabase):
             return self.data_source.query_latest_from_site()
 
     def _format_payload(self, payload):
         """Oranises payload into correct structure."""
         return payload
 
+    def _attach_swarm(self, swarm: object):
+        self.swarm = swarm
+
 
 class CR1000XDevice(BaseDevice):
     "Represents a CR1000X datalogger."