#! python
# -*- coding: utf-8 -*-


"""This is a module for using a SQLlite database as a collection of python objects.

For each SQL table, you derive a class from L{DBx}.   Each instance of that class corresponds to
one row of the table.
"""

import re
import sqlite3
import weakref
connect = sqlite3.connect

import die

DEBUG = 0

class DBMetaClass(type):
	"""This is used by all the SQL classes, and initializes
	the class.   Specifically, it makes sure that each class has
	a C{_cache} dictionary to cache instances of the class, and
	a C{idx} counter to be the integer index into the SQL database.
	"""
	
	def __new__(meta, classname, bases, cdict):
		cdict['idx'] = 0
		cdict['_cache'] = weakref.WeakValueDictionary()
		if 'SQL_name' not in cdict:
			cdict['SQL_name'] = classname
		# print 'Meta col types', cdict.get('SQL_name', '?'), "col-types=", cdict.get('COL_types', '?')
		if 'COL_types' in cdict:
			colpat = ','.join( ['?'] * (1+len(cdict['COL_types'])) )
			cdict['_ipat'] = "insert into %s values (%s)" % (cdict['SQL_name'], colpat)
			setpat = ','.join(["%s = ?"%c for (c,t) in cdict['COL_types']])
			cdict['_upat'] = "update %s set %s where id = ?" % (cdict['SQL_name'], setpat)
		return type.__new__(meta, classname, bases, cdict)



class SQLError(Exception):
	"""Any errors raised by this module.     This does not include errors raised by sqlite itself."""
	def __init__(self, *s):
		Exception.__init__(self, *s)

class NoSuchTable(SQLError):
	def __init__(self, *s):
		SQLError.__init__(self, *s)

class ColumnMismatchError(SQLError):
	def __init__(self, *s):
		SQLError.__init__(self, *s)

	


