From: mar77i Date: Wed, 14 Oct 2026 15:04:07 +0000 (+0200) Subject: broken design. Turns out a worker queue processor is not necessary, and the listener... X-Git-Url: https://git.mar77i.info/?a=commitdiff_plain;h=refs%2Fheads%2Fmaster;p=chat broken design. Turns out a worker queue processor is not necessary, and the listener task should manage its own connection --- diff --git a/activity/apps.py b/activity/apps.py deleted file mode 100644 index 34cfcc0..0000000 --- a/activity/apps.py +++ /dev/null @@ -1,5 +0,0 @@ -from django.apps import AppConfig - - -class ActivityConfig(AppConfig): - name = 'activity' diff --git a/activity/migrations/0002_initial.py b/activity/migrations/0002_initial.py deleted file mode 100644 index 2905784..0000000 --- a/activity/migrations/0002_initial.py +++ /dev/null @@ -1,33 +0,0 @@ -# Generated by Django 6.1.1 on 2026-09-27 10:49 - -import django.db.models.deletion -from django.conf import settings -from django.db import migrations, models - - -class Migration(migrations.Migration): - - initial = True - - dependencies = [ - ('activity', '0001_initial'), - migrations.swappable_dependency(settings.AUTH_USER_MODEL), - ] - - operations = [ - migrations.AddField( - model_name='websocket', - name='user', - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name='activity', - name='websockets', - field=models.ManyToManyField(related_name='activity', to='activity.websocket'), - ), - migrations.AddField( - model_name='websocket', - name='worker', - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='activity.worker'), - ), - ] diff --git a/activity/migrations/__init__.py b/activity/migrations/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/chat/asgi.py b/chat/asgi.py index 91bfbb8..8491522 100644 --- a/chat/asgi.py +++ b/chat/asgi.py @@ -24,7 +24,7 @@ class ChatASGIHandler(ASGIHandler): case "lifespan": await LifespanManager().serve(receive, send) case "websocket": - await Websocket().serve(receive, send) + await Websocket().serve(scope, receive, send) case "http": await super().__call__(scope, receive, send) diff --git a/chat/asyncpg_utils.py b/chat/asyncpg_utils.py index 1b3afc9..caeaa7d 100644 --- a/chat/asyncpg_utils.py +++ b/chat/asyncpg_utils.py @@ -1,8 +1,14 @@ from contextlib import asynccontextmanager +from typing import Any, AsyncIterator, TypeVar +from asyncpg import Pool +from django.core.exceptions import FieldDoesNotExist from django.db import DEFAULT_DB_ALIAS from django.db.models import Model, Field +pool: Pool | None = None +InstanceType = TypeVar("InstanceType") + class SQL: """ @@ -18,7 +24,7 @@ class SQL: return f"SQL({self.expression!r})" -class SQLHelper: +class StatementHelper: @staticmethod def quote(identifier: str): identifier = identifier.replace('"', '""') @@ -35,8 +41,10 @@ class SQLHelper: def comma_join(quoted_fields: list[str]) -> str: return ", ".join(quoted_fields) - def __init__(self, model: type[Model], using=None): - self.using = using + __slots__ = ("wrapper", "model", "attnames", "columns", "quoted_columns", "params") + + def __init__(self, wrapper: "AsyncPGWrapper", model: type[InstanceType]): + self.wrapper = wrapper self.model = model self.attnames = [] self.columns = [] @@ -47,6 +55,7 @@ class SQLHelper: if isinstance(field, str): field = self.model._meta.get_field(field) attname, column = field.get_attname_column() + assert isinstance(column, str) self.attnames.append(attname) self.columns.append(column) quoted_column = self.quote(column) @@ -65,7 +74,7 @@ class SQLHelper: def concrete_fields(self): return self.model._meta.concrete_fields - def __iter__(self): + def register_concrete_fields(self): for field in self.concrete_fields: attname, quoted_column = self.add(field) yield field, attname, quoted_column @@ -73,9 +82,9 @@ class SQLHelper: def get_field(self, field_name): return self.model._meta.get_field(field_name) - def hydrate(self, row) -> Model: + def hydrate(self, row: dict[str, Any]) -> InstanceType: return self.model.from_db( - self.using or DEFAULT_DB_ALIAS, + self.wrapper.using, self.attnames, [row[column] for column in self.columns], ) @@ -91,142 +100,199 @@ class SQLHelper: def joined_quoted_columns(self): return self.comma_join(self.quoted_columns) - def select_sql(self, filters: dict) -> str: - """Build SELECT of all concrete columns with equality filters.""" - # Work on a copy, any leftover keys are unknown fields - filters = filters.copy() - conditions = [] - for field, attname, quoted_column in self: - key = field.name if field.name in filters else attname - if key in filters: - value = filters.pop(key) - if value is None: - conditions.append(f"{quoted_column} IS NULL") - else: - conditions.append(f"{quoted_column}={self.bind(value)}") - if filters: + def resolve_filter_field(self, key: str) -> Field | None: + """Map a filter key (name, attname or 'pk') to a concrete field.""" + if key == "pk": + return self.model._meta.pk + try: + field = self.model._meta.get_field(key) + except FieldDoesNotExist: + return None + return field if field.concrete else None # excludes reverse rels / M2M + + def where_sql(self, filters: dict) -> str: + """ + Build an equality-only WHERE clause (with a leading space), or "" + if there are no filters. Values are bound into self.params. + """ + conditions, seen, unknown = [], {}, [] + for key, value in filters.items(): + field = self.resolve_filter_field(key) + if field is None: + unknown.append(key) + continue + if field in seen: + raise KeyError(f"{key}= conflicts with the {seen[field]}= kwarg.") + seen[field] = key + assert isinstance(field.column, str) + quoted_column = self.quote(field.column) + if value is None: + conditions.append(f"{quoted_column} IS NULL") + else: + conditions.append(f"{quoted_column}={self.bind(value)}") + if unknown: raise KeyError( f"Unknown field(s) for {self.model.__name__}: " - f"{self.comma_join(filters)}" - ) - query = f"SELECT {self.joined_quoted_columns} FROM {self.quoted_db_table}" - if conditions: - return f"{query} WHERE {' AND '.join(conditions)}" - return query - - -async def ainsert(con, instance: Model) -> Model: - if not instance._state.adding: - raise ValueError("Cannot insert a saved instance") - pk_is_none = instance.pk is None - helper = SQLHelper(type(instance), instance._state.db) - insert_columns, insert_values = [], [] - for field, attname, quoted_column in helper: - if field.primary_key and pk_is_none: - continue # let the database generate it - insert_columns.append(quoted_column) - insert_values.append( - helper.bind( - SQL("CURRENT_TIMESTAMP") - if helper.field_is_auto_now(field, add=True) - else getattr(instance, attname) + f"{self.comma_join(unknown)}" ) + return f" WHERE {' AND '.join(conditions)}" if conditions else "" + + def select_sql(self, filters: dict) -> str: + """Build SELECT of all concrete columns with equality filters.""" + where = self.where_sql(filters) + for _ in self.register_concrete_fields(): + pass + return ( + f"SELECT {self.joined_quoted_columns} " + f"FROM {self.quoted_db_table}{where}" ) - if insert_columns: - body = ( - f"({helper.comma_join(insert_columns)}) " - f"VALUES ({helper.comma_join(insert_values)})" + def delete_sql(self, filters: dict) -> str: + return f"DELETE FROM {self.quoted_db_table}{self.where_sql(filters)}" + + +@asynccontextmanager +async def guard_pool(new_pool: Pool): + global pool + pool = new_pool + try: + yield + finally: + pool = None + + +class AsyncPGWrapper: + def __init__(self, con, using=None): + self.con = con + self.using = using or DEFAULT_DB_ALIAS + + @classmethod + @asynccontextmanager + async def acquire(cls, using: str | None = None) -> AsyncIterator["AsyncPGWrapper"]: + async with pool.acquire() as con: + yield cls(con, using) + + async def ainsert(self, instance: Model) -> Model: + if not instance._state.adding: + raise ValueError("Cannot insert a saved instance") + pk_is_none = instance.pk is None + helper = StatementHelper(self, type(instance)) + insert_columns, insert_values = [], [] + for field, attname, quoted_column in helper.register_concrete_fields(): + if field.primary_key and pk_is_none: + continue # let the database generate it + insert_columns.append(quoted_column) + insert_values.append( + helper.bind( + SQL("CURRENT_TIMESTAMP") + if helper.field_is_auto_now(field, add=True) + else getattr(instance, attname) + ) + ) + + if insert_columns: + body = ( + f"({helper.comma_join(insert_columns)}) " + f"VALUES ({helper.comma_join(insert_values)})" + ) + else: + body = "DEFAULT VALUES" + row = await self.con.fetchrow( + ( + f"INSERT INTO {helper.quoted_db_table} {body} " + f"RETURNING {helper.joined_quoted_columns}" + ), + *helper.params, ) - else: - body = "DEFAULT VALUES" - row = await con.fetchrow( - ( - f"INSERT INTO {helper.quoted_db_table} {body} " - f"RETURNING {helper.joined_quoted_columns}" - ), - *helper.params, - ) - if row is None: - raise RuntimeError("INSERT did not return a row") - instance.__dict__.update(helper.hydrate(row).__dict__) - return instance - - -async def aupdate(con, instance: Model, **kwargs) -> Model: - if instance._state.adding: - raise ValueError("Cannot update an unsaved instance") - helper = SQLHelper(type(instance), instance._state.db) - for field in helper.concrete_fields: - if helper.field_is_auto_now(field, add=False): - kwargs[field.name] = SQL("CURRENT_TIMESTAMP") - if not kwargs: + if row is None: + raise RuntimeError("INSERT did not return a row") + instance.__dict__.update(helper.hydrate(row).__dict__) return instance - assignments = [] - for name, value in kwargs.items(): - quoted_column = helper.add(name)[1] - assignments.append(f"{quoted_column}={helper.bind(value)}") - row = await con.fetchrow( - ( - f"UPDATE {helper.quoted_db_table} " - f"SET {helper.comma_join(assignments)} " - f"WHERE {helper.quoted_pk_column}={helper.bind(instance.pk)} " - f"RETURNING {helper.joined_quoted_columns}" - ), - *helper.params, - ) - if row is None: - raise instance.DoesNotExist(f"{helper.model.__name__} {instance.pk!r}") - for attname, column in zip(helper.attnames, helper.columns): - setattr(instance, attname, row[column]) - return instance - - -async def adelete(con, instance: Model): - helper = SQLHelper(type(instance), instance._state.db) - if instance._state.adding: - raise ValueError("Cannot delete an unsaved instance") - await con.execute( - f"DELETE FROM {helper.quoted_db_table} WHERE {helper.quoted_pk_column}=$1", - instance.pk, - ) - instance._state.adding = True - instance.pk = None - - -async def aselect_one(con, model: type[Model], using=None, **filters) -> Model: - helper = SQLHelper(model, using) - query = helper.select_sql(filters) - rows = await con.fetch(f"{query} LIMIT 2", *helper.params) - if not rows: - raise model.DoesNotExist(f"{model.__name__} matching {filters!r}") - if len(rows) > 1: - raise model.MultipleObjectsReturned( - f"More than one {model.__name__} matching {filters!r}" + async def aupdate(self, instance: Model, **kwargs) -> Model: + if instance._state.adding: + raise ValueError("Cannot update an unsaved instance") + helper = StatementHelper(self, type(instance)) + for field in helper.concrete_fields: + if helper.field_is_auto_now(field, add=False): + kwargs[field.name] = SQL("CURRENT_TIMESTAMP") + if not kwargs: + return instance + + assignments = [] + for name, value in kwargs.items(): + quoted_column = helper.add(name)[1] + assignments.append(f"{quoted_column}={helper.bind(value)}") + row = await self.con.fetchrow( + ( + f"UPDATE {helper.quoted_db_table} " + f"SET {helper.comma_join(assignments)} " + f"WHERE {helper.quoted_pk_column}={helper.bind(instance.pk)} " + f"RETURNING {helper.joined_quoted_columns}" + ), + *helper.params, ) - return helper.hydrate(rows[0]) - + if row is None: + raise instance.DoesNotExist(f"{helper.model.__name__} {instance.pk!r}") + for attname, column in zip(helper.attnames, helper.columns): + setattr(instance, attname, row[column]) + instance._state.db = self.using + return instance -@asynccontextmanager -async def aselect_stream(con, model: type[Model], using=None, **filters): - """ - The transaction closes when the block exits, - so the iterator is only valid inside it. - """ - helper = SQLHelper(model, using) - query = helper.select_sql(filters) - async with con.transaction(): - yield ( - helper.hydrate(row) - async for row in con.cursor(query, *helper.params) + async def adelete(self, model: type[Model], **filters): + helper = StatementHelper(self, model) + if not filters: + raise ValueError( + "adelete() needs at least one filter; " + "refusing to delete every row" + ) + status = await self.con.execute(helper.delete_sql(filters), *helper.params) + return int(status.rpartition(" ")[2]) # "DELETE 3" -> 3 + + async def adelete_one(self, instance: Model): + helper = StatementHelper(self, type(instance)) + if instance._state.adding: + raise ValueError("Cannot delete an unsaved instance") + status = await self.con.execute( + f"DELETE FROM {helper.quoted_db_table} WHERE {helper.quoted_pk_column}=$1", + instance.pk, ) + if int(status.rpartition(" ")[2]): + instance._state.adding = True + instance._state.db = self.using + instance.pk = None + return instance + async def aselect_one(self, model: type[InstanceType], **filters) -> InstanceType: + helper = StatementHelper(self, model) + query = helper.select_sql(filters) + rows = await self.con.fetch(f"{query} LIMIT 2", *helper.params) + if not rows: + raise model.DoesNotExist(f"{model.__name__} matching {filters!r}") + if len(rows) > 1: + raise model.MultipleObjectsReturned( + f"More than one {model.__name__} matching {filters!r}" + ) + return helper.hydrate(rows[0]) + + @asynccontextmanager + async def aselect_stream(self, model: type[InstanceType], **filters) -> AsyncIterator[InstanceType]: + """ + The transaction closes when the block exits, + so the iterator is only valid inside it. + """ + helper = StatementHelper(self, model) + query = helper.select_sql(filters) + async with self.con.transaction(): + yield ( + helper.hydrate(row) + async for row in self.con.cursor(query, *helper.params) + ) -async def aselect_all(con, model: type[Model], using=None, **filters) -> list[Model]: - helper = SQLHelper(model, using) - query = helper.select_sql(filters) - return [ - helper.hydrate(row) - for row in await con.fetch(query, *helper.params) - ] + async def aselect_all(self, model: type[InstanceType], **filters) -> list[InstanceType]: + helper = StatementHelper(self, model) + query = helper.select_sql(filters) + return [ + helper.hydrate(row) + for row in await self.con.fetch(query, *helper.params) + ] diff --git a/chat/lifespan.py b/chat/lifespan.py index d1f775a..7723ba1 100644 --- a/chat/lifespan.py +++ b/chat/lifespan.py @@ -1,139 +1,133 @@ import os -from importlib import import_module +from asyncio import CancelledError, Queue, sleep +from contextlib import asynccontextmanager from socket import gethostname -from asyncio import CancelledError, Queue, sleep from asyncpg import create_pool -from contextlib import AsyncExitStack -from typing import TypedDict - from django.conf import settings -from django.contrib.auth import SESSION_KEY - -from .asyncpg_utils import adelete, ainsert -from .utils import CancellingTaskGroup -from .websocket import StopSend -from activity.models import Worker, Websocket from user.models import User -session_engine = import_module(settings.SESSION_ENGINE) - - -class LifespanResponse(TypedDict): - type: str | None - - -async def aget_user(scope): - headers = dict(scope.get("headers", [])) - session_key = None - cookie_str = headers.get(b"cookie", b"").decode() - cookie_name = "sessionid=" - pos = cookie_str.find(cookie_name) - if pos >= 0: - pos += len(cookie_name) - end_pos = cookie_str.find(";", pos) - if end_pos == -1: - end_pos = None - session_key = cookie_str[pos:end_pos] - - session = session_engine.SessionStore(session_key=session_key) - session_key = await session.aget(SESSION_KEY) - if session_key is None: - raise KeyError() - user_id = User._meta.pk.to_python(session_key) +from .asyncpg_utils import AsyncPGWrapper, guard_pool +from .models import Worker, Websocket +from .utils import CancellingTaskGroup, iterate_queue, log_exceptions class LifespanManager: queue = Queue() websockets = {} - def __init__(self): - self.worker_instance = None + __slots__ = ("pending_msg_type", "instance") - @staticmethod - async def send_pending_msg(send, pending_msg): - if pending_msg["type"] is None: + def __init__(self): + self.pending_msg_type: str | None = None + self.instance: Worker = None + + @asynccontextmanager + async def register_worker(self, pid): + async with AsyncPGWrapper.acquire() as wrapper: + self.instance = await wrapper.ainsert( + Worker( + hostname=gethostname(), pid=pid, starttime=Worker.get_starttime(pid) + ) + ) + try: + yield + finally: + async with AsyncPGWrapper.acquire() as wrapper: + await wrapper.adelete(Websocket, worker_id=self.instance.id) + await wrapper.adelete_one(self.instance) + self.instance = None + + @asynccontextmanager + async def background_tasks(self): + async with CancellingTaskGroup() as tg: + tg.create_task(self.consume_websocket_queue()) + tg.create_task(self.consume_channel("chat_channel")) + yield + + @asynccontextmanager + async def guard_pending_msg(self, send): + try: + yield + except (Exception, CancelledError): + if isinstance(self.pending_msg_type, str): + self.pending_msg_type = self.pending_msg_type.replace( + ".complete", ".failed" + ) + raise + finally: + await self.send_pending_msg(send) + + async def send_pending_msg(self, send): + if self.pending_msg_type is None: return - await send(pending_msg) - pending_msg["type"] = None + await send({"type": self.pending_msg_type}) + self.pending_msg_type = None + + async def serve_until_startup(self, receive): + while True: + match (await receive())["type"]: + case "lifespan.shutdown": + self.pending_msg_type = "lifespan.shutdown.complete" + return False + case "lifespan.startup": + return True + + async def serve_until_shutdown(self, receive): + while True: + match (await receive())["type"]: + case "lifespan.shutdown": + self.pending_msg_type = "lifespan.shutdown.complete" + return + case "lifespan.startup": + self.pending_msg_type = "lifespan.startup.failed" + raise ValueError("Unexpectedly repeating lifespan.startup") async def serve(self, receive, send): - pending_msg: LifespanResponse = {"type": None} creds = settings.DATABASES["default"] - async with create_pool( + async with log_exceptions(), create_pool( host=creds["HOST"], port=creds["PORT"], user=creds["USER"], password=creds["PASSWORD"], database=creds["NAME"], - ) as pool: - try: - async with AsyncExitStack() as exit_stack: - while (await receive())["type"] != "lifespan.startup": - pass - pending_msg["type"] = "lifespan.startup.complete" - async with pool.acquire() as con: - pid = os.getpid() - self.worker_instance = await ainsert( - con, - Worker( - hostname=gethostname(), - pid=pid, - starttime=Worker.get_starttime(pid), - ), - ) - task_group = await exit_stack.enter_async_context( - CancellingTaskGroup() - ) - task_group.create_task(self.consume_websocket_queue(pool)) - task_group.create_task(self.consume_channel(pool, "chat_channel")) - await self.send_pending_msg(send, pending_msg) - while msg := await receive(): - if msg["type"] == "lifespan.shutdown": - break - elif msg["type"] == "lifespan.startup": - pending_msg["type"] = "lifespan.startup.failed" - raise ValueError("Unexpectedly repeating lifespan.startup") - pending_msg["type"] = "lifespan.shutdown.complete" - except CancelledError: - raise - except Exception: - pending_msg_type = pending_msg["type"] - if isinstance(pending_msg_type, str): - pending_msg["type"] = pending_msg_type.replace( - ".complete", ".failed" - ) - raise - finally: - await self.send_pending_msg(send, pending_msg) - if isinstance(self.worker_instance, Worker): - async with pool.acquire() as con: - self.worker_instance = await adelete(con, self.worker_instance) - - async def consume_websocket_queue(self, pool): + ) as pool, guard_pool(pool), self.guard_pending_msg(send): + if not await self.serve_until_startup(receive): + return + self.pending_msg_type = "lifespan.startup.complete" + async with ( + self.register_worker(os.getpid()), + self.background_tasks(), + ): + await self.send_pending_msg(send) + await self.serve_until_shutdown(receive) + + async def consume_websocket_queue(self): # websockets -> pg rows (+ triggers) and activity - async with pool.acquire() as con: - while item := await self.queue.get(): + async with log_exceptions(), AsyncPGWrapper.acquire() as wrapper: + async for item in iterate_queue(self.queue): websocket = self.websockets[item["websocket"]] + if not websocket.registered and item["type"] != "register": + continue match item["type"]: case "register": - websocket.websocket_instance = await ainsert( - con, - Websocket( - worker=self.worker_instance, user=user - ), + user_id = item["user_id"] + websocket.instance = await wrapper.ainsert( + Websocket(worker=self.instance, user_id=user_id) ) + websocket.user = await wrapper.aselect_one(User, pk=user_id) + await websocket.registered.set() case "event": - pass + await websocket.send_queue.put(item["data"]) case "unregister": - websocket.websocket_instance = await adelete( - con, websocket.websocket_instance - ) - websocket.send_queue.put(StopSend) + try: + await wrapper.adelete_one(websocket.instance) + finally: + await websocket.registered.clear() - async def consume_channel(self, pool, channel_name): + async def consume_channel(self, channel_name): # pg channel -> websockets - async with pool.acquire() as con: + async with AsyncPGWrapper.acquire() as wrapper: while True: await sleep(1) diff --git a/activity/migrations/0001_initial.py b/chat/migrations/0001_initial.py similarity index 61% rename from activity/migrations/0001_initial.py rename to chat/migrations/0001_initial.py index cfea456..5c79f15 100644 --- a/activity/migrations/0001_initial.py +++ b/chat/migrations/0001_initial.py @@ -1,5 +1,7 @@ -# Generated by Django 6.1.1 on 2026-09-27 10:49 +# Generated by Django 6.1.2 on 2026-10-12 16:57 +import django.db.models.deletion +from django.conf import settings from django.db import migrations, models @@ -8,29 +10,25 @@ class Migration(migrations.Migration): initial = True dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), ] operations = [ migrations.CreateModel( - name='Activity', + name='Worker', fields=[ ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('status', models.CharField(max_length=256)), + ('hostname', models.CharField(max_length=256)), + ('pid', models.IntegerField()), + ('starttime', models.IntegerField()), ], ), migrations.CreateModel( name='Websocket', fields=[ ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ], - ), - migrations.CreateModel( - name='Worker', - fields=[ - ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('hostname', models.CharField(max_length=256)), - ('pid', models.IntegerField()), - ('starttime', models.IntegerField()), + ('user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='websockets', to=settings.AUTH_USER_MODEL)), + ('worker', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='websockets', to='chat.worker')), ], ), ] diff --git a/activity/__init__.py b/chat/migrations/__init__.py similarity index 100% rename from activity/__init__.py rename to chat/migrations/__init__.py diff --git a/activity/models.py b/chat/models.py similarity index 90% rename from activity/models.py rename to chat/models.py index d4ca9e0..6a99681 100644 --- a/activity/models.py +++ b/chat/models.py @@ -36,4 +36,6 @@ class Worker(Model): class Websocket(Model): worker = ForeignKey(Worker, on_delete=CASCADE, related_name="websockets") - user = ForeignKey(User, on_delete=CASCADE, related_name="websockets") + user = ForeignKey( + User, on_delete=CASCADE, related_name="websockets", blank=True, null=True + ) diff --git a/chat/settings.py b/chat/settings.py index 9200742..7458cbe 100644 --- a/chat/settings.py +++ b/chat/settings.py @@ -22,8 +22,8 @@ DEBUG = os.getenv("DEBUG", "1").lower() in ("1", "t", "true") ALLOWED_HOSTS = os.getenv("ALLOWED_HOSTS", "localhost").split(",") INSTALLED_APPS = [ - "activity", "chat", + "conversation", "user", "django.contrib.admin", "django.contrib.auth", @@ -31,6 +31,7 @@ INSTALLED_APPS = [ "django.contrib.sessions", "django.contrib.messages", "django.contrib.staticfiles", + "pgtrigger", ] MIDDLEWARE = [ @@ -81,7 +82,6 @@ AUTH_PASSWORD_VALIDATORS = [ }, {"NAME": "django.contrib.auth.password_validation.MinimumLengthValidator"}, {"NAME": "django.contrib.auth.password_validation.CommonPasswordValidator"}, - {"NAME": "django.contrib.auth.password_validation.NumericPasswordValidator"}, { "NAME": "django.contrib.auth.password_validation.MinimumLengthValidator", "OPTIONS": {"min_length": 9}, @@ -133,6 +133,10 @@ LOGGING = { }, } +SEQ_UPDATES_SLEEP_SECONDS = 10 +from django.conf import global_settings + + try: from .settings_local import * except ImportError: diff --git a/chat/static/chat/index.html b/chat/static/chat/index.html new file mode 100644 index 0000000..3dd93c2 --- /dev/null +++ b/chat/static/chat/index.html @@ -0,0 +1,26 @@ + + + + + + + Main Window + + + + Hello World. + + + + +
+ Requests: +
+ +
+ Replies: +
+
+ + diff --git a/chat/static/chat/js/index.js b/chat/static/chat/js/index.js new file mode 100644 index 0000000..3c6bc28 --- /dev/null +++ b/chat/static/chat/js/index.js @@ -0,0 +1,75 @@ +(function () { + "use strict"; + var socket = null; + + function get_url() { + var scheme = location.protocol === "https:" ? "wss:" : "ws:"; + return scheme + "//" + location.host + "/ws/chat/"; + } + + function toggle_buttons(connect_disabled) { + var i; + var buttons = document.getElementsByTagName("button"); + for (i = 0; i < buttons.length; i += 1) { + switch (buttons[i].textContent) { + case "Connect": + buttons[i].disabled = connect_disabled; + break; + case "Disconnect": + case "Send": + buttons[i].disabled = !connect_disabled; + break; + } + } + } + + function connect() { + var ws; + if (socket !== null) { + return; + } + ws = new WebSocket(get_url()); + socket = ws; + toggle_buttons(true); + ws.addEventListener( + "open", + function () { + console.log("Connected"); + document.getElementById("send_button").disabled = false; + } + ); + ws.addEventListener("message", function (event) { + var pre = document.createElement("pre"); + pre.innerHTML = event.data; + document.getElementById("replies").appendChild(pre); + }); + ws.addEventListener("close", function (event) { + console.log("Closed", event.code, event.reason); + document.getElementById("send_button").disabled = true; + if (socket === ws) { + toggle_buttons(false); + socket = null; + } + }); + ws.addEventListener("error", function (err) { + console.error("WebSocket error", err); + }); + } + + function send() { + if (socket === null) { + return; + } + socket.send(document.getElementsByTagName("textarea")[0].value); + } + + function disconnect() { + if (socket !== null) { + socket.close(); + } + } + + window.connect_ws = connect; + window.send_ws = send; + window.disconnect_ws = disconnect; +})(); diff --git a/chat/static/chat/js/login.js b/chat/static/chat/js/login.js new file mode 100644 index 0000000..41a27bd --- /dev/null +++ b/chat/static/chat/js/login.js @@ -0,0 +1,35 @@ +(function () { + "use strict"; + function set_form_element_value(form, element_name, value) { + form.elements[element_name].value = value; + } + + function setup_csrftoken() { + var buttons; + var cookie; + var cookie_prefix = "csrftoken="; + var cookies = document.cookie.split(";"); + var form = document.forms[0]; + var i; + + for (i = 0; i < cookies.length; i += 1) { + cookie = cookies[i].replace(/^\s+/, ""); + if (cookie.indexOf(cookie_prefix) === 0) { + set_form_element_value( + form, + "csrfmiddlewaretoken", + cookie.substring(cookie_prefix.length) + ); + return; + } + } + form.elements["username"].disabled = true; + form.elements["password"].disabled = true; + buttons = document.getElementsByTagName("button"); + for (i = 0; i < buttons.length; i += 1) { + buttons[i].disabled = true; + } + document.getElementsByTagName("h1")[0].textContent += " (csrf cookie missing)"; + } + window.addEventListener("load", setup_csrftoken); +})(); diff --git a/chat/static/chat/login.html b/chat/static/chat/login.html new file mode 100644 index 0000000..70ee15b --- /dev/null +++ b/chat/static/chat/login.html @@ -0,0 +1,29 @@ + + + + + + + Login + + + +

