Add read_stream and read_category special methods

Author: subhashb

.gitignore

@@ -145,3 +145,6 @@ poetry.lock
 
 # Postgres data
 pgdata/
+
+# Project-specific
+database/

message_db/client.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Dict
+from typing import Dict, List
 from uuid import uuid4
 
 from psycopg2.extensions import connection
@@ -98,7 +98,7 @@ def write(
 
         return position
 
-    def write_batch(self, stream_name, data, expected_version: int = None):
+    def write_batch(self, stream_name, data, expected_version: int = None) -> None:
         conn = self.connection_pool.get_connection()
 
         try:
@@ -117,14 +117,21 @@ def write_batch(self, stream_name, data, expected_version: int = None):
 
         return expected_version
 
-    def read(self, stream_name, position=0, no_of_messages=1000):
+    def read(
+        self,
+        stream_name: str,
+        sql: str = None,
+        position: int = 0,
+        no_of_messages: int = 1000,
+    ) -> List[Dict]:
         conn = self.connection_pool.get_connection()
         cursor = conn.cursor(cursor_factory=RealDictCursor)
 
-        if "-" in stream_name:
-            sql = "SELECT * FROM get_stream_messages(%(stream_name)s, %(position)s, %(batch_size)s);"
-        else:
-            sql = "SELECT * FROM get_category_messages(%(stream_name)s, %(position)s, %(batch_size)s);"
+        if not sql:
+            if "-" in stream_name:
+                sql = "SELECT * FROM get_stream_messages(%(stream_name)s, %(position)s, %(batch_size)s);"
+            else:
+                sql = "SELECT * FROM get_category_messages(%(stream_name)s, %(position)s, %(batch_size)s);"
 
         cursor.execute(
             sql,
@@ -142,6 +149,28 @@ def read(self, stream_name, position=0, no_of_messages=1000):
 
         return messages
 
+    def read_stream(self, stream_name, position=0, no_of_messages=1000) -> List[Dict]:
+        if "-" not in stream_name:
+            raise ValueError(f"{stream_name} is not a stream")
+
+        sql = "SELECT * FROM get_stream_messages(%(stream_name)s, %(position)s, %(batch_size)s);"
+
+        return self.read(
+            stream_name, sql=sql, position=position, no_of_messages=no_of_messages
+        )
+
+    def read_category(
+        self, category_name, position=0, no_of_messages=1000
+    ) -> List[Dict]:
+        if "-" in category_name:
+            raise ValueError(f"{category_name} is not a category")
+
+        sql = "SELECT * FROM get_category_messages(%(stream_name)s, %(position)s, %(batch_size)s);"
+
+        return self.read(
+            category_name, sql=sql, position=position, no_of_messages=no_of_messages
+        )
+
     def read_last_message(self, stream_name):
         conn = self.connection_pool.get_connection()
         cursor = conn.cursor(cursor_factory=RealDictCursor)

pyproject.toml

@@ -4,6 +4,10 @@ version = "0.1.0"
 description = "The Python interface to the MessageDB Event Store and Message Store"
 authors = ["Subhash Bhushan <[email protected]>"]
 
+packages = [
+    { include = "message_db" },
+]
+
 [tool.poetry.dependencies]
 python = "^3.7"
 psycopg2 = "^2.9.2"

tests/test_client.py

@@ -28,13 +28,13 @@ class TestMessageWrite:
     def test_write_to_store(self, client):
         client.write("testStream-123", "Event1", {"foo": "bar"})
 
-        messages = client.read("testStream-123", 0, 100)
+        messages = client.read("testStream-123")
         assert len(messages) == 1
 
     def test_writing_message_with_metadata(self, client):
         client.write("testStream-123", "Event1", {"foo": "bar"}, {"trace_id": "baz"})
 
-        messages = client.read("testStream-123", 0, 100)
+        messages = client.read("testStream-123")
         assert messages[0]["metadata"] == json.dumps({"trace_id": "baz"})
 
     def test_that_write_returns_position_of_message_written(self, client):
@@ -84,19 +84,19 @@ def test_that_write_fails_on_expected_version_mismatch(self, client):
         )
 
 
-class TestEventIO:
+class TestRead:
     def test_read_stream_from_store(self, client):
         client.write("testStream-123", "Event1", {"foo": "bar"})
 
-        messages = client.read("testStream-123", 0, 100)
+        messages = client.read("testStream-123")
         assert messages is not None
         assert messages[0]["data"] == json.dumps({"foo": "bar"})
 
     def test_read_multiple_stream_messages_from_store(self, client):
         for i in range(5):
             client.write("testStream-123", "Event1", {f"foo{i}": f"bar{i}"})
 
-        messages = client.read("testStream-123", 0, 100)
+        messages = client.read("testStream-123")
 
         assert messages is not None
         assert len(messages) == 5
@@ -106,7 +106,7 @@ def test_read_paginated_stream_messages_from_store(self, client):
         for i in range(5):
             client.write("testStream-123", "Event1", {f"foo{i}": f"bar{i}"})
 
-        messages = client.read("testStream-123", 0, 3)
+        messages = client.read("testStream-123", no_of_messages=3)
 
         assert messages is not None
         assert len(messages) == 3
@@ -126,7 +126,7 @@ def test_read_specific_stream_message(self, client):
         for i in range(5):
             client.write("testStream-456", "Event1", {"foo": f"baz{i}"})
 
-        messages = client.read("testStream-456", 0, 100)
+        messages = client.read("testStream-456")
 
         assert len(messages) == 5
         assert messages[4]["data"] == json.dumps({"foo": "baz4"})
@@ -135,7 +135,77 @@ def test_read_category_messages(self, client):
         for i in range(5):
             client.write("testStream-123", "Event1", {"foo": f"bar{i}"})
 
-        messages = client.read("testStream", 0, 100)
+        messages = client.read("testStream")
 
         assert len(messages) == 5
         assert messages[4]["data"] == json.dumps({"foo": "bar4"})
+
+
+class TestReadStream:
+    def test_reading_a_category_throws_error(self, client):
+        with pytest.raises(ValueError) as exc:
+            client.read_stream("testStream")
+
+        assert exc.value.args[0] == "testStream is not a stream"
+
+    def test_read_stream_from_store(self, client):
+        client.write("testStream-123", "Event1", {"foo": "bar"})
+
+        messages = client.read_stream("testStream-123")
+        assert messages is not None
+        assert messages[0]["data"] == json.dumps({"foo": "bar"})
+
+    def test_read_multiple_stream_messages_from_store(self, client):
+        for i in range(5):
+            client.write("testStream-123", "Event1", {f"foo{i}": f"bar{i}"})
+
+        messages = client.read_stream("testStream-123")
+
+        assert messages is not None
+        assert len(messages) == 5
+        assert messages[4]["data"] == json.dumps({"foo4": "bar4"})
+
+    def test_read_paginated_stream_messages_from_store(self, client):
+        for i in range(5):
+            client.write("testStream-123", "Event1", {f"foo{i}": f"bar{i}"})
+
+        messages = client.read_stream("testStream-123", no_of_messages=3)
+
+        assert messages is not None
+        assert len(messages) == 3
+        assert messages[2]["data"] == json.dumps({"foo2": "bar2"})
+
+
+class TestReadCategory:
+    def test_reading_a_stream_throws_error(self, client):
+        with pytest.raises(ValueError) as exc:
+            client.read_category("testStream-123")
+
+        assert exc.value.args[0] == "testStream-123 is not a category"
+
+    def test_read_category_messages_from_store(self, client):
+        client.write("testStream-123", "Event1", {"foo": "bar"})
+
+        messages = client.read_category("testStream")
+        assert messages is not None
+        assert messages[0]["data"] == json.dumps({"foo": "bar"})
+
+    def test_read_multiple_category_messages_from_store(self, client):
+        for i in range(5):
+            client.write("testStream-123", "Event1", {f"foo{i}": f"bar{i}"})
+
+        messages = client.read_category("testStream")
+
+        assert messages is not None
+        assert len(messages) == 5
+        assert messages[4]["data"] == json.dumps({"foo4": "bar4"})
+
+    def test_read_paginated_category_messages_from_store(self, client):
+        for i in range(5):
+            client.write("testStream-123", "Event1", {f"foo{i}": f"bar{i}"})
+
+        messages = client.read_category("testStream", no_of_messages=3)
+
+        assert messages is not None
+        assert len(messages) == 3
+        assert messages[2]["data"] == json.dumps({"foo2": "bar2"})