mercurial/wireprotov2server.py
changeset 37546 3a2367e6c6f2
parent 37545 93397c4633f6
child 37557 734515aca84d
--- a/mercurial/wireprotov2server.py	Mon Apr 09 19:35:04 2018 -0700
+++ b/mercurial/wireprotov2server.py	Mon Apr 09 19:35:39 2018 -0700
@@ -16,8 +16,10 @@
     interface as zi,
 )
 from . import (
+    encoding,
     error,
     pycompat,
+    util,
     wireproto,
     wireprotoframing,
     wireprototypes,
@@ -362,3 +364,111 @@
 
     def checkperm(self, perm):
         raise NotImplementedError
+
+def _capabilitiesv2(repo, proto):
+    """Obtain the set of capabilities for version 2 transports.
+
+    These capabilities are distinct from the capabilities for version 1
+    transports.
+    """
+    compression = []
+    for engine in wireproto.supportedcompengines(repo.ui, util.SERVERROLE):
+        compression.append({
+            b'name': engine.wireprotosupport().name,
+        })
+
+    caps = {
+        'commands': {},
+        'compression': compression,
+    }
+
+    for command, entry in wireproto.commandsv2.items():
+        caps['commands'][command] = {
+            'args': entry.args,
+            'permissions': [entry.permission],
+        }
+
+    return proto.addcapabilities(repo, caps)
+
+def wireprotocommand(*args, **kwargs):
+    def register(func):
+        return wireproto.wireprotocommand(
+            *args, transportpolicy=wireproto.POLICY_V2_ONLY, **kwargs)(func)
+
+    return register
+
+@wireprotocommand('branchmap', permission='pull')
+def branchmapv2(repo, proto):
+    branchmap = {encoding.fromlocal(k): v
+                 for k, v in repo.branchmap().iteritems()}
+
+    return wireprototypes.cborresponse(branchmap)
+
+@wireprotocommand('capabilities', permission='pull')
+def capabilitiesv2(repo, proto):
+    caps = _capabilitiesv2(repo, proto)
+
+    return wireprototypes.cborresponse(caps)
+
+@wireprotocommand('heads',
+                  args={
+                      'publiconly': False,
+                  },
+                  permission='pull')
+def headsv2(repo, proto, publiconly=False):
+    if publiconly:
+        repo = repo.filtered('immutable')
+
+    return wireprototypes.cborresponse(repo.heads())
+
+@wireprotocommand('known',
+                  args={
+                      'nodes': [b'deadbeef'],
+                  },
+                  permission='pull')
+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)
+
+@wireprotocommand('listkeys',
+                  args={
+                      'namespace': b'ns',
+                  },
+                  permission='pull')
+def listkeysv2(repo, proto, namespace=None):
+    keys = repo.listkeys(encoding.tolocal(namespace))
+    keys = {encoding.fromlocal(k): encoding.fromlocal(v)
+            for k, v in keys.iteritems()}
+
+    return wireprototypes.cborresponse(keys)
+
+@wireprotocommand('lookup',
+                  args={
+                      'key': b'foo',
+                  },
+                  permission='pull')
+def lookupv2(repo, proto, key):
+    key = encoding.tolocal(key)
+
+    # TODO handle exception.
+    node = repo.lookup(key)
+
+    return wireprototypes.cborresponse(node)
+
+@wireprotocommand('pushkey',
+                  args={
+                      'namespace': b'ns',
+                      'key': b'key',
+                      'old': b'old',
+                      'new': b'new',
+                  },
+                  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)