GCP Pub/Sub Publisher and Subscriber

Problem Statement

Cloud-native applications frequently need to decouple producers from consumers — a web frontend might publish order events, while a fleet of workers processes them asynchronously. Google Cloud Pub/Sub is the managed message bus for exactly this pattern, but the google-cloud-pubsub client library has enough moving parts (flow control, ack deadlines, streaming pull, futures, graceful shutdown) that a naive implementation often drops messages, leaks threads, or hangs on Ctrl+C.

This script provides a single, production-ready CLI that can either publish messages to a topic (from stdin, a file, or a generated stream) or subscribe to a topic/subscription and stream received messages to an NDJSON file for downstream processing. It handles batching, back-pressure, graceful shutdown, retries, and structured logging so you can use it as a debugging tool, a load-test generator, or the seed of a real consumer service.

How It Works

The script is organized around a PubSubClient class that wraps the PublisherClient and SubscriberClient from the official library. On publish, messages are batched by the client library (default 1 000 messages / 1 MB / 10 ms) and each publish() call returns a Future; we collect those futures and wait on them in bulk so throughput stays high while still surfacing per-message errors. On subscribe, we use the synchronous subscribe() callback API with a bounded flow-control window (max_messages) so the process doesn't pull more messages than it can write. Each message is written as a single NDJSON line containing ack_id, message_id, publish_time, attributes, and the base64-decoded data, then acked (or nacked if writing fails).

Shutdown is handled via a threading.Event set by SIGINT/SIGTERM handlers. The subscriber callback checks the event and calls future.cancel() to stop the streaming pull cleanly, while the publisher drains in-flight futures before closing. All log output goes to stderr so stdout (or the NDJSON file) stays clean for piping.

Full Script

#!/usr/bin/env python3
"""
GCP Pub/Sub Publisher and Subscriber

Publishes messages to a Pub/Sub topic or subscribes to a subscription and
streams received messages as NDJSON to stdout or a file.

Examples:
    # Publish lines from a file
    python3 pubsub_client.py --project my-proj --topic orders \
        --mode publish --input orders.txt

    # Subscribe and write to a file
    python3 pubsub_client.py --project my-proj --topic orders \
        --mode subscribe --output events.ndjson
"""

from __future__ import annotations

import argparse
import base64
import json
import logging
import os
import signal
import sys
import threading
import time
from concurrent.futures import TimeoutError as FuturesTimeoutError
from pathlib import Path
from typing import Iterable, Iterator, Optional, TextIO

from google.api_core import exceptions as gcp_exceptions
from google.cloud import pubsub_v1
from google.cloud.pubsub_v1.types import FlowControl

# ---------------------------------------------------------------------------
# Configuration (env vars with sensible defaults)
# ---------------------------------------------------------------------------
LOG_LEVEL = os.environ.get("PUBSUB_LOG_LEVEL", "INFO").upper()
BATCH_MAX_MESSAGES = int(os.environ.get("PUBSUB_BATCH_MAX_MESSAGES", "1000"))
BATCH_MAX_BYTES = int(os.environ.get("PUBSUB_BATCH_MAX_BYTES", str(1024 * 1024)))
BATCH_MAX_LATENCY = float(os.environ.get("PUBSUB_BATCH_MAX_LATENCY", "0.01"))
FLOW_MAX_MESSAGES = int(os.environ.get("PUBSUB_FLOW_MAX_MESSAGES", "100"))
FLOW_MAX_BYTES = int(os.environ.get("PUBSUB_FLOW_MAX_BYTES", str(10 * 1024 * 1024)))
ACK_DEADLINE = int(os.environ.get("PUBSUB_ACK_DEADLINE", "60"))
PUBLISH_TIMEOUT = float(os.environ.get("PUBSUB_PUBLISH_TIMEOUT", "30.0"))

logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s %(levelname)-8s %(name)s %(message)s",
    stream=sys.stderr,
)
log = logging.getLogger("pubsub_client")


