protocol: move basic ssh client commands to wirerepository

mercurial/sshrepo.py

@@ -1,293 +1,230 b''
 # sshrepo.py - ssh repository proxy class for mercurial
 #
 # Copyright 2005, 2006 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 from node import bin, hex
 from i18n import _
-import repo, util, error, encoding
+import repo, util, error, encoding, wireproto
 import re, urllib
 
 class remotelock(object):
     def __init__(self, repo):
         self.repo = repo
     def release(self):
         self.repo.unlock()
         self.repo = None
     def __del__(self):
         if self.repo:
             self.release()
 
-class sshrepository(repo.repository):
+class sshrepository(wireproto.wirerepository):
     def __init__(self, ui, path, create=0):
         self._url = path
         self.ui = ui
 
         m = re.match(r'^ssh://(([^@]+)@)?([^:/]+)(:(\d+))?(/(.*))?$', path)
         if not m:
             self.abort(error.RepoError(_("couldn't parse location %s") % path))
 
         self.user = m.group(2)
         self.host = m.group(3)
         self.port = m.group(5)
         self.path = m.group(7) or "."
 
         sshcmd = self.ui.config("ui", "ssh", "ssh")
         remotecmd = self.ui.config("ui", "remotecmd", "hg")
 
         args = util.sshargs(sshcmd, self.host, self.user, self.port)
 
         if create:
             cmd = '%s %s "%s init %s"'
             cmd = cmd % (sshcmd, args, remotecmd, self.path)
 
             ui.note(_('running %s\n') % cmd)
             res = util.system(cmd)
             if res != 0:
                 self.abort(error.RepoError(_("could not create remote repo")))
 
         self.validate_repo(ui, sshcmd, args, remotecmd)
 
     def url(self):
         return self._url
 
     def validate_repo(self, ui, sshcmd, args, remotecmd):
         # cleanup up previous run
         self.cleanup()
 
         cmd = '%s %s "%s -R %s serve --stdio"'
         cmd = cmd % (sshcmd, args, remotecmd, self.path)
 
         cmd = util.quotecommand(cmd)
         ui.note(_('running %s\n') % cmd)
         self.pipeo, self.pipei, self.pipee = util.popen3(cmd)
 
         # skip any noise generated by remote shell
         self.do_cmd("hello")
         r = self.do_cmd("between", pairs=("%s-%s" % ("0"*40, "0"*40)))
         lines = ["", "dummy"]
         max_noise = 500
         while lines[-1] and max_noise:
             l = r.readline()
             self.readerr()
             if lines[-1] == "1\n" and l == "\n":
                 break
             if l:
                 ui.debug("remote: ", l)
             lines.append(l)
             max_noise -= 1
         else:
             self.abort(error.RepoError(_("no suitable response from remote hg")))
 
         self.capabilities = set()
         for l in reversed(lines):
             if l.startswith("capabilities:"):
                 self.capabilities.update(l[:-1].split(":")[1].split())
                 break
 
     def readerr(self):
         while 1:
             size = util.fstat(self.pipee).st_size
             if size == 0:
                 break
             l = self.pipee.readline()
             if not l:
                 break
             self.ui.status(_("remote: "), l)
 
     def abort(self, exception):
         self.cleanup()
         raise exception
 
+    def _abort(self, exception):
+        self.cleanup()
+        raise exception
+
     def cleanup(self):
         try:
             self.pipeo.close()
             self.pipei.close()
             # read the error descriptor until EOF
             for l in self.pipee:
                 self.ui.status(_("remote: "), l)
             self.pipee.close()
         except:
             pass
 
     __del__ = cleanup
 
     def do_cmd(self, cmd, **args):
         self.ui.debug("sending %s command\n" % cmd)
         self.pipeo.write("%s\n" % cmd)
         for k, v in sorted(args.iteritems()):
             self.pipeo.write("%s %d\n" % (k, len(v)))
             self.pipeo.write(v)
         self.pipeo.flush()
 
         return self.pipei
 
     def call(self, cmd, **args):
         self.do_cmd(cmd, **args)
         return self._recv()
 
+    def _call(self, cmd, **args):
+        self.do_cmd(cmd, **args)
+        return self._recv()
+
     def _recv(self):
         l = self.pipei.readline()
         self.readerr()
         try:
             l = int(l)
         except:
             self.abort(error.ResponseError(_("unexpected response:"), l))
         return self.pipei.read(l)
 
     def _send(self, data, flush=False):
         self.pipeo.write("%d\n" % len(data))
         if data:
             self.pipeo.write(data)
         if flush:
             self.pipeo.flush()
         self.readerr()
 
     def lock(self):
         self.call("lock")
         return remotelock(self)
 
     def unlock(self):
         self.call("unlock")
 