class DB(object):
	"""Base class for all persistent objects in a sqlite database.   Not the normal API.
	Derived classes need to define a L{write}() method.
	"""
	__metaclass__ = DBMetaClass

	S_notinDB = 0
	S_inDB = 1
	S_modified = 2

	def __init__(self, connection):
		"""
		@type connection: sqlite3.Connection
		"""
		try:
			ct = self.COL_types
		except AttributeError:
			raise AttributeError, "Your class must define COL_types"
		if  len(ct)<=0 or len(ct[0])!=2 or not isinstance(ct[0][0], str) or not isinstance(ct[0][1], str):
			raise TypeError, "Your class must define COL_types as a tuple of (column_name, data_type_name) pairs."
		assert connection is not None and isinstance(connection, sqlite3.Connection)
		self.w = connection
		self.id = self.bump_id(1)
		self.state = self.S_notinDB
	
	@classmethod
	def bump_id(cls, delta):
		"""You can call this manually if you want to introduce a gap between ID numbers."""
		cls.idx += delta
		return cls.idx


	def _write(self, *stuff):
		"""Called by C{write} in the derived class.  This actually does the writing.
		This function gets called to insert new rows and also to update old rows.
		@param stuff: whatever should get written into the database table (neglecting "id").
		@type stuff: a L{tuple} of anything that L{sqlite3} can interpret.
		@rtype: int
		@return: the unique id number for the row that was just written.
		"""
		assert self.SQL_name is not None
		assert len(stuff) == len(self.COL_types), "len(stuff)=%d col_types=%s" % (len(stuff), self.COL_types)
		# die.info("Inserting: %s (%s) (%s)" % (self.SQL_name, col_pat, unicode((self.i,)+stuff)))
		if self.state == self.S_notinDB:
			if DEBUG > 1:
				die.info("sqlbase: %s %s" % (self._ipat, unicode((self.id,)+stuff)))
			self.w.cursor().execute(self._ipat, (self.id,) + stuff)
		else:
			if DEBUG > 1:
				die.info("sqlbase: %s %s" % (self._upat, unicode((self.id,)+stuff)))
			self.w.cursor().execute(self._upat, stuff + (self.id,))
		self.state = self.S_inDB
		self._cache[self.id] = self
		return self.id
	
	def write(self):
		"""This method writes an instance out to its database.
		This method should assemble its attributes to form an argument list
		for L{_write}, and then call L{_write} with that argument list.  It should be
		used like this::

			def write(self):
				return self._write(self.something, self.other, self.stuff, self.misc)

		@note: C{write()} forces an immediate write or an update of the database, whether or not you have
			called L{make_dirty}().   L{make_dirty}() simply arranges for a write to occur
			when the object is eventually destroyed.
		@return: the return value of L{_write}.
		@rtype: int
		@note: The instance remembers the database from which it came.
		@note: Do not pass C{self.id} to C{self.}L{_write()}.  That's handled automatically.
		"""
		raise RuntimeError, "Virtual method"
	
	def make_dirty(self):
		"""Changes to an instance are not written to the database unless you call
		L{make_dirty}() or L{write}().   C{write()} forces an immediate write,
		whereas C{make_dirty}() simply marks the instance so that it will eventually
		be written out when the destructor is called.
		@note: Global variables are destroyed very late, possibly after other necessary objects
			have already been destroyed, so you may not want to
			trust that C{make_Dirty()} will behave well on global variables.
		"""
		if self.state == self.S_inDB:
			self.state = self.S_modified

	def is_dirty(self):
		return self.state != self.S_inDB

	def __del__(self):
		if self.state != self.S_inDB:
			self.write()

	@classmethod
	def confirm(cls, connection):
		"""This can be called to confirm that the column names in the database matches what the class expects.
		@raise ColumnMismatchError: When something is wrong.
		@raise NoSuchTable: when the table doesn't exist.
		@return: None
		@rtype: None
		@param connection: A connection to a database.
		@type connection: L{sqlite3.Connection}
		"""
		assert isinstance(connection, sqlite3.Connection)
		dbc = connection.cursor()
		try:
			dbc.execute("select * from %s where id=0" % cls.SQL_name)
		except sqlite3.OperationalError, ex:
			if 'no such table' in str(ex):
				raise NoSuchTable(str(ex))
			raise
		ct = (('id', 'integer'),) + tuple(cls.COL_types)
		if len(dbc.description) != len(ct):
			raise ColumnMismatchError("Wrong number of columns: saw %d expected %d" % (len(dbc.description), len(ct)), [x[0] for x in dbc.description], ct )
		for (i,(x, ct)) in enumerate(zip(sorted(dbc.description), sorted(ct))):
			if x[0] != ct[0]:
				raise ColumnMismatchError( "Mismatched name column %d: saw %s expected %s" % (i, x[0], ct[0]), [x[0] for x in dbc.description] )

	@classmethod
	def create(cls, connection, if_not_exists=True):
		"""This creates the table in the database.   It is harmless if the table already exists.
		@return: None
		@rtype: None
		@param connection: A connection to a database.
		@type connection: L{sqlite3.Connection}
		"""
		assert cls.SQL_name is not None
		ct = [ ('id', 'INTEGER PRIMARY KEY ASC') ]
		cc = []
		for (nm,ty) in cls.COL_types:
			if ty.lower().startswith('fkey'):
				fk, table, col = ty.split()
				ct.append( (nm, 'INTEGER') )
				cc.append( ('FOREIGN KEY(%s)' % nm,
						'REFERENCES %s(%s)' % (table, col)) )
			else:
				ct.append( (nm, ty) )
		cstr = ','.join( [ "%s %s" % q for q in ct + cc] )
		ine = "if not exists" if if_not_exists else ""
		if DEBUG:
			die.info("sqlbase: create table %s %s (%s)" % (ine, cls.SQL_name, cstr))
		connection.execute("create table %s %s (%s)" % (ine, cls.SQL_name, cstr))


	@classmethod
	def set_ID(cls, connection):
		"""This should be called when you first open an existing database when
		you have the intent to write.  It sets the initial ID number
		to be larger than the largest ID in the database.
		"""
		tmp = connection.execute("select MAX(id) from %s" % cls.SQL_name)
		idx = tmp.fetchone()[0]
		if idx is not None:
			cls.idx = idx + 1
		
	@classmethod
	def index_unique(cls, connection, *columns):
		cls._index_guts(connection, "UNIQUE", columns)

	@classmethod
	def index_nonunique(cls, connection, *columns):
		cls._index_guts(connection, "", columns)

	@classmethod
	def _index_guts(cls, connection, u, columns):
		idx2 = ','.join(columns)
		idx1 = '_'.join([re.sub("_", "__", q) for q in columns])
		if DEBUG:
			die.info("sqlbase: %s INDEX if not exists %s_index_%s on %s (%s)"
					% (u, cls.SQL_name, idx1, cls.SQL_name, idx2)
					)
		connection.execute("create %s INDEX if not exists %s_index_%s on %s (%s)"
					% (u, cls.SQL_name, idx1, cls.SQL_name, idx2)
					)


	@classmethod
	def flush_cache(cls):
		cls._cache.clear()




