apt_dht_Khashmir/db.py

   1
   2 from datetime import datetime, timedelta
   3 from pysqlite2 import dbapi2 as sqlite
   4 import os
   5
   6 class DBExcept(Exception):
   7     pass
   8
   9 class DB:
  10     """Database access for storing persistent data."""
  11
  12     def __init__(self, db):
  13         self.db = db
  14         try:
  15             os.stat(db)
  16         except OSError:
  17             self._createNewDB(db)
  18         else:
  19             self._loadDB(db)
  20         self.conn.text_factory = str
  21
  22     def _loadDB(self, db):
  23         try:
  24             self.conn = sqlite.connect(database=db, detect_types=sqlite.PARSE_DECLTYPES)
  25         except:
  26             import traceback
  27             raise DBExcept, "Couldn't open DB", traceback.format_exc()
  28
  29     def _createNewDB(self, db):
  30         self.conn = sqlite.connect(database=db)
  31         c = self.conn.cursor()
  32         c.execute("CREATE TABLE kv (key TEXT, value TEXT, time TIMESTAMP, PRIMARY KEY (key, value))")
  33         c.execute("CREATE INDEX kv_key ON kv(key)")
  34         c.execute("CREATE INDEX kv_timestamp ON kv(time)")
  35         c.execute("CREATE TABLE nodes (id TEXT PRIMARY KEY, host TEXT, port NUMBER)")
  36         c.execute("CREATE TABLE self (num NUMBER PRIMARY KEY, id TEXT)")
  37         self.conn.commit()
  38
  39     def getSelfNode(self):
  40         c = self.conn.cursor()
  41         c.execute('SELECT id FROM self WHERE num = 0')
  42         if c.rowcount > 0:
  43             return c.fetchone()[0]
  44         else:
  45             return None
  46
  47     def saveSelfNode(self, id):
  48         c = self.conn.cursor()
  49         c.execute("INSERT OR REPLACE INTO self VALUES (0, ?)", (id,))
  50         self.conn.commit()
  51
  52     def dumpRoutingTable(self, buckets):
  53         """
  54             save routing table nodes to the database
  55         """
  56         c = self.conn.cursor()
  57         c.execute("DELETE FROM nodes WHERE id NOT NULL")
  58         for bucket in buckets:
  59             for node in bucket.l:
  60                 c.execute("INSERT INTO nodes VALUES (?, ?, ?)", (node.id, node.host, node.port))
  61         self.conn.commit()
  62
  63     def getRoutingTable(self):
  64         """
  65             load routing table nodes from database
  66             it's usually a good idea to call refreshTable(force=1) after loading the table
  67         """
  68         c = self.conn.cursor()
  69         c.execute("SELECT * FROM nodes")
  70         return c.fetchall()
  71
  72     def retrieveValues(self, key):
  73         c = self.conn.cursor()
  74         c.execute("SELECT value FROM kv WHERE key = ?", (key,))
  75         t = c.fetchone()
  76         l = []
  77         while t:
  78             l.append(t[0])
  79             t = c.fetchone()
  80         return l
  81
  82     def storeValue(self, key, value):
  83         """Store or update a key and value."""
  84         c = self.conn.cursor()
  85         c.execute("INSERT OR REPLACE INTO kv VALUES (?, ?, ?)", (key, value, datetime.now()))
  86         self.conn.commit()
  87
  88     def expireValues(self, expireAfter):
  89         """Expire older values after expireAfter seconds."""
  90         t = datetime.now() - timedelta(seconds=expireAfter)
  91         c = self.conn.cursor()
  92         c.execute("DELETE FROM kv WHERE time < ?", (t, ))
  93         self.conn.commit()
  94
  95     def close(self):
  96         self.conn.close()