dulwich/object_store.py

   1 # object_store.py -- Object store for git objects
   2 # Copyright (C) 2008-2009 Jelmer Vernooij <jelmer@samba.org>
   3 #
   4 # This program is free software; you can redistribute it and/or
   5 # modify it under the terms of the GNU General Public License
   6 # as published by the Free Software Foundation; either version 2
   7 # or (at your option) a later version of the License.
   8 #
   9 # This program is distributed in the hope that it will be useful,
  10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 # GNU General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License
  15 # along with this program; if not, write to the Free Software
  16 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
  17 # MA  02110-1301, USA.
  18
  19 import itertools
  20 import os
  21 import tempfile
  22 import urllib2
  23
  24 from dulwich.errors import (
  25     NotTreeError,
  26     )
  27 from dulwich.objects import (
  28     ShaFile,
  29     Tree,
  30     hex_to_sha,
  31     sha_to_hex,
  32     )
  33 from dulwich.pack import (
  34     Pack,
  35     PackData,
  36     iter_sha1,
  37     load_packs,
  38     load_pack_index,
  39     write_pack,
  40     write_pack_data,
  41     write_pack_index_v2,
  42     )
  43
  44 PACKDIR = 'pack'
  45
  46 class ObjectStore(object):
  47     """Object store."""
  48
  49     def __init__(self, path):
  50         """Open an object store.
  51
  52         :param path: Path of the object store.
  53         """
  54         self.path = path
  55         self._pack_cache = None
  56         self.pack_dir = os.path.join(self.path, PACKDIR)
  57
  58     def determine_wants_all(self, refs):
  59             return [sha for (ref, sha) in refs.iteritems() if not sha in self and not ref.endswith("^{}")]
  60
  61     def iter_shas(self, shas):
  62         """Iterate over the objects for the specified shas.
  63
  64         :param shas: Iterable object with SHAs
  65         """
  66         return ObjectStoreIterator(self, shas)
  67
  68     def __contains__(self, sha):
  69         """Check if a particular object is present by SHA1."""
  70         for pack in self.packs:
  71             if sha in pack:
  72                 return True
  73         ret = self._get_shafile(sha)
  74         if ret is not None:
  75             return True
  76         return False
  77
  78     def __iter__(self):
  79         """Iterate over the SHAs that are present in this store."""
  80         iterables = self.packs + [self._iter_shafile_shas()]
  81         return itertools.chain(*iterables)
  82
  83     @property
  84     def packs(self):
  85         """List with pack objects."""
  86         if self._pack_cache is None:
  87             self._pack_cache = list(load_packs(self.pack_dir))
  88         return self._pack_cache
  89
  90     def _add_known_pack(self, path):
  91         """Add a newly appeared pack to the cache by path.
  92
  93         """
  94         if self._pack_cache is not None:
  95             self._pack_cache.append(Pack(path))
  96
  97     def _get_shafile_path(self, sha):
  98         dir = sha[:2]
  99         file = sha[2:]
 100         # Check from object dir
 101         return os.path.join(self.path, dir, file)
 102
 103     def _iter_shafile_shas(self):
 104         for base in os.listdir(self.path):
 105             if len(base) != 2:
 106                 continue
 107             for rest in os.listdir(os.path.join(self.path, base)):
 108                 yield base+rest
 109
 110     def _get_shafile(self, sha):
 111         path = self._get_shafile_path(sha)
 112         if os.path.exists(path):
 113           return ShaFile.from_file(path)
 114         return None
 115
 116     def _add_shafile(self, sha, o):
 117         dir = os.path.join(self.path, sha[:2])
 118         if not os.path.isdir(dir):
 119             os.mkdir(dir)
 120         path = os.path.join(dir, sha[2:])
 121         f = open(path, 'w+')
 122         try:
 123             f.write(o.as_legacy_object())
 124         finally:
 125             f.close()
 126
 127     def get_raw(self, name):
 128         """Obtain the raw text for an object.
 129
 130         :param name: sha for the object.
 131         :return: tuple with object type and object contents.
 132         """
 133         if len(name) == 40:
 134             sha = hex_to_sha(name)
 135             hexsha = name
 136         elif len(name) == 20:
 137             sha = name
 138             hexsha = None
 139         else:
 140             raise AssertionError
 141         for pack in self.packs:
 142             try:
 143                 return pack.get_raw(sha)
 144             except KeyError:
 145                 pass
 146         if hexsha is None:
 147             hexsha = sha_to_hex(name)
 148         ret = self._get_shafile(hexsha)
 149         if ret is not None:
 150             return ret.type, ret.as_raw_string()
 151         raise KeyError(hexsha)
 152
 153     def __getitem__(self, sha):
 154         type, uncomp = self.get_raw(sha)
 155         return ShaFile.from_raw_string(type, uncomp)
 156
 157     def move_in_thin_pack(self, path):
 158         """Move a specific file containing a pack into the pack directory.
 159
 160         :note: The file should be on the same file system as the
 161             packs directory.
 162
 163         :param path: Path to the pack file.
 164         """
 165         data = PackData(path)
 166
 167         # Write index for the thin pack (do we really need this?)
 168         temppath = os.path.join(self.pack_dir,
 169             sha_to_hex(urllib2.randombytes(20))+".tempidx")
 170         data.create_index_v2(temppath, self.get_raw)
 171         p = Pack.from_objects(data, load_pack_index(temppath))
 172
 173         # Write a full pack version
 174         temppath = os.path.join(self.pack_dir,
 175             sha_to_hex(urllib2.randombytes(20))+".temppack")
 176         write_pack(temppath, ((o, None) for o in p.iterobjects(self.get_raw)),
 177                 len(p))
 178         pack_sha = load_pack_index(temppath+".idx").objects_sha1()
 179         newbasename = os.path.join(self.pack_dir, "pack-%s" % pack_sha)
 180         os.rename(temppath+".pack", newbasename+".pack")
 181         os.rename(temppath+".idx", newbasename+".idx")
 182         self._add_known_pack(newbasename)
 183
 184     def move_in_pack(self, path):
 185         """Move a specific file containing a pack into the pack directory.
 186
 187         :note: The file should be on the same file system as the
 188             packs directory.
 189
 190         :param path: Path to the pack file.
 191         """
 192         p = PackData(path)
 193         entries = p.sorted_entries()
 194         basename = os.path.join(self.pack_dir,
 195             "pack-%s" % iter_sha1(entry[0] for entry in entries))
 196         write_pack_index_v2(basename+".idx", entries, p.get_stored_checksum())
 197         os.rename(path, basename + ".pack")
 198         self._add_known_pack(basename)
 199
 200     def add_thin_pack(self):
 201         """Add a new thin pack to this object store.
 202
 203         Thin packs are packs that contain deltas with parents that exist
 204         in a different pack.
 205         """
 206         fd, path = tempfile.mkstemp(dir=self.pack_dir, suffix=".pack")
 207         f = os.fdopen(fd, 'w')
 208         def commit():
 209             os.fsync(fd)
 210             f.close()
 211             if os.path.getsize(path) > 0:
 212                 self.move_in_thin_pack(path)
 213         return f, commit
 214
 215     def add_pack(self):
 216         """Add a new pack to this object store.
 217
 218         :return: Fileobject to write to and a commit function to
 219             call when the pack is finished.
 220         """
 221         fd, path = tempfile.mkstemp(dir=self.pack_dir, suffix=".pack")
 222         f = os.fdopen(fd, 'w')
 223         def commit():
 224             os.fsync(fd)
 225             f.close()
 226             if os.path.getsize(path) > 0:
 227                 self.move_in_pack(path)
 228         return f, commit
 229
 230     def add_object(self, obj):
 231         """Add a single object to this object store.
 232
 233         """
 234         self._add_shafile(obj.id, obj)
 235
 236     def add_objects(self, objects):
 237         """Add a set of objects to this object store.
 238
 239         :param objects: Iterable over a list of objects.
 240         """
 241         if len(objects) == 0:
 242             return
 243         f, commit = self.add_pack()
 244         write_pack_data(f, objects, len(objects))
 245         commit()
 246
 247     def find_missing_objects(self, wants, graph_walker, progress):
 248         """Find the missing objects required for a set of revisions.
 249
 250         :param wants: Iterable over SHAs of objects to fetch.
 251         :param graph_walker: Object that can iterate over the list of revisions
 252             to fetch and has an "ack" method that will be called to acknowledge
 253             that a revision is present.
 254         :param progress: Simple progress function that will be called with
 255             updated progress strings.
 256         :return: Iterator over (sha, path) pairs.
 257         """
 258         return iter(MissingObjectFinder(self, wants, graph_walker, progress).next, None)
 259
 260
 261 class ObjectImporter(object):
 262     """Interface for importing objects."""
 263
 264     def __init__(self, count):
 265         """Create a new ObjectImporter.
 266
 267         :param count: Number of objects that's going to be imported.
 268         """
 269         self.count = count
 270
 271     def add_object(self, object):
 272         """Add an object."""
 273         raise NotImplementedError(self.add_object)
 274
 275     def finish(self, object):
 276         """Finish the imoprt and write objects to disk."""
 277         raise NotImplementedError(self.finish)
 278
 279
 280 class ObjectIterator(object):
 281     """Interface for iterating over objects."""
 282
 283     def iterobjects(self):
 284         raise NotImplementedError(self.iterobjects)
 285
 286
 287 class ObjectStoreIterator(ObjectIterator):
 288     """ObjectIterator that works on top of an ObjectStore."""
 289
 290     def __init__(self, store, sha_iter):
 291         self.store = store
 292         self.sha_iter = sha_iter
 293         self._shas = []
 294
 295     def __iter__(self):
 296         for sha, path in self.itershas():
 297             yield self.store[sha], path
 298
 299     def iterobjects(self):
 300         for o, path in self:
 301             yield o
 302
 303     def itershas(self):
 304         for sha in self._shas:
 305             yield sha
 306         for sha in self.sha_iter:
 307             self._shas.append(sha)
 308             yield sha
 309
 310     def __contains__(self, needle):
 311         """Check if an object is present.
 312
 313         :param needle: SHA1 of the object to check for
 314         """
 315         return needle in self.store
 316
 317     def __getitem__(self, key):
 318         """Find an object by SHA1."""
 319         return self.store[key]
 320
 321     def __len__(self):
 322         """Return the number of objects."""
 323         return len(list(self.itershas()))
 324
 325
 326 def tree_lookup_path(lookup_obj, root_sha, path):
 327     parts = path.split("/")
 328     sha = root_sha
 329     for p in parts:
 330         obj = lookup_obj(sha)
 331         if type(obj) is not Tree:
 332             raise NotTreeError(sha)
 333         if p == '':
 334             continue
 335         mode, sha = obj[p]
 336     return lookup_obj(sha)
 337
 338
 339 class MissingObjectFinder(object):
 340     """Find the objects missing from another object store.
 341
 342     :param object_store: Object store containing at least all objects to be
 343         sent
 344     :param wants: SHA1s of commits to send
 345     :param graph_walker: graph walker object used to see what the remote
 346         repo has and misses
 347     :param progress: Optional function to report progress to.
 348     """
 349
 350     def __init__(self, object_store, wants, graph_walker, progress=None):
 351         self.sha_done = set()
 352         self.objects_to_send = set([(w, None) for w in wants])
 353         self.object_store = object_store
 354         if progress is None:
 355             self.progress = lambda x: None
 356         else:
 357             self.progress = progress
 358         ref = graph_walker.next()
 359         while ref:
 360             if ref in self.object_store:
 361                 graph_walker.ack(ref)
 362             ref = graph_walker.next()
 363
 364     def add_todo(self, entries):
 365         self.objects_to_send.update([e for e in entries if not e in self.sha_done])
 366
 367     def parse_tree(self, tree):
 368         self.add_todo([(sha, name) for (mode, name, sha) in tree.entries()])
 369
 370     def parse_commit(self, commit):
 371         self.add_todo([(commit.tree, "")])
 372         self.add_todo([(p, None) for p in commit.parents])
 373
 374     def parse_tag(self, tag):
 375         self.add_todo([(tag.object[1], None)])
 376
 377     def next(self):
 378         if not self.objects_to_send:
 379             return None
 380         (sha, name) = self.objects_to_send.pop()
 381         o = self.object_store[sha]
 382         if isinstance(o, Commit):
 383             self.parse_commit(o)
 384         elif isinstance(o, Tree):
 385             self.parse_tree(o)
 386         elif isinstance(o, Tag):
 387             self.parse_tag(o)
 388         self.sha_done.add((sha, name))
 389         self.progress("counting objects: %d\r" % len(self.sha_done))
 390         return (sha, name)
 391
 392
 393