safeguard links

wmo-im · Feb 18, 2025 · d9dbc83 · d9dbc83
1 parent 17e2920
commit d9dbc83
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 5 deletions.
diff --git a/pywis_pubsub/message.py b/pywis_pubsub/message.py
@@ -64,7 +64,7 @@ def get_data(msg_dict: dict, verify_certs=True) -> bytes:
     :returns: `bytes` of data
     """
 
-    link = get_link(msg_dict['links'])
+    link = get_link(msg_dict.get('links', []))
 
     if link:
         LOGGER.debug(f'Found link: {link}')

diff --git a/pywis_pubsub/mqtt.py b/pywis_pubsub/mqtt.py
@@ -34,7 +34,7 @@
 
 class MQTTPubSubClient:
     """MQTT PubSub client"""
-    def __init__(self, broker: str, options: dict = {}) -> None:
+    def __init__(self, broker: str, options: dict = {}, conn=None) -> None:
         """
         PubSub initializer
 
@@ -62,6 +62,11 @@ def __init__(self, broker: str, options: dict = {}) -> None:
         if self.broker_url.scheme in ['ws', 'wss']:
             transport = 'websockets'
 
+        if conn is not None:
+            LOGGER.debug(f'Using existing connection {conn}')
+            self.conn = conn
+            return
+
         msg = f'Connecting to broker {self.broker_safe_url} with id {self.client_id}'  # noqa
         LOGGER.debug(msg)
         self.conn = mqtt_client.Client(mqtt_client.CallbackAPIVersion.VERSION2,

diff --git a/pywis_pubsub/subscribe.py b/pywis_pubsub/subscribe.py
@@ -70,7 +70,7 @@ def on_message_handler(client, userdata, msg):
             LOGGER.debug('Message geometry not within bbox; skipping')
             return
 
-    clink = get_link(msg_dict['links'])
+    clink = get_link(msg_dict.get('links', []))
     if not clink:
         LOGGER.warning('No valid data link found')
         return
@@ -106,7 +106,7 @@ def on_message_handler(client, userdata, msg):
         filepath = userdata['storage']['options'].get('filepath', 'data_id')
         LOGGER.debug(f'Using {filepath} for naming filepath')
 
-        link = get_link(msg_dict['links'])
+        link = get_link(msg_dict.get('links', []))
 
         if filepath == 'link':
             LOGGER.debug('Using link as filepath')

diff --git a/pywis_pubsub/verification.py b/pywis_pubsub/verification.py
@@ -97,7 +97,7 @@ def verify_data(instance: dict, verify_certs: bool = True) -> bool:
     if 'content' in instance['properties']:
         size = instance['properties']['content']['size']
     else:
-        size = get_link(instance['links'])['length']
+        size = get_link(instance.get('links', [])).get('length')
 
     LOGGER.debug(f'size: {size}')
     return data_verified(data, size, method, value)