From ed6b6b0a2e8b6d80c27069153214a7a28b524d7a Mon Sep 17 00:00:00 2001 From: mar77i Date: Tue, 6 Oct 2026 02:23:22 +0200 Subject: [PATCH] gather some basic interactions --- activity/models.py | 13 +- chat/async_db.py | 55 ------- chat/asyncpg_utils.py | 232 ++++++++++++++++++++++++++++ chat/lifespan.py | 68 +++++--- chat/triggers.py | 40 +++++ chat/websocket.py | 8 +- conversation/__init__.py | 0 conversation/apps.py | 5 + conversation/migrations/__init__.py | 0 conversation/models.py | 13 ++ user/models.py | 35 ++++- 11 files changed, 376 insertions(+), 93 deletions(-) delete mode 100644 chat/async_db.py create mode 100644 chat/asyncpg_utils.py create mode 100644 chat/triggers.py create mode 100644 conversation/__init__.py create mode 100644 conversation/apps.py create mode 100644 conversation/migrations/__init__.py create mode 100644 conversation/models.py diff --git a/activity/models.py b/activity/models.py index 3c9c7fd..d4ca9e0 100644 --- a/activity/models.py +++ b/activity/models.py @@ -1,9 +1,7 @@ from pathlib import Path from socket import gethostname -from django.db.models import ( - CASCADE, CharField, ForeignKey, IntegerField, ManyToManyField, Model -) +from django.db.models import CASCADE, CharField, ForeignKey, IntegerField, Model from user.models import User @@ -37,10 +35,5 @@ class Worker(Model): class Websocket(Model): - worker = ForeignKey(Worker, on_delete=CASCADE) - user = ForeignKey(User, on_delete=CASCADE) - - -class Activity(Model): - websockets = ManyToManyField(Websocket, related_name="activity") - status = CharField(max_length=256) + worker = ForeignKey(Worker, on_delete=CASCADE, related_name="websockets") + user = ForeignKey(User, on_delete=CASCADE, related_name="websockets") diff --git a/chat/async_db.py b/chat/async_db.py deleted file mode 100644 index b37e613..0000000 --- a/chat/async_db.py +++ /dev/null @@ -1,55 +0,0 @@ -from django.db import models - - -async def ainsert(con, instance: models.Model, using="default"): - model = type(instance) - concrete_fields = [ - field - for field in model._meta.concrete_fields - if not field.auto_created - ] - # Include the primary key only when the instance has an explicit value. - pk_is_none = getattr(instance, model._meta.pk.attname) is None - insert_columns = [] - insert_values = [] - returning_columns = [] - field_names = [] - for field in concrete_fields: - if not field.primary_key or not pk_is_none: - insert_columns.append(field.column) - insert_values.append(getattr(instance, field.attname)) - returning_columns.append(field.column) - field_names.append(field.attname) - - columns_sql = ", ".join(f'"{column}"' for column in insert_columns) - placeholders_sql = ", ".join(f"${index + 1}" for index in range(len(insert_values))) - returning_sql = ", ".join( - f'"{column}"' for column in returning_columns - ) - row = await con.fetchrow( - ( - f'INSERT INTO "{model._meta.db_table}" ({columns_sql}) ' - f'VALUES ({placeholders_sql}) RETURNING {returning_sql}' - ), - *insert_values, - ) - if row is None: - raise RuntimeError("INSERT did not return a row") - loaded_instance = model.from_db( - using, - field_names, - [row[field.column] for field in concrete_fields], - ) - instance.__dict__.update(loaded_instance.__dict__) - return instance - - -async def adelete(con, instance: models.Model): - model = type(instance) - if instance._state.adding: - raise ValueError("Cannot delete an unsaved instance") - await con.execute( - f'DELETE FROM "{model._meta.db_table}" WHERE "{model._meta.pk.column}"=$1', - instance.pk, - ) - return None diff --git a/chat/asyncpg_utils.py b/chat/asyncpg_utils.py new file mode 100644 index 0000000..1b3afc9 --- /dev/null +++ b/chat/asyncpg_utils.py @@ -0,0 +1,232 @@ +from contextlib import asynccontextmanager + +from django.db import DEFAULT_DB_ALIAS +from django.db.models import Model, Field + + +class SQL: + """ + Marks a trusted raw SQL expression, e.g. SQL("CURRENT_TIMESTAMP"). + Never build these from user input; plain values are parameterized instead. + """ + __slots__ = ("expression",) + + def __init__(self, expression: str): + self.expression = expression + + def __repr__(self): + return f"SQL({self.expression!r})" + + +class SQLHelper: + @staticmethod + def quote(identifier: str): + identifier = identifier.replace('"', '""') + return f'"{identifier}"' + + @staticmethod + def field_is_auto_now(field: Field, add: bool) -> bool: + """Mirror Django's pre_save logic for auto_now / auto_now_add.""" + return getattr(field, "auto_now", False) or ( + add and getattr(field, "auto_now_add", False) + ) + + @staticmethod + def comma_join(quoted_fields: list[str]) -> str: + return ", ".join(quoted_fields) + + def __init__(self, model: type[Model], using=None): + self.using = using + self.model = model + self.attnames = [] + self.columns = [] + self.quoted_columns = [] + self.params = [] + + def add(self, field: Field | str): + if isinstance(field, str): + field = self.model._meta.get_field(field) + attname, column = field.get_attname_column() + self.attnames.append(attname) + self.columns.append(column) + quoted_column = self.quote(column) + self.quoted_columns.append(quoted_column) + return attname, quoted_column + + @property + def quoted_db_table(self): + return self.quote(self.model._meta.db_table) + + @property + def quoted_pk_column(self): + return self.quote(self.model._meta.pk.column) + + @property + def concrete_fields(self): + return self.model._meta.concrete_fields + + def __iter__(self): + for field in self.concrete_fields: + attname, quoted_column = self.add(field) + yield field, attname, quoted_column + + def get_field(self, field_name): + return self.model._meta.get_field(field_name) + + def hydrate(self, row) -> Model: + return self.model.from_db( + self.using or DEFAULT_DB_ALIAS, + self.attnames, + [row[column] for column in self.columns], + ) + + def bind(self, value) -> str: + """Inline SQL(...) expressions; bind everything else as $n.""" + if isinstance(value, SQL): + return value.expression + self.params.append(value) + return f"${len(self.params)}" + + @property + 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: + 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) + ) + ) + + if insert_columns: + body = ( + f"({helper.comma_join(insert_columns)}) " + f"VALUES ({helper.comma_join(insert_values)})" + ) + 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: + 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}" + ) + return helper.hydrate(rows[0]) + + +@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 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) + ] diff --git a/chat/lifespan.py b/chat/lifespan.py index 4e1a6e9..d1f775a 100644 --- a/chat/lifespan.py +++ b/chat/lifespan.py @@ -1,4 +1,5 @@ import os +from importlib import import_module from socket import gethostname from asyncio import CancelledError, Queue, sleep @@ -7,18 +8,42 @@ from contextlib import AsyncExitStack from typing import TypedDict from django.conf import settings +from django.contrib.auth import SESSION_KEY -from .async_db import adelete, ainsert +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) + + class LifespanManager: queue = Queue() websockets = {} @@ -89,30 +114,23 @@ class LifespanManager: async def consume_websocket_queue(self, pool): # websockets -> pg rows (+ triggers) and activity async with pool.acquire() as con: - print("websocket queue") - try: - while item := await self.queue.get(): - websocket = self.websockets[item["websocket"]] - match item["type"]: - case "register": - websocket.websocket_instance = ( - await ainsert( - con, - Websocket( - worker=self.worker_instance, user=None - ), - ) - ) - case "unregister": - websocket.websocket_instance = await adelete( - con, websocket.websocket_instance - ) - websocket.send_queue.put(StopSend) - except CancelledError: - print("cancelled consume_websocket_queue") - raise - finally: - print("websocket queue done") + while item := await self.queue.get(): + websocket = self.websockets[item["websocket"]] + match item["type"]: + case "register": + websocket.websocket_instance = await ainsert( + con, + Websocket( + worker=self.worker_instance, user=user + ), + ) + case "event": + pass + case "unregister": + websocket.websocket_instance = await adelete( + con, websocket.websocket_instance + ) + websocket.send_queue.put(StopSend) async def consume_channel(self, pool, channel_name): # pg channel -> websockets diff --git a/chat/triggers.py b/chat/triggers.py new file mode 100644 index 0000000..6de5604 --- /dev/null +++ b/chat/triggers.py @@ -0,0 +1,40 @@ +from pgtrigger import After, Delete, Insert, Row, Statement, Trigger, Truncate, Update + + +def _notify(channel_name, row=None): + channel_name = channel_name.replace("'", "''") + id_part = f", 'id', {row}.id" if row else "" + return ( + "PERFORM pg_notify(" + f" '{channel_name}'," + f" jsonb_build_object(" + f" 'op', TG_OP, 'table', TG_TABLE_NAME::text{id_part}" + f" )::text" + "); RETURN NULL;" + ) + + +def model_triggers(channel_name): + return ( + Trigger( + name="notify_delete", + level=Row, + operation=Delete, + when=After, + func=_notify(channel_name, "OLD"), + ), + Trigger( + name="notify_insert_update", + level=Row, + operation=Insert | Update, + when=After, + func=_notify(channel_name, "NEW"), + ), + Trigger( + name="notify_truncate", + level=Statement, + operation=Truncate, + when=After, + func=_notify(channel_name), + ), + ) diff --git a/chat/websocket.py b/chat/websocket.py index c6318eb..d4f0a3a 100644 --- a/chat/websocket.py +++ b/chat/websocket.py @@ -29,7 +29,13 @@ class Websocket: if event["type"] == "websocket.receive": text = event.get("text") if text is not None: - await self.send_queue.put(json.loads(text)) + await self.send_queue.put( + { + "type": "event", + "websocket": id(self), + "data": json.loads(text), + } + ) elif event["type"] == "websocket.disconnect": disconnected = True break diff --git a/conversation/__init__.py b/conversation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/conversation/apps.py b/conversation/apps.py new file mode 100644 index 0000000..406f542 --- /dev/null +++ b/conversation/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class ConversationConfig(AppConfig): + name = 'conversation' diff --git a/conversation/migrations/__init__.py b/conversation/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/conversation/models.py b/conversation/models.py new file mode 100644 index 0000000..21866ef --- /dev/null +++ b/conversation/models.py @@ -0,0 +1,13 @@ +from django.db.models import BigIntegerField, CharField, ManyToManyField, Model + +from chat.triggers import model_triggers +from user.models import User + + +class Conversation(Model): + dm_key = CharField(max_length=64, blank=True, null=True) + users = ManyToManyField(User, related_name="conversations") + seq = BigIntegerField() + + class Meta: + triggers = [*model_triggers("chat_channel")] diff --git a/user/models.py b/user/models.py index 3d30525..20f62d1 100644 --- a/user/models.py +++ b/user/models.py @@ -1,5 +1,36 @@ -from django.contrib.auth.models import AbstractUser +from django.contrib.auth.models import AbstractUser, UserManager as DjangoUserManager +from django.db.models import ( + Case, CharField, DateTimeField, Exists, F, OuterRef, QuerySet, When, Value, +) +from chat.triggers import model_triggers -class User(AbstractUser): + +class UserQuerySet(QuerySet): + def with_effective_status(self): + from activity.models import Websocket + + return self.annotate( + effective_status=Case( + When( + Exists(Websocket.objects.filter(user_id=OuterRef("pk"))), + then=F("status"), + ), + default=Value("offline"), + output_field=CharField(), + ) + ) + + +class UserManager(DjangoUserManager): pass + + +class User(AbstractUser): + last_active = DateTimeField(auto_now=True) + status = CharField(max_length=256) + + objects = UserManager() + + class Meta: + triggers = [*model_triggers("chat_channel")] -- 2.55.0