straw/storage/SQLiteStorage.py

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