hgext/largefiles/proto.py
changeset 43077 687b865b95ad
parent 43076 2372284d9457
child 43085 eef9a2d67051
equal deleted inserted replaced
43076:2372284d9457 43077:687b865b95ad
    23 
    23 
    24 urlerr = util.urlerr
    24 urlerr = util.urlerr
    25 urlreq = util.urlreq
    25 urlreq = util.urlreq
    26 
    26 
    27 LARGEFILES_REQUIRED_MSG = (
    27 LARGEFILES_REQUIRED_MSG = (
    28     '\nThis repository uses the largefiles extension.'
    28     b'\nThis repository uses the largefiles extension.'
    29     '\n\nPlease enable it in your Mercurial config '
    29     b'\n\nPlease enable it in your Mercurial config '
    30     'file.\n'
    30     b'file.\n'
    31 )
    31 )
    32 
    32 
    33 eh = exthelper.exthelper()
    33 eh = exthelper.exthelper()
    34 
    34 
    35 # these will all be replaced by largefiles.uisetup
    35 # these will all be replaced by largefiles.uisetup
    48         try:
    48         try:
    49             for p in proto.getpayload():
    49             for p in proto.getpayload():
    50                 tmpfp.write(p)
    50                 tmpfp.write(p)
    51             tmpfp._fp.seek(0)
    51             tmpfp._fp.seek(0)
    52             if sha != lfutil.hexsha1(tmpfp._fp):
    52             if sha != lfutil.hexsha1(tmpfp._fp):
    53                 raise IOError(0, _('largefile contents do not match hash'))
    53                 raise IOError(0, _(b'largefile contents do not match hash'))
    54             tmpfp.close()
    54             tmpfp.close()
    55             lfutil.linktousercache(repo, sha)
    55             lfutil.linktousercache(repo, sha)
    56         except IOError as e:
    56         except IOError as e:
    57             repo.ui.warn(
    57             repo.ui.warn(
    58                 _('largefiles: failed to put %s into store: %s\n')
    58                 _(b'largefiles: failed to put %s into store: %s\n')
    59                 % (sha, e.strerror)
    59                 % (sha, e.strerror)
    60             )
    60             )
    61             return wireprototypes.pushres(
    61             return wireprototypes.pushres(
    62                 1, output.getvalue() if output else ''
    62                 1, output.getvalue() if output else b''
    63             )
    63             )
    64         finally:
    64         finally:
    65             tmpfp.discard()
    65             tmpfp.discard()
    66 
    66 
    67     return wireprototypes.pushres(0, output.getvalue() if output else '')
    67     return wireprototypes.pushres(0, output.getvalue() if output else b'')
    68 
    68 
    69 
    69 
    70 def getlfile(repo, proto, sha):
    70 def getlfile(repo, proto, sha):
    71     '''Server command for retrieving a largefile from the repository-local
    71     '''Server command for retrieving a largefile from the repository-local
    72     cache or user cache.'''
    72     cache or user cache.'''
    73     filename = lfutil.findfile(repo, sha)
    73     filename = lfutil.findfile(repo, sha)
    74     if not filename:
    74     if not filename:
    75         raise error.Abort(
    75         raise error.Abort(
    76             _('requested largefile %s not present in cache') % sha
    76             _(b'requested largefile %s not present in cache') % sha
    77         )
    77         )
    78     f = open(filename, 'rb')
    78     f = open(filename, b'rb')
    79     length = os.fstat(f.fileno())[6]
    79     length = os.fstat(f.fileno())[6]
    80 
    80 
    81     # Since we can't set an HTTP content-length header here, and
    81     # Since we can't set an HTTP content-length header here, and
    82     # Mercurial core provides no way to give the length of a streamres
    82     # Mercurial core provides no way to give the length of a streamres
    83     # (and reading the entire file into RAM would be ill-advised), we
    83     # (and reading the entire file into RAM would be ill-advised), we
    84     # just send the length on the first line of the response, like the
    84     # just send the length on the first line of the response, like the
    85     # ssh proto does for string responses.
    85     # ssh proto does for string responses.
    86     def generator():
    86     def generator():
    87         yield '%d\n' % length
    87         yield b'%d\n' % length
    88         for chunk in util.filechunkiter(f):
    88         for chunk in util.filechunkiter(f):
    89             yield chunk
    89             yield chunk
    90 
    90 
    91     return wireprototypes.streamreslegacy(gen=generator())
    91     return wireprototypes.streamreslegacy(gen=generator())
    92 
    92 
    98     The value 1 is reserved for mismatched checksum, but that is too expensive
    98     The value 1 is reserved for mismatched checksum, but that is too expensive
    99     to be verified on every stat and must be caught be running 'hg verify'
    99     to be verified on every stat and must be caught be running 'hg verify'
   100     server side.'''
   100     server side.'''
   101     filename = lfutil.findfile(repo, sha)
   101     filename = lfutil.findfile(repo, sha)
   102     if not filename:
   102     if not filename:
   103         return wireprototypes.bytesresponse('2\n')
   103         return wireprototypes.bytesresponse(b'2\n')
   104     return wireprototypes.bytesresponse('0\n')
   104     return wireprototypes.bytesresponse(b'0\n')
   105 
   105 
   106 
   106 
   107 def wirereposetup(ui, repo):
   107 def wirereposetup(ui, repo):
   108     class lfileswirerepository(repo.__class__):
   108     class lfileswirerepository(repo.__class__):
   109         def putlfile(self, sha, fd):
   109         def putlfile(self, sha, fd):
   110             # unfortunately, httprepository._callpush tries to convert its
   110             # unfortunately, httprepository._callpush tries to convert its
   111             # input file-like into a bundle before sending it, so we can't use
   111             # input file-like into a bundle before sending it, so we can't use
   112             # it ...
   112             # it ...
   113             if issubclass(self.__class__, httppeer.httppeer):
   113             if issubclass(self.__class__, httppeer.httppeer):
   114                 res = self._call(
   114                 res = self._call(
   115                     'putlfile',
   115                     b'putlfile',
   116                     data=fd,
   116                     data=fd,
   117                     sha=sha,
   117                     sha=sha,
   118                     headers={r'content-type': r'application/mercurial-0.1'},
   118                     headers={r'content-type': r'application/mercurial-0.1'},
   119                 )
   119                 )
   120                 try:
   120                 try:
   121                     d, output = res.split('\n', 1)
   121                     d, output = res.split(b'\n', 1)
   122                     for l in output.splitlines(True):
   122                     for l in output.splitlines(True):
   123                         self.ui.warn(_('remote: '), l)  # assume l ends with \n
   123                         self.ui.warn(_(b'remote: '), l)  # assume l ends with \n
   124                     return int(d)
   124                     return int(d)
   125                 except ValueError:
   125                 except ValueError:
   126                     self.ui.warn(_('unexpected putlfile response: %r\n') % res)
   126                     self.ui.warn(_(b'unexpected putlfile response: %r\n') % res)
   127                     return 1
   127                     return 1
   128             # ... but we can't use sshrepository._call because the data=
   128             # ... but we can't use sshrepository._call because the data=
   129             # argument won't get sent, and _callpush does exactly what we want
   129             # argument won't get sent, and _callpush does exactly what we want
   130             # in this case: send the data straight through
   130             # in this case: send the data straight through
   131             else:
   131             else:
   132                 try:
   132                 try:
   133                     ret, output = self._callpush("putlfile", fd, sha=sha)
   133                     ret, output = self._callpush(b"putlfile", fd, sha=sha)
   134                     if ret == "":
   134                     if ret == b"":
   135                         raise error.ResponseError(_('putlfile failed:'), output)
   135                         raise error.ResponseError(
       
   136                             _(b'putlfile failed:'), output
       
   137                         )
   136                     return int(ret)
   138                     return int(ret)
   137                 except IOError:
   139                 except IOError:
   138                     return 1
   140                     return 1
   139                 except ValueError:
   141                 except ValueError:
   140                     raise error.ResponseError(
   142                     raise error.ResponseError(
   141                         _('putlfile failed (unexpected response):'), ret
   143                         _(b'putlfile failed (unexpected response):'), ret
   142                     )
   144                     )
   143 
   145 
   144         def getlfile(self, sha):
   146         def getlfile(self, sha):
   145             """returns an iterable with the chunks of the file with sha sha"""
   147             """returns an iterable with the chunks of the file with sha sha"""
   146             stream = self._callstream("getlfile", sha=sha)
   148             stream = self._callstream(b"getlfile", sha=sha)
   147             length = stream.readline()
   149             length = stream.readline()
   148             try:
   150             try:
   149                 length = int(length)
   151                 length = int(length)
   150             except ValueError:
   152             except ValueError:
   151                 self._abort(
   153                 self._abort(
   152                     error.ResponseError(_("unexpected response:"), length)
   154                     error.ResponseError(_(b"unexpected response:"), length)
   153                 )
   155                 )
   154 
   156 
   155             # SSH streams will block if reading more than length
   157             # SSH streams will block if reading more than length
   156             for chunk in util.filechunkiter(stream, limit=length):
   158             for chunk in util.filechunkiter(stream, limit=length):
   157                 yield chunk
   159                 yield chunk
   159             # chunk of Chunked-Encoding so the connection can be reused.
   161             # chunk of Chunked-Encoding so the connection can be reused.
   160             if issubclass(self.__class__, httppeer.httppeer):
   162             if issubclass(self.__class__, httppeer.httppeer):
   161                 chunk = stream.read(1)
   163                 chunk = stream.read(1)
   162                 if chunk:
   164                 if chunk:
   163                     self._abort(
   165                     self._abort(
   164                         error.ResponseError(_("unexpected response:"), chunk)
   166                         error.ResponseError(_(b"unexpected response:"), chunk)
   165                     )
   167                     )
   166 
   168 
   167         @wireprotov1peer.batchable
   169         @wireprotov1peer.batchable
   168         def statlfile(self, sha):
   170         def statlfile(self, sha):
   169             f = wireprotov1peer.future()
   171             f = wireprotov1peer.future()
   170             result = {'sha': sha}
   172             result = {b'sha': sha}
   171             yield result, f
   173             yield result, f
   172             try:
   174             try:
   173                 yield int(f.value)
   175                 yield int(f.value)
   174             except (ValueError, urlerr.httperror):
   176             except (ValueError, urlerr.httperror):
   175                 # If the server returns anything but an integer followed by a
   177                 # If the server returns anything but an integer followed by a
   180 
   182 
   181     repo.__class__ = lfileswirerepository
   183     repo.__class__ = lfileswirerepository
   182 
   184 
   183 
   185 
   184 # advertise the largefiles=serve capability
   186 # advertise the largefiles=serve capability
   185 @eh.wrapfunction(wireprotov1server, '_capabilities')
   187 @eh.wrapfunction(wireprotov1server, b'_capabilities')
   186 def _capabilities(orig, repo, proto):
   188 def _capabilities(orig, repo, proto):
   187     '''announce largefile server capability'''
   189     '''announce largefile server capability'''
   188     caps = orig(repo, proto)
   190     caps = orig(repo, proto)
   189     caps.append('largefiles=serve')
   191     caps.append(b'largefiles=serve')
   190     return caps
   192     return caps
   191 
   193 
   192 
   194 
   193 def heads(orig, repo, proto):
   195 def heads(orig, repo, proto):
   194     '''Wrap server command - largefile capable clients will know to call
   196     '''Wrap server command - largefile capable clients will know to call
   198 
   200 
   199     return orig(repo, proto)
   201     return orig(repo, proto)
   200 
   202 
   201 
   203 
   202 def sshrepocallstream(self, cmd, **args):
   204 def sshrepocallstream(self, cmd, **args):
   203     if cmd == 'heads' and self.capable('largefiles'):
   205     if cmd == b'heads' and self.capable(b'largefiles'):
   204         cmd = 'lheads'
   206         cmd = b'lheads'
   205     if cmd == 'batch' and self.capable('largefiles'):
   207     if cmd == b'batch' and self.capable(b'largefiles'):
   206         args[r'cmds'] = args[r'cmds'].replace('heads ', 'lheads ')
   208         args[r'cmds'] = args[r'cmds'].replace(b'heads ', b'lheads ')
   207     return ssholdcallstream(self, cmd, **args)
   209     return ssholdcallstream(self, cmd, **args)
   208 
   210 
   209 
   211 
   210 headsre = re.compile(br'(^|;)heads\b')
   212 headsre = re.compile(br'(^|;)heads\b')
   211 
   213 
   212 
   214 
   213 def httprepocallstream(self, cmd, **args):
   215 def httprepocallstream(self, cmd, **args):
   214     if cmd == 'heads' and self.capable('largefiles'):
   216     if cmd == b'heads' and self.capable(b'largefiles'):
   215         cmd = 'lheads'
   217         cmd = b'lheads'
   216     if cmd == 'batch' and self.capable('largefiles'):
   218     if cmd == b'batch' and self.capable(b'largefiles'):
   217         args[r'cmds'] = headsre.sub('lheads', args[r'cmds'])
   219         args[r'cmds'] = headsre.sub(b'lheads', args[r'cmds'])
   218     return httpoldcallstream(self, cmd, **args)
   220     return httpoldcallstream(self, cmd, **args)