An SSE Relay Service: Perfect for Adding Real-Time Features to Your Django Project Without Introducing ASGI or Async Code.
from __future__ import annotations

import os

from dotenv import load_dotenv

load_dotenv()


class ConfigurationError(Exception):
    pass


def get_allowed_origins() -> list[str]:
    origins = os.getenv("ALLOWED_ORIGINS")
    return [] if origins is None else [v.strip() for v in origins.split(",")]


def get_debug_value() -> bool:
    return os.getenv("RELAY_SERVER_DEBUG") == "True"


def get_postgres_url() -> str | None:
    return os.getenv("DATABASE_URL")


def get_redis_url() -> str | None:
    return os.getenv("REDIS_URL")

src/sse_relay_server/config.py
Tobi-De/sse_relay_server
from typing import Protocol, AsyncGenerator

from sse_starlette import ServerSentEvent

from sse_relay_server.config import (
    ConfigurationError,
    get_redis_url,
    get_postgres_url,
)
from ._postgres import PostgresGateway
from ._redis import RedisGateway


def _select_gateway():
    if redis_url := get_redis_url():
        return RedisGateway(redis_url)
    if postgres_url := get_postgres_url():
        return PostgresGateway(postgres_url)
    else:
        raise ConfigurationError("Set either REDIS_URL or DATABASE_URL")


gateway = _select_gateway()


class Gateway(Protocol):
    async def listen(self, channel: str) -> AsyncGenerator[ServerSentEvent, None]:
        ...

    def notify(self, channel: str, sse_payload: dict) -> None:
        ...

src/sse_relay_server/gateways/__init__.py
Tobi-De/sse_relay_server
import json
from typing import AsyncGenerator

import dj_database_url
import psycopg
from loguru import logger
from sse_starlette import ServerSentEvent

from ..config import ConfigurationError


class PostgresGateway:
    def __init__(self, postgres_url: str) -> None:
        self.postgres_url = postgres_url
        parsed_params = dj_database_url.parse(postgres_url)
        if parsed_params["ENGINE"] not in [
            "django.db.backends.postgresql_psycopg2",
            "django.db.backends.postgresql",
        ]:
            raise ConfigurationError("A postgresql database is require")
        self.db_params = {
            "client_encoding": "UTF8",
            "dbname": parsed_params["NAME"],
            "user": parsed_params["USER"],
            "password": parsed_params["PASSWORD"],
            "host": parsed_params["HOST"],
        }

    async def listen(self, channel: str) -> AsyncGenerator[ServerSentEvent, None]:
        connection = await psycopg.AsyncConnection.connect(
            **self.db_params,
            autocommit=True,
        )

        async with connection.cursor() as cursor:
            await cursor.execute(f"LISTEN {channel}")
            generator = connection.notifies()
            async for notify_message in generator:
                payload = json.loads(notify_message.payload)
                logger.debug(f"Received from {channel}: {payload}")
                yield ServerSentEvent(**payload)

    def notify(self, channel: str, sse_payload: dict) -> None:
        connection = psycopg.Connection.connect(
            **self.db_params,
            autocommit=True,
        )
        with connection.cursor() as cursor:
            cursor.execute(f"NOTIFY {channel}, '{json.dumps(sse_payload)}'")

src/sse_relay_server/gateways/_postgres.py
Tobi-De/sse_relay_server
import json
from typing import AsyncGenerator

import redis
import redis.asyncio as async_redis
from loguru import logger
from sse_starlette import ServerSentEvent


class RedisGateway:
    def __init__(self, redis_url: str) -> None:
        self.redis_url = redis_url

    async def listen(self, channel: str) -> AsyncGenerator[ServerSentEvent, None]:
        r = async_redis.from_url(self.redis_url)
        async with r.pubsub() as pubsub:
            await pubsub.subscribe(channel)
            while True:
                message = await pubsub.get_message(ignore_subscribe_messages=True)
                if message is not None:
                    payload = json.loads(message["data"].decode())
                    logger.debug(f"Received from {channel}: {payload}")
                    yield ServerSentEvent(**payload)

    def notify(self, channel: str, sse_payload: dict) -> None:
        r = redis.from_url(self.redis_url)
        r.publish(channel=channel, message=json.dumps(sse_payload))

src/sse_relay_server/gateways/_redis.py
Tobi-De/sse_relay_server
from __future__ import annotations

import argparse

import uvicorn
from sse_starlette import ServerSentEvent
from sse_starlette.sse import EventSourceResponse
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.routing import Route

from .config import get_allowed_origins, get_debug_value
from .gateways import gateway

SSE_SERVER_DEBUG = get_debug_value()
ALLOWED_ORIGINS = get_allowed_origins()


async def generate_stop_event():
    # Set a high retry interval (e.g., 1 year in milliseconds)
    retry_interval = 365 * 24 * 60 * 60 * 1000
    yield ServerSentEvent(
        "Please stop, there is nothing here ;(",
        retry=retry_interval,
        comment="Shut down the connection",
    )


async def sse(request: Request):
    if channel := request.query_params.get("channel"):
        return EventSourceResponse(gateway.listen(channel))
    else:
        return EventSourceResponse(generate_stop_event())


routes = [Route("/", endpoint=sse)]

middleware = [
    Middleware(
        CORSMiddleware,
        allow_origins=ALLOWED_ORIGINS,
        allow_credentials=True,
        allow_methods=["GET"],
        allow_headers=["*"],
    )
]

app = Starlette(debug=SSE_SERVER_DEBUG, routes=routes, middleware=middleware)


def main():
    parser = argparse.ArgumentParser(
        prog="sse-relay-server",
        description="Relay server for SSE events",
    )
    parser.add_argument(
        "--host",
        type=str,
        default="0.0.0.0",
        help="server host",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8001,
        help="server port",
    )
    parser.add_argument(
        "--reload",
        action="store_true",
        help="enable auto-reload",
        default=False,
    )
    parser.add_argument("--workers", type=int)
    parser.add_argument(
        "--log-level",
        default="info",
        choices=["critical", "error", "warning", "info", "debug", "trace"],
    )

    args = parser.parse_args()

    uvicorn.run("sse_relay_server.main:app", **args.__dict__)


if __name__ == "__main__":
    main()

src/sse_relay_server/main.py
Tobi-De/sse_relay_server