~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/store/__init__.py

Merge in format-5 work - release bzr 0.1rc1.

Show diffs side-by-side

added added

removed removed

Lines of Context:
14
14
# along with this program; if not, write to the Free Software
15
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
16
16
 
 
17
# TODO: Could remember a bias towards whether a particular store is typically
 
18
# compressed or not.
 
19
 
17
20
"""
18
21
Stores are the main data-storage mechanism for Bazaar-NG.
19
22
 
21
24
unique ID.
22
25
"""
23
26
 
24
 
import os, tempfile, types, osutils, gzip, errno
25
 
from stat import ST_SIZE
26
 
from StringIO import StringIO
27
 
from bzrlib.errors import BzrError
 
27
from cStringIO import StringIO
 
28
 
 
29
from bzrlib.errors import BzrError, UnlistableStore, TransportNotPossible
28
30
from bzrlib.trace import mutter
29
 
import bzrlib.ui
 
31
import bzrlib.transport
 
32
from bzrlib.transport.local import LocalTransport
30
33
 
31
34
######################################################################
32
35
# stores
35
38
    pass
36
39
 
37
40
 
38
 
class ImmutableStore(object):
39
 
    """Store that holds files indexed by unique names.
40
 
 
 
41
class Store(object):
 
42
    """This class represents the abstract storage layout for saving information.
 
43
    
41
44
    Files can be added, but not modified once they are in.  Typically
42
45
    the hash is used as the name, or something else known to be unique,
43
46
    such as a UUID.
44
 
 
45
 
    >>> st = ImmutableScratchStore()
46
 
 
47
 
    >>> st.add(StringIO('hello'), 'aa')
48
 
    >>> 'aa' in st
49
 
    True
50
 
    >>> 'foo' in st
51
 
    False
52
 
 
53
 
    You are not allowed to add an id that is already present.
54
 
 
55
 
    Entries can be retrieved as files, which may then be read.
56
 
 
57
 
    >>> st.add(StringIO('goodbye'), '123123')
58
 
    >>> st['123123'].read()
59
 
    'goodbye'
60
 
 
61
 
    TODO: Atomic add by writing to a temporary file and renaming.
62
 
 
63
 
    In bzr 0.0.5 and earlier, files within the store were marked
64
 
    readonly on disk.  This is no longer done but existing stores need
65
 
    to be accomodated.
66
47
    """
67
48
 
68
 
    def __init__(self, basedir):
69
 
        self._basedir = basedir
70
 
 
71
 
    def _path(self, entry_id):
72
 
        if not isinstance(entry_id, basestring):
73
 
            raise TypeError(type(entry_id))
74
 
        if '\\' in entry_id or '/' in entry_id:
75
 
            raise ValueError("invalid store id %r" % entry_id)
76
 
        return os.path.join(self._basedir, entry_id)
77
 
 
78
 
    def __repr__(self):
79
 
        return "%s(%r)" % (self.__class__.__name__, self._basedir)
80
 
 
81
 
    def add(self, f, fileid, compressed=True):
82
 
        """Add contents of a file into the store.
83
 
 
84
 
        f -- An open file, or file-like object."""
85
 
        # FIXME: Only works on files that will fit in memory
86
 
        
87
 
        from bzrlib.atomicfile import AtomicFile
88
 
        
89
 
        mutter("add store entry %r" % (fileid))
90
 
        if isinstance(f, types.StringTypes):
91
 
            content = f
92
 
        else:
93
 
            content = f.read()
94
 
            
95
 
        p = self._path(fileid)
96
 
        if os.access(p, os.F_OK) or os.access(p + '.gz', os.F_OK):
97
 
            raise BzrError("store %r already contains id %r" % (self._basedir, fileid))
98
 
 
99
 
        fn = p
100
 
        if compressed:
101
 
            fn = fn + '.gz'
102
 
            
103
 
        af = AtomicFile(fn, 'wb')
104
 
        try:
105
 
            if compressed:
106
 
                gf = gzip.GzipFile(mode='wb', fileobj=af)
107
 
                gf.write(content)
108
 
                gf.close()
 
49
    def __len__(self):
 
50
        raise NotImplementedError('Children should define their length')
 
51
 
 
52
    def __getitem__(self, fileid):
 
53
        """Returns a file reading from a particular entry."""
 
54
        raise NotImplementedError
 
55
 
 
56
    def __contains__(self, fileid):
 
57
        """"""
 
58
        raise NotImplementedError
 
59
 
 
60
    def __iter__(self):
 
61
        raise NotImplementedError
 
62
 
 
63
    def add(self, f, fileid):
 
64
        """Add a file object f to the store accessible from the given fileid"""
 
65
        raise NotImplementedError('Children of Store must define their method of adding entries.')
 
66
 
 
67
    def add_multi(self, entries):
 
68
        """Add a series of file-like or string objects to the store with the given
 
