gpmq 0.4.0__py3-none-any.whl

gpmq/__init__.py

@@ -0,0 +1,110 @@
+"""gpmq - GPMQ (General Purpose Message Queue) - 面向个人项目的 Python 分布式消息队列包."""
+
+# Process policy
+from gpmq.process import use_spawn
+
+# Logging
+from gpmq.log import get_logger
+
+# Config
+from gpmq.config import (
+    GPMQConfig,
+    HandlerConfig,
+    PublisherConfig,
+    SubscriberConfig,
+    set_config_manager,
+    get_config_manager,
+    get_subscriber,
+    get_client,
+)
+
+# Models
+from gpmq.models import (
+    Message,
+    ProcessResult,
+    ProcessStatus,
+    WorkerState,
+    WorkerStatus,
+    WorkerInfo,
+    ConsumerGroupWorkers,
+)
+
+# Exceptions
+from gpmq.exceptions import (
+    DuplicateMessageTypeError,
+    GPMQError,
+    HandlerTypeMismatchError,
+    InvalidMessageError,
+    InvalidTargetSubscriberError,
+    RedisConnectionError,
+)
+
+# Storage
+from gpmq.storage import AuditStore, RedisAdapter
+
+# Publisher
+from gpmq.publisher import Publisher, ResultHandler
+
+# Client
+from gpmq.client import GPMQClient
+
+# Progress
+from gpmq.progress import GPMQProgress
+
+# Async Batch
+from gpmq.async_batch import GPMQAsyncBatchHandler
+
+# Subscriber
+from gpmq.subscriber import worker_context_processor
+from gpmq.subscriber.handler import HandlerRegistry, message_handler
+from gpmq.subscriber.subscriber import Subscriber
+
+__version__ = "0.4.0"
+
+__all__ = [
+    # Process policy
+    "use_spawn",
+    # Logging
+    "get_logger",
+    # Config
+    "GPMQConfig",
+    "HandlerConfig",
+    "PublisherConfig",
+    "SubscriberConfig",
+    "set_config_manager",
+    "get_config_manager",
+    "get_subscriber",
+    "get_client",
+    # Models
+    "Message",
+    "ProcessResult",
+    "ProcessStatus",
+    "WorkerState",
+    "WorkerStatus",
+    "WorkerInfo",
+    "ConsumerGroupWorkers",
+    # Exceptions
+    "GPMQError",
+    "RedisConnectionError",
+    "InvalidMessageError",
+    "InvalidTargetSubscriberError",
+    "DuplicateMessageTypeError",
+    "HandlerTypeMismatchError",
+    # Storage
+    "AuditStore",
+    "RedisAdapter",
+    # Publisher
+    "Publisher",
+    "ResultHandler",
+    # Client
+    "GPMQClient",
+    # Progress
+    "GPMQProgress",
+    # Async Batch
+    "GPMQAsyncBatchHandler",
+    # Subscriber
+    "message_handler",
+    "worker_context_processor",
+    "HandlerRegistry",
+    "Subscriber",
+]

gpmq/async_batch.py

