Initial commit
* Examples work * setup.py kinda updasted * Fork of txmongo but with new pymongo embedded
This commit is contained in:
Don Brown
191
asyncio_mongo/protocol.py
Normal file
191
asyncio_mongo/protocol.py
Normal file
@@ -0,0 +1,191 @@
|
||||
# coding: utf-8
|
||||
# Copyright 2009 Alexandre Fiori
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import logging
|
||||
|
||||
import struct
|
||||
import asyncio
|
||||
from asyncio_mongo.exceptions import ConnectionLostError
|
||||
import asyncio_mongo._bson as bson
|
||||
from asyncio_mongo.log import logger
|
||||
|
||||
_ONE = b"\x01\x00\x00\x00"
|
||||
_ZERO = b"\x00\x00\x00\x00"
|
||||
|
||||
"""Low level connection to Mongo."""
|
||||
|
||||
|
||||
class _MongoQuery(object):
|
||||
def __init__(self, id, collection, limit):
|
||||
self.id = id
|
||||
self.limit = limit
|
||||
self.collection = collection
|
||||
self.documents = []
|
||||
self.future = asyncio.Future()
|
||||
|
||||
|
||||
class MongoProtocol(asyncio.Protocol):
|
||||
def __init__(self):
|
||||
self.__id = 0
|
||||
self.__buffer = b""
|
||||
self.__queries = {}
|
||||
self.__datalen = None
|
||||
self.__response = 0
|
||||
self.__waiting_header = True
|
||||
self._pipelined_calls = set() # Set of all the pipelined calls.
|
||||
self.transport = None
|
||||
self._is_connected = False
|
||||
|
||||
def connection_made(self, transport):
|
||||
self.transport = transport
|
||||
self._is_connected = True
|
||||
logger.log(logging.INFO, 'Mongo connection made with %')
|
||||
|
||||
def connection_lost(self, exc):
|
||||
self._is_connected = False
|
||||
self.transport = None
|
||||
|
||||
# Raise exception on all waiting futures.
|
||||
for f in self.__queries:
|
||||
f.set_exception(ConnectionLostError(exc))
|
||||
|
||||
logger.log(logging.INFO, 'Mongo connection lost')
|
||||
|
||||
def data_received(self, data):
|
||||
while self.__waiting_header:
|
||||
self.__buffer += data
|
||||
if len(self.__buffer) < 16:
|
||||
break
|
||||
|
||||
# got full header, 16 bytes (or more)
|
||||
header, extra = self.__buffer[:16], self.__buffer[16:]
|
||||
self.__buffer = b""
|
||||
self.__waiting_header = False
|
||||
datalen, request, response, operation = struct.unpack("<iiii", header)
|
||||
self.__datalen = datalen - 16
|
||||
self.__response = response
|
||||
if extra:
|
||||
self.data_received(extra)
|
||||
break
|
||||
else:
|
||||
if self.__datalen is not None:
|
||||
data, extra = data[:self.__datalen], data[self.__datalen:]
|
||||
self.__datalen -= len(data)
|
||||
else:
|
||||
extra = b""
|
||||
|
||||
self.__buffer += data
|
||||
if self.__datalen == 0:
|
||||
self.message_received(self.__response, self.__buffer)
|
||||
self.__datalen = None
|
||||
self.__waiting_header = True
|
||||
self.__buffer = b""
|
||||
if extra:
|
||||
self.data_received(extra)
|
||||
|
||||
def message_received(self, request_id, packet):
|
||||
# Response Flags:
|
||||
# bit 0: Cursor Not Found
|
||||
# bit 1: Query Failure
|
||||
# bit 2: Shard Config Stale
|
||||
# bit 3: Await Capable
|
||||
# bit 4-31: Reserved
|
||||
QUERY_FAILURE = 1 << 1
|
||||
response_flag, cursor_id, start, length = struct.unpack("<iqii", packet[:20])
|
||||
if response_flag == QUERY_FAILURE:
|
||||
self.query_failure(request_id, cursor_id, response_flag, bson.BSON(packet[20:]).decode())
|
||||
return
|
||||
self.query_success(request_id, cursor_id, bson.decode_all(packet[20:]))
|
||||
|
||||
def send_message(self, operation, collection, message, query_opts=_ZERO):
|
||||
#print "sending %d to %s" % (operation, self)
|
||||
fullname = collection and bson._make_c_string(collection) or b""
|
||||
message = query_opts + fullname + message
|
||||
|
||||
# 16 is the size of the header in bytes
|
||||
header = struct.pack("<iiii", 16 + len(message), self.__id, 0, operation)
|
||||
self.transport.write(header + message)
|
||||
self.__id += 1
|
||||
|
||||
def OP_INSERT(self, collection, docs):
|
||||
docs = [bson.BSON.encode(doc) for doc in docs]
|
||||
self.send_message(2002, collection, b"".join(docs))
|
||||
|
||||
def OP_UPDATE(self, collection, spec, document, upsert=False, multi=False):
|
||||
options = 0
|
||||
if upsert:
|
||||
options += 1
|
||||
if multi:
|
||||
options += 2
|
||||
|
||||
message = struct.pack("<i", options) + \
|
||||
bson.BSON.encode(spec) + bson.BSON.encode(document)
|
||||
self.send_message(2001, collection, message)
|
||||
|
||||
def OP_DELETE(self, collection, spec):
|
||||
self.send_message(2006, collection, _ZERO + bson.BSON.encode(spec))
|
||||
|
||||
def OP_KILL_CURSORS(self, cursors):
|
||||
message = struct.pack("<i", len(cursors))
|
||||
for cursor_id in cursors:
|
||||
message += struct.pack("<q", cursor_id)
|
||||
self.send_message(2007, None, message)
|
||||
|
||||
def OP_GET_MORE(self, collection, limit, cursor_id):
|
||||
message = struct.pack("<iq", limit, cursor_id)
|
||||
self.send_message(2005, collection, message)
|
||||
|
||||
def OP_QUERY(self, collection, spec, skip, limit, fields=None):
|
||||
message = struct.pack("<ii", skip, limit) + bson.BSON.encode(spec)
|
||||
if fields:
|
||||
message += bson.BSON.encode(fields)
|
||||
|
||||
query = _MongoQuery(self.__id, collection, limit)
|
||||
self.__queries[self.__id] = query
|
||||
self.send_message(2004, collection, message)
|
||||
return query.future
|
||||
|
||||
def query_failure(self, request_id, cursor_id, response, raw_error):
|
||||
query = self.__queries.pop(request_id, None)
|
||||
if query:
|
||||
query.future.set_exception(ValueError("mongo error=%s" % repr(raw_error)))
|
||||
del query
|
||||
|
||||
def query_success(self, request_id, cursor_id, documents):
|
||||
try:
|
||||
query = self.__queries.pop(request_id)
|
||||
except KeyError:
|
||||
return
|
||||
if isinstance(documents, list):
|
||||
query.documents += documents
|
||||
else:
|
||||
query.documents.append(documents)
|
||||
if cursor_id:
|
||||
query.id = self.__id
|
||||
next_batch = 0
|
||||
if query.limit:
|
||||
next_batch = query.limit - len(query.documents)
|
||||
# Assert, because according to the protocol spec and my observations
|
||||
# there should be no problems with this, but who knows? At least it will
|
||||
# be noticed, if something unexpected happens. And it is definitely
|
||||
# better, than silently returning a wrong number of documents
|
||||
assert next_batch >= 0, "Unexpected number of documents received!"
|
||||
if not next_batch:
|
||||
self.OP_KILL_CURSORS([cursor_id])
|
||||
query.future.set_result(query.documents)
|
||||
return
|
||||
self.__queries[self.__id] = query
|
||||
self.OP_GET_MORE(query.collection, next_batch, cursor_id)
|
||||
else:
|
||||
query.future.set_result(query.documents)
|
||||
Reference in New Issue
Block a user