/hgext/largefiles/proto.py

https://bitbucket.org/mirror/mercurial/ · Python · 176 lines · 164 code · 7 blank · 5 comment · 6 complexity · cb527994636eb1466c85b03489b0e421 MD5 · raw file

  1. # Copyright 2011 Fog Creek Software
  2. #
  3. # This software may be used and distributed according to the terms of the
  4. # GNU General Public License version 2 or any later version.
  5. import os
  6. import urllib2
  7. import re
  8. from mercurial import error, httppeer, util, wireproto
  9. from mercurial.i18n import _
  10. import lfutil
  11. LARGEFILES_REQUIRED_MSG = ('\nThis repository uses the largefiles extension.'
  12. '\n\nPlease enable it in your Mercurial config '
  13. 'file.\n')
  14. # these will all be replaced by largefiles.uisetup
  15. capabilitiesorig = None
  16. ssholdcallstream = None
  17. httpoldcallstream = None
  18. def putlfile(repo, proto, sha):
  19. '''Put a largefile into a repository's local store and into the
  20. user cache.'''
  21. proto.redirect()
  22. path = lfutil.storepath(repo, sha)
  23. util.makedirs(os.path.dirname(path))
  24. tmpfp = util.atomictempfile(path, createmode=repo.store.createmode)
  25. try:
  26. try:
  27. proto.getfile(tmpfp)
  28. tmpfp._fp.seek(0)
  29. if sha != lfutil.hexsha1(tmpfp._fp):
  30. raise IOError(0, _('largefile contents do not match hash'))
  31. tmpfp.close()
  32. lfutil.linktousercache(repo, sha)
  33. except IOError, e:
  34. repo.ui.warn(_('largefiles: failed to put %s into store: %s\n') %
  35. (sha, e.strerror))
  36. return wireproto.pushres(1)
  37. finally:
  38. tmpfp.discard()
  39. return wireproto.pushres(0)
  40. def getlfile(repo, proto, sha):
  41. '''Retrieve a largefile from the repository-local cache or system
  42. cache.'''
  43. filename = lfutil.findfile(repo, sha)
  44. if not filename:
  45. raise util.Abort(_('requested largefile %s not present in cache') % sha)
  46. f = open(filename, 'rb')
  47. length = os.fstat(f.fileno())[6]
  48. # Since we can't set an HTTP content-length header here, and
  49. # Mercurial core provides no way to give the length of a streamres
  50. # (and reading the entire file into RAM would be ill-advised), we
  51. # just send the length on the first line of the response, like the
  52. # ssh proto does for string responses.
  53. def generator():
  54. yield '%d\n' % length
  55. for chunk in util.filechunkiter(f):
  56. yield chunk
  57. return wireproto.streamres(generator())
  58. def statlfile(repo, proto, sha):
  59. '''Return '2\n' if the largefile is missing, '0\n' if it seems to be in
  60. good condition.
  61. The value 1 is reserved for mismatched checksum, but that is too expensive
  62. to be verified on every stat and must be caught be running 'hg verify'
  63. server side.'''
  64. filename = lfutil.findfile(repo, sha)
  65. if not filename:
  66. return '2\n'
  67. return '0\n'
  68. def wirereposetup(ui, repo):
  69. class lfileswirerepository(repo.__class__):
  70. def putlfile(self, sha, fd):
  71. # unfortunately, httprepository._callpush tries to convert its
  72. # input file-like into a bundle before sending it, so we can't use
  73. # it ...
  74. if issubclass(self.__class__, httppeer.httppeer):
  75. res = None
  76. try:
  77. res = self._call('putlfile', data=fd, sha=sha,
  78. headers={'content-type':'application/mercurial-0.1'})
  79. d, output = res.split('\n', 1)
  80. for l in output.splitlines(True):
  81. self.ui.warn(_('remote: '), l) # assume l ends with \n
  82. return int(d)
  83. except (ValueError, urllib2.HTTPError):
  84. self.ui.warn(_('unexpected putlfile response: %r\n') % res)
  85. return 1
  86. # ... but we can't use sshrepository._call because the data=
  87. # argument won't get sent, and _callpush does exactly what we want
  88. # in this case: send the data straight through
  89. else:
  90. try:
  91. ret, output = self._callpush("putlfile", fd, sha=sha)
  92. if ret == "":
  93. raise error.ResponseError(_('putlfile failed:'),
  94. output)
  95. return int(ret)
  96. except IOError:
  97. return 1
  98. except ValueError:
  99. raise error.ResponseError(
  100. _('putlfile failed (unexpected response):'), ret)
  101. def getlfile(self, sha):
  102. """returns an iterable with the chunks of the file with sha sha"""
  103. stream = self._callstream("getlfile", sha=sha)
  104. length = stream.readline()
  105. try:
  106. length = int(length)
  107. except ValueError:
  108. self._abort(error.ResponseError(_("unexpected response:"),
  109. length))
  110. # SSH streams will block if reading more than length
  111. for chunk in util.filechunkiter(stream, 128 * 1024, length):
  112. yield chunk
  113. # HTTP streams must hit the end to process the last empty
  114. # chunk of Chunked-Encoding so the connection can be reused.
  115. if issubclass(self.__class__, httppeer.httppeer):
  116. chunk = stream.read(1)
  117. if chunk:
  118. self._abort(error.ResponseError(_("unexpected response:"),
  119. chunk))
  120. @wireproto.batchable
  121. def statlfile(self, sha):
  122. f = wireproto.future()
  123. result = {'sha': sha}
  124. yield result, f
  125. try:
  126. yield int(f.value)
  127. except (ValueError, urllib2.HTTPError):
  128. # If the server returns anything but an integer followed by a
  129. # newline, newline, it's not speaking our language; if we get
  130. # an HTTP error, we can't be sure the largefile is present;
  131. # either way, consider it missing.
  132. yield 2
  133. repo.__class__ = lfileswirerepository
  134. # advertise the largefiles=serve capability
  135. def capabilities(repo, proto):
  136. return capabilitiesorig(repo, proto) + ' largefiles=serve'
  137. def heads(repo, proto):
  138. if lfutil.islfilesrepo(repo):
  139. return wireproto.ooberror(LARGEFILES_REQUIRED_MSG)
  140. return wireproto.heads(repo, proto)
  141. def sshrepocallstream(self, cmd, **args):
  142. if cmd == 'heads' and self.capable('largefiles'):
  143. cmd = 'lheads'
  144. if cmd == 'batch' and self.capable('largefiles'):
  145. args['cmds'] = args['cmds'].replace('heads ', 'lheads ')
  146. return ssholdcallstream(self, cmd, **args)
  147. headsre = re.compile(r'(^|;)heads\b')
  148. def httprepocallstream(self, cmd, **args):
  149. if cmd == 'heads' and self.capable('largefiles'):
  150. cmd = 'lheads'
  151. if cmd == 'batch' and self.capable('largefiles'):
  152. args['cmds'] = headsre.sub('lheads', args['cmds'])
  153. return httpoldcallstream(self, cmd, **args)