Kiln » Kiln Extensions
Clone URL:  
bfproto.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import os import shutil import tempfile import urllib2 from mercurial.i18n import _ from mercurial import error, httprepo, util, wireproto import bfsetup import bfutil _heads_prefix = 'kbfiles\n' def putbfile(repo, proto, sha): """putbfile puts a bfile into a repository's local cache and into the system cache.""" f = None proto.redirect() try: try: f = tempfile.NamedTemporaryFile(mode='wb+', prefix='hg-putbfile-') proto.getfile(f) f.seek(0) if sha != bfutil.hexsha1(f): return wireproto.pushres(1) bfutil.copytocacheabsolute(repo, f.name, sha) except IOError: repo.ui.warn( _('error: could not put received data into bfile store')) return wireproto.pushres(1) finally: if f: f.close() return wireproto.pushres(0) def getbfile(repo, proto, sha): """getbfile retrieves a bfile from the repository-local cache or system cache.""" filename = bfutil.findfile(repo, sha) if not filename: raise util.Abort(_('requested bfile %s not present in cache') % sha) f = open(filename, 'rb') length = os.fstat(f.fileno())[6] # 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. def generator(): yield '%d\n' % length for chunk in f: yield chunk return wireproto.streamres(generator()) def statbfile(repo, proto, sha): """statbfile sends '2\n' if the bfile is missing, '1\n' if it has a mismatched checksum, or '0\n' if it is in good condition""" filename = bfutil.findfile(repo, sha) if not filename: return '2\n' fd = None try: fd = open(filename, 'rb') return bfutil.hexsha1(fd) == sha and '0\n' or '1\n' finally: if fd: fd.close() def wirereposetup(ui, repo): class kbfileswirerepository(repo.__class__): def putbfile(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 ... if issubclass(self.__class__, httprepo.httprepository): try: return int(self._call('putbfile', data=fd, sha=sha, headers={'content-type':'application/mercurial-0.1'})) except (ValueError, urllib2.HTTPError): 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: ret, output = self._callpush("putbfile", fd, sha=sha) if ret == "": raise error.ResponseError(_('putbfile failed:'), output) return int(ret) except IOError: return 1 except ValueError: raise error.ResponseError( _('putbfile failed (unexpected response):'), ret) def getbfile(self, sha): stream = self._callstream("getbfile", sha=sha) length = stream.readline() try: length = int(length) except ValueError: self._abort(error.ResponseError(_("unexpected response:"), l)) return (length, stream) def statbfile(self, sha): try: return int(self._call("statbfile", sha=sha)) except (ValueError, urllib2.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 bfile is present; either # way, consider it missing return 2 try: @wireproto.batchable def heads(self): f = wireproto.future() yield {}, f d = f.value if d[:len(_heads_prefix)] == _heads_prefix: d = d[len(_heads_prefix):] try: yield wireproto.decodelist(d[:-1]) except ValueError: self._abort(error.ResponseError(_("unexpected response:"), d)) except AttributeError: # Mercurial < 1.9 has no @batchable; define a normal wirerepo heads # command def heads(self): d = self._call('heads') if d[:len(_heads_prefix)] == _heads_prefix: d = d[len(_heads_prefix):] try: return wireproto.decodelist(d[:-1]) except ValueError: self._abort(error.ResponseError(_("unexpected response:"), d)) repo.__class__ = kbfileswirerepository # wrap dispatch to check for and remove the kbfiles argument so commands with # fixed argument lists don't complain def dispatch(repo, proto, command): func, spec = wireproto.commands[command] args = proto.getargs(spec) # remove the kbfiles argument and ignore it: it still needs to be sent # to avoid breaking compatibility with older versions of the extension, # but it is unused in current versions if len(args) > 0 and isinstance(args[-1], dict): args[-1].pop('kbfiles', None) return func(repo, proto, *args) # advertise the bfilestore=serve capability def capabilities(repo, proto): return capabilities_orig(repo, proto) + ' bfilestore=serve' def heads(repo, proto): if bfutil.iskbfilesrepo(repo): return _heads_prefix + wireproto.heads(repo, proto) return wireproto.heads(repo, proto)