from django.apps import AppConfig


class SseServerConfig(AppConfig):
    default_auto_field = "django.db.models.BigAutoField"
    name = "sse_server"

sse_server/apps.py
Tobi-De/sse_server_postgres_listen_notify
import json

import environ
import psycopg
from django.db import connection
from sse_starlette.sse import EventSourceResponse, ServerSentEvent
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

env = environ.Env()


async def event_publisher(req: Request):
    connection_params = connection.get_connection_params()
    connection_params.pop("cursor_factory")
    aconnection = await psycopg.AsyncConnection.connect(
        **connection_params,
        autocommit=True,
    )
    channel = req.query_params.get("channel")
    if not channel:
        return

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


async def sse(request: Request):
    return EventSourceResponse(event_publisher(request))


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


middleware = [
    Middleware(
        CORSMiddleware,
        allow_origins=env.list("SSE_ALLOWED_ORIGINS"),
        allow_credentials=True,
        allow_methods=["GET"],
        allow_headers=["*"],
    )
]

app = Starlette(debug=env("DJANGO_DEBUG"), routes=routes, middleware=middleware)

sse_server/asgi_app.py
Tobi-De/sse_server_postgres_listen_notify
import environ
from django.conf import settings


env = environ.Env()
environ.Env.read_env(settings.BASE_DIR / ".env")


def sse_server_url(_):
    return {"sse_server_url": env.str("SSE_SERVER_URL")}

sse_server/context_processors.py
Tobi-De/sse_server_postgres_listen_notify
import uvicorn
from django.conf import settings
from django.core.management.base import BaseCommand


class Command(BaseCommand):
    help_text = "Run uvicorn server to handle server sent events requests"

    def add_arguments(self, parser):
        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("--workers", type=int)

    def handle(self, *args, **options):
        uvicorn.run(
            "smartestate.sse_server.asgi_app:app",
            host=options.get("host"),
            port=options.get("port"),
            reload=settings.DEBUG,
            workers=options.get("workers"),
        )

sse_server/management/commands/serve_sse.py
Tobi-De/sse_server_postgres_listen_notify
import json

from django.db import connection


def postgres_notify(channel: str, sse_payload: dict):
    with connection.cursor() as cursor:
        cursor.execute(f"NOTIFY {channel}, '{json.dumps(sse_payload)}'")

sse_server/utils.py
Tobi-De/sse_server_postgres_listen_notify