-    def lookup(self, key):
-        self.requirecap('lookup', _('look up remote revision'))
-        d = self.call("lookup", key=key)
-        success, data = d[:-1].split(" ", 1)
-        if int(success):
-            return bin(data)
-        else:
-            self.abort(error.RepoError(data))
-
-    def heads(self):
-        d = self.call("heads")
-        try:
-            return map(bin, d[:-1].split(" "))
-        except:
-            self.abort(error.ResponseError(_("unexpected response:"), d))
-
-    def branchmap(self):
-        d = self.call("branchmap")
-        try:
-            branchmap = {}
-            for branchpart in d.splitlines():
-                branchheads = branchpart.split(' ')
-                branchname = urllib.unquote(branchheads[0])
-                # Earlier servers (1.3.x) send branch names in (their) local
-                # charset. The best we can do is assume it's identical to our
-                # own local charset, in case it's not utf-8.
-                try:
-                    branchname.decode('utf-8')
-                except UnicodeDecodeError:
-                    branchname = encoding.fromlocal(branchname)
-                branchheads = [bin(x) for x in branchheads[1:]]
-                branchmap[branchname] = branchheads
-            return branchmap
-        except:
-            raise error.ResponseError(_("unexpected response:"), d)
-
-    def branches(self, nodes):
-        n = " ".join(map(hex, nodes))
-        d = self.call("branches", nodes=n)
-        try:
-            br = [tuple(map(bin, b.split(" "))) for b in d.splitlines()]
-            return br
-        except:
-            self.abort(error.ResponseError(_("unexpected response:"), d))
-
-    def between(self, pairs):
-        n = " ".join(["-".join(map(hex, p)) for p in pairs])
-        d = self.call("between", pairs=n)
-        try:
-            p = [l and map(bin, l.split(" ")) or [] for l in d.splitlines()]
-            return p
-        except:
-            self.abort(error.ResponseError(_("unexpected response:"), d))
-
     def changegroup(self, nodes, kind):
         n = " ".join(map(hex, nodes))
         return self.do_cmd("changegroup", roots=n)
 
     def changegroupsubset(self, bases, heads, kind):
         self.requirecap('changegroupsubset', _('look up remote changes'))
         bases = " ".join(map(hex, bases))
         heads = " ".join(map(hex, heads))
         return self.do_cmd("changegroupsubset", bases=bases, heads=heads)
 
     def unbundle(self, cg, heads, source):
         '''Send cg (a readable file-like object representing the
         changegroup to push, typically a chunkbuffer object) to the
         remote server as a bundle. Return an integer indicating the
         result of the push (see localrepository.addchangegroup()).'''
         d = self.call("unbundle", heads=' '.join(map(hex, heads)))
         if d:
             # remote may send "unsynced changes"
             self.abort(error.RepoError(_("push refused: %s") % d))
 
         while 1:
             d = cg.read(4096)
             if not d:
                 break
             self._send(d)
 
         self._send("", flush=True)
 
         r = self._recv()
         if r:
             # remote may send "unsynced changes"
             self.abort(error.RepoError(_("push failed: %s") % r))
 
         r = self._recv()
         try:
             return int(r)
         except:
             self.abort(error.ResponseError(_("unexpected response:"), r))
 
     def addchangegroup(self, cg, source, url):
         '''Send a changegroup to the remote server.  Return an integer
         similar to unbundle(). DEPRECATED, since it requires locking the
         remote.'''
         d = self.call("addchangegroup")
         if d:
             self.abort(error.RepoError(_("push refused: %s") % d))
         while 1:
             d = cg.read(4096)
             if not d:
                 break
             self.pipeo.write(d)
             self.readerr()
 
         self.pipeo.flush()
 
         self.readerr()
         r = self._recv()
         if not r:
             return 1
         try:
             return int(r)
         except:
             self.abort(error.ResponseError(_("unexpected response:"), r))
 
     def stream_out(self):
         return self.do_cmd('stream_out')
 
-    def pushkey(self, namespace, key, old, new):
-        if not self.capable('pushkey'):
-            return False
-        d = self.call("pushkey",
-                      namespace=namespace, key=key, old=old, new=new)
-        return bool(int(d))
-
-    def listkeys(self, namespace):
-        if not self.capable('pushkey'):
-            return {}
-        d = self.call("listkeys", namespace=namespace)
-        r = {}
-        for l in d.splitlines():
-            k, v = l.split('\t')
-            r[k.decode('string-escape')] = v.decode('string-escape')
-        return r
-
 instance = sshrepository

mercurial/wireproto.py

@@ -1,97 +1,173 b''
 # wireproto.py - generic wire protocol support functions
 #
 # Copyright 2005-2010 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 from i18n import _
 from node import bin, hex
-import urllib, streamclone
+import urllib
+import streamclone, repo, error, encoding
 import pushkey as pushkey_
 
