<commit_before>import struct
from pymongo.bson import _make_c_string, BSON
from diesel import wait, fire, call, response, up
from diesel.protocols.mongodb import MongoProxy, Ops, MongoClient

OP_SUBSCRIBE = 1500

class SubscribingClient(MongoClient):
    @call
    def subscribe(self, col, spec):
        op = OP_SUBSCRIBE
        data = [
            "\x00\x00\x00\x00",
            _make_c_string(col), 
            struct.pack('<ii', 0, 0),
            BSON.from_dict(spec),
        ]
        msg = "".join(data)
        resp = yield self._put_request(op, msg)
        yield response('')

class SubscriptionProxy(MongoProxy):
    def handle_request(self, info, body):
        length, id, to, opcode = info
        if opcode in [OP_SUBSCRIBE, Ops.OP_UPDATE]:
            trimmed = body[4:]
            col, rest = trimmed.split('\0', 1)
            if opcode == OP_SUBSCRIBE:
                sl, raw_bson = rest[:8], rest[8:]
                spec = BSON(raw_bson).to_dict()
                yield wait(('update', col, str(spec)))
                yield up(('', info, body))
            elif opcode == Ops.OP_UPDATE:
                upsert, raw_bson = rest[:4], rest[4:]
                
                raw_spec, raw_doc = raw_bson.split('\x00\x00\x00\x00', 1)
                spec = BSON("%s\x00\x00\x00\x00" % raw_spec).to_dict()
                yield fire(('update', col, str(spec)))
                yield up((None, info, body))
        else:
            yield up((None, info, body))

if __name__ == '__main__':
    from diesel import Application, Service, sleep, Loop
    BACKEND_HOST = 'localhost'
    BACKEND_PORT = 27017
    FRONTEND_PORT = 27018
    def subber():
        c = SubscribingClient()
        c.connect(BACKEND_HOST, FRONTEND_PORT)
        print "subscribing ..."
        yield c.subscribe('sub.test', {'room':'general'})
        print "something published!"
        with (yield c.sub.test.find({'room':'general'})) as cursor:
            result = yield cursor.more()
            print "got", result
        a.halt()

    def pubber():
        c = MongoClient()
        c.connect(BACKEND_HOST, FRONTEND_PORT)
        print "pubber sleeping"
        yield sleep(5)
        yield c.sub.test.update({'room':'general'}, {'$set': {'users':['pubber']}}, upsert=1)
        print "pubber updated"

    a = Application()
    a.add_service(Service(SubscriptionProxy(BACKEND_HOST, BACKEND_PORT), FRONTEND_PORT))<commit_msg>mongodb: fixes to the proxy example application.<commit_after>import struct
from pymongo.bson import _make_c_string, BSON, _bson_to_dict
from diesel import wait, fire, call, response, up, bytes
from diesel.protocols.mongodb import MongoProxy, Ops, MongoClient

OP_SUBSCRIBE = 1500

class SubscribingClient(MongoClient):
    @call
    def subscribe(self, col, spec):
        op = OP_SUBSCRIBE
        data = [
            "\x00\x00\x00\x00",
            _make_c_string(col), 
            struct.pack('<ii', 0, 0),
            BSON.from_dict(spec),
        ]
        msg = "".join(data)
        yield self._put_request(op, msg)
        doclen = struct.unpack('<i', (yield bytes(4)))[0]
        doc = yield bytes(doclen)
        yield response(BSON(doc).to_dict())

class SubscriptionProxy(MongoProxy):
    def handle_request(self, info, body):
        length, id, to, opcode = info
        if opcode in [OP_SUBSCRIBE, Ops.OP_UPDATE]:
            trimmed = body[4:]
            col, rest = trimmed.split('\0', 1)
            if opcode == OP_SUBSCRIBE:
                sl, raw_bson = rest[:8], rest[8:]
                spec = BSON(raw_bson).to_dict()
                doc = yield wait(('update', col, str(spec)))
                doclen = len(doc)
                resp = "%s%s" % (struct.pack('<i', doclen), doc)
                yield up((resp, info, body))
            elif opcode == Ops.OP_UPDATE:
                upsert, raw_bson = rest[:4], rest[4:]
                spec, raw_doc = _bson_to_dict(raw_bson)
                yield fire(('update', col, str(spec)), raw_doc)
                yield up((None, info, body))
        else:
            yield up((None, info, body))

if __name__ == '__main__':
    from diesel import Application, Service, sleep, Loop
    BACKEND_HOST = 'localhost'
    BACKEND_PORT = 27017
    FRONTEND_PORT = 27018
    def subber():
        c = SubscribingClient()
        c.connect(BACKEND_HOST, FRONTEND_PORT)
        print "subber: subscribing ..."
        doc = yield c.subscribe('sub.test', {'room':'general'})
        print "subber: saw updated doc %r" % doc
        with (yield c.sub.test.find({'room':'general'})) as cursor:
            result = yield cursor.more()
            print "subber: got", result
        a.halt()

    def pubber():
        c = MongoClient()
        c.connect(BACKEND_HOST, FRONTEND_PORT)
        print "pubber: sleeping ..."
        yield sleep(5)
        print "pubber: updating ..."
        yield c.sub.test.update({'room':'general'}, {'$set': {'users':['pubber']}}, upsert=1)
        print "pubber: updated"

    a = Application()
    a.add_service(Service(SubscriptionProxy(BACKEND_HOST, BACKEND_PORT), FRONTEND_PORT))
{'fraction_non_alphanumeric': 0.1327329678935004, 'fraction_numerical': 0.01664056382145654, 'mean_word_length': 2.2051442910915933, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4502174', 'n_tokens_mistral': 1727, 'n_tokens_neox': 1610, 'n_words': 386}