# ---------------------------------------------------------------------------
# Client wrapper
# ---------------------------------------------------------------------------
class PubSubClient:
    """Thin wrapper around google-cloud-pubsub publisher/subscriber clients."""

    def __init__(self, project: str) -> None:
        self.project = project
        self._stop = threading.Event()
        self._publisher: Optional[pubsub_v1.PublisherClient] = None
        self._subscriber: Optional[pubsub_v1.SubscriberClient] = None

    # -- lifecycle ---------------------------------------------------------
    def install_signal_handlers(self) -> None:
        """Convert SIGINT/SIGTERM into a graceful stop request."""

        def _handler(signum: int, _frame: object) -> None:
            if self._stop.is_set():
                log.warning("Second signal (%s) received — forcing exit", signum)
                sys.exit(1)
            log.info("Signal %s received — shutting down gracefully", signum)
            self._stop.set()

        signal.signal(signal.SIGINT, _handler)
        signal.signal(signal.SIGTERM, _handler)

    @property
    def stopped(self) -> bool:
        return self._stop.is_set()

    # -- publish -----------------------------------------------------------
    def publish(self, topic: str, messages: Iterable[dict]) -> int:
        """Publish an iterable of {'data': str|bytes, 'attributes': dict} dicts.

        Returns the number of successfully published messages.
        """
        self._publisher = pubsub_v1.PublisherClient(
            batch_settings=pubsub_v1.types.BatchSettings(
                max_messages=BATCH_MAX_MESSAGES,
                max_bytes=BATCH_MAX_BYTES,
                max_latency=BATCH_MAX_LATENCY,
            )
        )
        topic_path = self._publisher.topic_path(self.project, topic)
        futures: list[tuple[str, object]] = []
        sent = 0

        for raw in messages:
            if self.stopped:
                break
            data = raw.get("data", "")
            if isinstance(data, str):
                data_bytes = data.encode("utf-8")
            else:
                data_bytes = bytes(data)
            attributes = {str(k): str(v) for k, v in (raw.get("attributes") or {}).items()}
            try:
                future = self._publisher.publish(topic_path, data_bytes, **attributes)
                futures.append((data_bytes.decode("utf-8", "replace"), future))
                sent += 1
            except gcp_exceptions.GoogleAPICallError as exc:
                log.error("Publish call failed: %s", exc)

        # Wait on all outstanding futures with a global deadline.
        ok, failed = 0, 0
        deadline = time.monotonic() + PUBLISH_TIMEOUT
        for preview, fut in futures:
            remaining = max(0.0, deadline - time.monotonic())
            try:
                fut.result(timeout=remaining)
                ok += 1
            except FuturesTimeoutError:
                failed += 1
                log.error("Publish timed out for message preview=%r", preview[:80])
            except gcp_exceptions.GoogleAPICallError as exc:
                failed += 1
                log.error("Publish failed for preview=%r: %s", preview[:80], exc)

        log.info("Publish complete: submitted=%d ok=%d failed=%d", sent, ok, failed)
        return ok

    # -- subscribe ---------------------------------------------------------
    def subscribe(
        self,
        topic: str,
        subscription: Optional[str],
        out: TextIO,
        max_messages: Optional[int] = None,
    ) -> int:
        """Subscribe and stream messages as NDJSON to `out`.

        Returns the number of messages written.
        """
        self._subscriber = pubsub_v1.SubscriberClient()
        sub_id = subscription or f"{topic}-sub"
        sub_path = self._subscriber.subscription_path(self.project, sub_id)

        flow = FlowControl(
            max_messages=FLOW_MAX_MESSAGES,
            max_bytes=FLOW_MAX_BYTES,
        )
        counter = {"written": 0}
        lock = threading.Lock()

        def _callback(message: pubsub_v1.subscriber.message.Message) -> None:
            if self.stopped:
                message.nack()
                return
            try:
                record = {
                    "message_id": message.message_id,
                    "ack_id": message.ack_id,
                    "publish_time": message.publish_time.isoformat()
                    if message.publish_time
                    else None,
                    "attributes": dict(message.attributes),
                    "data": message.data.decode("utf-8", "replace"),
                }
                out.write(json.dumps(record, ensure_ascii=False) + "\n")
                out.flush()
                message.ack()
                with lock:
                    counter["written"] += 1
                    if max_messages and counter["written"] >= max_messages:
                        self._stop.set()
            except Exception as exc:  # noqa: BLE001 — we must ack/nack every msg
                log.exception("Failed to process message %s: %s", message.message_id, exc)
                message.nack()

        streaming_pull = self._subscriber.subscribe(
            sub_path,
            callback=_callback,
            flow_control=flow,
            await_callbacks_on_shutdown=True,
        )
        log.info("Listening on %s (Ctrl+C to stop)", sub_path)

        # Poll the stop event while the pull runs in a background thread.
        try:
            while not self._stop.wait(timeout=1.0):
                pass
        finally:
            streaming_pull.cancel()
            try:
                streaming_pull.result(timeout=30)
            except (FuturesTimeoutError, gcp_exceptions.GoogleAPICallError):
                pass
            self._subscriber.close()

        log.info("Subscribe complete: wrote=%d", counter["written"])
        return counter["written"]


