1##############################################################################
2#
3# Copyright (c) Zope Foundation and Contributors.
4# All Rights Reserved.
5#
6# This software is subject to the provisions of the Zope Public License,
7# Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.
8# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
9# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
10# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
11# FOR A PARTICULAR PURPOSE
12#
13##############################################################################
14"""A simple in-memory mapping-based ZODB storage
15
16This storage provides an example implementation of a fairly full
17storage without distracting storage details.
18"""
19
20import BTrees
21import time
22import ZODB.BaseStorage
23import ZODB.interfaces
24import ZODB.POSException
25import ZODB.TimeStamp
26import ZODB.utils
27import zope.interface
28
29
30@zope.interface.implementer(
31        ZODB.interfaces.IStorage,
32        ZODB.interfaces.IStorageIteration,
33        )
34class MappingStorage(object):
35    """In-memory storage implementation
36
37    Note that this implementation is somewhat naive and inefficient
38    with regard to locking.  Its implementation is primarily meant to
39    be a simple illustration of storage implementation. It's also
40    useful for testing and exploration where scalability and efficiency
41    are unimportant.
42    """
43
44    def __init__(self, name='MappingStorage'):
45        """Create a mapping storage
46
47        The name parameter is used by the
48        :meth:`~ZODB.interfaces.IStorage.getName` and
49        :meth:`~ZODB.interfaces.IStorage.sortKey` methods.
50        """
51        self.__name__ = name
52        self._data = {}                               # {oid->{tid->pickle}}
53        self._transactions = BTrees.OOBTree.OOBTree() # {tid->TransactionRecord}
54        self._ltid = ZODB.utils.z64
55        self._last_pack = None
56        self._lock = ZODB.utils.RLock()
57        self._commit_lock = ZODB.utils.Lock()
58        self._opened = True
59        self._transaction = None
60        self._oid = 0
61
62    ######################################################################
63    # Preconditions:
64
65    def opened(self):
66        """The storage is open
67        """
68        return self._opened
69
70    def not_in_transaction(self):
71        """The storage is not committing a transaction
72        """
73        return self._transaction is None
74
75    #
76    ######################################################################
77
78    # testing framework (lame)
79    def cleanup(self):
80        pass
81
82    # ZODB.interfaces.IStorage
83    @ZODB.utils.locked
84    def close(self):
85        self._opened = False
86
87    # ZODB.interfaces.IStorage
88    def getName(self):
89        return self.__name__
90
91    # ZODB.interfaces.IStorage
92    @ZODB.utils.locked(opened)
93    def getSize(self):
94        size = 0
95        for oid, tid_data in self._data.items():
96            size += 50
97            for tid, pickle in tid_data.items():
98                size += 100+len(pickle)
99        return size
100
101    # ZEO.interfaces.IServeable
102    @ZODB.utils.locked(opened)
103    def getTid(self, oid):
104        tid_data = self._data.get(oid)
105        if tid_data:
106            return tid_data.maxKey()
107        raise ZODB.POSException.POSKeyError(oid)
108
109    # ZODB.interfaces.IStorage
110    @ZODB.utils.locked(opened)
111    def history(self, oid, size=1):
112        tid_data = self._data.get(oid)
113        if not tid_data:
114            raise ZODB.POSException.POSKeyError(oid)
115
116        tids = tid_data.keys()[-size:]
117        tids.reverse()
118        return [
119            dict(
120                time = ZODB.TimeStamp.TimeStamp(tid).timeTime(),
121                tid = tid,
122                serial = tid,
123                user_name = self._transactions[tid].user,
124                description = self._transactions[tid].description,
125                extension = self._transactions[tid].extension,
126                size = len(tid_data[tid])
127                )
128            for tid in tids]
129
130    # ZODB.interfaces.IStorage
131    def isReadOnly(self):
132        return False
133
134    # ZODB.interfaces.IStorageIteration
135    def iterator(self, start=None, end=None):
136        for transaction_record in self._transactions.values(start, end):
137            yield transaction_record
138
139    # ZODB.interfaces.IStorage
140    @ZODB.utils.locked(opened)
141    def lastTransaction(self):
142        return self._ltid
143
144    # ZODB.interfaces.IStorage
145    @ZODB.utils.locked(opened)
146    def __len__(self):
147        return len(self._data)
148
149    load = ZODB.utils.load_current
150
151    # ZODB.interfaces.IStorage
152    @ZODB.utils.locked(opened)
153    def loadBefore(self, oid, tid):
154        tid_data = self._data.get(oid)
155        if tid_data:
156            before = ZODB.utils.u64(tid)
157            if not before:
158                return None
159            before = ZODB.utils.p64(before-1)
160            tids_before = tid_data.keys(None, before)
161            if tids_before:
162                tids_after = tid_data.keys(tid, None)
163                tid = tids_before[-1]
164                return (tid_data[tid], tid,
165                        (tids_after and tids_after[0] or None)
166                        )
167        else:
168            raise ZODB.POSException.POSKeyError(oid)
169
170
171    # ZODB.interfaces.IStorage
172    @ZODB.utils.locked(opened)
173    def loadSerial(self, oid, serial):
174        tid_data = self._data.get(oid)
175        if tid_data:
176            try:
177                return tid_data[serial]
178            except KeyError:
179                pass
180
181        raise ZODB.POSException.POSKeyError(oid, serial)
182
183    # ZODB.interfaces.IStorage
184    @ZODB.utils.locked(opened)
185    def new_oid(self):
186        self._oid += 1
187        return ZODB.utils.p64(self._oid)
188
189    # ZODB.interfaces.IStorage
190    @ZODB.utils.locked(opened)
191    def pack(self, t, referencesf, gc=True):
192        if not self._data:
193            return
194
195        stop = ZODB.TimeStamp.TimeStamp(*time.gmtime(t)[:5]+(t%60,)).raw()
196        if self._last_pack is not None and self._last_pack >= stop:
197            if self._last_pack == stop:
198                return
199            raise ValueError("Already packed to a later time")
200
201        self._last_pack = stop
202        transactions = self._transactions
203
204        # Step 1, remove old non-current records
205        for oid, tid_data in self._data.items():
206            tids_to_remove = tid_data.keys(None, stop)
207            if tids_to_remove:
208                tids_to_remove.pop()    # Keep the last, if any
209
210                if tids_to_remove:
211                    for tid in tids_to_remove:
212                        del tid_data[tid]
213                        if transactions[tid].pack(oid):
214                            del transactions[tid]
215
216        if gc:
217            # Step 2, GC.  A simple sweep+copy
218            new_data = BTrees.OOBTree.OOBTree()
219            to_copy = set([ZODB.utils.z64])
220            while to_copy:
221                oid = to_copy.pop()
222                tid_data = self._data.pop(oid)
223                new_data[oid] = tid_data
224                for pickle in tid_data.values():
225                    for oid in referencesf(pickle):
226                        if oid in new_data:
227                            continue
228                        to_copy.add(oid)
229
230            # Remove left over data from transactions
231            for oid, tid_data in self._data.items():
232                for tid in tid_data:
233                    if transactions[tid].pack(oid):
234                        del transactions[tid]
235
236            self._data.clear()
237            self._data.update(new_data)
238
239    # ZODB.interfaces.IStorage
240    def registerDB(self, db):
241        pass
242
243    # ZODB.interfaces.IStorage
244    def sortKey(self):
245        return self.__name__
246
247    # ZODB.interfaces.IStorage
248    @ZODB.utils.locked(opened)
249    def store(self, oid, serial, data, version, transaction):
250        assert not version, "Versions are not supported"
251        if transaction is not self._transaction:
252            raise ZODB.POSException.StorageTransactionError(self, transaction)
253
254        old_tid = None
255        tid_data = self._data.get(oid)
256        if tid_data:
257            old_tid = tid_data.maxKey()
258            if serial != old_tid:
259                raise ZODB.POSException.ConflictError(
260                    oid=oid, serials=(old_tid, serial), data=data)
261
262        self._tdata[oid] = data
263
264    checkCurrentSerialInTransaction = (
265        ZODB.BaseStorage.checkCurrentSerialInTransaction)
266
267    # ZODB.interfaces.IStorage
268    @ZODB.utils.locked(opened)
269    def tpc_abort(self, transaction):
270        if transaction is not self._transaction:
271            return
272        self._transaction = None
273        self._commit_lock.release()
274
275    # ZODB.interfaces.IStorage
276    def tpc_begin(self, transaction, tid=None):
277        with self._lock:
278
279            ZODB.utils.check_precondition(self.opened)
280
281            # The tid argument exists to support testing.
282            if transaction is self._transaction:
283                raise ZODB.POSException.StorageTransactionError(
284                    "Duplicate tpc_begin calls for same transaction")
285
286        self._commit_lock.acquire()
287
288        with self._lock:
289            self._transaction = transaction
290            self._tdata = {}
291            if tid is None:
292                if self._transactions:
293                    old_tid = self._transactions.maxKey()
294                else:
295                    old_tid = None
296                tid = ZODB.utils.newTid(old_tid)
297            self._tid = tid
298
299    # ZODB.interfaces.IStorage
300    @ZODB.utils.locked(opened)
301    def tpc_finish(self, transaction, func = lambda tid: None):
302        if (transaction is not self._transaction):
303            raise ZODB.POSException.StorageTransactionError(
304                "tpc_finish called with wrong transaction")
305
306        tid = self._tid
307        func(tid)
308
309        tdata = self._tdata
310        for oid in tdata:
311            tid_data = self._data.get(oid)
312            if tid_data is None:
313                tid_data = BTrees.OOBTree.OOBucket()
314                self._data[oid] = tid_data
315            tid_data[tid] = tdata[oid]
316
317        self._ltid = tid
318        self._transactions[tid] = TransactionRecord(tid, transaction, tdata)
319        self._transaction = None
320        del self._tdata
321        self._commit_lock.release()
322        return tid
323
324    # ZEO.interfaces.IServeable
325    @ZODB.utils.locked(opened)
326    def tpc_transaction(self):
327        return self._transaction
328
329    # ZODB.interfaces.IStorage
330    def tpc_vote(self, transaction):
331        if transaction is not self._transaction:
332            raise ZODB.POSException.StorageTransactionError(
333                "tpc_vote called with wrong transaction")
334
335class TransactionRecord(object):
336
337    status = ' '
338
339    def __init__(self, tid, transaction, data):
340        self.tid = tid
341        self.user = transaction.user
342        self.description = transaction.description
343        extension = transaction.extension
344        self.extension = extension
345        self.data = data
346
347    _extension = property(lambda self: self.extension,
348                          lambda self, v: setattr(self, 'extension', v),
349                          )
350
351    def __iter__(self):
352        for oid, data in self.data.items():
353            yield DataRecord(oid, self.tid, data)
354
355    def pack(self, oid):
356        self.status = 'p'
357        del self.data[oid]
358        return not self.data
359
360@zope.interface.implementer(ZODB.interfaces.IStorageRecordInformation)
361class DataRecord(object):
362    """Abstract base class for iterator protocol"""
363
364
365    version = ''
366    data_txn = None
367
368    def __init__(self, oid, tid, data):
369        self.oid = oid
370        self.tid = tid
371        self.data = data
372
373def DB(*args, **kw):
374    return ZODB.DB(MappingStorage(), *args, **kw)
375