+# client side
+
+class wirerepository(repo.repository):
+    def lookup(self, key):
+        self.requirecap('lookup', _('look up remote revision'))
+        d = self._call("lookup", key=key)
+        success, data = d[:-1].split(" ", 1)
+        if int(success):
+            return bin(data)
+        self._abort(error.RepoError(data))
+
+    def heads(self):
+        d = self._call("heads")
+        try:
+            return map(bin, d[:-1].split(" "))
+        except:
+            self.abort(error.ResponseError(_("unexpected response:"), d))
+
+    def branchmap(self):
+        d = self._call("branchmap")
+        try:
+            branchmap = {}
+            for branchpart in d.splitlines():
+                branchheads = branchpart.split(' ')
+                branchname = urllib.unquote(branchheads[0])
+                # Earlier servers (1.3.x) send branch names in (their) local
+                # charset. The best we can do is assume it's identical to our
+                # own local charset, in case it's not utf-8.
+                try:
+                    branchname.decode('utf-8')
+                except UnicodeDecodeError:
+                    branchname = encoding.fromlocal(branchname)
+                branchheads = [bin(x) for x in branchheads[1:]]
+                branchmap[branchname] = branchheads
+            return branchmap
+        except TypeError:
+            self._abort(error.ResponseError(_("unexpected response:"), d))
+
+    def branches(self, nodes):
+        n = " ".join(map(hex, nodes))
+        d = self._call("branches", nodes=n)
+        try:
+            br = [tuple(map(bin, b.split(" "))) for b in d.splitlines()]
+            return br
+        except:
+            self._abort(error.ResponseError(_("unexpected response:"), d))
+
+    def between(self, pairs):
+        n = " ".join(["-".join(map(hex, p)) for p in pairs])
+        d = self._call("between", pairs=n)
+        try:
+            p = [l and map(bin, l.split(" ")) or [] for l in d.splitlines()]
+            return p
+        except:
+            self._abort(error.ResponseError(_("unexpected response:"), d))
+
+    def pushkey(self, namespace, key, old, new):
+        if not self.capable('pushkey'):
+            return False
+        d = self._call("pushkey",
+                      namespace=namespace, key=key, old=old, new=new)
+        return bool(int(d))
+
+    def listkeys(self, namespace):
+        if not self.capable('pushkey'):
+            return {}
+        d = self._call("listkeys", namespace=namespace)
+        r = {}
+        for l in d.splitlines():
+            k, v = l.split('\t')
+            r[k.decode('string-escape')] = v.decode('string-escape')
+        return r
+
+# server side
+
 def dispatch(repo, proto, command):
     if command not in commands:
         return False
     func, spec = commands[command]
     args = proto.getargs(spec)
     r = func(repo, proto, *args)
     if r != None:
         proto.respond(r)
     return True
 
 def between(repo, proto, pairs):
     pairs = [map(bin, p.split("-")) for p in pairs.split(" ")]
     r = []
     for b in repo.between(pairs):
         r.append(" ".join(map(hex, b)) + "\n")
     return "".join(r)
 
 def branchmap(repo, proto):
     branchmap = repo.branchmap()
     heads = []
     for branch, nodes in branchmap.iteritems():
         branchname = urllib.quote(branch)
         branchnodes = [hex(node) for node in nodes]
         heads.append('%s %s' % (branchname, ' '.join(branchnodes)))
     return '\n'.join(heads)
 
 def branches(repo, proto, nodes):
     nodes = map(bin, nodes.split(" "))
     r = []
     for b in repo.branches(nodes):
         r.append(" ".join(map(hex, b)) + "\n")
     return "".join(r)
 
 def changegroup(repo, proto, roots):
     nodes = map(bin, roots.split(" "))
     cg = repo.changegroup(nodes, 'serve')
     proto.sendchangegroup(cg)
 
 def changegroupsubset(repo, proto, bases, heads):
     bases = [bin(n) for n in bases.split(' ')]
     heads = [bin(n) for n in heads.split(' ')]
     cg = repo.changegroupsubset(bases, heads, 'serve')
     proto.sendchangegroup(cg)
 
 def heads(repo, proto):
     h = repo.heads()
     return " ".join(map(hex, h)) + "\n"
 
 def listkeys(repo, proto, namespace):
     d = pushkey_.list(repo, namespace).items()
     t = '\n'.join(['%s\t%s' % (k.encode('string-escape'),
                                v.encode('string-escape')) for k, v in d])
     return t
 
 def lookup(repo, proto, key):
     try:
         r = hex(repo.lookup(key))
         success = 1
     except Exception, inst:
         r = str(inst)
         success = 0
     return "%s %s\n" % (success, r)
 
 def pushkey(repo, proto, namespace, key, old, new):
     r = pushkey_.push(repo, namespace, key, old, new)
     return '%s\n' % int(r)
 
 def stream(repo, proto):
     try:
         proto.sendstream(streamclone.stream_out(repo))
     except streamclone.StreamException, inst:
         return str(inst)
 
 commands = {
     'between': (between, 'pairs'),
     'branchmap': (branchmap, ''),
     'branches': (branches, 'nodes'),
     'changegroup': (changegroup, 'roots'),
     'changegroupsubset': (changegroupsubset, 'bases heads'),
     'heads': (heads, ''),
     'listkeys': (listkeys, 'namespace'),
     'lookup': (lookup, 'key'),
     'pushkey': (pushkey, 'namespace key old new'),
     'stream_out': (stream, ''),
 }