Login

+
+ + +
+ + +
+ + +
+ + + diff --git a/chat/urls.py b/chat/urls.py index 69ffe9a..e79573c 100644 --- a/chat/urls.py +++ b/chat/urls.py @@ -1,22 +1,26 @@ -""" -URL configuration for chat project. - -The `urlpatterns` list routes URLs to views. For more information please see: - https://docs.djangoproject.com/en/6.1/topics/http/urls/ -Examples: -Function views - 1. Add an import: from my_app import views - 2. Add a URL to urlpatterns: path('', views.home, name='home') -Class-based views - 1. Add an import: from other_app.views import Home - 2. Add a URL to urlpatterns: path('', Home.as_view(), name='home') -Including another URLconf - 1. Import the include() function: from django.urls import include, path - 2. Add a URL to urlpatterns: path('blog/', include('blog.urls')) -""" +from django.conf import settings from django.contrib import admin -from django.urls import path +from django.contrib.staticfiles.views import serve +from django.urls import path, re_path +from django.views.decorators.csrf import ensure_csrf_cookie + +from .views import async_serve, login_auth urlpatterns = [ - path('admin/', admin.site.urls), + path("admin/", admin.site.urls), + path("login-auth/", login_auth, name="login-auth"), ] + +if settings.DEBUG: + urlpatterns.extend( + ( + path("", async_serve, {"path": "chat/index.html"}, "main-window"), + path( + "login/", + ensure_csrf_cookie(async_serve), + {"path": "chat/login.html"}, + "login-form", + ), + re_path("^static/(?P.*)$", serve), + ) + ) diff --git a/chat/utils.py b/chat/utils.py index 4b34279..ff51e6f 100644 --- a/chat/utils.py +++ b/chat/utils.py @@ -1,5 +1,7 @@ -from asyncio import TaskGroup -from contextlib import contextmanager +from asyncio import Condition, QueueShutDown, TaskGroup +from contextlib import asynccontextmanager, contextmanager +from functools import partialmethod +from traceback import print_exc class CancellingTaskGroup: @@ -32,3 +34,54 @@ def replace_attr(obj, attr, new_value): delattr(obj, attr) else: setattr(obj, attr, old_value) + + +@asynccontextmanager +async def log_exceptions(): + try: + yield + except Exception: + print_exc() + raise + + +async def iterate_queue(queue): + while True: + try: + item = await queue.get() + except QueueShutDown: + return + try: + yield item + finally: + queue.task_done() + + +class BooleanCondition: + __slots__ = ("_cond", "_state") + + def __init__(self, initial_state = False): + self._cond = Condition() + self._state = bool(initial_state) + + async def _set(self, value): + async with self._cond: + if self._state != value: + self._state = value + self._cond.notify_all() + + set = partialmethod(_set, True) + clear = partialmethod(_set, False) + + async def _wait(self, value): + async with self._cond: + await self._cond.wait_for(lambda: self._state == value) + + wait_set = partialmethod(_wait, True) + wait_clear = partialmethod(_wait, False) + + def is_set(self): + return self._state + + def __bool__(self): + return self._state diff --git a/chat/views.py b/chat/views.py new file mode 100644 index 0000000..cd05f45 --- /dev/null +++ b/chat/views.py @@ -0,0 +1,49 @@ +import mimetypes +from asyncio import to_thread +from pathlib import Path +from urllib.parse import urlunsplit + +from django.contrib.auth import alogin, acheck_password_with_timing_attack_mitigation +from django.contrib.staticfiles import finders +from django.http import Http404, HttpResponse, HttpResponseRedirect +from django.urls import reverse +from django.views.decorators.http import require_POST + +from .asyncpg_utils import AsyncPGWrapper +from user.models import User + + +async def async_serve(_, path): + """ + Simple drop-in for django.contrib.staticfiles.serve, + but does not print a warning when used in an asgi environment + """ + abspath = await to_thread(finders.find, path) + if abspath is None: + raise Http404(path) + data = await to_thread(Path(abspath).read_bytes) + content_type = mimetypes.guess_type(abspath)[0] + return HttpResponse(data, content_type=content_type or "application/octet-stream") + + +@require_POST +async def login_auth(request): + username = request.POST.get("username") + password = request.POST.get("password") + if username is not None and password is not None: + async with AsyncPGWrapper.acquire() as wrapper: + for kwargs in ({"username": username}, {"email": username}): + try: + user: User = await wrapper.aselect_one( + User, is_active=True, **kwargs + ) + if await acheck_password_with_timing_attack_mitigation( + user, password + ): + await alogin(request, user) + return HttpResponseRedirect(reverse("main-window")) + except (User.DoesNotExist, User.MultipleObjectsReturned): + continue + return HttpResponseRedirect( + urlunsplit(("", "", reverse("login-form"), "error=invalid", "")) + ) diff --git a/chat/websocket.py b/chat/websocket.py index d4f0a3a..120fd02 100644 --- a/chat/websocket.py +++ b/chat/websocket.py @@ -1,61 +1,75 @@ import json -from asyncio import CancelledError, Queue, create_task, shield, wait_for -from contextlib import suppress +from asyncio import Queue, create_task, shield, wait_for +from contextlib import asynccontextmanager, suppress +from importlib import import_module +from django.conf import settings +from django.contrib.auth import SESSION_KEY from django.contrib.auth.models import AnonymousUser -StopSend = object() +from .utils import BooleanCondition, iterate_queue + +session_engine = import_module(settings.SESSION_ENGINE) + + +async def aget_user_id(headers) -> str | None: + for key, value in headers: + if key.lower() != b"cookie": + continue + for cookie in value.decode().split(";"): + name, value = cookie.split("=", 1) + if name.lstrip() != "sessionid": + continue + return await session_engine.SessionStore(session_key=value).aget( + SESSION_KEY + ) + return None class Websocket: - # keep it dumb, we'd like to run many of these concurrently + __slots__ = ("close_code", "instance", "registered", "send_queue", "user") + def __init__(self): - self.websocket_instance = None - self.user = AnonymousUser + self.close_code = 1000 self.send_queue = Queue() + self.instance = None + self.user = AnonymousUser + self.registered = BooleanCondition() - async def serve(self, send, receive): + @asynccontextmanager + async def register_with_worker(self, user_id): from .lifespan import LifespanManager - disconnected = False - send_task = None + self_id = id(self) + LifespanManager.websockets[self_id] = self try: - await send({"type": "websocket.accept"}) - send_task = create_task(self.serve_send_queue(send)) - LifespanManager.websockets[id(self)] = self - await LifespanManager.queue.put({"type": "register", "websocket": id(self)}) - while True: - event = await receive() - if event["type"] == "websocket.receive": - text = event.get("text") - if text is not None: - await self.send_queue.put( - { - "type": "event", - "websocket": id(self), - "data": json.loads(text), - } - ) - elif event["type"] == "websocket.disconnect": - disconnected = True - break - except CancelledError: - if disconnected: - raise - if send_task: - send_task.cancel() - raise - finally: await LifespanManager.queue.put( - {"type": "unregister", "websocket": id(self)} + {"type": "register", "user_id": user_id, "websocket": self_id} ) - if send_task: - await wait_for(send_task, 10) + await self.registered.wait_set() + yield LifespanManager.queue + finally: + if self.registered: + await LifespanManager.queue.put( + {"type": "unregister", "websocket": self_id} + ) + await self.registered.wait_clear() + self.instance, self.user = None, AnonymousUser + self.send_queue.shutdown() + LifespanManager.websockets.pop(self_id, None) + + @asynccontextmanager + async def guard_send_task(self, send): + send_task = create_task(self.serve_send_queue(send)) + try: + yield + finally: + self.send_queue.shutdown() + await wait_for(send_task, 10) async def serve_send_queue(self, send): try: - while True: - item = await self.send_queue.get() + async for item in iterate_queue(self.send_queue): if isinstance(item, dict): await send( { @@ -63,11 +77,59 @@ class Websocket: "text": json.dumps(item, ensure_ascii=False), } ) - self.send_queue.task_done() - if item is StopSend: - break finally: - with suppress(Exception): - await wait_for( - shield(send({"type": "websocket.close", "code": 1013})), 10 - ) + if self.close_code is not None: + with suppress(Exception): + await wait_for( + shield( + send({"type": "websocket.close", "code": self.close_code}) + ), + 10, + ) + + async def serve_until_connect(self, receive): + while True: + match (await receive())["type"]: + case "websocket.connect": + return True + case "websocket.disconnect": + return False + + async def process_event(self, worker_queue, text): + if not text: + return + try: + data = json.loads(text) + except ValueError: + self.close_code = 1007 + raise + await worker_queue.put( + { + "type": "event", + "websocket": id(self), + "data": data, + } + ) + + async def serve_until_disconnect(self, receive, worker_queue): + while True: + match (event := await receive())["type"]: + case "websocket.receive": + await self.process_event(worker_queue, event.get("text")) + case "websocket.disconnect": + self.close_code = None + break + + async def serve(self, scope, receive, send): + if not await self.serve_until_connect(receive): + return + await send({"type": "websocket.accept"}) + async with self.guard_send_task(send): + if ( + scope["path"] != "/ws/chat/" + or (user_id := await aget_user_id(scope["headers"])) is None + ): + self.close_code = 1008 + return + async with self.register_with_worker(int(user_id)) as worker_queue: + await self.serve_until_disconnect(receive, worker_queue) diff --git a/conversation/migrations/0001_initial.py b/conversation/migrations/0001_initial.py new file mode 100644 index 0000000..20c660c --- /dev/null +++ b/conversation/migrations/0001_initial.py @@ -0,0 +1,22 @@ +# Generated by Django 6.1.2 on 2026-10-12 16:56 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ] + + operations = [ + migrations.CreateModel( + name='Conversation', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('dm_key', models.CharField(blank=True, max_length=64, null=True)), + ('seq', models.BigIntegerField()), + ], + ), + ] diff --git a/conversation/migrations/0002_initial.py b/conversation/migrations/0002_initial.py new file mode 100644 index 0000000..5eb5490 --- /dev/null +++ b/conversation/migrations/0002_initial.py @@ -0,0 +1,36 @@ +# Generated by Django 6.1.2 on 2026-10-12 16:56 + +import pgtrigger.compiler +import pgtrigger.migrations +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ('conversation', '0001_initial'), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.AddField( + model_name='conversation', + name='users', + field=models.ManyToManyField(related_name='conversations', to=settings.AUTH_USER_MODEL), + ), + pgtrigger.migrations.AddTrigger( + model_name='conversation', + trigger=pgtrigger.compiler.Trigger(name='notify_delete', sql=pgtrigger.compiler.UpsertTriggerSql(func="PERFORM pg_notify( 'chat_channel', jsonb_build_object( 'op', TG_OP, 'table', TG_TABLE_NAME::text, 'id', OLD.id )::text); RETURN NULL;", hash='bcf7a86be12aa014c97be65f3800c5324d552627', operation='DELETE', pgid='pgtrigger_notify_delete_c9ca6', table='conversation_conversation', when='AFTER')), + ), + pgtrigger.migrations.AddTrigger( + model_name='conversation', + trigger=pgtrigger.compiler.Trigger(name='notify_insert_update', sql=pgtrigger.compiler.UpsertTriggerSql(func="PERFORM pg_notify( 'chat_channel', jsonb_build_object( 'op', TG_OP, 'table', TG_TABLE_NAME::text, 'id', NEW.id )::text); RETURN NULL;", hash='7645a8ddb34827fffb4256bc2d053a3617cb2ba3', operation='INSERT OR UPDATE', pgid='pgtrigger_notify_insert_update_c3422', table='conversation_conversation', when='AFTER')), + ), + pgtrigger.migrations.AddTrigger( + model_name='conversation', + trigger=pgtrigger.compiler.Trigger(name='notify_truncate', sql=pgtrigger.compiler.UpsertTriggerSql(func="PERFORM pg_notify( 'chat_channel', jsonb_build_object( 'op', TG_OP, 'table', TG_TABLE_NAME::text )::text); RETURN NULL;", hash='8eacbe262f80a66ffc527f150972eddad34df819', level='STATEMENT', operation='TRUNCATE', pgid='pgtrigger_notify_truncate_ddf6b', table='conversation_conversation', when='AFTER')), + ), + ] diff --git a/gunicorn.conf.py b/gunicorn.conf.py index 060fea5..04f6b3c 100644 --- a/gunicorn.conf.py +++ b/gunicorn.conf.py @@ -1,9 +1,18 @@ -from activity.models import Worker +from django.db import connections + +from chat.models import Worker def on_starting(_): assert not Worker.cleanup() + connections.close_all() def on_exit(_): Worker.cleanup() + connections.close_all() + + +def child_exit(_, __): + Worker.cleanup() + connections.close_all() diff --git a/setup_venv.sh b/setup_venv.sh index cc3a7ff..7c9f4e1 100644 --- a/setup_venv.sh +++ b/setup_venv.sh @@ -20,4 +20,5 @@ if [[ -z "$(ip route show default)" ]]; then pip_args+=(--no-index --no-build-isolation) fi PIP_RETRIES=2 PIP_TIMEOUT=2 "${PYTHON%%*/}" -m pip install "${pip_args[@]}" \ - asyncpg django django-pgtrigger gunicorn pip 'psycopg[c]' uvicorn-worker + asyncpg django django-stubs django-pgtrigger gunicorn pip 'psycopg[c]' \ + uvicorn-worker websockets diff --git a/user/migrations/0001_initial.py b/user/migrations/0001_initial.py index c138ea5..c98bec7 100644 --- a/user/migrations/0001_initial.py +++ b/user/migrations/0001_initial.py @@ -1,8 +1,10 @@ -# Generated by Django 6.1.1 on 2026-09-27 10:49 +# Generated by Django 6.1.2 on 2026-10-12 16:56 -import django.contrib.auth.models import django.contrib.auth.validators import django.utils.timezone +import pgtrigger.compiler +import pgtrigger.migrations +import user.models from django.db import migrations, models @@ -29,16 +31,25 @@ class Migration(migrations.Migration): ('is_staff', models.BooleanField(default=False, help_text='Designates whether the user can log into this admin site.', verbose_name='staff status')), ('is_active', models.BooleanField(default=True, help_text='Designates whether this user should be treated as active. Unselect this instead of deleting accounts.', verbose_name='active')), ('date_joined', models.DateTimeField(default=django.utils.timezone.now, verbose_name='date joined')), + ('last_active', models.DateTimeField(auto_now=True)), + ('status', models.CharField(max_length=256)), ('groups', models.ManyToManyField(blank=True, help_text='The groups this user belongs to. A user will get all permissions granted to each of their groups.', related_name='user_set', related_query_name='user', to='auth.group', verbose_name='groups')), ('user_permissions', models.ManyToManyField(blank=True, help_text='Specific permissions for this user.', related_name='user_set', related_query_name='user', to='auth.permission', verbose_name='user permissions')), ], - options={ - 'verbose_name': 'user', - 'verbose_name_plural': 'users', - 'abstract': False, - }, managers=[ - ('objects', django.contrib.auth.models.UserManager()), + ('objects', user.models.UserManager()), ], ), + pgtrigger.migrations.AddTrigger( + model_name='user', + trigger=pgtrigger.compiler.Trigger(name='notify_delete', sql=pgtrigger.compiler.UpsertTriggerSql(func="PERFORM pg_notify( 'chat_channel', jsonb_build_object( 'op', TG_OP, 'table', TG_TABLE_NAME::text, 'id', OLD.id )::text); RETURN NULL;", hash='3eb6b22cb07ed7ed6cb71de3ca66e1ef5f6be480', operation='DELETE', pgid='pgtrigger_notify_delete_9e81d', table='user_user', when='AFTER')), + ), + pgtrigger.migrations.AddTrigger( + model_name='user', + trigger=pgtrigger.compiler.Trigger(name='notify_insert_update', sql=pgtrigger.compiler.UpsertTriggerSql(func="PERFORM pg_notify( 'chat_channel', jsonb_build_object( 'op', TG_OP, 'table', TG_TABLE_NAME::text, 'id', NEW.id )::text); RETURN NULL;", hash='0f814337c836aad989d061f7a737d43d10cf72b0', operation='INSERT OR UPDATE', pgid='pgtrigger_notify_insert_update_98106', table='user_user', when='AFTER')), + ), + pgtrigger.migrations.AddTrigger( + model_name='user', + trigger=pgtrigger.compiler.Trigger(name='notify_truncate', sql=pgtrigger.compiler.UpsertTriggerSql(func="PERFORM pg_notify( 'chat_channel', jsonb_build_object( 'op', TG_OP, 'table', TG_TABLE_NAME::text )::text); RETURN NULL;", hash='5ed1b4a60771276933ae5065b2cfed19c41597c8', level='STATEMENT', operation='TRUNCATE', pgid='pgtrigger_notify_truncate_9c110', table='user_user', when='AFTER')), + ), ] diff --git a/user/password_validation.py b/user/password_validation.py new file mode 100644 index 0000000..26e2622 --- /dev/null +++ b/user/password_validation.py @@ -0,0 +1,27 @@ +from django.core.exceptions import ValidationError +from django.utils.translation import gettext_lazy as _ + + +class ComplexityValidator: + def validate(self, password, user=None): + if not any(char.isupper() for char in password): + raise ValidationError( + _("The password must contain at least 1 uppercase letter."), + code="password_no_upper", + ) + if not any(char.islower() for char in password): + raise ValidationError( + _("The password must contain at least 1 lowercase letter."), + code="password_no_lower", + ) + if not any(char.isdigit() for char in password): + raise ValidationError( + _("The password must contain at least 1 digit."), + code="password_no_number", + ) + + def get_help_text(self): + return _( + "Your password must contain at least 1 uppercase letter, " + "1 lowercase letter, and 1 digit." + )