Add Redis stream discovery helpers

src/braingeneers/iot/messaging.py

@@ -112,7 +112,12 @@ class MessageBroker:
         https://awslabs.github.io/aws-crt-python/
     """
 
-    def __init__(self, name: str = None, credentials_file: (str, io.IOBase) = None, logger: logging.Logger = None):
+    def __init__(
+            self,
+            name: str = None,
+            credentials_file: (str, io.IOBase) = None,
+            logger: logging.Logger = None,
+    ):
         """
         Typical usage example:
             mb = MessageBroker()
@@ -128,8 +133,8 @@ def __init__(self, name: str = None, credentials_file: (str, io.IOBase) = None,
         :param endpoint: optional AWS endpoint, defaults to Braingeneers standard us-west-2
         :param credentials_file: optional file path string or file-like object containing the
             standard `~/.aws/credentials` file. See https://github.com/braingeneers/wiki/blob/main/shared/permissions.md
-            defaults to looking in `~/.aws/credentials` if left as None. This file expects to find profiles named
-            'aws-braingeneers-iot' and 'redis' in it.
+            defaults to looking in `~/.aws/credentials` if left as None. Each backing service validates its
+            own credentials lazily when first used.
         :param logger: optional logger object, defaults to a new logger with the name of this class.
         """
         self.logger = logger if logger is not None else logging.getLogger(__name__)
@@ -145,36 +150,19 @@ def __init__(self, name: str = None, credentials_file: (str, io.IOBase) = None,
             assert hasattr(credentials_file, 'read'), 'credentials_file parameter must be a filename string or file-like object.'
             self._credentials = credentials_file.read()
 
-        config = configparser.ConfigParser()
-        config.read_file(io.StringIO(self._credentials))
-
-        assert 'braingeneers-mqtt' in config, \
-            'Your AWS credentials_file file is missing a section [braingeneers-mqtt], you may have the wrong ' \
-            'version of the credentials_file file.'
-        assert 'profile-id' in config['braingeneers-mqtt'], \
-            'Your AWS credentials_file file is malformed, profile-id is missing from the [braingeneers-mqtt] section.'
-        assert 'profile-key' in config['braingeneers-mqtt'], \
-            'Your AWS credentials_file file is malformed, profile-key was not found under the [braingeneers-mqtt] section.'
-        assert 'endpoint' in config['braingeneers-mqtt'], \
-            'Your AWS credentials_file file is malformed, endpoint was not found under the [braingeneers-mqtt] section.'
-        assert 'port' in config['braingeneers-mqtt'], \
-            'Your AWS credentials_file file is malformed, ' \
-                                                      'port was not found under the [braingeneers-mqtt] section.'
-
         self.certs_temp_dir = None
         self._mqtt_connection = None
         self._mqtt_loop_running = False
         self._mqtt_client_id = f"braingeneerspy-{uuid.uuid4()}"
-        self._mqtt_profile_id = config['braingeneers-mqtt']['profile-id']
-        self._mqtt_profile_key = config['braingeneers-mqtt']['profile-key']
-        self._mqtt_endpoint = config['braingeneers-mqtt']['endpoint']
-        self._mqtt_port = int(config['braingeneers-mqtt']['port'])
+        self._mqtt_profile_id = None
+        self._mqtt_profile_key = None
+        self._mqtt_endpoint = None
+        self._mqtt_port = None
         self._boto_iot_client = None
         self._boto_iot_data_client = None
         self._redis_client = None
         self._jwt_service_account_token = None
-
-        self.shadow_interface = sh.DatabaseInteractor(jwt_service_token=self.jwt_service_account_token)
+        self._shadow_interface = None
 
         self._subscribed_data_streams = set()  # keep track of subscribed data streams
         self._recent_duplicate_mqtt_messages = OrderedDict()
@@ -543,6 +531,66 @@ def get_metadata_for_stream(self, stream_name: str) -> dict:
             value = value.decode('utf-8')
         return json.loads(value)
 
+    @staticmethod
+    def _decode_redis_text(value: Union[str, bytes]) -> str:
+        return value.decode('utf-8') if isinstance(value, bytes) else value
+
+    def get_data_stream_info(self, stream_name: str) -> dict:
+        """
+        Retrieve Redis stream metadata for a data stream.
+
+        This wraps Redis XINFO STREAM and returns a small JSON-friendly subset
+        of the stream info that is useful for discovery and diagnostics.
+
+        :param stream_name: Name of the data stream.
+        :return: A dictionary with name, length, first_entry_id, and last_entry_id.
+        """
+        info = self.redis_client.xinfo_stream(stream_name)
+        first_entry = info.get("first-entry")
+        last_entry = info.get("last-entry")
+        return {
+            "name": stream_name,
+            "length": info.get("length"),
+            "first_entry_id": self._decode_redis_text(first_entry[0]) if first_entry else None,
+            "last_entry_id": self._decode_redis_text(last_entry[0]) if last_entry else None,
+        }
+
+    def list_data_streams(
+        self,
+        pattern: str = "*",
+        *,
+        include_metadata_keys: bool = False,
+        scan_count: int = 2000,
+    ) -> List[dict]:
+        """
+        Discover Redis data streams matching a key pattern.
+
+        This scans Redis keys, keeps only keys that are Redis Streams, and
+        returns a JSON-friendly summary for each stream. Non-stream keys that
+        match the pattern are ignored.
+
+        :param pattern: Redis key pattern, for example "braindance/*".
+        :param include_metadata_keys: Include sorted metadata key names from
+            get_metadata_for_stream() when True.
+        :param scan_count: Redis SCAN count hint.
+        :return: A list of stream summary dictionaries sorted by stream name.
+        """
+        streams = []
+        seen_stream_names = set()
+        for raw_key in self.redis_client.scan_iter(match=pattern, count=scan_count):
+            stream_name = self._decode_redis_text(raw_key)
+            if stream_name in seen_stream_names:
+                continue
+            seen_stream_names.add(stream_name)
+            try:
+                stream_info = self.get_data_stream_info(stream_name)
+            except redis.exceptions.ResponseError:
+                continue
+            if include_metadata_keys:
+                stream_info["metadata_keys"] = sorted(self.get_metadata_for_stream(stream_name).keys())
+            streams.append(stream_info)
+        return sorted(streams, key=lambda stream: stream["name"])
+
     def publish_data_stream(self, stream_name: str, data: Dict[Union[str, bytes], bytes], stream_size: int) -> None:
         """
         Publish (potentially large) data to a stream.
@@ -914,6 +962,7 @@ def _is_loop_alive(client) -> bool:
             thread = getattr(client, "_thread", None)
             return bool(thread and thread.is_alive())
         if self._mqtt_connection is None:
+            self._load_mqtt_credentials()
             '''
             root certs only required for https connection our current mqtt broker does not have this yet
             '''
@@ -980,6 +1029,31 @@ def on_log(client, userdata, level, buf):
 
         return self._mqtt_connection
 
+    def _load_mqtt_credentials(self) -> None:
+        if self._mqtt_profile_id is not None:
+            return
+
+        config = configparser.ConfigParser()
+        config.read_file(io.StringIO(self._credentials))
+
+        assert 'braingeneers-mqtt' in config, \
+            'Your AWS credentials_file file is missing a section [braingeneers-mqtt], you may have the wrong ' \
+            'version of the credentials_file file.'
+        assert 'profile-id' in config['braingeneers-mqtt'], \
+            'Your AWS credentials_file file is malformed, profile-id is missing from the [braingeneers-mqtt] section.'
+        assert 'profile-key' in config['braingeneers-mqtt'], \
+            'Your AWS credentials_file file is malformed, profile-key was not found under the [braingeneers-mqtt] section.'
+        assert 'endpoint' in config['braingeneers-mqtt'], \
+            'Your AWS credentials_file file is malformed, endpoint was not found under the [braingeneers-mqtt] section.'
+        assert 'port' in config['braingeneers-mqtt'], \
+            'Your AWS credentials_file file is malformed, ' \
+            'port was not found under the [braingeneers-mqtt] section.'
+
+        self._mqtt_profile_id = config['braingeneers-mqtt']['profile-id']
+        self._mqtt_profile_key = config['braingeneers-mqtt']['profile-key']
+        self._mqtt_endpoint = config['braingeneers-mqtt']['endpoint']
+        self._mqtt_port = int(config['braingeneers-mqtt']['port'])
+
     @property
     def redis_client(self) -> redis.Redis:
         """ Lazy initialization of the redis client. """
@@ -997,6 +1071,17 @@ def redis_client(self) -> redis.Redis:
 
         return self._redis_client
 
+    @property
+    def shadow_interface(self):
+        """ Lazy initialization of the Braingeneers device-shadow interface. """
+        if self._shadow_interface is None:
+            self._shadow_interface = sh.DatabaseInteractor(jwt_service_token=self.jwt_service_account_token)
+        return self._shadow_interface
+
+    @shadow_interface.setter
+    def shadow_interface(self, value):
+        self._shadow_interface = value
+
     @property
     def jwt_service_account_token(self) -> str:
         """ Lazy initialization of the JWT service account token. """

tests/test_messaging.py

@@ -1,5 +1,6 @@
 """ Unit test for BraingeneersMqttClient, assumes Braingeneers ~/.aws/credentials file exists """
 import logging
+import io
 import queue
 import threading
 import time
@@ -450,8 +451,15 @@ def test_shutdown_stops_mqtt_client_and_clears_duplicate_cache(self):
 
 
 class FakeRedisMetadataStore:
-    def __init__(self):
+    def __init__(self, *args, **kwargs):
+        del args, kwargs
         self.data = {}
+        self.streams = {}
+        self.duplicate_scan_keys = []
+        self.xinfo_stream_calls = {}
+
+    def config_set(self, *args, **kwargs):
+        pass
 
     def set(self, key, value):
         self.data[key] = value
@@ -460,6 +468,22 @@ def get(self, key):
         value = self.data.get(key)
         return None if value is None else value.encode("utf-8")
 
+    def scan_iter(self, match="*", count=2000):
+        del count
+        prefix = match.removesuffix("*")
+        for key in sorted([*self.data.keys(), *self.streams.keys()]):
+            if key.startswith(prefix):
+                yield key.encode("utf-8")
+        for key in self.duplicate_scan_keys:
+            if key.startswith(prefix):
+                yield key.encode("utf-8")
+
+    def xinfo_stream(self, key):
+        self.xinfo_stream_calls[key] = self.xinfo_stream_calls.get(key, 0) + 1
+        if key not in self.streams:
+            raise messaging.redis.exceptions.ResponseError("no such key")
+        return self.streams[key]
+
 
 class TestStreamMetadata(unittest.TestCase):
     def setUp(self):
@@ -499,6 +523,144 @@ def test_set_stream_metadata_requires_dict(self):
                 "ephys/experiment_123", ["not", "a", "dict"]
             )
 
