kubernetes/scripts/setup_jetstream.py

#!/usr/bin/env python3
"""
Setup JetStream streams and consumers for NATS
Run this after NATS is deployed and running

Durable consumer names are derived with the exact same logic worker.py uses
at runtime (NATS_CONSUMER + sanitized subject). WORKER_DOMAINS below must be
kept in sync with the NATS_STREAM / NATS_CONSUMER / FILTER_SUBJECT env vars
in manifests/core/workers.yaml - if they drift apart, the consumers created
here (with their max_deliver/ack_wait settings) will never be the ones the
workers actually bind to, and this script's config becomes a no-op.
"""
import asyncio
import nats
import os
import sys

WORKER_DOMAINS = [
    {"stream": "ORDERS", "consumer": "orders-worker", "subjects": [
        "api.v1.mob.orders.createorder",
    ]},
    {"stream": "DELIVERIES", "consumer": "deliveries-worker", "subjects": [
        "api.v1.deliveries.createdeliveries",
        "api.v1.deliveries.updatedelivery",
        "api.v2.deliveries.createdeliverylog",
    ]},
    {"stream": "CUSTOMERS", "consumer": "customers-worker", "retention": "work", "subjects": [
        "api.v1.mob.customers.login",
        "api.v1.mob.customers.create",
    ]},
    {"stream": "RIDER", "consumer": "rider-logs-worker", "subjects": [
        "api.v2.partners.createriderlog",
        "api.v2.partners.createbreaklog",
        "api.v2.partners.updatebreaklog",
    ]},
    {"stream": "PRODUCTS", "consumer": "products-worker", "subjects": [
        "api.v1.web.products.create",
    ]},
]


def durable_name(consumer: str, subject: str) -> str:
    """Mirrors worker.py's durable-name derivation exactly."""
    suffix = subject.replace(".", "_").replace("*", "all").replace(">", "all")
    return f"{consumer}_{suffix}"


async def setup_jetstream():
    nats_url = os.getenv("NATS_URL", "nats://nats.workolik.com:4222")
    nats_user = os.getenv("NATS_USER", "admin")
    nats_password = os.getenv("NATS_PASSWORD", "")

    try:
        print(f"Connecting to NATS at {nats_url}...")
        nc = await nats.connect(
            servers=[nats_url],
            user=nats_user,
            password=nats_password,
        )
        print("✅ Connected to NATS")

        js = nc.jetstream()

        # Create / recreate streams
        for domain in WORKER_DOMAINS:
            stream_name = domain["stream"]
            subjects = domain["subjects"]
            retention_policy = domain.get("retention", "limits")
            try:
                info = await js.stream_info(stream_name)
                print(f"⚠️  Stream '{stream_name}' already exists with subjects={info.config.subjects}, updating...")
                await js.update_stream(
                    name=stream_name,
                    subjects=subjects,
                    storage="memory",
                    retention=retention_policy,
                    max_age=24 * 60 * 60,
                    max_msgs=50000,
                    max_bytes=512 * 1024 * 1024,
                )
                print(f"✅ Stream '{stream_name}' updated")
            except Exception as e:
                if "not found" in str(e).lower() or "404" in str(e).lower():
                    print(f"Creating stream '{stream_name}'...")
                    await js.add_stream(
                        name=stream_name,
                        subjects=subjects,
                        storage="memory",
                        retention=retention_policy,
                        max_age=24 * 60 * 60,
                        max_msgs=50000,
                        max_bytes=512 * 1024 * 1024,
                    )
                    print(f"✅ Stream '{stream_name}' created")
                else:
                    print(f"⚠️  Could not inspect stream '{stream_name}': {e}")

        # Create durable consumers per subject - names match what worker.py
        # computes at runtime, so max_deliver/ack_wait here actually apply.
        print("\nConfiguring consumers...")
        for domain in WORKER_DOMAINS:
            stream_name = domain["stream"]
            consumer = domain["consumer"]
            for subject in domain["subjects"]:
                durable = durable_name(consumer, subject)
                try:
                    print(f"Creating consumer '{durable}' on stream '{stream_name}' for subject '{subject}'...")
                    await js.add_consumer(
                        stream_name,
                        durable_name=durable,
                        filter_subject=subject,
                        ack_policy="explicit",
                        deliver_policy="all",
                        max_deliver=5,
                        ack_wait=60,
                    )
                    print(f"✅ Consumer '{durable}' created")
                except Exception as e:
                    if "already in use" in str(e).lower() or "already exists" in str(e).lower():
                        print(f"⚠️  Consumer '{durable}' already exists, skipping...")
                    else:
                        raise

        print("\n✅ JetStream setup complete!")
        print("   Streams:")
        for domain in WORKER_DOMAINS:
            print(f"     - {domain['stream']}: {', '.join(domain['subjects'])}")
        print("   Consumers:")
        for domain in WORKER_DOMAINS:
            for subject in domain["subjects"]:
                print(f"     - {durable_name(domain['consumer'], subject)}: stream={domain['stream']}, subject={subject}")

        await nc.close()
        sys.exit(0)

    except Exception as e:
        print(f"❌ Error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(setup_jetstream())