Mercurial > hg
view hgext/infinitepush/store.py @ 49541:976648e20856
tests: remove non-python3 line matching and tests block
We don't support Python2 anymore
author | Pierre-Yves David <pierre-yves.david@octobus.net> |
---|---|
date | Thu, 22 Sep 2022 16:27:17 +0200 |
parents | 642e31cb55f0 |
children |
line wrap: on
line source
# 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> import abc import os import subprocess from mercurial.node import hex from mercurial.pycompat import open from mercurial import pycompat from mercurial.utils import ( hashutil, procutil, ) class BundleWriteException(Exception): pass class BundleReadException(Exception): pass class abstractbundlestore: # pytype: disable=ignored-metaclass """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: """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(b'scratchbranch', b'storepath') if not self.storepath: self.storepath = self.repo.vfs.join( b"scratchbranches", b"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 = hex(hashutil.sha1(data).digest()) dirpath = self._dirpath(filename) if not os.path.exists(dirpath): os.makedirs(dirpath) with open(self._filepath(filename), b'wb') as f: f.write(data) return filename def read(self, key): try: with open(self._filepath(key), b'rb') as f: return f.read() except IOError: return None 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 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( pycompat.rapply(procutil.tonativestr, 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 # TODO: rewrite without str.format() and replace NamedTemporaryFile() # with pycompat.namedtempfile() with pycompat.namedtempfile() as temp: temp.write(data) temp.flush() temp.seek(0) formatted_args = format_placeholders_args( self.put_args, filename=temp.name ) returncode, stdout, stderr = self._call_binary( [self.put_binary] + formatted_args ) if returncode != 0: raise BundleWriteException( b'Failed to upload to external store: %s' % stderr ) stdout_lines = stdout.splitlines() if len(stdout_lines) == 1: return stdout_lines[0] else: raise BundleWriteException( b'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 pycompat.namedtempfile() as temp: formatted_args = format_placeholders_args( self.get_args, filename=temp.name, handle=handle ) returncode, stdout, stderr = self._call_binary( [self.get_binary] + formatted_args ) if returncode != 0: raise BundleReadException( b'Failed to download from external store: %s' % stderr ) return temp.read()