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): 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 or c == sender_id) for waiter in waiters: 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) 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 = html.escape(message) message = message.replace('\n', '
') 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 team_id and 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 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 += '

' + docs_links; 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\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()