@@ -0,0 +1,173 @@
+"""GPMQAsyncBatchHandler - batch async message processing."""
+
+import time
+from typing import Optional, Union
+
+from gpmq.client import GPMQClient
+from gpmq.models.result import ProcessResult, ProcessStatus
+from gpmq.progress import GPMQProgress
+from gpmq.publisher.result_handler import ResultHandler
+
+
+class GPMQAsyncBatchHandler:
+    """Batch async message processor.
+
+    Collects messages from send_message() calls and provides
+    unified wait_all() with sliding window concurrency control,
+    progress reporting and timeout control.
+    """
+
+    def __init__(
+        self,
+        client: GPMQClient,
+        progress: Optional[GPMQProgress] = None,
+        timeout: Optional[float] = None,
+        queue_size: int = 32,
+    ) -> None:
+        if queue_size < 1:
+            raise ValueError(f"queue_size must be >= 1, got {queue_size}")
+        self._client = client
+        self._progress = progress
+        self._timeout = timeout
+        self._queue_size = queue_size
+        self._pending_messages: list[
+            tuple[str, dict, Optional[str], Optional[Union[str, list[str]]]]
+        ] = []
+        self._all_results: list[dict[str, ProcessResult]] = []
+        self._waited = False
+
+    def send_message(
+        self,
+        message_type: str,
+        payload: dict,
+        correlation_id: Optional[str] = None,
+        target_subscribers: Optional[Union[str, list[str]]] = None,
+    ) -> None:
+        """Store message for deferred publishing.
+
+        Args:
+            message_type: Type of message to publish.
+            payload: Message payload dict.
+            correlation_id: Optional correlation ID.
+            target_subscribers: Optional subscriber name or list of names.
+                Stored as-is and forwarded to ``client.publish_async`` during
+                ``wait_all``; normalization and validation happen there.
+        """
+        self._pending_messages.append(
+            (message_type, payload, correlation_id, target_subscribers)
+        )
+
+    def wait_all(self) -> list[dict[str, ProcessResult]]:
+        """Wait for all messages to be sent and results collected.
+
+        Uses a sliding window: sends up to queue_size messages initially,
+        then replenishes one-for-one as results complete. This limits
+        concurrent threads to queue_size regardless of total message count.
+
+        Returns:
+            List of result dicts, one per message, in send_message() order.
+        """
+        total = len(self._pending_messages)
+        if total == 0:
+            return []
+
+        if self._waited:
+            return self._all_results
+
+        if self._progress:
+            self._progress.create(total)
+            self._progress.start()
+
+        deadline = time.monotonic() + self._timeout if self._timeout else None
+
+        # Send initial batch
+        window_size = min(self._queue_size, total)
+        active: list[tuple[int, ResultHandler]] = []  # (original_index, handler)
+        next_idx = 0
+
+        for i in range(window_size):
+            msg_type, payload, corr_id, target = self._pending_messages[i]
+            handler = self._client.publish_async(msg_type, payload, corr_id, target)
+            active.append((i, handler))
+            next_idx = i + 1
+
+        result_map: dict[int, dict[str, ProcessResult]] = {}
+
+        while True:
+            if deadline is not None and time.monotonic() >= deadline:
+                break
+
+            if not active:
+                break
+
+            # Wait on current active window
+            active_handlers = [h for _, h in active]
+            wait_results = self._client.wait_all(
+                active_handlers, timeout=1.0, allow_timeout_retry=True
+            )
+
+            # Separate completed vs still-pending
+            # With allow_timeout_retry=True, TIMEOUT in results means
+            # "subscriber hasn't responded yet" — so "no TIMEOUT" = complete
+            new_active = []
+            for (orig_idx, _handler), result in zip(active, wait_results):
+                if all(pr.status != ProcessStatus.TIMEOUT for pr in result.values()):
+                    result_map[orig_idx] = result
+                    # Refill: send next pending message
+                    if next_idx < total:
+                        msg_type, payload, corr_id, target = self._pending_messages[next_idx]
+                        try:
+                            new_handler = self._client.publish_async(
+                                msg_type, payload, corr_id, target
+                            )
+                            new_active.append((next_idx, new_handler))
+                        except Exception as exc:
+                            raise RuntimeError(
+                                f"async_batch aborted: publish_async failed for "
+                                f"message index {next_idx} "
+                                f"(type={msg_type!r}, payload={payload!r})"
+                            ) from exc
+                        next_idx += 1
+                else:
+                    new_active.append((orig_idx, _handler))
+
+            active = new_active
+
+            total_completed = len(result_map)
+            if self._progress:
+                self._progress.update(total_completed)
+
+            if not active:
+                break
+
+        # Final cleanup: PubSub unsubscribe and save TIMEOUT results to audit
+        if active:
+            active_handlers = [h for _, h in active]
+            active_indices = [idx for idx, _ in active]
+            final_results = self._client.wait_all(
+                active_handlers, timeout=1.0, allow_timeout_retry=False
+            )
+            for orig_idx, result in zip(active_indices, final_results):
+                result_map[orig_idx] = result
+
+        self._all_results = [result_map.get(i, {}) for i in range(total)]
+
+        if self._progress:
+            self._progress.finish()
+
+        self._waited = True
+        return self._all_results
+
+    def get_all_result(self) -> list[dict[str, ProcessResult]]:
+        """Get all message processing results.
+
+        Available after wait_all() or context manager exit.
+        Returns empty list if called before wait_all().
+        """
+        return self._all_results
+
+    def __enter__(self) -> "GPMQAsyncBatchHandler":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.wait_all()

gpmq/cli.py

