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
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")
+++ /dev/null
-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
--- /dev/null
+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)
+ ]
import os
+from importlib import import_module
from socket import gethostname
from asyncio import CancelledError, Queue, sleep
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 = {}
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
--- /dev/null
+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),
+ ),
+ )
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
--- /dev/null
+from django.apps import AppConfig
+
+
+class ConversationConfig(AppConfig):
+ name = 'conversation'
--- /dev/null
+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")]
-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")]
projects / chat / commitdiff