#Author: James Henstridge <james@daa.com.au>
"""This module implements a `shelf' like class, which stores the information
in a MySQL table.  I am not sure how fast this is in comparison to the dbm
shelves (I'm sure this can be optimised).  It was mainly written as an example
of how to use the new mysqldb interface.  Of course it may actually be useful
to someone.

To use it, do some initialisation like this:
  >>> import mysqlshelve
  >>> s = mysqlshelve.MysqlShelf('shelf', 'test')

This will set up a shelf that stores its data in table shelf of database test
on the local machine.  The second argument is passed directly to mysqldb, and
has the form 'db@host user pass', 'db@host user', 'db@host', 'db',
'db user pass' or 'db user'.

After this assignment, s can be used as a dictionary, provided you only use
strings for keys.

If the table doesn't exist, it is created.  If it exists, it is expected to
have a varchar(127) column named `name' and a longblob column named `value'.
Created tables set these columns to be not null and default to ''. It also
sets name to be a primary key.
"""

import Mysqldb
import pickle, StringIO

error = 'mysqlshelve.error'

def _quote(s):
	quote = {'\0': '\\0', '\n': '\\n', '\t': '\\t', '\b': '\\b',
		"'": "\\'", '"': '\\"', '\\': '\\\\', '%': '\\%',
		'_': '\\_'}
	r = range(len(s))
	r.reverse()       # iterate backwards, so as not to destroy indexing
	for i in r:
		if quote.has_key(s[i]):
			s = s[:i] + quote[s[i]] + s[i+1:]
	return s

class MysqlShelf:
	def __init__(self, table, connect_string):
		self.conn = Mysqldb.mysqldb(connect_string)
		curs = self.conn.cursor()
		self.table = table
		curs.execute("show tables like '%s'" % (_quote(table),))
		if len(curs.fetchall()) == 0:
			if not curs.execute("create table %s (" % (table,) +
			"name varchar(127) not null default '' primary key," +
			"value longblob not null default '')"):
				raise error, "Cannot create table"
		curs.close()

	def keys(self):
		self.conn.execute("select name from %s" % (self.table,))
		return map(lambda (x,): x, self.conn.fetchall())

	def __len__(self):
		self.conn.execute("select name from %s" % (self.table,))
		return len(self.conn.fetchall())
	
	def has_key(self, key):
		self.conn.execute("select name from %s where name = '%s'" %
			(self.table, _quote(key)))
		return self.conn.fetchone() != None

	def __getitem__(self, key):
		self.conn.execute("select value from %s where name = '%s'" %
			(self.table, _quote(key)))
		val = self.conn.fetchone()
		if val == None: raise KeyError, key
		f = StringIO.StringIO(val[0])
		return pickle.Unpickler(f).load()
	
	def __setitem__(self, key, value):
		f = StringIO.StringIO()
		p = pickle.Pickler(f)
		p.dump(value)
		if not self.conn.execute(   # check how many records updated
			"update %s set value='%s' where name='%s'" %
			(self.table, _quote(f.getvalue()), _quote(key))):
			self.conn.execute(
				"insert into %s (name,value) values ('%s','%s')"
				% (self.table, _quote(key),
				_quote(f.getvalue())))

	def __delitem__(self, key):
		if not self.conn.execute("delete from %s where name='%s'" %
				(self.table, _quote(key))):
			raise KeyError, key

	def close(self):
		self.conn.close()
		del self.conn