pypsyc/mjacob2/pypsyc/server/db.py

106 lines
2.9 KiB
Python

"""
pypsyc.server.db
~~~~~~~~~~~~~~~~
:copyright: 2010 by Manuel Jacob
:license: MIT
"""
import sqlite3
from threading import Thread
from Queue import Queue
from greenlet import getcurrent
from twisted.internet import reactor
class Database(object):
"""
Open a sqlite database that queues the database operations and executes
them in a thread.
"""
def __init__(self, filename):
self.filename = filename
self.sync = (filename == ':memory:')
if self.sync:
self._connect()
else:
self.queue = Queue()
self.thread = Thread(target=self.run_async)
self.thread.start()
self.main_greenlet = getcurrent()
def execute(self, *args):
"""
Execute a query.
:param query: the sqlite query
:type query: string
:param \*params: additional query parameters
:type \*params: one or more strings
"""
if self.sync:
return self._execute(*args)
self.queue.put(('execute', args, getcurrent()))
assert getcurrent() is not self.main_greenlet
return self.main_greenlet.switch()
def fetch(self, *args):
"""
Fetch a result from the database asynchronously.
:param query: the sqlite query
:type query: string
:param \*params: additional query parameters
:type \*params: one or more strings
:rtype: result rows
"""
if self.sync:
return self._fetch(*args)
self.queue.put(('fetch', args, getcurrent()))
assert getcurrent() is not self.main_greenlet
return self.main_greenlet.switch()
def stop(self):
"""Close the database and stop the worker thread."""
if self.sync:
return self._close()
del self.sync # make database unusable
self.queue.put(('stop', (), None))
self.thread.join()
def run_async(self):
self._connect()
while 1:
cmd, args, gl = self.queue.get()
if cmd == 'stop':
self._close()
break
try:
if cmd == 'execute':
ret = self._execute(*args)
elif cmd == 'fetch':
ret = self._fetch(*args)
except Exception, e:
reactor.callFromThread(gl.throw, e)
else:
reactor.callFromThread(gl.switch, ret)
def _connect(self):
self.conn = sqlite3.connect(self.filename)
# return bytestrings instead of unicode strings
self.conn.text_factory = str
def _close(self):
self.conn.close()
def _execute(self, query, *params):
cur = self.conn.cursor()
cur.execute(query, params)
self.conn.commit()
def _fetch(self, query, *params):
cur = self.conn.cursor()
cur.execute(query, params)
return cur.fetchall()