Gram Pubsub Python

Gram Pub/Sub in Python (pystreams)

pystreams is the Python home for Pub/Sub stream consumers in the Gram monorepo. It runs alongside — not instead of — the Go gram streams process, and it consumes the same proto-declared Pub/Sub topology. This skill covers the Python runtime: how a subscriber is written, registered, and run. For declaring the topic/subscription itself (the (gcp.pubsub.v1.topic) / (gcp.pubsub.v1.subscription) proto options and infra/gen/kcc.yaml), use the gram-pubsub skill — that half is language-agnostic and shared.

Why Python exists here, and when to reach for it

The original and default stream-processing home is the Go streams component (server/cmd/gram/streams.go). Go's concurrency model and the GCP Pub/Sub Go SDK's use of it are substantially better suited to high-throughput message fan-out than the Python equivalent, async Python included. So the rule of thumb is: build new consumers in Go streams by default.

pystreams exists for the cases where Python's ecosystem is the deciding factor — overwhelmingly working with language models, transformers, NLP/ML use cases. Presidio (PII detection), prompt-injection classifier models, spaCy, and the transformer tooling these depend on are Python-first and have no comparable Go story. Reaching into that ecosystem from Go would mean reimplementing or shelling out; running the consumer in Python is the honest path. So the decision is not "which language do I prefer" but "does this consumer need a Python-only library to do its job?" If yes, pystreams; if no, Go streams.

The first real pystreams consumer is the Presidio PII scanner (pystreams/src/pystreams/risk/handler.py); ping is a heartbeat used to keep the publish→subscribe path exercised, mirroring the Go ping.

How the pieces fit