+    def test_message_broker_lazily_initializes_non_redis_interfaces(self):
+        credentials = io.StringIO(
+            "[redis]\n"
+            "redis_password = test\n"
+        )
+        redis_store = FakeRedisMetadataStore()
+        redis_store.streams["ephys/stream_1"] = {
+            "length": 1,
+            "first-entry": (b"1-0", {}),
+            "last-entry": (b"1-0", {}),
+        }
+
+        with patch.object(messaging.redis, "Redis", return_value=redis_store), \
+             patch.object(messaging.sh, "DatabaseInteractor") as database_interactor:
+            broker = messaging.MessageBroker(
+                name="redis-only-test",
+                credentials_file=credentials,
+            )
+
+            self.assertIsNone(broker._shadow_interface)
+            self.assertIsNone(broker._mqtt_connection)
+            database_interactor.assert_not_called()
+
+            self.assertEqual(
+                broker.list_data_streams("ephys/*"),
+                [
+                    {
+                        "name": "ephys/stream_1",
+                        "length": 1,
+                        "first_entry_id": "1-0",
+                        "last_entry_id": "1-0",
+                    },
+                ],
+            )
+            database_interactor.assert_not_called()
+
+            broker._jwt_service_account_token = {
+                "access_token": "test-token",
+                "expires_at": "2099-01-01 00:00:00 UTC",
+            }
+            shadow_interface = broker.shadow_interface
+
+            database_interactor.assert_called_once_with(
+                jwt_service_token=broker._jwt_service_account_token
+            )
+            self.assertIs(shadow_interface, database_interactor.return_value)
+
+    def test_mqtt_credentials_are_validated_when_mqtt_is_first_used(self):
+        broker = messaging.MessageBroker(
+            name="redis-only-test",
+            credentials_file=io.StringIO(
+                "[redis]\n"
+                "redis_password = test\n"
+            ),
+        )
+
+        with self.assertRaisesRegex(AssertionError, r"\[braingeneers-mqtt\]"):
+            broker.mqtt_connection
+
+    def test_get_data_stream_info_returns_json_friendly_subset(self):
+        self.broker.redis_client.streams["ephys/stream_1"] = {
+            "length": 2,
+            "first-entry": (b"1-0", {}),
+            "last-entry": (b"2-0", {}),
+        }
+
+        self.assertEqual(
+            self.broker.get_data_stream_info("ephys/stream_1"),
+            {
+                "name": "ephys/stream_1",
+                "length": 2,
+                "first_entry_id": "1-0",
+                "last_entry_id": "2-0",
+            },
+        )
+
+    def test_list_data_streams_filters_non_stream_keys_and_includes_metadata_keys(self):
+        self.broker.redis_client.streams["ephys/stream_1"] = {
+            "length": 2,
+            "first-entry": (b"1-0", {}),
+            "last-entry": (b"2-0", {}),
+        }
+        self.broker.redis_client.streams["ephys/stream_2"] = {
+            "length": 1,
+            "first-entry": (b"3-0", {}),
+            "last-entry": (b"3-0", {}),
+        }
+        self.broker.set_metadata_for_stream("ephys/stream_1", {"sample_rate": 25000})
+        self.broker.redis_client.data["ephys/not_stream"] = "{}"
+
+        self.assertEqual(
+            self.broker.list_data_streams("ephys/*", include_metadata_keys=True),
+            [
+                {
+                    "name": "ephys/stream_1",
+                    "length": 2,
+                    "first_entry_id": "1-0",
+                    "last_entry_id": "2-0",
+                    "metadata_keys": ["sample_rate"],
+                },
+                {
+                    "name": "ephys/stream_2",
+                    "length": 1,
+                    "first_entry_id": "3-0",
+                    "last_entry_id": "3-0",
+                    "metadata_keys": [],
+                },
+            ],
+        )
+
+    def test_list_data_streams_deduplicates_scan_results_before_metadata_lookup(self):
+        self.broker.redis_client.streams["ephys/stream_1"] = {
+            "length": 2,
+            "first-entry": (b"1-0", {}),
+            "last-entry": (b"2-0", {}),
+        }
+        self.broker.set_metadata_for_stream("ephys/stream_1", {"sample_rate": 25000})
+        self.broker.redis_client.duplicate_scan_keys = [
+            "ephys/stream_1",
+            "ephys/stream_1",
+        ]
+
+        self.assertEqual(
+            self.broker.list_data_streams("ephys/*", include_metadata_keys=True),
+            [
+                {
+                    "name": "ephys/stream_1",
+                    "length": 2,
+                    "first_entry_id": "1-0",
+                    "last_entry_id": "2-0",
+                    "metadata_keys": ["sample_rate"],
+                },
+            ],
+        )
+        self.assertEqual(
+            self.broker.redis_client.xinfo_stream_calls["ephys/stream_1"], 1
+        )
+
 
 class TestInterprocessQueue(unittest.TestCase):
     def setUp(self) -> None: