~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/store.py

  • Committer: Martin Pool
  • Date: 2005-09-22 07:04:14 UTC
  • Revision ID: mbp@sourcefrog.net-20050922070414-06a503a1d47b5ce4
- avoid converting inventories to/from StringIO 

Show diffs side-by-side

added added

removed removed

Lines of Context:
21
21
unique ID.
22
22
"""
23
23
 
24
 
import bzrlib
 
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
25
28
from bzrlib.trace import mutter
26
29
import bzrlib.ui
27
 
import bzrlib.transport
28
30
 
29
31
######################################################################
30
32
# stores
32
34
class StoreError(Exception):
33
35
    pass
34
36
 
35
 
class Store(object):
36
 
    """This class represents the abstract storage layout for saving information.
 
37
 
 
38
class ImmutableStore(object):
 
39
    """Store that holds files indexed by unique names.
 
40
 
 
41
    Files can be added, but not modified once they are in.  Typically
 
42
    the hash is used as the name, or something else known to be unique,
 
43
    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.
37
66
    """
38
 
    _transport = None
39
 
    _max_buffered_requests = 10
40
 
 
41
 
    def __init__(self, transport):
42
 
        assert isinstance(transport, bzrlib.transport.Transport)
43
 
        self._transport = transport
 
67
 
 
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)
44
77
 
45
78
    def __repr__(self):
46
 
        if self._transport is None:
47
 
            return "%s(None)" % (self.__class__.__name__)
 
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
48
92
        else:
49
 
            return "%s(%r)" % (self.__class__.__name__, self._transport.base)
50
 
 
51
 
    __str__ = __repr__
52
 
 
53
 
    def __len__(self):
54
 
        raise NotImplementedError('Children should define their length')
55
 
 
56
 
    def __getitem__(self, fileid):
57
 
        """Returns a file reading from a particular entry."""
58
 
        raise NotImplementedError
59
 
 
60
 
    def __contains__(self, fileid):
61
 
        """"""
62
 
        raise NotImplementedError
63
 
 
64
 
    def __iter__(self):
65
 
        raise NotImplementedError
66
 
 
67
 
    def add(self, f, fileid):
68
 
        """Add a file object f to the store accessible from the given fileid"""
69
 
        raise NotImplementedError('Children of Store must define their method of adding entries.')
70
 
 
71
 
    def add_multi(self, entries):
72
 
        """Add a series of file-like or string objects to the store with the given
73
 
        identities.
74
 
        
75
 
        :param entries: A list of tuples of file,id pairs [(file1, id1), (file2, id2), ...]
76
 
                        This could also be a generator yielding (file,id) pairs.
77
 
        """
78
 
        for f, fileid in entries:
79
 
            self.add(f, fileid)
80
 
 
81
 
    def has(self, fileids):
82
 
        """Return True/False for each entry in fileids.
83
 
 
84
 
        :param fileids: A List or generator yielding file ids.
85
 
        :return: A generator or list returning True/False for each entry.
86
 
        """
87
 
        for fileid in fileids:
88
 
            if fileid in self:
89
 
                yield True
 
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()
90
109
            else:
91
 
                yield False
92
 
 
93
 
    def get(self, fileids, ignore_missing=False, pb=None):
94
 
        """Return a set of files, one for each requested entry.
95
 
        
96
 
        :param ignore_missing: 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 ignore_missing:
105
 
                    yield None
106
 
                else:
107
 
                    raise
108
 
 
109
 
    def copy_multi(self, other, ids, pb=None, permit_failure=False):
 
110
                af.write(content)
 
111
            af.commit()
 
112
        finally:
 
113
            af.close()
 
114
 
 
115
 
 
116
    def copy_multi(self, other, ids, permit_failure=False):
110
117
        """Copy texts for ids from other into self.
111
118
 
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).
 
119
        If an id is present in self, it is skipped.
114
120
 
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)
 
121
        Returns (count_copied, failed), where failed is a collection of ids
 
122
        that could not be copied.
122
123
        """
123
 
        if pb is None:
124
 
            pb = bzrlib.ui.ui_factory.progress_bar()
125
 
 
126
 
        ids = list(ids) # Make sure we don't have a generator, since we iterate 2 times
 
124
        pb = bzrlib.ui.ui_factory.progress_bar()
 
125
        
127
126
        pb.update('preparing to copy')
128
 
        to_copy = []
129
 
        for file_id, has in zip(ids, self.has(ids)):
130
 
            if not has:
131
 
                to_copy.append(file_id)
132
 
        return self._do_copy(other, to_copy, pb, permit_failure=permit_failure)
133
 
 
134
 
    def _do_copy(self, other, to_copy, pb, permit_failure=False):
135
 
        """This is the standard copying mechanism, just get them one at
