mirror of https://github.com/borgbackup/borg.git
271 lines
9.7 KiB
Python
271 lines
9.7 KiB
Python
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 darc store' % path)
|
|
db_path = os.path.join(path, 'darcstore.db')
|
|
if not os.path.exists(db_path):
|
|
raise Exception('%s Does not look like a darc store')
|
|
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, 'darcstore.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()
|