mercurial/wireprotov2server.py
changeset 39559 07b58266bce3
parent 39448 660879e49b46
child 39630 9c2c77c73f23
--- a/mercurial/wireprotov2server.py	Mon Aug 27 13:30:44 2018 -0700
+++ b/mercurial/wireprotov2server.py	Wed Sep 05 09:06:40 2018 -0700
@@ -19,7 +19,6 @@
     wireprototypes,
 )
 from .utils import (
-    cborutil,
     interfaceutil,
 )
 
@@ -295,31 +294,19 @@
             res.setbodybytes(_('command in frame must match command in URL'))
             return True
 
-    rsp = dispatch(repo, proto, command['command'])
-
     res.status = b'200 OK'
     res.headers[b'Content-Type'] = FRAMINGTYPE
 
-    # TODO consider adding a type to represent an iterable of values to
-    # be CBOR encoded.
-    if isinstance(rsp, wireprototypes.cborresponse):
-        # TODO consider calling oncommandresponsereadygen().
-        encoded = b''.join(cborutil.streamencode(rsp.value))
-        action, meta = reactor.oncommandresponseready(outstream,
-                                                      command['requestid'],
-                                                      encoded)
-    elif isinstance(rsp, wireprototypes.v2streamingresponse):
-        action, meta = reactor.oncommandresponsereadygen(outstream,
-                                                         command['requestid'],
-                                                         rsp.gen)
-    elif isinstance(rsp, wireprototypes.v2errorresponse):
-        action, meta = reactor.oncommanderror(outstream,
-                                              command['requestid'],
-                                              rsp.message,
-                                              rsp.args)
-    else:
+    try:
+        objs = dispatch(repo, proto, command['command'])
+
+        action, meta = reactor.oncommandresponsereadyobjects(
+            outstream, command['requestid'], objs)
+
+    except Exception as e:
         action, meta = reactor.onservererror(
-            _('unhandled response type from wire proto command'))
+            outstream, command['requestid'],
+            _('exception when invoking command: %s') % e)
 
     if action == 'sendframes':
         res.setbodygen(meta['framegen'])
@@ -430,6 +417,12 @@
     respectively. Default is to assume command requires ``push`` permissions
     because otherwise commands not declaring their permissions could modify
     a repository that is supposed to be read-only.
+
+    Wire protocol commands are generators of objects to be serialized and
+    sent to the client.
+
+    If a command raises an uncaught exception, this will be translated into
+    a command error.
     """
     transports = {k for k, v in wireprototypes.TRANSPORTS.items()
                   if v['version'] == 2}
@@ -460,16 +453,12 @@
 
 @wireprotocommand('branchmap', permission='pull')
 def branchmapv2(repo, proto):
-    branchmap = {encoding.fromlocal(k): v
-                 for k, v in repo.branchmap().iteritems()}
-
-    return wireprototypes.cborresponse(branchmap)
+    yield {encoding.fromlocal(k): v
+           for k, v in repo.branchmap().iteritems()}
 
 @wireprotocommand('capabilities', permission='pull')
 def capabilitiesv2(repo, proto):
-    caps = _capabilitiesv2(repo, proto)
-
-    return wireprototypes.cborresponse(caps)
+    yield _capabilitiesv2(repo, proto)
 
 @wireprotocommand('heads',
                   args={
@@ -480,7 +469,7 @@
     if publiconly:
         repo = repo.filtered('immutable')
 
-    return wireprototypes.cborresponse(repo.heads())
+    yield repo.heads()
 
 @wireprotocommand('known',
                   args={
@@ -490,7 +479,7 @@
 def knownv2(repo, proto, nodes=None):
     nodes = nodes or []
     result = b''.join(b'1' if n else b'0' for n in repo.known(nodes))
-    return wireprototypes.cborresponse(result)
+    yield result
 
 @wireprotocommand('listkeys',
                   args={
@@ -502,7 +491,7 @@
     keys = {encoding.fromlocal(k): encoding.fromlocal(v)
             for k, v in keys.iteritems()}
 
-    return wireprototypes.cborresponse(keys)
+    yield keys
 
 @wireprotocommand('lookup',
                   args={
@@ -515,7 +504,7 @@
     # TODO handle exception.
     node = repo.lookup(key)
 
-    return wireprototypes.cborresponse(node)
+    yield node
 
 @wireprotocommand('pushkey',
                   args={
@@ -527,9 +516,7 @@
                   permission='push')
 def pushkeyv2(repo, proto, namespace, key, old, new):
     # TODO handle ui output redirection
-    r = repo.pushkey(encoding.tolocal(namespace),
-                     encoding.tolocal(key),
-                     encoding.tolocal(old),
-                     encoding.tolocal(new))
-
-    return wireprototypes.cborresponse(r)
+    yield repo.pushkey(encoding.tolocal(namespace),
+                       encoding.tolocal(key),
+                       encoding.tolocal(old),
+                       encoding.tolocal(new))