69
        identities.
 
70
        
 
71
        :param entries: A list of tuples of file,id pairs [(file1, id1), (file2, id2), ...]
 
72
                        This could also be a generator yielding (file,id) pairs.
 
73
        """
 
74
        for f, fileid in entries:
 
75
            self.add(f, fileid)
 
76
 
 
77
    def has(self, fileids):
 
78
        """Return True/False for each entry in fileids.
 
79
 
 
80
        :param fileids: A List or generator yielding file ids.
 
81
        :return: A generator or list returning True/False for each entry.
 
82
        """
 
83
        for fileid in fileids:
 
84
            if fileid in self:
 
85
                yield True
109
86
            else:
110
 
                af.write(content)
111
 
            af.commit()
112
 
        finally:
113
 
            af.close()
114
 
 
115
 
 
116
 
    def copy_multi(self, other, ids, permit_failure=False):
 
87
                yield False
 
88
 
 
89
    def listable(self):
 
90
        """Return True if this store is able to be listed."""
 
91
        return hasattr(self, "__iter__")
 
92
 
 
93
    def get(self, fileids, permit_failure=False, pb=None):
 
94
        """Return a set of files, one for each requested entry.
 
95
        
 
96
        :param permit_failure: If true, return None for entries which do not 
 
97
                               exist.
 
98
        :return: A list or generator of file-like objects, one for each id.
 
99
        """
 
100
        for fileid in fileids:
 
101
            try:
 
102
                yield self[fileid]
 
103
            except KeyError:
 
104
                if permit_failure:
 
105
                    yield None
 
106
                else:
 
107
                    raise
 
108
 
 
109
    def copy_multi(self, other, ids, pb=None, permit_failure=False):
117
110
        """Copy texts for ids from other into self.
118
111
 
119
 
        If an id is present in self, it is skipped.
 
112
        If an id is present in self, it is skipped.  A count of copied
 
113
        ids is returned, which may be less than len(ids).
120
114
 
121
 
        Returns (count_copied, failed), where failed is a collection of ids
122
 
        that could not be copied.
 
115
        :param other: Another Store object
 
116
        :param ids: A list of entry ids to be copied
 
117
        :param pb: A ProgressBar object, if none is given, the default will be created.
 
118
        :param permit_failure: Allow missing entries to be ignored
 
119
        :return: (n_copied, [failed]) The number of entries copied successfully,
 
120
            followed by a list of entries which could not be copied (because they
 
121
            were missing)
123
122
        """
124
 
        pb = bzrlib.ui.ui_factory.progress_bar()
125
 
        
 
123
        if pb is None:
 
124
            pb = bzrlib.ui.ui_factory.progress_bar()
 
125
 
 
126
        # XXX: Is there any reason why we couldn't make this accept a generator
 
127
        # and build a list as it finds things to copy?
 
128
        ids = list(ids) # Make sure we don't have a generator, since we iterate 2 times
126
129
        pb.update('preparing to copy')
127
 
        to_copy = [id for id in ids if id not in self]
128
 
        if isinstance(other, ImmutableStore):
129
 
            return self.copy_multi_immutable(other, to_copy, pb, 
130
 
                                             permit_failure=permit_failure)
131
 
        count = 0
 
130
        to_copy = []
 
131
        for file_id, has in zip(ids, self.has(ids)):
 
132
            if not has:
 
133
                to_copy.append(file_id)
 
134
        return self._do_copy(other, to_copy, pb, permit_failure=permit_failure)
 
135
 
 
136
    def _do_copy(self, other, to_copy, pb, permit_failure=False):
 
137
        """This is the standard copying mechanism, just get them one at
 
138
        a time from remote, and store them locally.
 
139
 
 
140
        :param other: Another Store object
 
141
        :param to_copy: A list of entry ids to copy
 
142
        :param pb: A ProgressBar object to display completion status.
 
143
        :param permit_failure: Allow missing entries to be ignored
 
144
        :return: (n_copied, [failed])
 
145
            The number of entries copied, and a list of failed entries.
 
