borg/dedupestore/store.py

import os
import tempfile
import shutil
import unittest
import sqlite3
import fcntl

Binary = sqlite3.Binary


class Store(object):
    """
    """

    class DoesNotExist(KeyError):
        """"""

    class AlreadyExists(KeyError):
        """"""

    IDLE = 'Idle'
    OPEN = 'Open'
    ACTIVE = 'Active'
    BAND_LIMIT = 1 * 1024 * 1024

    def __init__(self, path, create=False):
        self.read_fd = None
        self.write_fd = None
        if create:
            self.create(path)
        self.open(path)

    def get_option(self, key):
        return self.cursor.execute('SELECT value FROM system WHERE key=?', (key,)) \
            .fetchone()[0]

    def set_option(self, key, value):
        return self.cursor.execute('UPDATE system SET value=? WHERE key=?',
                                   (value, key))

    def open(self, path):
        if not os.path.isdir(path):
            raise Exception('%s Does not look like a store' % path)
        db_path = os.path.join(path, 'dedupestore.db')
        if not os.path.exists(db_path):
            raise Exception('%s Does not look like a store2')
        self.lock_fd = open(os.path.join(path, 'lock'), 'w')
        fcntl.flock(self.lock_fd, fcntl.LOCK_EX)
        self.path = path
        self.cnx = sqlite3.connect(db_path)
        self.cnx.text_factory = str
        self.cursor = self.cnx.cursor()
        self._begin()

    def _begin(self):
        if self.read_fd:
            self.read_fd.close()
            self.read_fd = None
        if self.write_fd:
            self.write_fd.close()
            self.write_fd = None
        self.version = self.get_option('version')
        self.id = self.get_option('id').decode('hex')
        self.tid = self.get_option('tid')
        self.nextband = self.get_option('nextband')
        self.bandlimit = self.get_option('bandlimit')
        assert self.version == 1
        self.state = self.OPEN
        self.read_band = None
        self.write_band = None
        self.to_delete = set()
        band = self.nextband
        while os.path.exists(self.band_filename(band)):
            os.unlink(self.band_filename(band))
            band += 1

    def create(self, path):
        if os.path.exists(path) and (not os.path.isdir(path) or os.listdir(path)):
            raise Exception('Path "%s" already exists' % path)
        if not os.path.exists(path):
            os.mkdir(path)
        os.mkdir(os.path.join(path, 'bands'))
        cnx = sqlite3.connect(os.path.join(path, 'dedupestore.db'))
        cnx.execute('CREATE TABLE objects(ns BINARY NOT NULL, id BINARY NOT NULL, '
                    'band NOT NULL, offset NOT NULL, size NOT NULL)')
        cnx.execute('CREATE UNIQUE INDEX objects_pk ON objects(ns, id)')
        cnx.execute('CREATE TABLE system(key UNIQUE NOT NULL, value)')
        cnx.executemany('INSERT INTO system VALUES(?, ?)',
                        (('id', os.urandom(32).encode('hex')),
                         ('version', 1),
                         ('tid', 0),
                         ('nextband', 0),
                         ('bandlimit', self.BAND_LIMIT)))
        cnx.commit()

    def close(self):
        self.rollback()
        self.cnx.close()
        self.lock_fd.close()
        os.unlink(os.path.join(self.path, 'lock'))

    def commit(self):
        """
        """
        self.band = None
        self.delete_bands()
        self.tid += 1
        self._begin()

    def delete_bands(self):
        for b in self.to_delete:
            objects = self.cursor.execute('SELECT ns, id, offset, size '
                                          'FROM objects WHERE band=? ORDER BY offset',
                                          (b,)).fetchall()
            for o in objects:
                band, offset, size = self.store_data(self.retrieve_data(b, *o[2:]))
                self.cursor.execute('UPDATE objects SET band=?, offset=?, size=? '
                                    'WHERE ns=? AND id=?', (band, offset, size,
                                    Binary(o[0]), Binary(o[1])))
        self.set_option('tid', self.tid + 1)
        self.set_option('nextband', self.nextband)
        self.cnx.commit()
        for b in self.to_delete:
            os.unlink(self.band_filename(b))

    def rollback(self):
        """
        """
        self.cnx.rollback()
        self._begin()

    def get(self, ns, id):
        """
        """
        self.cursor.execute('SELECT band, offset, size FROM objects WHERE ns=? and id=?',
                            (Binary(ns), Binary(id)))
        row = self.cursor.fetchone()
        if row:
            return self.retrieve_data(*row)
        else:
            raise self.DoesNotExist

    def band_filename(self, band):
        return os.path.join(self.path, 'bands', str(band / 1000), str(band))

    def retrieve_data(self, band, offset, size):
        if self.write_band == band:
            self.write_fd.flush()
        if self.read_band != band:
            self.read_band = band
            if self.read_fd:
                self.read_fd.close()
            self.read_fd = open(self.band_filename(band), 'rb')
        self.read_fd.seek(offset)
        return self.read_fd.read(size)

    def store_data(self, data):
        if self.write_band is None:
            self.write_band = self.nextband
            self.nextband += 1
            if self.write_band % 1000 == 0:
                path = os.path.join(self.path, 'bands', str(self.write_band / 1000))
                if not os.path.exists(path):
                    os.mkdir(path)
            assert not os.path.exists(self.band_filename(self.write_band))
            self.write_fd = open(self.band_filename(self.write_band), 'ab')
        band = self.write_band
        offset = self.write_fd.tell()
        self.write_fd.write(data)
        if offset + len(data) > self.bandlimit:
            self.write_band = None
        return band, offset, len(data)

    def put(self, ns, id, data):
        """
        """
        try:
            band, offset, size = self.store_data(data)
            self.cursor.execute('INSERT INTO objects (ns, id, band, offset, size) '
                                'VALUES(?, ?, ?, ?, ?)',
                                (Binary(ns), Binary(id), band, offset, size))
        except sqlite3.IntegrityError:
            raise self.AlreadyExists

    def delete(self, ns, id):
        """
        """
        self.cursor.execute('SELECT band FROM objects WHERE ns=? and id=?',
                            (Binary(ns), Binary(id)))
        row = self.cursor.fetchone()
        if not row:
            raise self.DoesNotExist
        self.cursor.execute('DELETE FROM objects WHERE ns=? AND id=?',
                           (Binary(ns), Binary(id)))
        self.to_delete.add(row[0])

    def list(self, ns, prefix='', marker=None, max_keys=1000000):
        """
        """
        sql = 'SELECT id FROM objects WHERE ns=:ns'
        args = dict(ns=Binary(ns))
        if prefix:
            args['prefix'] = Binary(prefix)
            args['end'] = Binary(prefix + chr(255))
            sql += ' AND id >= :prefix AND id < :end'
        if marker:
            sql += ' AND id >= :marker'
            args['marker'] = Binary(marker)
        for row in self.cursor.execute(sql + ' LIMIT ' + str(max_keys), args):
            yield str(row[0])


