def open_stream(self, user_id):
"""Opens a new message stream for the given user_id.
WARNING: This function returns a *blocking* generator which yields new
messages as they are received from the server.
"""
stream_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
stream_socket.connect(self._broadcast_address)
request = protocol.OpenStreamRequest(user_id)
socket_lib.send_request('OpenStreamRequest',
request.to_json(),
sock=stream_socket,
keep_alive=True)
while True:
message = socket_lib.recv_message(stream_socket)
message = json.loads(message)['result']
yield message
def receive_one_message(self, user_id, timeout=None):
"""Waits up to 'timeout' seconds to receive a single message.
If 'timeout' is 'None', blocks indefinetly until a message is received.
Returns an empty object if 'timeout' is not 'None' and no message is
received within the 'timeout' interval.
"""
stream_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
stream_socket.settimeout(timeout)
try:
stream_socket.connect(self._broadcast_address)
request = protocol.OpenStreamRequest(user_id)
socket_lib.send_request('OpenStreamRequest',
request.to_json(),
sock=stream_socket,
keep_alive=True)
message = socket_lib.recv_message(stream_socket)
return json.loads(message)['result']
except socket.timeout:
return {}
finally:
stream_socket.close()
def handle_request(self, client_socket):
"""See the base class."""
db_client = database_client.DatabaseClient(self._db_path)
request = socket_lib.recv_message(client_socket)
request = json.loads(request)
method, params, id_ = request['method'], request['params'], request[
'id']
if method == 'GetUser':
request = protocol.GetUserRequest.from_json(params)
user = db_client.get_user(request.user_id)
response = protocol.GetUserResponse(user)
elif method == 'InsertUser':
request = protocol.InsertUserRequest.from_json(params)
user = db_client.insert_user(request.user_name)
response = protocol.InsertUserResponse(user)
elif method == 'GetChats':
# TODO(eugen): This is horrible!!! We need to be smarter about
# how we fill out the chat objects. Either via joins, or by bulk
# requesting users and messages.
request = protocol.GetChatsRequest.from_json(params)
chats = []
for chat in db_client.get_chats(request.user_id):
users = {}
for p in db_client.get_participants(chat.chat_id):
users[p.user_id] = protocol.User(p.user_id, p.user_name)
messages = []
for m in db_client.get_messages(chat.chat_id):
message = protocol.Message(m.message_id, m.chat_id,
users[m.user_id],
m.message_text, m.message_ts)
messages.append(message)
# TODO(eugenhotaj): Move chat_name logic to a helper function.
users = list(users.values())
chat_name = chat.chat_name
if len(users) == 2:
chat_name = [
user.user_name for user in users
if user.user_id != request.user_id
][0]
chat = protocol.Chat(chat.chat_id, chat_name, users, messages)
chats.append(chat)
chats = sorted(chats,
key=lambda chat: chat.messages[-1].message_ts)
response = protocol.GetChatsResponse(chats)
elif method == 'InsertChat':
request = protocol.InsertChatRequest.from_json(params)
user_ids = request.user_ids
chat_id = None
if len(user_ids) == 2:
chat = db_client.get_private_chat_id(*user_ids)
if not chat_id:
chat = db_client.insert_chat(request.chat_name, user_ids)
response = protocol.InsertChatResponse(chat)
elif method == 'GetMessages':
# TODO(eugen): This is horrible!!! We need to be smarter about
# how we fill out the chat objects. Either via joins, or by bulk
# requesting users and messages.
request = protocol.GetMessagesRequest.from_json(params)
users = {}
for p in db_client.get_participants(request.chat_id):
users[p.user_id] = protocol.User(p.user_id, p.user_name)
messages = []
for m in db_client.get_messages(request.chat_id):
message = protocol.Message(m.message_id, m.chat_id,
users[m.user_id], m.message_text,
m.message_ts)
messages.append(message)
response = protocol.GetMessagesResponse(messages)
elif method == 'InsertMessage':
request = protocol.InsertMessageRequest.from_json(params)
message_ts = int(time.time() * constants.MILLIS_PER_SEC)
message = db_client.insert_message(request.chat_id,
request.user_id,
request.message_text,
message_ts)
user = db_client.get_user(request.user_id)
user = protocol.User(user.user_id, user.user_name)
message = protocol.Message(message.message_id, message.chat_id,
user, message.message_text,
message.message_ts)
receivers = db_client.get_participants(request.chat_id)
receiver_ids = [
receiver.user_id for receiver in receivers
if receiver.user_id != request.user_id
]
request = protocol.BroadcastRequest(receiver_ids, message)
socket_lib.send_request('BroadcastRequest',
request.to_json(),
address=self._broadcast_address)
response = protocol.InsertMessageResponse(message)
else:
# TODO(eugenhotaj): Return back a malformed request response.
raise NotImplementedError()
response = {'result': response.to_json(), 'id': id_}
socket_lib.send_message(client_socket, json.dumps(response))
def insert_message(self, chat_id, user_id, message_text):
request = protocol.InsertMessageRequest(chat_id, user_id, message_text)
response = socket_lib.send_request('InsertMessage',
request.to_json(),
address=self._data_address)
return response
def get_messages(self, chat_id):
request = protocol.GetMessagesRequest(chat_id)
response = socket_lib.send_request('GetMessages',
request.to_json(),
address=self._data_address)
return response
def get_chats(self, user_id):
request = protocol.GetChatsRequest(user_id)
response = socket_lib.send_request('GetChats',
request.to_json(),
address=self._data_address)
return response