class DBx(DB):
	"""This class is the main interface.   Normally, you derive a class from this to represent a table
	in the database.   Each of your derived classes MUST contain an attribute C{COL_type} which
	specifies what columns are in that table.
	Each derived class will have a unique ID number called C{id}; this will be a column in the database,
	and it will be stored in each instance object.  The L{DBx} class manages that ID number for you.

	Beyond that, a derived class MUST redefine L{_fromtuple}.
	
	@cvar COL_type: a sequence of C{(column_name, data_type_name)} pairs.   Those name the columns of the table and tell what
		type of data is in each.   C{data_type_name} must be the name of a legal sqlite3 data type OR,
		it must be C{"FKEY tablename column_name"} where I{tablename} is the name of the table into which the
		foreign key points and I{column_name} is the name of the column in that table.
		Typical tuples in C{COL_type} are C{("Number_of_dogs", "INTEGER")} or C{("dog_id", "FKEY dog_info id")}.
		Note that C{COL_type} should not include the ID number, which is always column C{id}; it is added automatically.

	@cvar SQL_name: the name of the corresponding table in the database.   This defaults to the name of the leaf class.
	@ivar id: the instance's ID number.
		Note that an ID number is allocated for every instance you create.  Instances you look up with C{select}
		won't necesarily be new: they bay just be a pointer to an existing, perhaps cached instance.
		The ID number of an instance is persistant and is written into the database.
	@cvar idx: the lowest unused ID number.
	"""
	def __init__(self, connection):
		"""@param connection: a database connection.
		@type connection: L{sqlite3.Connection}
		"""
		DB.__init__(self, connection)


	def get_id(self):
		if self.state == self.S_notinDB:
			self.write()
		return self.id

	@classmethod
	def select(cls, db, where, *args):
		assert isinstance(db, sqlite3.Connection)
		dbc = db.cursor()
		cln = cls.SQL_name
		if DEBUG > 2:
			die.info("sqlbase: select * from %s %s (%s)" % (cln, where, unicode(args)))
		dbc.execute("select * from %s %s" % (cln, where), args)
		for x in dbc.fetchall():
			idx = x[0]
			try:
				# The intent here is that there should be only one
				# in-memory copy of each object in the database.
				# Unfortunately, this means we will not see
				# externally-caused changes to the database,
				# but that seems better than having two objects in
				# memory where you might accidentally change one and
				# look at the other.
				yield cls._cache[idx]
			except KeyError:
				rv = cls._fromtuple_id(db, *x)
				cls._cache[idx] = rv
				yield rv



	@classmethod
	def get_by_id(cls, db, idx):
		assert isinstance(db, sqlite3.Connection)
		dbc = db.cursor()
		if idx is None:
			return None
		assert isinstance(idx, int)

		try:
			return cls._cache[idx]
		except KeyError:
			pass

		# die.info("select * from %s where id = %s" % (cls.SQL_name, idx))
		if DEBUG > 2:
			die.info("select * from %s where id = ? (%s)" % (cls.SQL_name, unicode(idx)))
		dbc.execute("select * from %s where id = ?" % cls.SQL_name, (idx,))
		x = dbc.fetchone()
		if x is None:
			raise KeyError, "No %s with id=%d" % (cls.SQL_name, idx)
		rv = cls._fromtuple_id(db, *x)
		cls._cache[idx] = rv
		return rv

	@classmethod
	def _fromtuple(cls, db, *rest):
		"""This method takes a database connection and a tuple of data from a row of the database and must
		return an instance of the class.   Normally, it checks its input arguments, converts them from
		one format to another, and then calls the constructor and returns its result.
		@param db: connection to a sqlite3 database.
		@type db: L{sqlite3.Connection}
		@param rest: a tuple derived by reading a row of data from the database.
			This tuple does NOT contain the row's unique ID number, which is the first column in the database.
		@type rest: tuple of stuff.
		@return: a constructed instance of the class.
		@rtype: an instance of a suitable class derived from L{DBx}.
		"""
		raise RuntimeError, "Virtual Method"

	@classmethod
	def _fromtuple_id(cls, db, idx, *rest):
		"""This wraps L{_fromtuple} and adds in the ID number information.
		"""
		rv = cls._fromtuple(db, *rest)
		rv.state = cls.S_inDB
		rv.id = idx
		# Somewhere inside the call to cls._fromtuple(), we called the constructor which
		# called bump_id().   But, this isn't a new entry: it's one that already exists in
		# the database, so we don't actually want idx incremented.   Thus, we bump it back down.
		cls.bump_id(-1)
		return rv

	def __repr__(self):
		contents = []
		for (nm, ty) in self.__class__.COL_types:
			try:
				v = getattr(self, nm)
			except AttributeError:
				continue
			if ty.startswith('FKEY'):
				fkey, ttype, rest = ty.split(None, 2)
				contents.append("%s=<%s %d>" % (nm, ttype, v.id))
			else:
				contents.append('%s=%s' % (nm, v))
		return "<%s %s>" % (self.__class__.SQL_name, ', '.join(contents))



