apt_dht_Khashmir/db.py

   1
   2 from datetime import datetime, timedelta
   3 from pysqlite2 import dbapi2 as sqlite
   4 from binascii import a2b_base64, b2a_base64
   5 from time import sleep
   6 import os
   7
   8 from twisted.trial import unittest
   9
  10 class DBExcept(Exception):
  11     pass
  12
  13 class khash(str):
  14     """Dummy class to convert all hashes to base64 for storing in the DB."""
  15
  16 class dht_value(str):
  17     """Dummy class to convert all DHT values to base64 for storing in the DB."""
  18
  19 sqlite.register_adapter(khash, b2a_base64)
  20 sqlite.register_converter("KHASH", a2b_base64)
  21 sqlite.register_converter("khash", a2b_base64)
  22 sqlite.register_adapter(dht_value, b2a_base64)
  23 sqlite.register_converter("DHT_VALUE", a2b_base64)
  24 sqlite.register_converter("dht_value", a2b_base64)
  25
  26 class DB:
  27     """Database access for storing persistent data."""
  28
  29     def __init__(self, db):
  30         self.db = db
  31         try:
  32             os.stat(db)
  33         except OSError:
  34             self._createNewDB(db)
  35         else:
  36             self._loadDB(db)
  37         if sqlite.version_info < (2, 1):
  38             sqlite.register_converter("TEXT", str)
  39             sqlite.register_converter("text", str)
  40         else:
  41             self.conn.text_factory = str
  42
  43     def _loadDB(self, db):
  44         try:
  45             self.conn = sqlite.connect(database=db, detect_types=sqlite.PARSE_DECLTYPES)
  46         except:
  47             import traceback
  48             raise DBExcept, "Couldn't open DB", traceback.format_exc()
  49
  50     def _createNewDB(self, db):
  51         self.conn = sqlite.connect(database=db, detect_types=sqlite.PARSE_DECLTYPES)
  52         c = self.conn.cursor()
  53         c.execute("CREATE TABLE kv (key KHASH, value DHT_VALUE, last_refresh TIMESTAMP, PRIMARY KEY (key, value))")
  54         c.execute("CREATE INDEX kv_key ON kv(key)")
  55         c.execute("CREATE INDEX kv_last_refresh ON kv(last_refresh)")
  56         c.execute("CREATE TABLE nodes (id KHASH PRIMARY KEY, host TEXT, port NUMBER)")
  57         c.execute("CREATE TABLE self (num NUMBER PRIMARY KEY, id KHASH)")
  58         self.conn.commit()
  59
  60     def getSelfNode(self):
  61         c = self.conn.cursor()
  62         c.execute('SELECT id FROM self WHERE num = 0')
  63         id = c.fetchone()
  64         if id:
  65             return id[0]
  66         else:
  67             return None
  68
  69     def saveSelfNode(self, id):
  70         c = self.conn.cursor()
  71         c.execute("INSERT OR REPLACE INTO self VALUES (0, ?)", (khash(id),))
  72         self.conn.commit()
  73
  74     def dumpRoutingTable(self, buckets):
  75         """
  76             save routing table nodes to the database
  77         """
  78         c = self.conn.cursor()
  79         c.execute("DELETE FROM nodes WHERE id NOT NULL")
  80         for bucket in buckets:
  81             for node in bucket.l:
  82                 c.execute("INSERT INTO nodes VALUES (?, ?, ?)", (khash(node.id), node.host, node.port))
  83         self.conn.commit()
  84
  85     def getRoutingTable(self):
  86         """
  87             load routing table nodes from database
  88             it's usually a good idea to call refreshTable(force=1) after loading the table
  89         """
  90         c = self.conn.cursor()
  91         c.execute("SELECT * FROM nodes")
  92         return c.fetchall()
  93
  94     def retrieveValues(self, key):
  95         """Retrieve values from the database."""
  96         c = self.conn.cursor()
  97         c.execute("SELECT value FROM kv WHERE key = ?", (khash(key),))
  98         l = []
  99         rows = c.fetchall()
 100         for row in rows:
 101             l.append(row[0])
 102         return l
 103
 104     def storeValue(self, key, value):
 105         """Store or update a key and value."""
 106         c = self.conn.cursor()
 107         c.execute("INSERT OR REPLACE INTO kv VALUES (?, ?, ?)",
 108                   (khash(key), dht_value(value), datetime.now()))
 109         self.conn.commit()
 110
 111     def expireValues(self, expireAfter):
 112         """Expire older values after expireAfter seconds."""
 113         t = datetime.now() - timedelta(seconds=expireAfter)
 114         c = self.conn.cursor()
 115         c.execute("DELETE FROM kv WHERE last_refresh < ?", (t, ))
 116         self.conn.commit()
 117
 118     def close(self):
 119         self.conn.close()
 120
 121 class TestDB(unittest.TestCase):
 122     """Tests for the khashmir database."""
 123
 124     timeout = 5
 125     db = '/tmp/khashmir.db'
 126     key = '\xca\xec\xb8\x0c\x00\xe7\x07\xf8~])\x8f\x9d\xe5_B\xff\x1a\xc4!'
 127
 128     def setUp(self):
 129         self.store = DB(self.db)
 130
 131     def test_selfNode(self):
 132         self.store.saveSelfNode(self.key)
 133         self.failUnlessEqual(self.store.getSelfNode(), self.key)
 134
 135     def test_Value(self):
 136         self.store.storeValue(self.key, self.key)
 137         val = self.store.retrieveValues(self.key)
 138         self.failUnlessEqual(len(val), 1)
 139         self.failUnlessEqual(val[0], self.key)
 140
 141     def test_expireValues(self):
 142         self.store.storeValue(self.key, self.key)
 143         sleep(2)
 144         self.store.storeValue(self.key, self.key+self.key)
 145         self.store.expireValues(1)
 146         val = self.store.retrieveValues(self.key)
 147         self.failUnlessEqual(len(val), 1)
 148         self.failUnlessEqual(val[0], self.key+self.key)
 149
 150     def test_RoutingTable(self):
 151         class dummy:
 152             id = self.key
 153             host = "127.0.0.1"
 154             port = 9977
 155             def contents(self):
 156                 return (self.id, self.host, self.port)
 157         dummy2 = dummy()
 158         dummy2.id = '\xaa\xbb\xcc\x0c\x00\xe7\x07\xf8~])\x8f\x9d\xe5_B\xff\x1a\xc4!'
 159         dummy2.host = '205.23.67.124'
 160         dummy2.port = 12345
 161         class bl:
 162             def __init__(self):
 163                 self.l = []
 164         bl1 = bl()
 165         bl1.l.append(dummy())
 166         bl2 = bl()
 167         bl2.l.append(dummy2)
 168         buckets = [bl1, bl2]
 169         self.store.dumpRoutingTable(buckets)
 170         rt = self.store.getRoutingTable()
 171         self.failUnlessIn(dummy().contents(), rt)
 172         self.failUnlessIn(dummy2.contents(), rt)
 173
 174     def tearDown(self):
 175         self.store.close()
 176         os.unlink(self.db)