fix(pubsub): read entire message before resetting

Author: danielebriggi

pyproject.toml

@@ -6,6 +6,7 @@
 profile = "black"
 
 [tool.pytest.ini_options]
+log_level = "INFO"
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
     "serial",

src/sqlitecloud/driver.py

@@ -207,10 +207,10 @@ def _internal_setup_pubsub(
     def _internal_pubsub_thread(self, connection: SQLiteCloudConnect) -> None:
         blen = 2048
         buffer: bytes = b""
+        tread = 0
 
         try:
             while True:
-                tread = 0
 
                 try:
                     if not connection.pubsub_socket:
@@ -240,7 +240,6 @@ def _internal_pubsub_thread(self, connection: SQLiteCloudConnect) -> None:
 
                 nread = len(data)
                 tread += nread
-                blen -= nread
                 buffer += data
 
                 sqlitecloud_number = self._internal_parse_number(buffer)
@@ -262,11 +261,16 @@ def _internal_pubsub_thread(self, connection: SQLiteCloudConnect) -> None:
                 connection.pubsub_callback(
                     connection, SQLiteCloudResultSet(result), connection.pubsub_data
                 )
+
+                # reset after having read the message
+                tread = 0
+                buffer: bytes = b""
         except Exception as e:
             logging.error(f"An error occurred while parsing data: {e}.")
 
         finally:
-            connection.pubsub_callback(connection, None, connection.pubsub_data)
+            if connection and connection.pubsub_callback:
+                connection.pubsub_callback(connection, None, connection.pubsub_data)
 
     def upload_database(
         self,

src/tests/integration/test_pubsub.py

@@ -47,6 +47,38 @@ def assert_callback(conn, result, data):
 
         assert callback_called
 
+    def test_notify_multiple_messages(self, sqlitecloud_connection):
+        connection, _ = sqlitecloud_connection
+
+        called_times = 3
+        flag = threading.Event()
+
+        def assert_callback(conn, result, data):
+            nonlocal called_times
+            nonlocal flag
+
+            if isinstance(result, SQLiteCloudResultSet):
+                assert data == ["somedataX"]
+                called_times -= 1
+                if called_times == 0:
+                    flag.set()
+
+        pubsub = SQLiteCloudPubSub()
+        subject_type = SQLITECLOUD_PUBSUB_SUBJECT.CHANNEL
+        channel = "channel" + str(uuid.uuid4())
+
+        pubsub.create_channel(connection, channel)
+        pubsub.listen(connection, subject_type, channel, assert_callback, ["somedataX"])
+
+        pubsub.notify_channel(connection, channel, "somedataX")
+        pubsub.notify_channel(connection, channel, "somedataX")
+        pubsub.notify_channel(connection, channel, "somedataX")
+
+        # wait for callback to be called
+        flag.wait(30)
+
+        assert called_times == 0
+
     def test_unlisten_channel(self, sqlitecloud_connection):
         connection, _ = sqlitecloud_connection