straw/storage/SQLiteStorage.py

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