class DbHandler(Thread):
running = True
subscribed_event = [
'gcp_state_changed', 'device_state_changed', 'iot_traffic',
'host_health', 'device_sensor_data', 'digital_twin'
]
def __init__(self, queue: Queue, thread_event: Event) -> None:
Thread.__init__(self)
self.mongo = MongoHandler(db_name='iot_db')
self.observer_publish_queue = queue
self._thread_ready = thread_event
self.observer_notify_queue = Queue(maxsize=100)
self.log = Logging(owner=__file__, config=True)
def __del__(self) -> None:
self.running = False
def run(self) -> None:
self._thread_ready.set()
while self.running:
item = self.observer_notify_queue.get()
action = self.action_selector(event=item.event)
action(msg=item)
def notify(self, msg: ObserverMessage) -> None:
self.observer_notify_queue.put(msg)
def action_selector(self, event: str) -> Callable:
action_map = {
'gcp_state_changed': self.store_state_data,
'device_state_changed': self.store_state_data,
'iot_traffic': self.add_document_row,
'host_health': self.add_document_row,
'device_sensor_data': self.add_document_row,
'digital_twin': self.handle_digital_twin
}
return action_map.get(event, self.action_skip)
@staticmethod
def action_skip():
pass
def get_document_data(self, document: str) -> list:
return self.mongo.get(collection_name=document)
def store_state_data(self, msg: ObserverMessage) -> None:
document_data = {
'device_id': msg.data.get('device_id'),
'event': msg.event,
'change_source': msg.data.get('event'),
'state': msg.data.get('state')
}
self.mongo.write(collection_name="states",
data=document_data,
key="device_id")
def add_document_row(self, msg: ObserverMessage) -> None:
self.mongo.insert(collection_name=msg.event, data=msg.data)
def handle_digital_twin(self, msg: ObserverMessage) -> None:
if msg.subject == "fetch_digital_twin":
self._fetch_digital_twin()
elif msg.subject == "save_digital_twin":
self._save_digital_twin(twin=msg.data)
else:
self.action_skip()
def _fetch_digital_twin(self) -> None:
self.log.info("Fetching digital twin from DB")
twin_data = self.get_document_data(document="digital_twin")
digital_twin = self._outbound_adapter(data=twin_data)
msg = ObserverMessage(event="digital_twin",
data=digital_twin,
subject="retrieved_digital_twin")
self.observer_publish_queue.put(msg)
def _save_digital_twin(self, twin: list) -> None:
self.log.info("Uploading updated digital twin")
self.mongo.write(collection_name='digital_twin',
data=twin,
key='device_name')
@staticmethod
def _outbound_adapter(data: list) -> list:
""" Removes the object_id field from each data entry, preping the data for transportation """
for entry in data:
entry.pop("_id", None)
return data
def test_get(self, mock_connect):
mock_connect.return_value = MockMongoClient
mongo = MongoHandler("db_name")
result = mongo.get(collection_name="mock_collection",
query={"abc": 123})
self.assertListEqual(result, [1, 2, 3])