pgmq-python: adding support for Transaction

tembo-pgmq-python/example/example_app_async.py

@@ -0,0 +1,87 @@
+import asyncio
+from tembo_pgmq_python.async_queue import PGMQueue
+from tembo_pgmq_python.decorators import async_transaction as transaction
+
+
+async def main():
+    # Initialize the queue
+    queue = PGMQueue(
+        host="0.0.0.0",
+        port="5432",
+        username="postgres",
+        password="postgres",
+        database="postgres",
+        verbose=True,
+        log_filename="pgmq_async.log",
+    )
+    await queue.init()
+
+    test_queue = "transactional_queue_async"
+
+    # Clean up if the queue already exists
+    queues = await queue.list_queues()
+    if test_queue in queues:
+        await queue.drop_queue(test_queue)
+    await queue.create_queue(test_queue)
+
+    # Example messages
+    message1 = {"id": 1, "content": "First message"}
+    message2 = {"id": 2, "content": "Second message"}
+
+    # Transactional operation: send messages within a transaction
+    @transaction
+    async def transactional_operation(queue: PGMQueue, conn=None):
+        # Perform multiple queue operations within a transaction
+        await queue.send(test_queue, message1, conn=conn)
+        await queue.send(test_queue, message2, conn=conn)
+        # If an exception occurs here, all previous operations will be rolled back
+
+    # Execute the transactional function (Success Case)
+    try:
+        await transactional_operation(queue)
+        print("Transaction committed successfully.")
+    except Exception as e:
+        print(f"Transaction failed: {e}")
+
+    # Read messages outside of the transaction
+    read_message1 = await queue.read(test_queue)
+    read_message2 = await queue.read(test_queue)
+    print("Messages read after transaction commit:")
+    if read_message1:
+        print(f"Message 1: {read_message1.message}")
+    if read_message2:
+        print(f"Message 2: {read_message2.message}")
+
+    # Purge the queue for the failure case
+    await queue.purge(test_queue)
+
+    # Transactional operation: simulate failure
+    @transaction
+    async def transactional_operation_failure(queue: PGMQueue, conn=None):
+        await queue.send(test_queue, message1, conn=conn)
+        await queue.send(test_queue, message2, conn=conn)
+        # Simulate an error to trigger rollback
+        raise Exception("Simulated failure")
+
+    # Execute the transactional function (Failure Case)
+    try:
+        await transactional_operation_failure(queue)
+    except Exception as e:
+        print(f"Transaction failed: {e}")
+
+    # Attempt to read messages after failed transaction
+    read_message = await queue.read(test_queue)
+    if read_message:
+        print("Message read after failed transaction (should not exist):")
+        print(read_message.message)
+    else:
+        print("No messages found after transaction rollback.")
+
+    # Clean up
+    await queue.drop_queue(test_queue)
+    await queue.pool.close()
+
+
+# Run the main function
+if __name__ == "__main__":
+    asyncio.run(main())

tembo-pgmq-python/example/example_app_sync.py

@@ -0,0 +1,108 @@
+from tembo_pgmq_python.queue import PGMQueue
+from tembo_pgmq_python.decorators import transaction
+
+queue = PGMQueue(
+    host="localhost",
+    port="5432",
+    username="postgres",
+    password="postgres",
+    database="postgres",
+    verbose=True,
+    log_filename="pgmq_sync.log",
+)
+
+test_queue = "transaction_queue_sync"
+
+# Clean up if the queue already exists
+queues = queue.list_queues()
+if test_queue in queues:
+    queue.drop_queue(test_queue)  # Pass queue name as positional argument
+queue.create_queue(test_queue)  # Pass queue name as positional argument
+
+# Example messages
+messages = [
+    {"id": 1, "content": "First message"},
+    {"id": 2, "content": "Second message"},
+    {"id": 3, "content": "Third message"},
+]
+
+
+# Transactional operation: send multiple messages and perform additional operations within a transaction
+@transaction
+def transactional_operations(queue, conn=None):
+    # Send multiple messages
+    msg_ids = queue.send_batch(
+        test_queue,  # Positional argument
+        messages=messages,
+        conn=conn,
+    )
+    print(f"Messages sent with IDs: {msg_ids}")
+
+    # Read messages within the transaction
+    internal_messages = queue.read_batch(
+        test_queue,  # Positional argument
+        batch_size=10,
+        conn=conn,
+    )
+    print(f"Messages read within transaction: {internal_messages}")
+
+    # Perform additional operations
+    if msg_ids:
+        queue.delete(
+            test_queue,  # Positional argument
+            msg_id=msg_ids[0],
+            conn=conn,
+        )
+        print(f"Deleted message ID: {msg_ids[0]} within transaction")
+
+
+# Execute the transactional operations (Success Case)
+print("=== Executing Transactional Operations (Success Case) ===")
+try:
+    transactional_operations(queue)
+except Exception as e:
+    print(f"Transaction failed: {e}")
+
+# Read messages after transaction commit
+external_messages = queue.read_batch(test_queue, batch_size=10)
+print("Messages read after transaction commit:")
+for msg in external_messages:
+    print(f"ID: {msg.msg_id}, Content: {msg.message}")
+
+# Clean up for failure case
+queue.purge(test_queue)
+
+
+# Transactional operation: simulate failure
+@transaction
+def transactional_operations_failure(queue, conn=None):
+    # Send multiple messages
+    msg_ids = queue.send_batch(
+        test_queue,  # Positional argument
+        messages=messages,
+        conn=conn,
+    )
+    print(f"Messages sent with IDs: {msg_ids}")
+
+    # Simulate an error to trigger a rollback
+    raise Exception("Simulated failure in transactional operations")
+
+
+# Execute the transactional operations (Failure Case)
+print("\n=== Executing Transactional Operations (Failure Case) ===")
+try:
+    transactional_operations_failure(queue)
+except Exception as e:
+    print(f"Transaction failed: {e}")
+
+# Read messages after transaction rollback
+external_messages = queue.read_batch(test_queue, batch_size=10)
+if external_messages:
+    print("Messages read after transaction rollback:")
+    for msg in external_messages:
+        print(f"ID: {msg.msg_id}, Content: {msg.message}")
+else:
+    print("No messages found after transaction rollback.")
+
+# Clean up
+queue.drop_queue(test_queue)

tembo-pgmq-python/tembo_pgmq_python/__init__.py

@@ -1,3 +1,4 @@
 from tembo_pgmq_python.queue import Message, PGMQueue  # type: ignore
+from tembo_pgmq_python.decorators import transaction, async_transaction
 
-__all__ = ["Message", "PGMQueue"]
+__all__ = ["Message", "PGMQueue", "transaction", "async_transaction"]