store.py
198 lines
| 6.2 KiB
| text/x-python
|
PythonLexer
|
r37204 | # This software may be used and distributed according to the terms of the | ||
# GNU General Public License version 2 or any later version. | ||||
# based on bundleheads extension by Gregory Szorc <gps@mozilla.com> | ||||
from __future__ import absolute_import | ||||
import abc | ||||
import os | ||||
import subprocess | ||||
import tempfile | ||||
|
r43355 | from mercurial.pycompat import open | ||
|
r39851 | from mercurial import ( | ||
|
r40288 | node, | ||
|
r39851 | pycompat, | ||
) | ||||
|
r44519 | from mercurial.utils import ( | ||
hashutil, | ||||
procutil, | ||||
) | ||||
|
r39851 | |||
|
r43346 | |||
|
r37204 | class BundleWriteException(Exception): | ||
pass | ||||
|
r43346 | |||
|
r37204 | class BundleReadException(Exception): | ||
pass | ||||
|
r43346 | |||
|
r43775 | class abstractbundlestore(object): # pytype: disable=ignored-metaclass | ||
|
r37204 | """Defines the interface for bundle stores. | ||
A bundle store is an entity that stores raw bundle data. It is a simple | ||||
key-value store. However, the keys are chosen by the store. The keys can | ||||
be any Python object understood by the corresponding bundle index (see | ||||
``abstractbundleindex`` below). | ||||
""" | ||||
|
r43346 | |||
|
r37204 | __metaclass__ = abc.ABCMeta | ||
@abc.abstractmethod | ||||
def write(self, data): | ||||
"""Write bundle data to the store. | ||||
This function receives the raw data to be written as a str. | ||||
Throws BundleWriteException | ||||
The key of the written data MUST be returned. | ||||
""" | ||||
@abc.abstractmethod | ||||
def read(self, key): | ||||
"""Obtain bundle data for a key. | ||||
Returns None if the bundle isn't known. | ||||
Throws BundleReadException | ||||
The returned object should be a file object supporting read() | ||||
and close(). | ||||
""" | ||||
|
r43346 | |||
|
r37204 | class filebundlestore(object): | ||
"""bundle store in filesystem | ||||
meant for storing bundles somewhere on disk and on network filesystems | ||||
""" | ||||
|
r43346 | |||
|
r37204 | def __init__(self, ui, repo): | ||
self.ui = ui | ||||
self.repo = repo | ||||
|
r43347 | self.storepath = ui.configpath(b'scratchbranch', b'storepath') | ||
|
r37204 | if not self.storepath: | ||
|
r43346 | self.storepath = self.repo.vfs.join( | ||
|
r43347 | b"scratchbranches", b"filebundlestore" | ||
|
r43346 | ) | ||
|
r37204 | if not os.path.exists(self.storepath): | ||
os.makedirs(self.storepath) | ||||
def _dirpath(self, hashvalue): | ||||
"""First two bytes of the hash are the name of the upper | ||||
level directory, next two bytes are the name of the | ||||
next level directory""" | ||||
return os.path.join(self.storepath, hashvalue[0:2], hashvalue[2:4]) | ||||
def _filepath(self, filename): | ||||
return os.path.join(self._dirpath(filename), filename) | ||||
def write(self, data): | ||||
|
r44519 | filename = node.hex(hashutil.sha1(data).digest()) | ||
|
r37204 | dirpath = self._dirpath(filename) | ||
if not os.path.exists(dirpath): | ||||
os.makedirs(dirpath) | ||||
|
r43347 | with open(self._filepath(filename), b'wb') as f: | ||
|
r37204 | f.write(data) | ||
return filename | ||||
def read(self, key): | ||||
try: | ||||
|
r43347 | with open(self._filepath(key), b'rb') as f: | ||
|
r37814 | return f.read() | ||
|
r37204 | except IOError: | ||
return None | ||||
|
r43346 | |||
|
r45753 | def format_placeholders_args(args, filename=None, handle=None): | ||
"""Formats `args` with Infinitepush replacements. | ||||
Hack to get `str.format()`-ed strings working in a BC way with | ||||
bytes. | ||||
""" | ||||
formatted_args = [] | ||||
for arg in args: | ||||
if filename and arg == b'{filename}': | ||||
formatted_args.append(filename) | ||||
elif handle and arg == b'{handle}': | ||||
formatted_args.append(handle) | ||||
else: | ||||
formatted_args.append(arg) | ||||
return formatted_args | ||||
|
r37204 | class externalbundlestore(abstractbundlestore): | ||
def __init__(self, put_binary, put_args, get_binary, get_args): | ||||
""" | ||||
`put_binary` - path to binary file which uploads bundle to external | ||||
storage and prints key to stdout | ||||
`put_args` - format string with additional args to `put_binary` | ||||
{filename} replacement field can be used. | ||||
`get_binary` - path to binary file which accepts filename and key | ||||
(in that order), downloads bundle from store and saves it to file | ||||
`get_args` - format string with additional args to `get_binary`. | ||||
{filename} and {handle} replacement field can be used. | ||||
""" | ||||
self.put_args = put_args | ||||
self.get_args = get_args | ||||
self.put_binary = put_binary | ||||
self.get_binary = get_binary | ||||
def _call_binary(self, args): | ||||
p = subprocess.Popen( | ||||
|
r39851 | pycompat.rapply(procutil.tonativestr, args), | ||
|
r43346 | stdout=subprocess.PIPE, | ||
stderr=subprocess.PIPE, | ||||
close_fds=True, | ||||
) | ||||
|
r37204 | stdout, stderr = p.communicate() | ||
returncode = p.returncode | ||||
return returncode, stdout, stderr | ||||
def write(self, data): | ||||
# Won't work on windows because you can't open file second time without | ||||
# closing it | ||||
|
r38184 | # TODO: rewrite without str.format() and replace NamedTemporaryFile() | ||
# with pycompat.namedtempfile() | ||||
|
r45752 | with pycompat.namedtempfile() as temp: | ||
|
r37204 | temp.write(data) | ||
temp.flush() | ||||
temp.seek(0) | ||||
|
r45753 | formatted_args = format_placeholders_args( | ||
self.put_args, filename=temp.name | ||||
) | ||||
|
r37204 | returncode, stdout, stderr = self._call_binary( | ||
|
r43346 | [self.put_binary] + formatted_args | ||
) | ||||
|
r37204 | |||
if returncode != 0: | ||||
raise BundleWriteException( | ||||
|
r43347 | b'Failed to upload to external store: %s' % stderr | ||
|
r43346 | ) | ||
|
r37204 | stdout_lines = stdout.splitlines() | ||
if len(stdout_lines) == 1: | ||||
return stdout_lines[0] | ||||
else: | ||||
raise BundleWriteException( | ||||
|
r43347 | b'Bad output from %s: %s' % (self.put_binary, stdout) | ||
|
r43346 | ) | ||
|
r37204 | |||
def read(self, handle): | ||||
# Won't work on windows because you can't open file second time without | ||||
# closing it | ||||
|
r45752 | with pycompat.namedtempfile() as temp: | ||
|
r45753 | formatted_args = format_placeholders_args( | ||
self.get_args, filename=temp.name, handle=handle | ||||
) | ||||
|
r37204 | returncode, stdout, stderr = self._call_binary( | ||
|
r43346 | [self.get_binary] + formatted_args | ||
) | ||||
|
r37204 | |||
if returncode != 0: | ||||
raise BundleReadException( | ||||
|
r43347 | b'Failed to download from external store: %s' % stderr | ||
|
r43346 | ) | ||
|
r37204 | return temp.read() | ||