class StoreTestCase(unittest.TestCase):

    def setUp(self):
        self.tmppath = tempfile.mkdtemp()
        self.store = Store(os.path.join(self.tmppath, 'store'), create=True)

    def tearDown(self):
        shutil.rmtree(self.tmppath)

    def test1(self):
        self.assertEqual(self.store.tid, 0)
        self.assertEqual(self.store.state, self.store.OPEN)
        self.store.put('SOMENS', 'SOMEID', 'SOMEDATA')
        self.assertRaises(self.store.AlreadyExists, lambda: self.store.put('SOMENS', 'SOMEID', 'SOMEDATA'))
        self.assertEqual(self.store.get('SOMENS', 'SOMEID'), 'SOMEDATA')
        self.store.rollback()
        self.assertRaises(self.store.DoesNotExist, lambda: self.store.get('SOMENS', 'SOMEID'))
        self.assertEqual(self.store.tid, 0)

    def test2(self):
        self.assertEqual(self.store.tid, 0)
        self.assertEqual(self.store.state, self.store.OPEN)
        self.store.put('SOMENS', 'SOMEID', 'SOMEDATA')
        self.assertEqual(self.store.get('SOMENS', 'SOMEID'), 'SOMEDATA')
        self.store.commit()
        self.assertEqual(self.store.tid, 1)
        self.assertEqual(self.store.get('SOMENS', 'SOMEID'), 'SOMEDATA')
        self.store.delete('SOMENS', 'SOMEID')
        self.assertRaises(self.store.DoesNotExist, lambda: self.store.get('SOMENS', 'SOMEID'))
        self.store.rollback()
        self.assertEqual(self.store.get('SOMENS', 'SOMEID'), 'SOMEDATA')
        self.store.delete('SOMENS', 'SOMEID')
        self.assertRaises(self.store.DoesNotExist, lambda: self.store.get('SOMENS', 'SOMEID'))
        self.store.commit()
        self.assertEqual(self.store.tid, 2)
        self.assertRaises(self.store.DoesNotExist, lambda: self.store.get('SOMENS', 'SOMEID'))

    def test_list(self):
        self.store.put('SOMENS', 'SOMEID12', 'SOMEDATA')
        self.store.put('SOMENS', 'SOMEID', 'SOMEDATA')
        self.store.put('SOMENS', 'SOMEID1', 'SOMEDATA')
        self.store.put('SOMENS', 'SOMEID123', 'SOMEDATA')
        self.store.commit()
        self.assertEqual(list(self.store.list('SOMENS', max_keys=3)),
            ['SOMEID', 'SOMEID1', 'SOMEID12'])
        self.assertEqual(list(self.store.list('SOMENS', marker='SOMEID12')),
            ['SOMEID12', 'SOMEID123'])
        self.assertEqual(list(self.store.list('SOMENS', prefix='SOMEID1', max_keys=2)),
            ['SOMEID1', 'SOMEID12'])
        self.assertEqual(list(self.store.list('SOMENS', prefix='SOMEID1', marker='SOMEID12')),
            ['SOMEID12', 'SOMEID123'])


def suite():
    return unittest.TestLoader().loadTestsFromTestCase(StoreTestCase)

if __name__ == '__main__':
    unittest.main()