def multiselect(cll, db, where, *args):
	"""This is used for SQL select operations on a join of several tables.
	@type cll: a sequence of C{(class, str)} where C{class} is derived from L{DBx}
	@param cll: This a sequence of C{(class, name)} pairs, where C{class} matches
		a table in the database, and C{name} is used to refer to that table in
		the SQL select statement.
	@type db: C{sqlite3.Connection}
	@type where: str
	@param where: This is the body of the select statement.   Everything after
		C{select * from tablename}.
	@param args: All the arguments for the C{select} statement.
	@note: As an example, if you have a table C{X} in the database and a class C{x},
		and C{x.SQL_name=="X"}, and each row in C{X} has an C{id} number, and
		also the ID number of the C{next} item,
		then
		you can find pairs of adjacent items with this call::
		
			multiselect(((x, "x1"), (x, "x2")), db, "x1.next = x2.id")
	"""
	assert isinstance(db, sqlite3.Connection)
	dbc = db.cursor()
	cn = []
	frt = []
	ioff = 0
	for c in cll:
		if isinstance(c, tuple):
			if len(c) != 2:
				raise TypeError, "cll is ((class, 'name_in_select'),...)"
			cn.append("%s %s" % (c[0].SQL_name, c[1]))
			c = c[0]
		else:
			cn.append(c.SQL_name)
		frt.append((c, ioff))
		ioff += len(c.COL_types)+1
	# print 'cn=', cn
	if DEBUG > 2:
		die.info("select * from %s %s (%s)" % (', '.join(cn), where, unicode(args)))
	try:
		dbc.execute("select * from %s %s" % (', '.join(cn), where), args)
	except sqlite3.OperationalError, ex:
		die.info("multiselect: from: %s" % ', '.join(cn))
		die.info("multiselect: where: %s" % where)
		die.info("multiselect: args: %s" % unicode(args))
		die.warn("multiselect: sqlite3.OperationalError: %s" % ex)
		raise

	for x in dbc.fetchall():
		# print "x=", x
		rv = []
		assert len(x) == ioff
		for (cls, i) in frt:
			# print "cls=", cls, "i=", i
			idx = x[i]
			try:
				tmp = cls._cache[idx]
				# print "cached=", tmp
			except KeyError:
				tmp = cls._fromtuple_id(db, *(x[i:i+len(cls.COL_types)+1]))
				cls._cache[idx] = tmp
				# print "created=", tmp
			rv.append(tmp)
		# print "rv=", rv
		yield tuple(rv)


def get_version(db_connection, list_of_classes):
	"""This finds which of several variants of a table actually exists in the database.
	The intent is to let you upgrade from one version of table format to another.
	@param db_connection: A L{sqlite3.Connection} to a database.
	@param list_of_classes: Possible classes to try.   Each will be checked for consistency with the database.
	@type list_of_classes: sequence of some class derived from sqlbase.L{DBx}.   These are the classes themselves,
		not class instances.
	@rtype: One of the classes in C{list_of_classes}.
	@return: The first class in the list that is consistent with the structure of the database.
	"""
	errs = []
	evtype = None
        for tmp in list_of_classes:
                try:
                        tmp.confirm(db_connection)
                        evtype = tmp
                        break
                except (ColumnMismatchError, NoSuchTable), ex:
			errs.append( "%s: %s" % (str(tmp), str(ex)) )
                        pass
        if evtype is None:
		raise NoSuchTable("None of the supplied tables match.", *errs)
	return evtype


