hyperboria/idm/api/services/chat_manager.py

import sys

from grpc import StatusCode
from idm.api.proto.chat_manager_service_pb2 import Chat as ChatPb
from idm.api.proto.chat_manager_service_pb2 import Chats as ChatsPb
from idm.api.proto.chat_manager_service_pb2_grpc import (
    ChatManagerServicer,
    add_ChatManagerServicer_to_server,
)
from library.aiogrpctools.base import (
    BaseService,
    aiogrpc_request_wrapper,
)
from psycopg.rows import dict_row
from pypika import (
    PostgreSQLQuery,
    Table,
)


class ChatManagerService(ChatManagerServicer, BaseService):
    chats_table = Table('chats')

    async def start(self):
        add_ChatManagerServicer_to_server(self, self.application.server)

    @aiogrpc_request_wrapper()
    async def create_chat(self, request, context, metadata):
        chat = ChatPb(
            chat_id=request.chat_id,
            language=request.language,
            username=request.username,
            is_system_messaging_enabled=True,
            is_discovery_enabled=True,
            is_connectome_enabled=False,
        )
        query = (
            PostgreSQLQuery
            .into(self.chats_table)
            .columns(
                self.chats_table.chat_id,
                self.chats_table.language,
                self.chats_table.username,
                self.chats_table.is_system_messaging_enabled,
                self.chats_table.is_discovery_enabled,
                self.chats_table.is_connectome_enabled,
            )
            .insert(
                chat.chat_id,
                chat.language,
                chat.username,
                chat.is_system_messaging_enabled,
                chat.is_discovery_enabled,
                chat.is_connectome_enabled,
            )
            .on_conflict('chat_id')
            .do_nothing()
        ).get_sql()
        await self.application.pool_holder['idm'].execute(query)
        return await self._get_chat(chat_id=request.chat_id, context=context)

    async def _get_chat(self, chat_id, context):
        sql = (
            PostgreSQLQuery
            .from_(self.chats_table)
            .select(
                self.chats_table.chat_id,
                self.chats_table.username,
                self.chats_table.language,
                self.chats_table.is_system_messaging_enabled,
                self.chats_table.is_discovery_enabled,
                self.chats_table.is_connectome_enabled,
                self.chats_table.ban_until,
                self.chats_table.ban_message,
                self.chats_table.is_admin,
                self.chats_table.created_at,
                self.chats_table.updated_at,
            )
            .where(self.chats_table.chat_id == chat_id)
        ).get_sql()

        chats = [ChatPb(**row) async for row in self.application.pool_holder['idm'].iterate(sql, row_factory=dict_row)]
        if not chats:
            await context.abort(StatusCode.NOT_FOUND, 'not_found')
        return chats[0]

    @aiogrpc_request_wrapper(log=False)
    async def get_chat(self, request, context, metadata):
        return await self._get_chat(chat_id=request.chat_id, context=context)

    @aiogrpc_request_wrapper(log=False)
    async def list_chats(self, request, context, metadata):
        sql = (
            PostgreSQLQuery
            .from_(self.chats_table)
            .select(
                self.chats_table.chat_id,
                self.chats_table.username,
                self.chats_table.language,
                self.chats_table.is_system_messaging_enabled,
                self.chats_table.is_discovery_enabled,
                self.chats_table.is_connectome_enabled,
                self.chats_table.ban_until,
                self.chats_table.ban_message,
                self.chats_table.is_admin,
                self.chats_table.created_at,
                self.chats_table.updated_at,
            )
            .where(self.chats_table.ban_until > request.banned_at_moment)
            .limit(10)
        ).get_sql()
        return ChatsPb(chats=[ChatPb(**row) async for row in self.application.pool_holder['idm'].iterate(sql, row_factory=dict_row)])

    @aiogrpc_request_wrapper()
    async def update_chat(self, request, context, metadata):
        sql = PostgreSQLQuery.update(self.chats_table)
        for field in request.DESCRIPTOR.fields:
            if field.containing_oneof and request.HasField(field.name):
                field_value = getattr(request, field.name)
                sql = sql.set(field.name, field_value)
        sql = sql.where(self.chats_table.chat_id == request.chat_id).returning(
            self.chats_table.chat_id,
            self.chats_table.username,
            self.chats_table.language,
            self.chats_table.is_system_messaging_enabled,
            self.chats_table.is_discovery_enabled,
            self.chats_table.is_connectome_enabled,
            self.chats_table.ban_until,
            self.chats_table.ban_message,
            self.chats_table.is_admin,
            self.chats_table.created_at,
            self.chats_table.updated_at,
        ).get_sql()
        rows = []
        async for row in self.application.pool_holder['idm'].iterate(sql, row_factory=dict_row):
            rows.append(row)
        if not rows:
            return await context.abort(StatusCode.NOT_FOUND, 'not_found')
        return ChatPb(**rows[0])