store.py
155 lines
| 5.3 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 hashlib | ||||
import os | ||||
import subprocess | ||||
import tempfile | ||||
NamedTemporaryFile = tempfile.NamedTemporaryFile | ||||
class BundleWriteException(Exception): | ||||
pass | ||||
class BundleReadException(Exception): | ||||
pass | ||||
class abstractbundlestore(object): | ||||
"""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). | ||||
""" | ||||
__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(). | ||||
""" | ||||
class filebundlestore(object): | ||||
"""bundle store in filesystem | ||||
meant for storing bundles somewhere on disk and on network filesystems | ||||
""" | ||||
def __init__(self, ui, repo): | ||||
self.ui = ui | ||||
self.repo = repo | ||||
self.storepath = ui.configpath('scratchbranch', 'storepath') | ||||
if not self.storepath: | ||||
self.storepath = self.repo.vfs.join("scratchbranches", | ||||
"filebundlestore") | ||||
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): | ||||
filename = hashlib.sha1(data).hexdigest() | ||||
dirpath = self._dirpath(filename) | ||||
if not os.path.exists(dirpath): | ||||
os.makedirs(dirpath) | ||||
with open(self._filepath(filename), 'w') as f: | ||||
f.write(data) | ||||
return filename | ||||
def read(self, key): | ||||
try: | ||||
f = open(self._filepath(key), 'r') | ||||
except IOError: | ||||
return None | ||||
return f.read() | ||||
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( | ||||
args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, | ||||
close_fds=True) | ||||
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 | ||||
with NamedTemporaryFile() as temp: | ||||
temp.write(data) | ||||
temp.flush() | ||||
temp.seek(0) | ||||
formatted_args = [arg.format(filename=temp.name) | ||||
for arg in self.put_args] | ||||
returncode, stdout, stderr = self._call_binary( | ||||
[self.put_binary] + formatted_args) | ||||
if returncode != 0: | ||||
raise BundleWriteException( | ||||
'Failed to upload to external store: %s' % stderr) | ||||
stdout_lines = stdout.splitlines() | ||||
if len(stdout_lines) == 1: | ||||
return stdout_lines[0] | ||||
else: | ||||
raise BundleWriteException( | ||||
'Bad output from %s: %s' % (self.put_binary, stdout)) | ||||
def read(self, handle): | ||||
# Won't work on windows because you can't open file second time without | ||||
# closing it | ||||
with NamedTemporaryFile() as temp: | ||||
formatted_args = [arg.format(filename=temp.name, handle=handle) | ||||
for arg in self.get_args] | ||||
returncode, stdout, stderr = self._call_binary( | ||||
[self.get_binary] + formatted_args) | ||||
if returncode != 0: | ||||
raise BundleReadException( | ||||
'Failed to download from external store: %s' % stderr) | ||||
return temp.read() | ||||