146
        """
 
147
        # This should be updated to use add_multi() rather than
 
148
        # the current methods of buffering requests.
 
149
        # One question, is it faster to queue up 1-10 and then copy 1-10
 
150
        # then queue up 11-20, copy 11-20
 
151
        # or to queue up 1-10, copy 1, queue 11, copy 2, etc?
 
152
        # sort of pipeline versus batch.
 
153
 
 
154
        # We can't use self._transport.copy_to because we don't know
 
155
        # whether the local tree is in the same format as other
132
156
        failed = set()
133
 
        for id in to_copy:
134
 
            count += 1
135
 
            pb.update('copy', count, len(to_copy))
136
 
            if not permit_failure:
137
 
                self.add(other[id], id)
138
 
            else:
 
157
        def buffer_requests():
 
158
            count = 0
 
159
            buffered_requests = []
 
160
            for fileid in to_copy:
139
161
                try:
140
 
                    entry = other[id]
 
162
                    f = other[fileid]
141
163
                except KeyError:
142
 
                    failed.add(id)
143
 
                    continue
144
 
                self.add(entry, id)
145
 
                
146
 
        if not permit_failure:
 
164
                    if permit_failure:
 
165
                        failed.add(fileid)
 
166
                        continue
 
167
                    else:
 
168
                        raise
 
169
 
 
170
                buffered_requests.append((f, fileid))
 
171
                if len(buffered_requests) > self._max_buffered_requests:
 
172
                    yield buffered_requests.pop(0)
 
173
                    count += 1
 
174
                    pb.update('copy', count, len(to_copy))
 
175
 
 
176
            for req in buffered_requests:
 
177
                yield req
 
178
                count += 1
 
179
                pb.update('copy', count, len(to_copy))
 
180
 
147
181
            assert count == len(to_copy)
148
 
        pb.clear()
149
 
        return count, failed
150
 
 
151
 
    def copy_multi_immutable(self, other, to_copy, pb, permit_failure=False):
152
 
        from shutil import copyfile
153
 
        count = 0
154
 
        failed = set()
155
 
        for id in to_copy:
156
 
            p = self._path(id)
157
 
            other_p = other._path(id)
158
 
            try:
159
 
                copyfile(other_p, p)
160
 
            except IOError, e:
161
 
                if e.errno == errno.ENOENT:
162
 
                    if not permit_failure:
163
 
                        copyfile(other_p+".gz", p+".gz")
164
 
                    else:
165
 
                        try:
166
 
                            copyfile(other_p+".gz", p+".gz")
167
 
                        except IOError, e:
168
 
                            if e.errno == errno.ENOENT:
169
 
                                failed.add(id)
170
 
                            else:
171
 
                                raise
172
 
                else:
173
 
                    raise
174
 
            
175
 
            count += 1
176
 
            pb.update('copy', count, len(to_copy))
177
 
        assert count == len(to_copy)
178
 
        pb.clear()
179
 
        return count, failed
180
 
    
 
182
 
 
183
        self.add_multi(buffer_requests())
 
184
 
 
185
        pb.clear()
 
186
        return len(to_copy), failed
 
187
 
 
188
 
 
189
class TransportStore(Store):
 
190
    """A TransportStore is a Store superclass for Stores that use Transports."""
 
191
 
 
192
    _max_buffered_requests = 10
 
193
 
 
194
    def __init__(self, transport):
 
195
        assert isinstance(transport, bzrlib.transport.Transport)
 
196
        super(TransportStore, self).__init__()
 
197
        self._transport = transport
 
198
 
 
199
    def __repr__(self):
 
200
        if self._transport is None:
 
201
            return "%s(None)" % (self.__class__.__name__)
 
202
        else:
 
203
            return "%s(%r)" % (self.__class__.__name__, self._transport.base)
 
204
 
 
205
    __str__ = __repr__
 
206
 
 
207
    def listable(self):
 
208
        """Return True if this store is able to be listed."""
 
209
        return self._transport.listable()
 
210
 
 
211
 
 
212
class ImmutableMemoryStore(Store):
 
213
    """A memory only store."""
181
214
 
182
215
    def __contains__(self, fileid):
183
 
        """"""
184
 
        p = self._path(fileid)
185
 
        return (os.access(p, os.R_OK)
186
 
                or os.access(p + '.gz', os.R_OK))
187
 
 
188
 
    # TODO: Guard against the same thing being stored twice, compressed and uncompresse
189
 
 
190
 
    def __iter__(self):
191
 
        for f in os.listdir(self._basedir):
192
 
            if f[-3:] == '.gz':
193
 
                # TODO: case-insensitive?
194
 
                yield f[:-3]
195
 
            else:
196
 
                yield f
197
 
 
198
 
    def __len__(self):
199
 
        return len(os.listdir(self._basedir))
200
 
 
 
216
        return self._contents.has_key(fileid)
 
217
 
 
218
    def __init__(self):
 
219
        super(ImmutableMemoryStore, self).__init__()
 
220
        self._contents = {}
 
221
 
 
222
    def add(self, stream, fileid, compressed=True):
 
223
        if self._contents.has_key(fileid):
 
224
            raise StoreError("fileid %s already in the store" % fileid)
 
225
        self._contents[fileid] = stream.read()
201
226
 
202
227
    def __getitem__(self, fileid):
203
228
        """Returns a file reading from a particular entry."""
204
 
        p = self._path(fileid)
205
 
        try:
206
 
            return gzip.GzipFile(p + '.gz', 'rb')
207
 
        except IOError, e:
208
 
            if e.errno != errno.ENOENT:
209
 
                raise
210
 
 
211
 
        try:
212
 
            return file(p, 'rb')
213
 
        except IOError, e:
214
 
            if e.errno != errno.ENOENT:
215
 
                raise
216
 
 
217
 
        raise KeyError(fileid)
218
 
 
 
229
        if not self._contents.has_key(fileid):
 
230
            raise IndexError
 
231
        return StringIO(self._contents[fileid])
 
232
 
 
233
    def _item_size(self, fileid):
 
234
        return len(self._contents[fileid])
 
235
 
 
236
    def __iter__(self):
 
237
        return iter(self._contents.keys())
219
238
 
220
239
    def total_size(self):
221
 
        """Return (count, bytes)