# ---------------------------------------------------------------------------
# Input helpers
# ---------------------------------------------------------------------------
def read_messages(path: Optional[str]) -> Iterator[dict]:
    """Yield message dicts from a file or stdin.

    Each line may be either plain text (published as-is) or a JSON object
    with keys `data` and optional `attributes`.
    """
    stream: TextIO
    if path and path != "-":
        stream = Path(path).open("r", encoding="utf-8")
        close = True
    else:
        stream = sys.stdin
        close = False
    try:
        for line in stream:
            line = line.rstrip("\n")
            if not line:
                continue
            if line.lstrip().startswith("{"):
                try:
                    yield json.loads(line)
                    continue
                except json.JSONDecodeError:
                    pass  # fall through and treat as raw text
            yield {"data": line}
    finally:
        if close:
            stream.close()


def open_output(path: Optional[str]) -> tuple[TextIO, bool]:
    if path and path != "-":
        return Path(path).open("a", encoding="utf-8"), True
    return sys.stdout, False


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Publish to or subscribe from GCP Pub/Sub.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--project", required=True, help="GCP project ID")
    p.add_argument("--topic", required=True, help="Pub/Sub topic name")
    p.add_argument(
        "--subscription",
        default=None,
        help="Subscription name (default: <topic>-sub)",
    )
    p.add_argument(
        "--mode",
        choices=("publish", "subscribe"),
        required=True,
        help="Operation to perform",
    )
    p.add_argument(
        "--input",
        default="-",
        help="Publish mode: file with messages, or '-' for stdin",
    )
    p.add_argument(
        "--output",
        default="-",
        help="Subscribe mode: NDJSON output file, or '-' for stdout",
    )
    p.add_argument(
        "--count",
        type=int,
        default=0,
        help="Subscribe mode: stop after N messages (0 = unlimited)",
    )
    return p