infra/proto/**.proto            topic/subscription declarations (shared, see gram-pubsub)
  └─ buf generate ─┬─ infra/gen/**        Go types  + infra/gen/kcc.yaml
                   └─ infra/gen_py/**      Python types (gram.*.v1.*_pb2)

infra/gram_infra/pubsub/        the Python publisher/subscriber library
  ├─ broker.py                  PubSubBroker / EmulatedPubSubBroker (mirrors infra/pkg/gcp)
  ├─ publisher.py               Publisher[M], pubsub_publisher_for_message
  ├─ subscriber.py              Subscriber[M], pubsub_subscriber_for_message
  └─ discover.py                proto-option → resource-name resolution

pystreams/                      the runnable service ("multi" command)
  └─ src/pystreams/cmd/multi.py the entrypoint — START HERE

infra/gram_infra (the gram-infra package) is a uv workspace member shared by pystreams; it holds both the generated Python protobufs and the hand-written Pub/Sub convenience layer. pystreams is the deployable that wires handlers onto subscriptions and runs the receive loops. Read the library as the Python counterpart of infra/pkg/gcp/ — the docstrings deliberately point back to the Go files they mirror, so the two layers stay behaviorally aligned.

Read multi.py first

pystreams/src/pystreams/cmd/multi.py is the whole service in one screen. It is a Click command (multi) that anyio.runs an async multi(...) coroutine. Walking it top to bottom teaches the runtime:

1. Configure logging (structlog) with service attributes.
2. Pick a project id: real --gcp-project-id, or a throwaway when an emulator host is set (the emulator doesn't care about the project).
3. Build a broker (_build_broker): EmulatedPubSubBroker when --pubsub-emulator-host is set (it reconciles topics/subscriptions on demand, since the emulator has no Config Connector), else PubSubBroker (assumes Config Connector already provisioned everything). This is the exact prod-vs-local split the Go layer makes.
4. Enter the broker as a context manager (with broker:) — it owns the publisher/subscriber clients and flushes + closes them on exit, including a clean Ctrl-C teardown.
5. Open one anyio task group and start, in order: a signal handler that cancels the group on SIGINT/SIGTERM, the event-loop-lag monitor, the health server (awaited so it's bound before consuming begins), then the ReceiverGroup.
6. Register receivers, one receivers.receive(...) call per subscription.
7. Flip readiness on (health_state.set_ready()) only after receivers are wired, so /readyz doesn't go green before the service can actually consume.

The structured-concurrency shape matters: everything runs as a child task of that single task group, so any fatal error or a shutdown signal tears the whole process down together for a clean restart — the Python analogue of the Go errgroup in streams.go.

Adding a subscriber — the two steps

Exactly like the Go side: write a handler, then register it.

Step 1 — write an async handler

A handler is an async callable (message, meta) -> None. The proto message type is the topic's payload; meta is gram_infra.pubsub.subscriber.MessageMetadata (id, attributes, delivery_attempt). The return/raise is the ack/nack signal, identical to Go: returning normally acks; raising nacks (triggering redelivery and eventual dead-lettering if the subscription declares a dead_letter policy). You never call ack()/nack() yourself in the callback form — the library does it from your handler's outcome.

The minimal shape is PingHandler (pystreams/src/pystreams/ping/handler.py): a class holding its dependencies, with an async def handle(self, message, meta). Register handle (the bound method), not the class.

The real reference is PresidioHandler (pystreams/src/pystreams/risk/handler.py). It demonstrates the patterns that matter for ML/NLP consumers:

• Load the model once, reuse across messages. Constructing Presidio's AnalyzerEngine loads a spaCy model — expensive. Do it in __init__, not per delivery.
• Run CPU-bound work off the event loop. This is the Python-specific hazard. anyio runs everything on one event-loop thread; a synchronous, CPU-bound call (model inference, regex over large text) that doesn't await will stall every other subscription until it returns. Wrap such work with asyncer.asyncify(...) so it runs in a worker thread and the loop stays responsive. The loop-lag monitor exists precisely to make this kind of stall visible before it becomes an outage.
• Make heavy dependencies injectable behind a Protocol. PresidioHandler depends on a narrow Analyzer protocol, so tests inject a lightweight fake instead of loading the NLP model (see pystreams/tests/test_risk_handler.py).
• Decide ack/nack deliberately, especially when no DLQ policy is declared. The PresidioScanner subscription declares no dead_letter policy, so a raised exception would nack and redeliver the same message for the full retention window — one poison input could loop for 30 days. Because this is best-effort shadow processing, the handler swallows Exception and returns (acking) rather than poisoning the subscription, logging the error type only. Catch Exception, never BaseException, so cancellation (graceful shutdown) still propagates. Match the policy to intent: raise only when you genuinely want the message retried.
• Never log PII or sensitive data. PII/security handlers log entity types and counts, request ids, and delivery attempts — never the text or the matches. An error string or traceback can echo the input, which is why the failure path logs type(exc).__name__ and not the exception message.

Step 2 — register it in multi.py

Add a receivers.receive(...) call in the marked block, mirroring the topic and subscription proto declarations:

receivers.receive(
    presidio_request_pb2.PresidioRequest,  # the TOPIC message type → fixes the handler's `message` type
    presidio_scanner_pb2.PresidioScanner,  # the SUBSCRIPTION marker message (carries the subscription option)
    PresidioHandler(logger).handle,        # your async handler callback
)

The three arguments line up one-to-one with the proto options the gram-pubsub skill describes: the topic-declaring message, the subscription-declaring marker, and the handler. ReceiverGroup.receive (pystreams/src/pystreams/cmd/receiver.py) resolves a Subscriber via pubsub_subscriber_for_message, wraps your handler in per-message tracing (deps/tracing.py — a stream.handleMessage span tagged with the topic/subscription proto names, continuing the publisher's trace via W3C context), and starts the receive loop as a child task. It's the direct analogue of mustReceive/receiverGroup in streams.go, so each subscription is one line at the call site.

Go and Python can consume the same topic through separate subscriptions: gram.ping.v2.Message has both a Processor marker (consumed by Go) and a PyProcessor marker (consumed by pystreams). If you want a Python consumer of a topic the Go side already reads, declare a new subscription marker rather than stealing the existing one. Declaration is a gram-pubsub task.