proekton.com/chat/chat.py

import html
from tornado import gen, web, websocket, escape, options
from tornado.ioloop import IOLoop
from tornado.httpserver import HTTPServer
from tornado.options import parse_command_line
from settings import settings, PORT, DATABASE_DSN

import psycopg2
import momoko
import json


class BaseHandler(web.RequestHandler):
    @property
    def db(self):
        return self.application.db


class ChatHandler(websocket.WebSocketHandler):
    @property
    def db(self):
        return self.application.db

    waiters = set()

    def open(self, *args, **kwargs):
        self.user_id = kwargs.get('user_id', 1)
        self.waiters.add((self.user_id, self))

    # @gen.coroutine
    def on_message(self, message):
        parsed = escape.json_decode(message)

        if 'dummy' in parsed:
            return

        if 'approve_stages' in parsed['format_type']:
            self.approve_stages(parsed)
        else:
            self.add_message(parsed)

    def on_close(self):
        self.waiters.remove((self.user_id, self))

    @gen.coroutine
    def approve_stages(self, data):
        data['data']['chat_message'] = data['data']['msg']
        self.add_message(data)
        # sender_id = data['data']['sender_id']
        # recipent_id = data['data']['recipent_id']
        # order_id = data['data'].get('order_id')
        # message = data['data'].get('msg', 'Этапы обновлены')
        # answer_type = data['format_type']
        # waiters = tuple(w for c, w in self.waiters if c == recipent_id)
        # for waiter in waiters:
        #     print(waiter)
        #     # waiter.write_message({'msg': message, 'order_id': order_id, 'answer_type': answer_type})

    @gen.coroutine
    def add_message(self, message_data):
        sender_id = message_data['data']['sender_id']
        recipent_id = message_data['data'].get('recipent_id', None)
        order_id = message_data['data'].get('order_id', None)
        team_id = message_data['data'].get('team_id', None)
        team_ids_raw = message_data['data'].get('team_ids', None)
        message = message_data['data'].get('chat_message', None)
        docs_send_links = message_data['data'].get('document_send_links', None)
        if 'document_data' in message_data['data']:
            docs_links = message_data['data']['document_data'].get('document_links', "");
            docs_attach = message_data['data']['document_data'].get('document_attach_files', "")
        else:
            docs_links = ''
            docs_attach = ''

        message_type = message_data.get('message_type', None)
        if not message_type:
            message = html.escape(message)
            message = message.replace('\n', '<br />')

        answer_type = message_data['format_type']
        private_type = 'true' if not order_id and not team_id else 'false'
        is_new = 'true'
        is_delete = 'false'

        team_value = "NULL" if team_id is None or not team_id else team_id

        if not recipent_id:
            recipent_id = sender_id

        order_value = "NULL" if order_id is None or not order_id else order_id

        insert_sql = "INSERT INTO chat_message (id,text,created, sender_id,recipent_id," \
                     " private_type,team_id, order_id,is_delete,is_new) " \
                     "VALUES (DEFAULT,'{0}',NOW(),{1},{2},{3},{4},{5},{6},{7}) RETURNING id".\
            format(message, sender_id, recipent_id, private_type, team_value,order_value, is_delete, is_new)

        cursor_list = yield dict(cursor=self.db.execute(insert_sql))
        cursor = cursor_list.get('cursor')
        result = cursor.fetchone()
        message_id = result[0]
        if team_ids_raw:
            team_ids = [t for t in team_ids_raw.rstrip(';').split(';')]
            values_str = '';
            for t in team_ids:
                values_str +='(DEFAULT,{0},{1}),'.format(message_id,t)
            values_str = values_str.rstrip(',')
            insert_new_messages = "INSERT INTO chat_newmessage (id,message_id, user_id) VALUES{0}".\
                format(values_str)
        else:
            insert_new_messages = "INSERT INTO chat_newmessage (id,message_id, user_id) VALUES(DEFAULT,{0},{1})".\
                format(message_id, recipent_id)
        yield self.db.execute(insert_new_messages)

        if docs_send_links:
            is_send = 'true'
            docs_send_ids = docs_send_links.rstrip(';').replace(';', ',')
            update_sql_documents = "UPDATE chat_documents SET message_id={0},is_send={1} WHERE id IN({2})".format(message_id, is_send, docs_send_ids)
            yield self.db.execute(update_sql_documents)

        select_last_sql = "SELECT chat_message.id, chat_message.text, chat_message.created, chat_message.sender_id," \
                          "users_user.id, users_user.username FROM chat_message" \
                          " INNER JOIN users_user ON (chat_message.sender_id = users_user.id)" \
                          " WHERE chat_message.id = {0}".format(message_id)

        cursor_msg = yield self.db.execute(select_last_sql)
        msg_data = cursor_msg.fetchone()
        sender_name = msg_data[5]
        msg_time = msg_data[2].strftime("%Y-%m-%d %H:%M:%S")
        if docs_links:
            message += '<br><br>' + docs_links;

        if message_type:
            waiters = tuple(w for c, w in self.waiters if c == recipent_id)
        else:
            waiters = tuple(w for c, w in self.waiters if c == recipent_id or c == sender_id)

        for waiter in waiters:
            waiter.write_message({'msg': message,
                                  'msg_time': msg_time,
                                  'order_id': order_id,
                                  'recipent_id': recipent_id,
                                  'sender_id': sender_id,
                                  'sender_name': sender_name,
                                  'answer_type': answer_type,
                                  'docs_attach': docs_attach,
            })

    def check_origin(self, origin):
        return True


class Application(web.Application):
    def __init__(self):
        handlers = [
            (r"/chat/(?P<user_id>\d+)/", ChatHandler),
        ]
        super().__init__(handlers=handlers, **settings)


if __name__ == '__main__':
    parse_command_line()
    application = Application()
    ioloop = IOLoop.instance()

    application.db = momoko.Pool(
        dsn=DATABASE_DSN,
        size=1,
        ioloop=ioloop,
    )
    future = application.db.connect()
    ioloop.add_future(future, lambda f: ioloop.stop())
    ioloop.start()
    future.result()

    http_server = HTTPServer(application)
    http_server.listen(PORT, '127.0.0.1')
    ioloop.start()