def main(argv: Optional[list[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    client = PubSubClient(args.project)
    client.install_signal_handlers()

    try:
        if args.mode == "publish":
            count = client.publish(args.topic, read_messages(args.input))
            return 0 if count >= 0 else 1
        else:
            out, close_out = open_output(args.output)
            try:
                client.subscribe(
                    args.topic,
                    args.subscription,
                    out,
                    max_messages=args.count or None,
                )
            finally:
                if close_out:
                    out.close()
            return 0
    except gcp_exceptions.NotFound as exc:
        log.error("Resource not found: %s", exc)
        return 2
    except gcp_exceptions.PermissionDenied as exc:
        log.error("Permission denied: %s", exc)
        return 3
    except gcp_exceptions.GoogleAPICallError as exc:
        log.error("Pub/Sub API error: %s", exc)
        return 4
    except KeyboardInterrupt:
        log.info("Interrupted by user")
        return 130


if __name__ == "__main__":
    sys.exit(main())

Configuration

Setting	Env var	Default	Purpose
Log level	`PUBSUB_LOG_LEVEL`	`INFO`	Standard `logging` level. All logs go to stderr.
Batch max messages	`PUBSUB_BATCH_MAX_MESSAGES`	`1000`	Max messages per publish batch.
Batch max bytes	`PUBSUB_BATCH_MAX_BYTES`	`1048576`	Max bytes per publish batch.
Batch max latency	`PUBSUB_BATCH_MAX_LATENCY`	`0.01`	Seconds to wait before flushing a partial batch.
Flow max messages	`PUBSUB_FLOW_MAX_MESSAGES`	`100`	Subscriber back-pressure: max unacked messages.
Flow max bytes	`PUBSUB_FLOW_MAX_BYTES`	`10485760`	Subscriber back-pressure: max unacked bytes.
Ack deadline	`PUBSUB_ACK_DEADLINE`	`60`	Informational; the client auto-extends acks.
Publish timeout	`PUBSUB_PUBLISH_TIMEOUT`	`30.0`	Seconds to wait for all publish futures.

Authentication: the script uses Application Default Credentials. Run gcloud auth application-default login locally, or set GOOGLE_APPLICATION_CREDENTIALS=/path/to/sa.json for a service account. The identity needs roles/pubsub.publisher on the topic and roles/pubsub.subscriber on the subscription.

Usage Examples

1. Publish every line of a file as a message:

$ python3 pubsub_client.py --project my-proj --topic orders \
    --mode publish --input orders.txt
2024-06-01 10:00:00 INFO     pubsub_client Publish complete: submitted=42 ok=42 failed=0

2. Subscribe and stream to NDJSON (stop after 5 messages):

$ python3 pubsub_client.py --project my-proj --topic orders \
    --mode subscribe --output events.ndjson --count 5
2024-06-01 10:00:05 INFO     pubsub_client Listening on projects/my-proj/subscriptions/orders-sub (Ctrl+C to stop)
2024-06-01 10:00:07 INFO     pubsub_client Subscribe complete: wrote=5
$ head -n1 events.ndjson
{"message_id":"123","ack_id":"...","publish_time":"2024-06-01T10:00:06.123Z","attributes":{"source":"web"},"data":"{\"order_id\":42}"}

3. Publish JSON-encoded messages from stdin with attributes:

$ echo '{"data":"hello","attributes":{"env":"prod"}}' | \
    python3 pubsub_client.py --project my-proj --topic orders --mode publish
2024-06-01 10:01:00 INFO     pubsub_client Publish complete: submitted=1 ok=1 failed=0

Error Handling

NotFound (exit 2) — topic or subscription doesn't exist. Create them with gcloud pubsub topics create orders / gcloud pubsub subscriptions create orders-sub --topic orders.
PermissionDenied (exit 3) — missing IAM role or wrong credentials. Verify with gcloud auth application-default print-access-token.
GoogleAPICallError (exit 4) — transient API errors, quota issues, or invalid arguments. The client library retries idempotent publishes automatically; we surface the final failure.
Publish timeouts — if a future doesn't resolve within PUBSUB_PUBLISH_TIMEOUT, we log it and continue. ok/failed counts make partial failures visible.
Callback exceptions — any error while writing or decoding a message triggers message.nack(), so Pub/Sub redelivers it (up to the subscription's dead_letter_policy).
Signal handling — SIGINT/SIGTERM set a stop flag; the subscriber cancels the streaming pull and waits up to 30 s for in-flight callbacks. A second signal forces exit with code 1.

Extending Ideas

Dead-letter routing: add --dead-letter-topic and, on repeated callback failures, publish the raw message there with an error attribute before acking.
Avro/Protobuf schemas: swap the JSON body for a schema-registry-backed encoder and validate on publish; Pub/Sub supports schema enforcement at the topic level.
Async publisher: replace the sync publisher with pubsub_v1.PublisherClient wrapped in an asyncio loop (or use google-cloud-pubsub's async client) for very high throughput with fewer threads.
Prometheus metrics: expose pubsub_messages_published_total, pubsub_messages_acked_total, and pubsub_publish_latency_seconds via prometheus_client.start_http_server() for Kubernetes-friendly observability.
Exactly-once delivery: enable --enable-exactly-once-delivery on the subscription and pass enable_exactly_once_delivery=True to subscribe(); note the throughput trade-off.