space_worker/supabase.py

import json
import logging
import time
from urllib import parse, request

import websocket


logger = logging.getLogger("celery-space.supabase")


def json_bytes(payload):
    return json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def realtime_topic(topic):
    if topic.startswith("realtime:"):
        return topic
    return f"realtime:{topic}"


def extract_broadcast(message):
    if isinstance(message, (bytes, bytearray)):
        message = decode_binary_broadcast(message)
        if message is None:
            return None

    if message.get("event") != "broadcast":
        payload = message.get("payload")
        if isinstance(payload, dict) and ("user_id" in payload or "message" in payload):
            return message.get("event"), payload
        return None

    payload = message.get("payload") or {}
    event = payload.get("event") or payload.get("type")
    body = payload.get("payload", payload)
    return event, body


def decode_binary_broadcast(frame):
    try:
        if not frame or frame[0] != 4 or len(frame) < 6:
            return None

        topic_size = frame[1]
        event_size = frame[2]
        metadata_size = frame[3]
        payload_encoding = frame[4]
        cursor = 5

        topic = frame[cursor : cursor + topic_size].decode("utf-8")
        cursor += topic_size
        event = frame[cursor : cursor + event_size].decode("utf-8")
        cursor += event_size
        metadata = frame[cursor : cursor + metadata_size]
        cursor += metadata_size
        raw_payload = frame[cursor:]

        if payload_encoding == 1:
            payload = json.loads(raw_payload.decode("utf-8"))
        else:
            payload = raw_payload.decode("utf-8", errors="replace")

        return {
            "topic": topic,
            "event": "broadcast",
            "payload": {
                "event": event,
                "metadata": metadata.decode("utf-8", errors="replace"),
                "payload": payload,
            },
        }
    except Exception as exc:
        logger.warning("failed to decode binary realtime frame: %s", exc)
        return None


def normalize_realtime_message(message):
    if isinstance(message, list) and len(message) >= 5:
        join_ref, ref, topic, event, payload = message[:5]
        return {
            "join_ref": join_ref,
            "ref": ref,
            "topic": topic,
            "event": event,
            "payload": payload,
        }
    if isinstance(message, dict):
        return message
    return {"event": "unknown", "payload": message}


class RefCounter:
    def __init__(self):
        self.value = 0

    def next(self):
        self.value += 1
        return str(self.value)


class SupabaseBroadcaster:
    def __init__(self, settings):
        self.settings = settings

    def send(self, topic, event, payload, private=True):
        self.settings.require_supabase()
        encoded_topic = parse.quote(topic, safe="")
        encoded_event = parse.quote(event, safe="")
        private_query = "true" if private else "false"
        url = (
            f"{self.settings.supabase_url}/realtime/v1/api/broadcast/"
            f"{encoded_topic}/events/{encoded_event}?private={private_query}"
        )
        req = request.Request(
            url,
            data=json_bytes(payload),
            method="POST",
            headers={
                "apikey": self.settings.supabase_key,
                "Authorization": f"Bearer {self.settings.supabase_key}",
                "Content-Type": "application/json",
            },
        )

        with request.urlopen(req, timeout=10) as response:
            logger.info("broadcast delivered topic=%s event=%s status=%s", topic, event, response.status)


class RealtimeConnection:
    def __init__(self, settings):
        self.settings = settings
        self.refs = RefCounter()
        self.ws = None

    def __enter__(self):
        self.settings.require_supabase()
        self.ws = websocket.create_connection(self.websocket_url(), timeout=30)
        self.ws.settimeout(1)
        self.join_inbox()
        return self

    def __exit__(self, exc_type, exc, traceback):
        if self.ws is not None:
            self.ws.close()

    def websocket_url(self):
        parsed = parse.urlparse(self.settings.supabase_url)
        if parsed.scheme == "https":
            scheme = "wss"
        elif parsed.scheme == "http":
            scheme = "ws"
        else:
            raise RuntimeError("SUPABASE_URL must start with http:// or https://")

        query = parse.urlencode({"apikey": self.settings.supabase_key, "vsn": self.settings.realtime_vsn})
        return parse.urlunparse((scheme, parsed.netloc, "/realtime/v1/websocket", "", query, ""))

    def send(self, topic, event, payload, join_ref=None, ref=None):
        if ref is None:
            ref = self.refs.next()
        if self.settings.realtime_vsn == "2.0.0":
            message = [join_ref, ref, topic, event, payload]
        else:
            message = {
                "topic": topic,
                "event": event,
                "payload": payload,
                "ref": ref,
                "join_ref": join_ref,
            }
        self.ws.send(json.dumps(message, separators=(",", ":"), ensure_ascii=False))
        return ref

    def join_inbox(self):
        topic = realtime_topic(self.settings.inbox_topic)
        join_ref = self.refs.next()
        payload = {
            "config": {
                "broadcast": {"ack": True, "self": self.settings.broadcast_self},
                "presence": {"enabled": False},
                "postgres_changes": [],
                "private": self.settings.inbox_private,
            },
            "access_token": self.settings.supabase_key,
        }
        self.send(topic, "phx_join", payload, join_ref=join_ref, ref=join_ref)
        logger.info(
            "joining supabase topic=%s private=%s vsn=%s self=%s",
            self.settings.inbox_topic,
            self.settings.inbox_private,
            self.settings.realtime_vsn,
            self.settings.broadcast_self,
        )

    def send_heartbeat(self):
        self.send("phoenix", "heartbeat", {})

    def iter_messages(self):
        last_heartbeat = 0
        while True:
            now = time.monotonic()
            if now - last_heartbeat >= self.settings.heartbeat_seconds:
                self.send_heartbeat()
                last_heartbeat = now

            try:
                raw_message = self.ws.recv()
            except websocket.WebSocketTimeoutException:
                continue

            if isinstance(raw_message, (bytes, bytearray)):
                logger.info("received binary realtime frame bytes=%s", len(raw_message))
                yield raw_message
                continue

            message = normalize_realtime_message(json.loads(raw_message))
            if message.get("event") == "phx_reply":
                payload = message.get("payload") or {}
                topic = message.get("topic")
                status = payload.get("status")
                if not (topic == "phoenix" and status == "ok"):
                    logger.info(
                        "received realtime reply topic=%s status=%s ref=%s",
                        topic,
                        status,
                        message.get("ref"),
                    )
            else:
                logger.info("received realtime event=%s topic=%s", message.get("event"), message.get("topic"))
            yield message