@@ -0,0 +1,432 @@
+"""GPMQ CLI - Command-line interface for GPMQ message queue management."""
+
+import json
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Any
+
+import click
+from gpconfig import GPConfigManager
+
+from gpmq.config import (
+    GPMQConfig,
+    get_config_manager,
+    get_subscriber,
+    set_config_manager,
+)
+from gpmq.models.result import ProcessStatus
+from gpmq.storage.audit_store import AuditStore
+
+
+def _init_cfg_mgr(config_path: str | None) -> GPConfigManager:
+    """Initialize GPConfigManager and set it as global config manager.
+
+    Args:
+        config_path: Optional path to config folder containing global_env.yaml.
+            If not provided, GPConfigManager will search via env variable
+            GPMQ_CLI_CFG_PATH or user home directory.
+
+    Returns:
+        Initialized GPConfigManager instance.
+    """
+    cfg_mgr = GPConfigManager(
+        project_name="gpmq_cli",
+        cfg_folder=config_path,
+    )
+    set_config_manager(cfg_mgr)
+    return cfg_mgr
+
+
+def _lazy_init_cfg_mgr(ctx: click.Context, param: click.Parameter, value: str | None) -> None:
+    """Callback for --config option on subcommands: init GPConfigManager if not done."""
+    if get_config_manager() is not None:
+        return
+    _init_cfg_mgr(value)
+
+
+def _config_option(func):
+    """Decorator that adds --config option to a subcommand for lazy init."""
+    return click.option(
+        "--config",
+        default=None,
+        expose_value=False,
+        callback=_lazy_init_cfg_mgr,
+        help="Configuration folder path (must contain global_env.yaml). "
+        "If not provided, searches GPMQ_CLI_CFG_PATH env var or user home.",
+    )(func)
+
+
+def _format_output(data: Any, fmt: str) -> str:
+    """Format data for CLI output.
+
+    Args:
+        data: Data to format (list of dicts or dict).
+        fmt: Output format - "json" or "table".
+
+    Returns:
+        Formatted string.
+    """
+    if fmt == "json":
+        return json.dumps(data, indent=2, default=str)
+
+    if fmt == "table":
+        if not data:
+            return ""
+
+        # Handle dict (stats) vs list of dicts (query results)
+        if isinstance(data, dict):
+            lines = []
+            for key, value in data.items():
+                if isinstance(value, dict):
+                    lines.append(f"{key}:")
+                    for k, v in value.items():
+                        lines.append(f"  {k}: {v}")
+                else:
+                    lines.append(f"{key}: {value}")
+            return "\n".join(lines)
+
+        # List of dicts - column-aligned table
+        if isinstance(data, list) and len(data) > 0:
+            columns = list(data[0].keys())
+            # Calculate column widths
+            widths = {}
+            for col in columns:
+                widths[col] = len(str(col))
+                for row in data:
+                    widths[col] = max(widths[col], len(str(row.get(col, ""))))
+
+            # Build header
+            header = "  ".join(str(col).ljust(widths[col]) for col in columns)
+            separator = "  ".join("-" * widths[col] for col in columns)
+            # Build rows
+            rows = []
+            for row in data:
+                line = "  ".join(
+                    str(row.get(col, "")).ljust(widths[col]) for col in columns
+                )
+                rows.append(line)
+
+            return "\n".join([header, separator] + rows)
+
+        return ""
+
+    raise ValueError(f"Unsupported output format: {fmt!r} (expected 'json' or 'table')")
+
+
+def _resolve_audit_config(cfg_path: str) -> tuple[str, bool]:
+    """Resolve audit_db_path and enable_audit for a given component config path.
+
+    Merges common config with the specific component config, following the
+    same pattern as get_subscriber() and get_client().
+
+    Args:
+        cfg_path: Dot-notation config path, e.g. "gpmq.subscribers.data_loader"
+            or "gpmq.publisher.main".
+
+    Returns:
+        Tuple of (audit_db_path, enable_audit).
+    """
+    cfg_mgr = get_config_manager()
+    common_dict = cfg_mgr.get_config(GPMQConfig.default_cfg_path).model_dump()
+    specific = cfg_mgr.get_config(cfg_path)
+    specific_set = specific.model_dump(exclude_unset=True)
+    merged = {**common_dict, **specific_set}
+
+    return merged.get("audit_db_path", "gpmq_audit.db"), merged.get("enable_audit", True)
+
+
+@contextmanager
+def _get_audit_store(db_path: str):
+    """Context manager that creates and closes an AuditStore.
+
+    Args:
+        db_path: Path to the SQLite audit database file.
+
+    Yields:
+        AuditStore instance.
+    """
+    store = AuditStore(db_path)
+    store.connect()
+    try:
+        yield store
+    finally:
+        store.close()
+
+
+def _parse_iso_datetime(value: str) -> datetime:
+    """Parse ISO format datetime string.
+
+    Args:
+        value: ISO datetime string (e.g., '2026-01-01T00:00:00').
+
+    Returns:
+        datetime object.
+
+    Raises:
+        click.BadParameter: If the string cannot be parsed.
+    """
+    try:
+        return datetime.fromisoformat(value)
+    except ValueError:
+        raise click.BadParameter(
+            f"Invalid datetime format: {value}. Use ISO format like 2026-01-01T00:00:00"
+        )
+
+
+@contextmanager
+def _create_redis_adapter():
+    """Context manager that creates, connects, and disconnects a RedisAdapter.
+
+    Uses the global GPConfigManager to load configuration.
+
+    Yields:
+        Connected RedisAdapter instance.
+    """
+    from gpmq.storage.redis_adapter import RedisAdapter
+
+    cfg_mgr = get_config_manager()
+    config = cfg_mgr.get_config(GPMQConfig.default_cfg_path)
+    adapter = RedisAdapter(config)
+    adapter.connect()
+    try:
+        yield adapter
+    finally:
+        adapter.disconnect()
+
+
+@click.group()
+@click.option(
+    "--config",
+    default=None,
+    help="Configuration folder path (must contain global_env.yaml). "
+    "If not provided, searches GPMQ_CLI_CFG_PATH env var or user home.",
+)
+@click.pass_context
+def cli(ctx: click.Context, config: str | None) -> None:
+    """GPMQ - General Purpose Message Queue CLI."""
+    from gpmq import use_spawn
+    use_spawn()
+    ctx.ensure_object(dict)
+    if config is not None:
+        _init_cfg_mgr(config)
+
+
+@cli.group()
+@click.pass_context
+def audit(ctx: click.Context) -> None:
+    """Audit message and result records."""
+    pass
+
+
+@audit.command("query")
+@_config_option
+@click.argument("component", metavar="COMPONENT_NAME")
+@click.option("--type", "message_type", default=None, help="Filter by message type.")
+@click.option(
+    "--status",
+    "status_filter",
+    default=None,
+    help="Filter by processing status (success/failure/exception/timeout).",
+)
+@click.option("--correlation-id", default=None, help="Filter by correlation ID.")
+@click.option("--from", "from_time", default=None, help="Start time (ISO format).")
+@click.option("--to", "to_time", default=None, help="End time (ISO format).")
+@click.option("--limit", default=100, type=int, help="Max results (default: 100).")
+@click.option(
+    "--format",
+    "fmt",
+    type=click.Choice(["json", "table"]),
+    default="table",
+    help="Output format.",
+)
+@click.pass_context
+def audit_query(
+    ctx: click.Context,
+    component: str,
+    message_type: str | None,
+    status_filter: str | None,
+    correlation_id: str | None,
+    from_time: str | None,
+    to_time: str | None,
+    limit: int,
+    fmt: str,
+) -> None:
+    """Query audit records with optional filters."""
+    audit_db_path, enable_audit = _resolve_audit_config(component)
+    if not enable_audit:
+        click.echo(f"Audit store is disabled for '{component}'.")
+        return
+
+    click.echo(f"Audit store: {audit_db_path}")
+
+    with _get_audit_store(audit_db_path) as store:
+        results = []
+
+        if correlation_id:
+            results = store.query_by_correlation_id(correlation_id)
+        elif status_filter:
+            try:
+                status_enum = ProcessStatus(status_filter)
+            except ValueError:
+                valid = ", ".join(s.value for s in ProcessStatus)
+                raise click.BadParameter(
+                    f"Invalid status: {status_filter}. Valid values: {valid}"
+                )
+            results = store.query_by_status(status_enum, limit=limit)
+        elif message_type:
+            start = _parse_iso_datetime(from_time) if from_time else None
+            end = _parse_iso_datetime(to_time) if to_time else None
+            results = store.query_by_type(
+                message_type, start_time=start, end_time=end, limit=limit
+            )
+        elif from_time and to_time:
+            start = _parse_iso_datetime(from_time)
+            end = _parse_iso_datetime(to_time)
+            results = store.query_by_time_range(start, end, limit=limit)
+        else:
+            # Fallback: no filters, get recent
+            click.echo(
+                "No filters specified. Use --type, --status, --correlation-id, or --from/--to."
+            )
+            return
+
+        click.echo(_format_output(results, fmt))
+
+
+@audit.command("cleanup")
+@_config_option
+@click.argument("component", metavar="COMPONENT_NAME")
+@click.option(
+    "--before",
+    required=True,
+    help="Delete records before this timestamp (ISO format).",
+)
+@click.option("--type", "message_type", default=None, help="Delete by message type.")
+@click.option(
+    "--format",
+    "fmt",
+    type=click.Choice(["json", "table"]),
+    default="table",
+    help="Output format.",
+)
+@click.pass_context
+def audit_cleanup(
+    ctx: click.Context, component: str, before: str, message_type: str | None, fmt: str
+) -> None:
+    """Delete old audit records."""
+    audit_db_path, enable_audit = _resolve_audit_config(component)
+    if not enable_audit:
+        click.echo(f"Audit store is disabled for '{component}'.")
+        return
+
+    click.echo(f"Audit store: {audit_db_path}")
+    before_dt = _parse_iso_datetime(before)
+
+    with _get_audit_store(audit_db_path) as store:
+        if message_type:
+            count = store.cleanup_by_type_before(message_type, before_dt)
+        else:
+            count = store.cleanup_by_time(before_dt)
+
+        result = {"deleted_count": count}
+        click.echo(_format_output(result, fmt))
+
+
+@audit.command("stats")
+@_config_option
+@click.argument("component", metavar="COMPONENT_NAME")
+@click.option(
+    "--format",
+    "fmt",
+    type=click.Choice(["json", "table"]),
+    default="table",
+    help="Output format.",
+)
+@click.pass_context
+def audit_stats(ctx: click.Context, component: str, fmt: str) -> None:
+    """Show audit statistics."""
+    audit_db_path, enable_audit = _resolve_audit_config(component)
+    if not enable_audit:
+        click.echo(f"Audit store is disabled for '{component}'.")
+        return
+
+    click.echo(f"Audit store: {audit_db_path}")
+
+    with _get_audit_store(audit_db_path) as store:
+        stats = store.stats()
+        click.echo(_format_output(stats, fmt))
+
+
+@cli.command()
+@_config_option
+@click.option("--stream", required=True, help="Stream name to clear.")
+@click.pass_context
+def clear(ctx: click.Context, stream: str) -> None:
+    """Clear a Redis Stream (requires confirmation)."""
+    with _create_redis_adapter() as adapter:
+        if click.confirm(f"Are you sure you want to clear stream '{stream}'?"):
+            adapter.clear_stream(stream)
+            click.echo(f"Stream '{stream}' cleared.")
+        else:
+            click.echo("Cancelled.")
+
+
+@cli.command()
+@_config_option
+@click.option(
+    "--format",
+    "fmt",
+    type=click.Choice(["json", "table"]),
+    default="table",
+    help="Output format.",
+)
+@click.pass_context
+def status(ctx: click.Context, fmt: str) -> None:
+    """Show system status: active subscribers and their info."""
+    with _create_redis_adapter() as adapter:
+        names = adapter.get_active_subscribers()
+        subscribers = []
+        for name in names:
+            info = adapter.get_subscriber_info(name)
+            if info:
+                entry = {"name": name}
+                entry.update(info)
+                subscribers.append(entry)
+        click.echo(_format_output(subscribers, fmt))
+
+
+def _run_worker(subscriber_name: str) -> None:
+    """Common worker startup logic.
+
+    Args:
+        subscriber_name: Dot-notation config path, e.g. "gpmq.subscribers.data_loader".
+    """
+    sub = get_subscriber(subscriber_name)
+    click.echo(f"Subscriber '{sub.name}' started. Press Ctrl+C to stop.")
+    sub.run_forever()
+    click.echo("Subscriber stopped.")
+
+
+@cli.command()
+@_config_option
+@click.argument("subscriber_name")
+@click.pass_context
+def worker(ctx: click.Context, subscriber_name: str) -> None:
+    """Start a GPMQ subscriber worker process."""
+    _run_worker(subscriber_name)
+
+
+@click.command()
+@click.argument("subscriber_name")
+@click.option(
+    "--config",
+    default=None,
+    help="Configuration folder path (must contain global_env.yaml). "
+    "If not provided, searches GPMQ_CLI_CFG_PATH env var or user home.",
+)
+def gpmq_worker(subscriber_name: str, config: str | None) -> None:
+    """Start a GPMQ subscriber worker process."""
+    from gpmq import use_spawn
+    use_spawn()
+    _init_cfg_mgr(config)
+    _run_worker(subscriber_name)