136
 
        a time from remote, and store them locally.
137
 
 
138
 
        :param other: Another Store object
139
 
        :param to_copy: A list of entry ids to copy
140
 
        :param pb: A ProgressBar object to display completion status.
141
 
        :param permit_failure: Allow missing entries to be ignored
142
 
        :return: (n_copied, [failed])
143
 
            The number of entries copied, and a list of failed entries.
144
 
        """
145
 
        # This should be updated to use add_multi() rather than
146
 
        # the current methods of buffering requests.
147
 
        # One question, is it faster to queue up 1-10 and then copy 1-10
148
 
        # then queue up 11-20, copy 11-20
149
 
        # or to queue up 1-10, copy 1, queue 11, copy 2, etc?
150
 
        # sort of pipeline versus batch.
151
 
 
152
 
        # We can't use self._transport.copy_to because we don't know
153
 
        # whether the local tree is in the same format as other
 
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
        count = 0
154
131
        failed = set()
155
 
        def buffer_requests():
156
 
            count = 0
157
 
            buffered_requests = []
158
 
            for fileid in to_copy:
 
132
        for id in to_copy:
 
133
            count += 1
 
134
            pb.update('copy', count, len(to_copy))
 
135
            if not permit_failure:
 
136
                self.add(other[id], id)
 
137
            else:
159
138
                try:
160
 
                    f = other[fileid]
161
 
                except KeyError:
162
 
                    if permit_failure:
163
 
                        failed.add(fileid)
164
 
                        continue
 
139
                    entry = other[id]
 
140
                except IndexError:
 
141
                    failed.add(id)
 
142
                    continue
 
143
                self.add(entry, id)
 
144
                
 
145
        if not permit_failure:
 
146
            assert count == len(to_copy)
 
147
        pb.clear()
 
148
        return count, failed
 
149
 
 
150
    def copy_multi_immutable(self, other, to_copy, pb, permit_failure=False):
 
151
        from shutil import copyfile
 
152
        count = 0
 
153
        failed = set()
 
154
        for id in to_copy:
 
155
            p = self._path(id)
 
156
            other_p = other._path(id)
 
157
            try:
 
158
                copyfile(other_p, p)
 
159
            except IOError, e:
 
160
                if e.errno == errno.ENOENT:
 
161
                    if not permit_failure:
 
162
                        copyfile(other_p+".gz", p+".gz")
165
163
                    else:
166
 
                        raise
167
 
 
168
 
                buffered_requests.append((f, fileid))
169
 
                if len(buffered_requests) > self._max_buffered_requests:
170
 
                    yield buffered_requests.pop(0)
171
 
                    count += 1
172
 
                    pb.update('copy', count, len(to_copy))
173
 
 
174
 
            for req in buffered_requests:
175
 
                yield req
176
 
                count += 1
177
 
                pb.update('copy', count, len(to_copy))
178
 
 
179
 
            assert count == len(to_copy)
180
 
 
181
 
        self.add_multi(buffer_requests())
182
 
 
 
164
                        try:
 
165
                            copyfile(other_p+".gz", p+".gz")
 
166
                        except IOError, e:
 
167
                            if e.errno == errno.ENOENT:
 
168
                                failed.add(id)
 
169
                            else:
 
170
                                raise
 
171
                else:
 
172
                    raise
 
173
            
 
174
            count += 1
 
175
            pb.update('copy', count, len(to_copy))
 
176
        assert count == len(to_copy)
183
177
        pb.clear()
184
 
        return len(to_copy), failed
185
 
 
 
178
        return count, failed
 
179
    
 
180
 
 
181
    def __contains__(self, fileid):
 
182
        """"""
 
183
        p = self._path(fileid)
 
184
        return (os.access(p, os.R_OK)
 
185
                or os.access(p + '.gz', os.R_OK))
 
186
 
 
187
    # TODO: Guard against the same thing being stored twice,
 
188
    # compressed and uncompressed
 
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
 
 
201
 
 
202
    def __getitem__(self, fileid):
 
203
        """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 IndexError(fileid)
 
218
 
 
219
 
 
220
    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
 
226
        count = 0
 
227
        for fid in self:
 
228
            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.
 
245
    """
 
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)