diff --git a/client/djrandom_client/upload.py b/client/djrandom_client/upload.py index cb3826a3ae8ed2dabc6000e200d555f0dd72d8a1..206f06e0ebeae71277d02b866b8b2f5340778b71 100644 --- a/client/djrandom_client/upload.py +++ b/client/djrandom_client/upload.py @@ -6,6 +6,7 @@ import os import optparse import platform import socket +import sqlite3 import threading import time import urllib2 @@ -19,20 +20,42 @@ log = logging.getLogger(__name__) class FileDatabase(object): + """SQLite-backed database of seen hashes.""" def __init__(self, dbpath=None): if not dbpath: dbpath = os.path.join(os.getenv('HOME'), '.djrandom') - self.db = anydbm.open(dbpath, 'c', 0600) + if not dbpath.endswith('.db'): + dbpath += '.db' + + is_new = not os.path.exist(dbpath) + self.conn = sqlite3.connect(dbpath) + self.conn.text_factory = str + + if is_new: + self.conn.execute( + 'create table seen (sha1 text(40), stamp integer, ' + 'primary key (sha1))') + self.conn.commit() def close(self): - self.db.close() + self.conn.close() def has(self, key): - return key in self.db + c = self.conn.cursor() + row = c.execute( + 'select stamp from seen where sha1 = ?', (key,) + ).fetchone() + if row: + return True + else: + return False def add(self, key): - self.db[key] = '%i' % time.time() + c = self.conn.cursor() + c.execute('insert or replace into seen (sha1, stamp) values (?, ?)', + (key, int(time.time()))) + c.commit() class Uploader(object):