largefiles: redo heads interception
authorJoerg Sonnenberger <joerg@bec.de>
Mon, 28 Dec 2020 03:18:16 +0100
changeset 46200 bd31462a86a2
parent 46199 8e7ce6555ea7
child 46201 b986e3342827
largefiles: redo heads interception The largefiles extension replaces the "heads" wire command and tries to redirect all uses towards the custom "lheads" wire command. As seen in issue6384, this doesn't currently work for ssh. Instead of hooking into the _callstream interface, properly register the command for the peer instance and monkeypatch the executor to do the redirection. This works transparently for both all kinds of peers and both for the batch and non-batch case. Differential Revision: https://phab.mercurial-scm.org/D9663
hgext/largefiles/__init__.py
hgext/largefiles/proto.py
--- a/hgext/largefiles/__init__.py	Thu Jan 07 11:51:18 2021 +0100
+++ b/hgext/largefiles/__init__.py	Mon Dec 28 03:18:16 2020 +0100
@@ -190,13 +190,6 @@
     )
     # TODO also wrap wireproto.commandsv2 once heads is implemented there.
 
-    # can't do this in reposetup because it needs to have happened before
-    # wirerepo.__init__ is called
-    proto.ssholdcallstream = sshpeer.sshv1peer._callstream
-    proto.httpoldcallstream = httppeer.httppeer._callstream
-    sshpeer.sshv1peer._callstream = proto.sshrepocallstream
-    httppeer.httppeer._callstream = proto.httprepocallstream
-
     # override some extensions' stuff as well
     for name, module in extensions.extensions():
         if name == b'rebase':
--- a/hgext/largefiles/proto.py	Thu Jan 07 11:51:18 2021 +0100
+++ b/hgext/largefiles/proto.py	Mon Dec 28 03:18:16 2020 +0100
@@ -33,10 +33,6 @@
 
 eh = exthelper.exthelper()
 
-# these will all be replaced by largefiles.uisetup
-ssholdcallstream = None
-httpoldcallstream = None
-
 
 def putlfile(repo, proto, sha):
     """Server command for putting a largefile into a repository's local store
@@ -106,7 +102,27 @@
 
 
 def wirereposetup(ui, repo):
+    orig_commandexecutor = repo.commandexecutor
+
     class lfileswirerepository(repo.__class__):
+        def commandexecutor(self):
+            executor = orig_commandexecutor()
+            if self.capable(b'largefiles'):
+                orig_callcommand = executor.callcommand
+
+                class lfscommandexecutor(executor.__class__):
+                    def callcommand(self, command, args):
+                        if command == b'heads':
+                            command = b'lheads'
+                        return orig_callcommand(command, args)
+
+                executor.__class__ = lfscommandexecutor
+            return executor
+
+        @wireprotov1peer.batchable
+        def lheads(self):
+            return self.heads.batchable(self)
+
         def putlfile(self, sha, fd):
             # unfortunately, httprepository._callpush tries to convert its
             # input file-like into a bundle before sending it, so we can't use
@@ -200,22 +216,3 @@
         return wireprototypes.ooberror(LARGEFILES_REQUIRED_MSG)
 
     return orig(repo, proto)
-
-
-def sshrepocallstream(self, cmd, **args):
-    if cmd == b'heads' and self.capable(b'largefiles'):
-        cmd = b'lheads'
-    if cmd == b'batch' and self.capable(b'largefiles'):
-        args['cmds'] = args[r'cmds'].replace(b'heads ', b'lheads ')
-    return ssholdcallstream(self, cmd, **args)
-
-
-headsre = re.compile(br'(^|;)heads\b')
-
-
-def httprepocallstream(self, cmd, **args):
-    if cmd == b'heads' and self.capable(b'largefiles'):
-        cmd = b'lheads'
-    if cmd == b'batch' and self.capable(b'largefiles'):
-        args['cmds'] = headsre.sub(b'lheads', args['cmds'])
-    return httpoldcallstream(self, cmd, **args)