proto.py
218 lines
| 7.6 KiB
| text/x-python
|
PythonLexer
|
r15168 | # Copyright 2011 Fog Creek Software | ||
# | ||||
# This software may be used and distributed according to the terms of the | ||||
# GNU General Public License version 2 or any later version. | ||||
|
r29312 | from __future__ import absolute_import | ||
|
r15168 | |||
import os | ||||
from mercurial.i18n import _ | ||||
|
r43355 | from mercurial.pycompat import open | ||
|
r15168 | |||
|
r29312 | from mercurial import ( | ||
error, | ||||
|
r41092 | exthelper, | ||
|
r29312 | httppeer, | ||
util, | ||||
|
r36091 | wireprototypes, | ||
|
r37632 | wireprotov1peer, | ||
|
r41092 | wireprotov1server, | ||
|
r29312 | ) | ||
|
r43346 | from . import lfutil | ||
|
r29312 | |||
|
r28883 | urlerr = util.urlerr | ||
urlreq = util.urlreq | ||||
|
r43346 | LARGEFILES_REQUIRED_MSG = ( | ||
|
r43347 | b'\nThis repository uses the largefiles extension.' | ||
b'\n\nPlease enable it in your Mercurial config ' | ||||
b'file.\n' | ||||
|
r43346 | ) | ||
|
r15168 | |||
|
r41092 | eh = exthelper.exthelper() | ||
|
r43346 | |||
|
r15168 | def putlfile(repo, proto, sha): | ||
|
r46554 | """Server command for putting a largefile into a repository's local store | ||
and into the user cache.""" | ||||
|
r36084 | with proto.mayberedirectstdio() as output: | ||
path = lfutil.storepath(repo, sha) | ||||
util.makedirs(os.path.dirname(path)) | ||||
tmpfp = util.atomictempfile(path, createmode=repo.store.createmode) | ||||
|
r16594 | |||
|
r36084 | try: | ||
|
r37432 | for p in proto.getpayload(): | ||
tmpfp.write(p) | ||||
|
r36084 | tmpfp._fp.seek(0) | ||
if sha != lfutil.hexsha1(tmpfp._fp): | ||||
|
r43347 | raise IOError(0, _(b'largefile contents do not match hash')) | ||
|
r36084 | tmpfp.close() | ||
lfutil.linktousercache(repo, sha) | ||||
except IOError as e: | ||||
|
r43346 | repo.ui.warn( | ||
|
r43347 | _(b'largefiles: failed to put %s into store: %s\n') | ||
|
r43346 | % (sha, e.strerror) | ||
) | ||||
|
r37309 | return wireprototypes.pushres( | ||
|
r43347 | 1, output.getvalue() if output else b'' | ||
|
r43346 | ) | ||
|
r36084 | finally: | ||
tmpfp.discard() | ||||
|
r15168 | |||
|
r43347 | return wireprototypes.pushres(0, output.getvalue() if output else b'') | ||
|
r15168 | |||
|
r43346 | |||
|
r15168 | def getlfile(repo, proto, sha): | ||
|
r46554 | """Server command for retrieving a largefile from the repository-local | ||
cache or user cache.""" | ||||
|
r15168 | filename = lfutil.findfile(repo, sha) | ||
if not filename: | ||||
|
r43346 | raise error.Abort( | ||
|
r43347 | _(b'requested largefile %s not present in cache') % sha | ||
|
r43346 | ) | ||
|
r43347 | f = open(filename, b'rb') | ||
|
r15168 | length = os.fstat(f.fileno())[6] | ||
|
r15252 | |||
# Since we can't set an HTTP content-length header here, and | ||||
# Mercurial core provides no way to give the length of a streamres | ||||
# (and reading the entire file into RAM would be ill-advised), we | ||||
# just send the length on the first line of the response, like the | ||||
# ssh proto does for string responses. | ||||
|
r15168 | def generator(): | ||
|
r43347 | yield b'%d\n' % length | ||
|
r19009 | for chunk in util.filechunkiter(f): | ||
|
r15168 | yield chunk | ||
|
r43346 | |||
|
r37309 | return wireprototypes.streamreslegacy(gen=generator()) | ||
|
r15168 | |||
|
r43346 | |||
|
r15168 | def statlfile(repo, proto, sha): | ||
|
r46554 | """Server command for checking if a largefile is present - returns '2\n' if | ||
|
r28576 | the largefile is missing, '0\n' if it seems to be in good condition. | ||
|
r18488 | |||
The value 1 is reserved for mismatched checksum, but that is too expensive | ||||
to be verified on every stat and must be caught be running 'hg verify' | ||||
|
r46554 | server side.""" | ||
|
r15168 | filename = lfutil.findfile(repo, sha) | ||
if not filename: | ||||
|
r43347 | return wireprototypes.bytesresponse(b'2\n') | ||
return wireprototypes.bytesresponse(b'0\n') | ||||
|
r15168 | |||
|
r43346 | |||
|
r15168 | def wirereposetup(ui, repo): | ||
|
r46816 | orig_commandexecutor = repo.commandexecutor | ||
|
r15168 | class lfileswirerepository(repo.__class__): | ||
|
r46816 | 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) | ||||
|
r15168 | 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 | ||||
# it ... | ||||
|
r17192 | if issubclass(self.__class__, httppeer.httppeer): | ||
|
r43346 | res = self._call( | ||
|
r43347 | b'putlfile', | ||
|
r43346 | data=fd, | ||
sha=sha, | ||||
|
r43906 | headers={'content-type': 'application/mercurial-0.1'}, | ||
|
r43346 | ) | ||
|
r15168 | try: | ||
|
r43347 | d, output = res.split(b'\n', 1) | ||
|
r15778 | for l in output.splitlines(True): | ||
|
r43347 | self.ui.warn(_(b'remote: '), l) # assume l ends with \n | ||
|
r15778 | return int(d) | ||
|
r26825 | except ValueError: | ||
|
r43347 | self.ui.warn(_(b'unexpected putlfile response: %r\n') % res) | ||
|
r15168 | return 1 | ||
# ... but we can't use sshrepository._call because the data= | ||||
# argument won't get sent, and _callpush does exactly what we want | ||||
# in this case: send the data straight through | ||||
else: | ||||
try: | ||||
|
r43347 | ret, output = self._callpush(b"putlfile", fd, sha=sha) | ||
if ret == b"": | ||||
raise error.ResponseError( | ||||
_(b'putlfile failed:'), output | ||||
) | ||||
|
r15168 | return int(ret) | ||
except IOError: | ||||
return 1 | ||||
except ValueError: | ||||
raise error.ResponseError( | ||||
|
r43347 | _(b'putlfile failed (unexpected response):'), ret | ||
|
r43346 | ) | ||
|
r15168 | |||
def getlfile(self, sha): | ||||
|
r19004 | """returns an iterable with the chunks of the file with sha sha""" | ||
|
r43347 | stream = self._callstream(b"getlfile", sha=sha) | ||
|
r15168 | length = stream.readline() | ||
try: | ||||
length = int(length) | ||||
except ValueError: | ||||
|
r43346 | self._abort( | ||
|
r43347 | error.ResponseError(_(b"unexpected response:"), length) | ||
|
r43346 | ) | ||
|
r19004 | |||
|
r19005 | # SSH streams will block if reading more than length | ||
|
r30181 | for chunk in util.filechunkiter(stream, limit=length): | ||
|
r19004 | yield chunk | ||
|
r19006 | # HTTP streams must hit the end to process the last empty | ||
# chunk of Chunked-Encoding so the connection can be reused. | ||||
if issubclass(self.__class__, httppeer.httppeer): | ||||
chunk = stream.read(1) | ||||
if chunk: | ||||
|
r43346 | self._abort( | ||
|
r43347 | error.ResponseError(_(b"unexpected response:"), chunk) | ||
|
r43346 | ) | ||
|
r15168 | |||
|
r37632 | @wireprotov1peer.batchable | ||
|
r15168 | def statlfile(self, sha): | ||
|
r48687 | def decode(d): | ||
try: | ||||
return int(d) | ||||
except (ValueError, urlerr.httperror): | ||||
# If the server returns anything but an integer followed by a | ||||
# newline, newline, it's not speaking our language; if we get | ||||
# an HTTP error, we can't be sure the largefile is present; | ||||
# either way, consider it missing. | ||||
return 2 | ||||
|
r43347 | result = {b'sha': sha} | ||
|
r48687 | return result, decode | ||
|
r15168 | |||
repo.__class__ = lfileswirerepository | ||||
|
r43346 | |||
|
r15168 | # advertise the largefiles=serve capability | ||
|
r43347 | @eh.wrapfunction(wireprotov1server, b'_capabilities') | ||
|
r35523 | def _capabilities(orig, repo, proto): | ||
'''announce largefile server capability''' | ||||
caps = orig(repo, proto) | ||||
|
r43347 | caps.append(b'largefiles=serve') | ||
|
r35523 | return caps | ||
|
r15168 | |||
|
r43346 | |||
|
r37502 | def heads(orig, repo, proto): | ||
|
r46554 | """Wrap server command - largefile capable clients will know to call | ||
lheads instead""" | ||||
|
r15168 | if lfutil.islfilesrepo(repo): | ||
|
r37309 | return wireprototypes.ooberror(LARGEFILES_REQUIRED_MSG) | ||
|
r37502 | |||
return orig(repo, proto) | ||||