straw/storage/SQLiteStorage.py

   1 from pysqlite2 import dbapi2 as sqlite
   2 import threading
   3
   4 DATABASE_FILE_NAME = "data.db"
   5
   6 class SQLiteStorage:
   7     """
   8     Provides an entry point to the database.
   9
  10     """
  11
  12     def __init__(self, db_path):
  13         self._connections = {}
  14         self._txs = {}
  15         self._db_path = db_path
  16         self._init_db()
  17
  18     def _get_sql(self):
  19         f = open('data/sql/create_01.sql', 'r')
  20         sql = f.read()
  21         f.close()
  22         return sql.split("--")
  23
  24     def _init_db(self):
  25         do_init = False
  26
  27         try:
  28             # Check if we need init.
  29             # TODO: use smarter strategy maybe?
  30             result = self.query("SELECT COUNT(*) FROM feeds")
  31         except Exception:
  32             # DB seems not initialized.
  33             do_init = True
  34
  35         if not do_init:
  36             return
  37
  38         c = self._connect()
  39
  40         try:
  41             sql = self._get_sql()
  42             for statement in sql:
  43                 c.execute(statement)
  44             self._tx_commit()
  45         except Exception, e:
  46             print "DB init failed -- %s" % e
  47
  48     def _connect(self):
  49         key = threading.currentThread()
  50
  51         if not self._connections.has_key(key):
  52             self._connections[key] = sqlite.connect(self._db_path)
  53             self._connections[key].execute("PRAGMA cache_size = 20000;")
  54             self._connections[key].execute("PRAGMA synchronous = NORMAL;")
  55             self._connections[key].row_factory = sqlite.Row
  56
  57         return self._connections[key]
  58
  59     def _tx_begin(self):
  60         key = threading.currentThread()
  61
  62         if not self._txs.has_key(key):
  63             self._txs[key] = True
  64
  65     def _tx_commit(self):
  66         self._connect().commit()
  67
  68         key = threading.currentThread()
  69         self._txs[key] = False
  70
  71     def _in_tx(self):
  72         key = threading.currentThread()
  73         return self._txs.has_key(key) and self._txs[key]
  74
  75     def query(self, query, params = None):
  76         if params == None:
  77             params = ()
  78         cursor = self._connect()
  79         #print query
  80         res = cursor.execute(query, params)
  81         return res.fetchall()
  82
  83     def insert(self, table, data):
  84         """
  85         Inserts some data into the database.
  86
  87         @param table: a name of the table to insert to
  88         @param data: a dictionary where keys are field names and values
  89             are field values in the given table
  90
  91         """
  92
  93         if len(data) == 0:
  94             return None
  95
  96         cursor = self._connect().cursor()
  97         query = "INSERT INTO %s (%s) VALUES (?%s)" % (table, ", ".join(data.keys()),
  98                                                       ", ?" * (len(data.keys()) - 1))
  99         #print query
 100         #print data
 101         cursor.execute(query, data.values())
 102
 103         #print self._in_tx()
 104         if not self._in_tx():
 105             self._tx_commit()
 106
 107         return cursor.lastrowid
 108
 109     def update(self, table, id, data):
 110         """
 111         Updates single row identified by a primary key in the table.
 112
 113         @param table: a name of the table to do the update in
 114         @param id: primary key of the item to update
 115         @param data: a dictionary where keys are field names and values
 116             are field values in the given table
 117
 118         """
 119
 120         params = data.values()
 121         cursor = self._connect().cursor()
 122         assignments = ", ".join([("%s = ?" % field_name) for field_name in data.keys()])
 123         query = "UPDATE %s SET %s WHERE id = ?" % (table, assignments)
 124         #print query
 125         params.append(id)
 126         cursor.execute(query, params)
 127
 128         if not self._in_tx():
 129             self._tx_commit()
 130
 131         return cursor.lastrowid