1 # This software may be used and distributed according to the terms of the |
|
2 # GNU General Public License version 2 or any later version. |
|
3 |
|
4 # based on bundleheads extension by Gregory Szorc <gps@mozilla.com> |
|
5 |
|
6 |
|
7 import abc |
|
8 import os |
|
9 import subprocess |
|
10 |
|
11 from mercurial.node import hex |
|
12 from mercurial.pycompat import open |
|
13 from mercurial import pycompat |
|
14 from mercurial.utils import ( |
|
15 hashutil, |
|
16 procutil, |
|
17 ) |
|
18 |
|
19 |
|
20 class BundleWriteException(Exception): |
|
21 pass |
|
22 |
|
23 |
|
24 class BundleReadException(Exception): |
|
25 pass |
|
26 |
|
27 |
|
28 class abstractbundlestore: # pytype: disable=ignored-metaclass |
|
29 """Defines the interface for bundle stores. |
|
30 |
|
31 A bundle store is an entity that stores raw bundle data. It is a simple |
|
32 key-value store. However, the keys are chosen by the store. The keys can |
|
33 be any Python object understood by the corresponding bundle index (see |
|
34 ``abstractbundleindex`` below). |
|
35 """ |
|
36 |
|
37 __metaclass__ = abc.ABCMeta |
|
38 |
|
39 @abc.abstractmethod |
|
40 def write(self, data): |
|
41 """Write bundle data to the store. |
|
42 |
|
43 This function receives the raw data to be written as a str. |
|
44 Throws BundleWriteException |
|
45 The key of the written data MUST be returned. |
|
46 """ |
|
47 |
|
48 @abc.abstractmethod |
|
49 def read(self, key): |
|
50 """Obtain bundle data for a key. |
|
51 |
|
52 Returns None if the bundle isn't known. |
|
53 Throws BundleReadException |
|
54 The returned object should be a file object supporting read() |
|
55 and close(). |
|
56 """ |
|
57 |
|
58 |
|
59 class filebundlestore: |
|
60 """bundle store in filesystem |
|
61 |
|
62 meant for storing bundles somewhere on disk and on network filesystems |
|
63 """ |
|
64 |
|
65 def __init__(self, ui, repo): |
|
66 self.ui = ui |
|
67 self.repo = repo |
|
68 self.storepath = ui.configpath(b'scratchbranch', b'storepath') |
|
69 if not self.storepath: |
|
70 self.storepath = self.repo.vfs.join( |
|
71 b"scratchbranches", b"filebundlestore" |
|
72 ) |
|
73 if not os.path.exists(self.storepath): |
|
74 os.makedirs(self.storepath) |
|
75 |
|
76 def _dirpath(self, hashvalue): |
|
77 """First two bytes of the hash are the name of the upper |
|
78 level directory, next two bytes are the name of the |
|
79 next level directory""" |
|
80 return os.path.join(self.storepath, hashvalue[0:2], hashvalue[2:4]) |
|
81 |
|
82 def _filepath(self, filename): |
|
83 return os.path.join(self._dirpath(filename), filename) |
|
84 |
|
85 def write(self, data): |
|
86 filename = hex(hashutil.sha1(data).digest()) |
|
87 dirpath = self._dirpath(filename) |
|
88 |
|
89 if not os.path.exists(dirpath): |
|
90 os.makedirs(dirpath) |
|
91 |
|
92 with open(self._filepath(filename), b'wb') as f: |
|
93 f.write(data) |
|
94 |
|
95 return filename |
|
96 |
|
97 def read(self, key): |
|
98 try: |
|
99 with open(self._filepath(key), b'rb') as f: |
|
100 return f.read() |
|
101 except IOError: |
|
102 return None |
|
103 |
|
104 |
|
105 def format_placeholders_args(args, filename=None, handle=None): |
|
106 """Formats `args` with Infinitepush replacements. |
|
107 |
|
108 Hack to get `str.format()`-ed strings working in a BC way with |
|
109 bytes. |
|
110 """ |
|
111 formatted_args = [] |
|
112 for arg in args: |
|
113 if filename and arg == b'{filename}': |
|
114 formatted_args.append(filename) |
|
115 elif handle and arg == b'{handle}': |
|
116 formatted_args.append(handle) |
|
117 else: |
|
118 formatted_args.append(arg) |
|
119 return formatted_args |
|
120 |
|
121 |
|
122 class externalbundlestore(abstractbundlestore): |
|
123 def __init__(self, put_binary, put_args, get_binary, get_args): |
|
124 """ |
|
125 `put_binary` - path to binary file which uploads bundle to external |
|
126 storage and prints key to stdout |
|
127 `put_args` - format string with additional args to `put_binary` |
|
128 {filename} replacement field can be used. |
|
129 `get_binary` - path to binary file which accepts filename and key |
|
130 (in that order), downloads bundle from store and saves it to file |
|
131 `get_args` - format string with additional args to `get_binary`. |
|
132 {filename} and {handle} replacement field can be used. |
|
133 """ |
|
134 |
|
135 self.put_args = put_args |
|
136 self.get_args = get_args |
|
137 self.put_binary = put_binary |
|
138 self.get_binary = get_binary |
|
139 |
|
140 def _call_binary(self, args): |
|
141 p = subprocess.Popen( |
|
142 pycompat.rapply(procutil.tonativestr, args), |
|
143 stdout=subprocess.PIPE, |
|
144 stderr=subprocess.PIPE, |
|
145 close_fds=True, |
|
146 ) |
|
147 stdout, stderr = p.communicate() |
|
148 returncode = p.returncode |
|
149 return returncode, stdout, stderr |
|
150 |
|
151 def write(self, data): |
|
152 # Won't work on windows because you can't open file second time without |
|
153 # closing it |
|
154 # TODO: rewrite without str.format() and replace NamedTemporaryFile() |
|
155 # with pycompat.namedtempfile() |
|
156 with pycompat.namedtempfile() as temp: |
|
157 temp.write(data) |
|
158 temp.flush() |
|
159 temp.seek(0) |
|
160 formatted_args = format_placeholders_args( |
|
161 self.put_args, filename=temp.name |
|
162 ) |
|
163 returncode, stdout, stderr = self._call_binary( |
|
164 [self.put_binary] + formatted_args |
|
165 ) |
|
166 |
|
167 if returncode != 0: |
|
168 raise BundleWriteException( |
|
169 b'Failed to upload to external store: %s' % stderr |
|
170 ) |
|
171 stdout_lines = stdout.splitlines() |
|
172 if len(stdout_lines) == 1: |
|
173 return stdout_lines[0] |
|
174 else: |
|
175 raise BundleWriteException( |
|
176 b'Bad output from %s: %s' % (self.put_binary, stdout) |
|
177 ) |
|
178 |
|
179 def read(self, handle): |
|
180 # Won't work on windows because you can't open file second time without |
|
181 # closing it |
|
182 with pycompat.namedtempfile() as temp: |
|
183 formatted_args = format_placeholders_args( |
|
184 self.get_args, filename=temp.name, handle=handle |
|
185 ) |
|
186 returncode, stdout, stderr = self._call_binary( |
|
187 [self.get_binary] + formatted_args |
|
188 ) |
|
189 |
|
190 if returncode != 0: |
|
191 raise BundleReadException( |
|
192 b'Failed to download from external store: %s' % stderr |
|
193 ) |
|
194 return temp.read() |
|