222
 
 
223
 
        This is the (compressed) size stored on disk, not the size of
224
 
        the content."""
225
 
        total = 0
 
240
        result = 0
226
241
        count = 0
227
 
        for fid in self:
 
242
        for fileid in self:
228
243
            count += 1
229
 
            p = self._path(fid)
230
 
            try:
231
 
                total += os.stat(p)[ST_SIZE]
232
 
            except OSError:
233
 
                total += os.stat(p + '.gz')[ST_SIZE]
234
 
                
235
 
        return count, total
236
 
 
237
 
 
238
 
 
239
 
 
240
 
class ImmutableScratchStore(ImmutableStore):
241
 
    """Self-destructing test subclass of ImmutableStore.
242
 
 
243
 
    The Store only exists for the lifetime of the Python object.
244
 
 Obviously you should not put anything precious in it.
 
244
            result += self._item_size(fileid)
 
245
        return count, result
 
246
        
 
247
 
 
248
class CachedStore(Store):
 
249
    """A store that caches data locally, to avoid repeated downloads.
 
250
    The precacache method should be used to avoid server round-trips for
 
251
    every piece of data.
245
252
    """
246
 
    def __init__(self):
247
 
        ImmutableStore.__init__(self, tempfile.mkdtemp())
248
 
 
249
 
    def __del__(self):
250
 
        for f in os.listdir(self._basedir):
251
 
            fpath = os.path.join(self._basedir, f)
252
 
            # needed on windows, and maybe some other filesystems
253
 
            os.chmod(fpath, 0600)
254
 
            os.remove(fpath)
255
 
        os.rmdir(self._basedir)
256
 
        mutter("%r destroyed" % self)
 
253
 
 
254
    def __init__(self, store, cache_dir):
 
255
        super(CachedStore, self).__init__()
 
256
        self.source_store = store
 
257
        # This clones the source store type with a locally bound
 
258
        # transport. FIXME: it assumes a constructor is == cloning.
 
259
        # clonable store - it might be nicer to actually have a clone()
 
260
        # or something. RBC 20051003
 
261
        self.cache_store = store.__class__(LocalTransport(cache_dir))
 
262
 
 
263
    def __getitem__(self, id):
 
264
        mutter("Cache add %s" % id)
 
265
        if id not in self.cache_store:
 
266
            self.cache_store.add(self.source_store[id], id)
 
267
        return self.cache_store[id]
 
268
 
 
269
    def __contains__(self, fileid):
 
270
        if fileid in self.cache_store:
 
271
            return True
 
272
        if fileid in self.source_store:
 
273
            # We could copy at this time
 
274
            return True
 
275
        return False
 
276
 
 
277
    def get(self, fileids, permit_failure=False, pb=None):
 
278
        fileids = list(fileids)
 
279
        hasids = self.cache_store.has(fileids)
 
280
        needs = set()
 
281
        for has, fileid in zip(hasids, fileids):
 
282
            if not has:
 
283
                needs.add(fileid)
 
284
        if needs:
 
285
            self.cache_store.copy_multi(self.source_store, needs,
 
286
                    permit_failure=permit_failure)
 
287
        return self.cache_store.get(fileids,
 
288
                permit_failure=permit_failure, pb=pb)
 
289
 
 
290
    def prefetch(self, ids):
 
291
        """Copy a series of ids into the cache, before they are used.
 
292
        For remote stores that support pipelining or async downloads, this can
 
293
        increase speed considerably.
 
294
 
 
295
        Failures while prefetching are ignored.
 
296
        """
 
297
        mutter("Prefetch of ids %s" % ",".join(ids))
 
298
        self.cache_store.copy_multi(self.source_store, ids, 
 
299
                                    permit_failure=True)
 
300
 
 
301
 
 
302
def copy_all(store_from, store_to):
 
303
    """Copy all ids from one store to another."""
 
304
    # TODO: Optional progress indicator
 
305
    if not store_from.listable():
 
306
        raise UnlistableStore(store_from)
 
307
    ids = [f for f in store_from]
 
308
    store_to.copy_multi(store_from, ids)
 
309