|
# retoor <retoor@molodetz.nl>
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import APIRouter, Request, WebSocket, WebSocketDisconnect
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from devplacepy.services.manager import service_manager
|
|
from devplacepy.services.pubsub import policy
|
|
from devplacepy.services.pubsub.hub import pubsub
|
|
|
|
from devplacepy.routers.dbapi._shared import error, read_body, require_dbapi_caller
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter()
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def _message(topic: str, data) -> dict:
|
|
return {"type": "message", "topic": topic, "data": data, "ts": _now()}
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def pubsub_ws(websocket: WebSocket):
|
|
await websocket.accept()
|
|
svc = service_manager.get_service("pubsub")
|
|
if svc is None or not svc.is_enabled():
|
|
await websocket.close(code=1013)
|
|
return
|
|
if not service_manager.owns_lock():
|
|
await websocket.close(code=4013)
|
|
return
|
|
actor = policy.resolve_actor(websocket)
|
|
if actor.kind == "guest" and not policy.guests_enabled():
|
|
await websocket.close(code=1008)
|
|
return
|
|
await websocket.send_json({"type": "ready", "actor": actor.kind})
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
kind = data.get("type")
|
|
topic = str(data.get("topic", "")).strip()
|
|
if kind in ("subscribe", "unsubscribe", "publish") and not policy.valid_topic(
|
|
topic
|
|
):
|
|
await websocket.send_json({"type": "error", "message": "Invalid topic name."})
|
|
continue
|
|
if kind == "subscribe":
|
|
if policy.can_subscribe(actor, topic):
|
|
pubsub.subscribe(topic, websocket)
|
|
await websocket.send_json({"type": "subscribed", "topic": topic})
|
|
else:
|
|
await websocket.send_json(
|
|
{"type": "error", "message": f"Not allowed to subscribe to {topic}."}
|
|
)
|
|
elif kind == "unsubscribe":
|
|
pubsub.unsubscribe(topic, websocket)
|
|
await websocket.send_json({"type": "unsubscribed", "topic": topic})
|
|
elif kind == "publish":
|
|
if not policy.can_publish(actor, topic):
|
|
await websocket.send_json(
|
|
{"type": "error", "message": f"Not allowed to publish to {topic}."}
|
|
)
|
|
continue
|
|
payload = data.get("data")
|
|
if len(json.dumps(payload, default=str)) > policy.MAX_PAYLOAD_BYTES:
|
|
await websocket.send_json({"type": "error", "message": "Payload too large."})
|
|
continue
|
|
delivered = await pubsub.publish(topic, _message(topic, payload))
|
|
await websocket.send_json(
|
|
{"type": "ack", "topic": topic, "delivered": delivered}
|
|
)
|
|
else:
|
|
await websocket.send_json({"type": "error", "message": "Unknown frame type."})
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception:
|
|
logger.exception("pubsub websocket loop failed")
|
|
finally:
|
|
pubsub.drop_socket(websocket)
|
|
|
|
|
|
@router.post("/publish")
|
|
async def pubsub_http_publish(request: Request):
|
|
caller = require_dbapi_caller(request)
|
|
if not service_manager.owns_lock():
|
|
return error(409, "Pub/sub is served by the lock-owner worker. Retry.")
|
|
body = await read_body(request)
|
|
topic = str(body.get("topic", "")).strip()
|
|
if not policy.valid_topic(topic):
|
|
return error(400, "Invalid topic name.")
|
|
payload = body.get("data")
|
|
delivered = await pubsub.publish(topic, _message(topic, payload))
|
|
from devplacepy.services.audit import record as audit
|
|
|
|
audit.record(
|
|
request,
|
|
"pubsub.publish",
|
|
target_type="topic",
|
|
target_uid=topic,
|
|
summary=f"{caller.username or caller.kind} published to {topic}",
|
|
metadata={"topic": topic, "delivered": delivered, "caller": caller.kind},
|
|
)
|
|
return JSONResponse({"topic": topic, "delivered": delivered})
|
|
|
|
|
|
@router.get("/topics")
|
|
async def pubsub_list_topics(request: Request):
|
|
require_dbapi_caller(request)
|
|
return JSONResponse({"topics": pubsub.topics()})
|