def multiclass_get_by_id(possible_classes, db, idx):
	"""This reads in a row from the database when you don't know exactly what information is available
	The idea is that there is a one of several possible types of objects stored in the database, and
	you will be happy to take whichever one is available.   This makes most sense when they are
	different versions of the same class.   You access the objects by their ID number.
	@type possible_classes: a sequence of classes, where each class is derived from L{DBx}
	@param possible_classes: possible classes to match a table in the database.
		Classes must be derived from L{DBx}, must have a name that matches a table in the database,
		must have the right number of columns, and the raw data must be convertible to an object
		of that class.
	@type db: a L{sqlite3.Connection} to a database.
	@type idx: int
	@param idx: the C{id} of a row in the database.
	@note: Ideally, there should be a 1:1 relationship between the names of tables in the database and
		the C{SQL_name} attributes of classes.   So, when you have a new version of a class, you
		really should have a new name for the corresponding table in the database.  However, if
		you wish to have several classes that have the same C{SQL_name} attribute, a class can raise
		C{ColumnMismatchError} if the class and the data row are incompatible.
	"""
	errs = []
	for c in possible_classes:
		try:
			return c.get_by_id(db, idx)
		except sqlite3.OperationalError, ex:
			if 'no such table' not in str(ex):
				raise
			errs.append("%s: %s" % (str(c), str(ex)))
		except ColumnMismatchError, ex:
			errs.append("%s: %s" % (str(c), str(ex)))
	raise NoSuchTable("Cannot find a matching class", *errs)



class _SampleTable(DBx):
	"""Used for L{test}()."""
	SQL_name = 'SampleTable'
	COL_types = (('name', 'text'), ('some_value', 'integer'))

	def __init__(self, name, some_value, db=None):
		assert name is not None
		DBx.__init__(self, db)
		self.name = name
		self.some_value = int(some_value)
	
	def write(self):
		return self._write(self.name, self.some_value)
		
	@classmethod
	def _fromtuple(cls, db, name, some_value):
		return cls(name, some_value, db=db)


class _InconsistentTable(DBx):
	"""Used for L{test}()."""
	SQL_name = 'SampleTable'
	COL_types = (('name', 'text'), ('foo', 'text'))

	def __init__(self, name, some_value, db=None):
		assert name is not None
		DBx.__init__(self, db)
		self.name = name
		self.some_value = some_value
	
	def write(self):
		return self._write(self.name, self.some_value)
		
	@classmethod
	def _fromtuple(cls, db, name, some_value):
		return cls(name, some_value, db=db)


def test():
	global DEBUG
	# DEBUG=10
	db = sqlite3.connect(":memory:")
	_SampleTable.create(db)
	_SampleTable.confirm(db)
	_SampleTable.index_nonunique(db, 'name')
	ok = False
	try:
		_InconsistentTable.confirm(db)
	except ColumnMismatchError:
		ok = True
	assert ok
	_SampleTable.set_ID(db);
	assert _SampleTable.idx == 0, "Expected idx=0, got %d" % _SampleTable.idx
	# Create rows
	_SampleTable('my name', 3, db=db).write()
	assert _SampleTable.idx == 1, "Expected idx=1, got %d" % _SampleTable.idx
	_SampleTable('another name', 2, db=db).write()
	assert _SampleTable.idx == 2, "Expected idx=2, got %d" % _SampleTable.idx
	# Select a row:
	n = 0
	for x in _SampleTable.select(db, "where name = ?", "my name"):
		n += 1
		assert x.some_value == 3
	assert n == 1
	assert _SampleTable.idx == 2, "Expected idx=2, got %d" % _SampleTable.idx
	# Get an object and modify it:
	tmp = list(_SampleTable.select(db, "where name = ?", "my name"))[0]
	tmp.some_value = 7
	# Check that the value is correct:
	tmp2 = list(_SampleTable.select(db, "where name = ?", "my name"))[0]
	assert tmp2.some_value == 7
	# Force a write to the database:
	tmp.make_dirty()
	tmp.write()
	del tmp
	del tmp2
	# Check that the value is correct:
	tmp2 = list(_SampleTable.select(db, "where name = ?", "my name"))[0]
	assert tmp2.some_value == 7

if __name__ == '__main__':
	test()
