Merge branch 'main' into feature/add-transaction-class

Author: jason810496

.readthedocs.yml

@@ -32,4 +32,6 @@ sphinx:
 # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
 python:
   install:
-    - requirements: doc/requirements.txt
+    - requirements: doc/requirements.txt
+    - method: pip
+      path: .

README.md

@@ -116,10 +116,3 @@ print(metrics.total_messages)
 
 Welcome to open an issue or pull request ! <br>
 See [`Development` on Online Document](https://pgmq-sqlalchemy.readthedocs.io/en/latest/) or [CONTRIBUTING.md](.github/CONTRIBUTING.md) for more information.
-
-## TODO 
-
-- [ ] Add **time-based** partition option and validation to `create_partitioned_queue` method.
-- [ ] Read(single/batch) Archive Table ( `read_archive` method )
-- [ ] Detach Archive Table ( `detach_archive` method )
-- [ ] Add `set_vt` utils method.

doc/conf.py

@@ -6,10 +6,27 @@
 import os
 import sys
 
+# Handle Python 3.9+ compatibility for tomllib
+try:
+    import tomllib
+except ModuleNotFoundError:
+    import tomli as tomllib
+
 # path setup
 sys.path.insert(0, os.path.abspath(".."))
 sys.path.insert(0, os.path.abspath("../pgmq_sqlalchemy"))
 
+# Read version from pyproject.toml
+_pyproject_path = os.path.join(os.path.dirname(__file__), "..", "pyproject.toml")
+try:
+    with open(_pyproject_path, "rb") as f:
+        _pyproject_data = tomllib.load(f)
+        _version = _pyproject_data["project"]["version"]
+except (FileNotFoundError, KeyError, OSError) as e:
+    # Fallback to a default version if pyproject.toml is missing or invalid
+    print(f"Warning: Could not read version from pyproject.toml: {e}")
+    _version = "0.0.0"
+
 extensions = [
     "sphinx_copybutton",
     "sphinx.ext.autodoc",
@@ -22,6 +39,13 @@
 project = "pgmq-sqlalchemy"
 copyright = f'2024-{time.strftime("%Y")}, the pgmq-sqlalchemy developers'
 
+# Version information
+# The short X.Y version
+_version_parts = _version.split(".")
+version = ".".join(_version_parts[:2]) if len(_version_parts) >= 2 else _version
+# The full version, including alpha/beta/rc tags
+release = _version
+
 source_suffix = {
     ".rst": "restructuredtext",
 }

doc/requirements.txt

@@ -1,4 +1,5 @@
 sphinx
 sphinx-rtd-theme
 sphinx-copybutton
-SQLAlchemy
+SQLAlchemy
+tomli; python_version < "3.11"

pgmq_sqlalchemy/queue.py

@@ -1,5 +1,6 @@
 import asyncio
-from typing import List, Optional
+import re
+from typing import List, Optional, Union
 
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
@@ -190,6 +191,35 @@ async def _create():
             queue_name, unlogged, session=session, commit=commit
         )
 
+    def _validate_partition_interval(self, interval: Union[int, str]) -> str:
+        """Validate partition interval format.
+
+        Args:
+            interval: Either an integer for numeric partitioning or a string for time-based partitioning
+                     (e.g., '1 day', '1 hour', '7 days')
+
+        Returns:
+            The validated interval as a string
+
+        Raises:
+            ValueError: If the interval format is invalid
+        """
+        if isinstance(interval, int):
+            if interval <= 0:
+                raise ValueError("Numeric partition interval must be positive")
+            return str(interval)
+
+        # Validate time-based interval format
+        # Valid PostgreSQL interval formats: '1 day', '7 days', '1 hour', '1 month', etc.
+        time_pattern = r"^\d+\s+(microsecond|millisecond|second|minute|hour|day|week|month|year)s?$"
+        if not re.match(time_pattern, interval.strip(), re.IGNORECASE):
+            raise ValueError(
+                f"Invalid time-based partition interval: '{interval}'. "
+                "Expected format: '<number> <unit>' where unit is one of: "
+                "microsecond, millisecond, second, minute, hour, day, week, month, year"
+            )
+        return interval.strip()
+
     def create_partitioned_queue(
         self,
         queue_name: str,
@@ -206,16 +236,23 @@ def create_partitioned_queue(
 
         .. code-block:: python
 
+                # Numeric partitioning (by msg_id)
                 pgmq_client.create_partitioned_queue('my_partitioned_queue', partition_interval=10000, retention_interval=100000)
 
+                # Time-based partitioning (by enqueued_at)
+                pgmq_client.create_partitioned_queue('my_time_queue', partition_interval='1 day', retention_interval='7 days')
+
         Args:
             queue_name (str): The name of the queue, should be less than 48 characters.
-            partition_interval (int): Will create a new partition every ``partition_interval`` messages.
-            retention_interval (int): The interval for retaining partitions. Any messages that have a `msg_id` less than ``max(msg_id)`` - ``retention_interval`` will be dropped.
+            partition_interval (Union[int, str]): For numeric partitioning, the number of messages per partition.
+                                                   For time-based partitioning, a PostgreSQL interval string (e.g., '1 day', '1 hour').
+            retention_interval (Union[int, str]): For numeric partitioning, messages with msg_id less than max(msg_id) - retention_interval will be dropped.
+                                                   For time-based partitioning, a PostgreSQL interval string (e.g., '7 days').
 
                 .. note::
-                    | Currently, only support for partitioning by **msg_id**.
-                    | Will add **time-based partitioning** in the future ``pgmq-sqlalchemy`` release.
+                    | Supports both **numeric** (by ``msg_id``) and **time-based** (by ``enqueued_at``) partitioning.
+                    | For time-based partitioning, use interval strings like '1 day', '1 hour', '7 days', etc.
+                    | For numeric partitioning, use integer values.
 
         .. important::
             | You must make sure that the ``pg_partman`` extension already **installed** in the Postgres.
@@ -227,6 +264,10 @@ def create_partitioned_queue(
         # check if the pg_partman extension exists before creating a partitioned queue at runtime
         self._check_pg_partman_ext()
 
+        # Validate partition intervals
+        partition_interval = self._validate_partition_interval(partition_interval)
+        retention_interval = self._validate_partition_interval(retention_interval)
+
         if self.is_async:
             if session is None:
 

tests/test_queue.py

@@ -2,6 +2,7 @@
 import pytest
 import time
 
+from sqlalchemy import text
 from sqlalchemy.exc import ProgrammingError
 from filelock import FileLock
 from pgmq_sqlalchemy import PGMQueue
@@ -286,6 +287,12 @@ def test_set_vt_to_smaller_value(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
     assert pgmq.read(queue_name) is not None
 
 
+def test_set_vt_not_exist(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
+    pgmq, queue_name = pgmq_setup_teardown
+    msg_updated = pgmq.set_vt(queue_name, 999, 20)
+    assert msg_updated is None
+
+
 def test_pop(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
     pgmq, queue_name = pgmq_setup_teardown
     msg = MSG
@@ -427,3 +434,158 @@ def test_metrics_all_queues(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
         assert queue_2.queue_length == 2
         assert queue_1.total_messages == 3
         assert queue_2.total_messages == 2
+
+
+# Tests for detach_archive method
+@pgmq_deps
+def test_detach_archive(pgmq_fixture, db_session):
+    """Test detach_archive method - detaches archive table from queue."""
+    pgmq: PGMQueue = pgmq_fixture
+    queue_name = f"test_queue_{uuid.uuid4().hex}"
+    pgmq.create_queue(queue_name)
+    msg = MSG
+    msg_id = pgmq.send(queue_name, msg)
+    pgmq.archive(queue_name, msg_id)
+
+    # Detach archive should not raise an error
+    pgmq.detach_archive(queue_name)
+
+    # Read the archive to ensure it still exists after detaching
+    archived_msg = pgmq.read_archive(queue_name)
+    assert archived_msg is not None
+    assert archived_msg.msg_id == msg_id
+
+    # Cleanup: Drop the archive and queue tables
+    # After detaching, the archive is no longer part of the extension
+    # We need to drop both tables manually by first removing them from the extension
+    if pgmq.is_async:
+
+        async def cleanup():
+            async with pgmq.session_maker() as session:
+                # Drop archive table (already detached)
+                await session.execute(
+                    text(f"DROP TABLE IF EXISTS pgmq.a_{queue_name} CASCADE;")
+                )
+                # Detach and drop queue table
+                await session.execute(
+                    text(f"ALTER EXTENSION pgmq DROP TABLE pgmq.q_{queue_name};")
+                )
+                await session.execute(
+                    text(f"DROP TABLE IF EXISTS pgmq.q_{queue_name} CASCADE;")
+                )
+                await session.commit()
+
+        pgmq.loop.run_until_complete(cleanup())
+    else:
+        with pgmq.session_maker() as session:
+            # Drop archive table (already detached)
+            session.execute(text(f"DROP TABLE IF EXISTS pgmq.a_{queue_name} CASCADE;"))
+            # Detach and drop queue table
+            session.execute(
+                text(f"ALTER EXTENSION pgmq DROP TABLE pgmq.q_{queue_name};")
+            )
+            session.execute(text(f"DROP TABLE IF EXISTS pgmq.q_{queue_name} CASCADE;"))
+            session.commit()
+
+
+# Tests for read_archive methods
+def test_read_archive(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
+    pgmq, queue_name = pgmq_setup_teardown
+    msg = MSG
+    msg_ids = pgmq.send_batch(queue_name, [msg, msg, msg])
+    pgmq.archive(queue_name, msg_ids[0])
+    archived_msg = pgmq.read_archive(queue_name)
+    assert archived_msg is not None
+    assert archived_msg.msg_id == msg_ids[0]
+    assert archived_msg.message == msg
+
+
+def test_read_archive_empty(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
+    pgmq, queue_name = pgmq_setup_teardown
+    archived_msg = pgmq.read_archive(queue_name)
+    assert archived_msg is None
+
+
+def test_read_archive_batch(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
+    pgmq, queue_name = pgmq_setup_teardown
+    msg = MSG
+    msg_ids = pgmq.send_batch(queue_name, [msg, msg, msg])
+    pgmq.archive_batch(queue_name, msg_ids)
+    archived_msgs = pgmq.read_archive_batch(queue_name, batch_size=10)
+    assert archived_msgs is not None
+    assert len(archived_msgs) == 3
+    assert [m.msg_id for m in archived_msgs] == msg_ids
+    for m in archived_msgs:
+        assert m.message == msg
+
+
+def test_read_archive_batch_empty(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
+    pgmq, queue_name = pgmq_setup_teardown
+    archived_msgs = pgmq.read_archive_batch(queue_name, batch_size=10)
+    assert archived_msgs is None
+
+
+def test_read_archive_batch_limit(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
+    pgmq, queue_name = pgmq_setup_teardown
+    msg = MSG
+    msg_ids = pgmq.send_batch(queue_name, [msg, msg, msg, msg, msg])
+    pgmq.archive_batch(queue_name, msg_ids)
+    archived_msgs = pgmq.read_archive_batch(queue_name, batch_size=3)
+    assert archived_msgs is not None
+    assert len(archived_msgs) == 3
+
+
+# Tests for time-based partitioned queues
+@pgmq_deps
+def test_create_time_based_partitioned_queue(pgmq_fixture, db_session):
+    pgmq: PGMQueue = pgmq_fixture
+    queue_name = f"test_queue_{uuid.uuid4().hex}"
+    pgmq.create_partitioned_queue(
+        queue_name, partition_interval="1 day", retention_interval="7 days"
+    )
+    assert check_queue_exists(db_session, queue_name) is True
+
+
+@pgmq_deps
+def test_create_time_based_partitioned_queue_various_intervals(
+    pgmq_fixture, db_session
+):
+    pgmq: PGMQueue = pgmq_fixture
+
+    # Test with hour
+    queue_name_hour = f"test_queue_{uuid.uuid4().hex}"
+    pgmq.create_partitioned_queue(
+        queue_name_hour, partition_interval="1 hour", retention_interval="24 hours"
+    )
+    assert check_queue_exists(db_session, queue_name_hour) is True
+
+    # Test with week
+    queue_name_week = f"test_queue_{uuid.uuid4().hex}"
+    pgmq.create_partitioned_queue(
+        queue_name_week, partition_interval="1 week", retention_interval="4 weeks"
+    )
+    assert check_queue_exists(db_session, queue_name_week) is True
+
+
+@pgmq_deps
+def test_create_partitioned_queue_invalid_time_interval(pgmq_fixture):
+    pgmq: PGMQueue = pgmq_fixture
+    queue_name = f"test_queue_{uuid.uuid4().hex}"
+    with pytest.raises(ValueError) as e:
+        pgmq.create_partitioned_queue(
+            queue_name,
+            partition_interval="invalid interval",
+            retention_interval="7 days",
+        )
+    assert "Invalid time-based partition interval" in str(e.value)
+
+
+@pgmq_deps
+def test_create_partitioned_queue_invalid_numeric_interval(pgmq_fixture):
+    pgmq: PGMQueue = pgmq_fixture
+    queue_name = f"test_queue_{uuid.uuid4().hex}"
+    with pytest.raises(ValueError) as e:
+        pgmq.create_partitioned_queue(
+            queue_name, partition_interval=-100, retention_interval=100000
+        )
+    assert "Numeric partition interval must be positive" in str(e.value)