# Bender Lab Animal Colony Database Management System - Python Database Interaction Functions

import sqlite3 as lite
import sys
import debug

class Table_cur(object):
	"""Object representing a table, with following methods:

	--searching through tables (get)
	--adding (add_row) or updating (update_row) rows for the table

	--NOTE-- : search_param only handles equality tests for now
	"""
	def __init__(self, db, table):
		"""Tell it which database file to open
		and which table in that file to work with.
		"""

		self.db = db
		self.table = table
		self.col_list = []

		if debug.level > 0:
			print "\nIN Table_cur.__init__:\n\tself.db =", self.db
			print "\tself.table =", self.table
			print "\tself.col_list =", self.col_list

		# make self.col_list an ordered list of columns in self.table
		con = lite.connect(self.db)
		with con:
			cur = con.cursor()
			if debug.level > 0:
				print '\nexecuting.... PRAGMA table_info('+self.table+')'
			cur.execute('PRAGMA table_info('+self.table+')')
			data = cur.fetchall()
			print 'data =', data
			for d in data:
				self.col_list.append(d[1])
		con.close()
		if debug.level > 0:
			print "self.col_list =", self.col_list

	def get(self, table=None, columns="*", search_param=None):
		"""Retrieves table data based on search parameters.

		-accepts table, columns, and search_param dictionary as parameters
		-default retrieves entire table
		-returns data in form of a list of tuples
			-list entries are rows, tuple positions are columns
			-order is order of result of select sqlite command
			-None if there was a problem with parameter type

		(make sure all parameters are strings or dictionaries of strings, 
		as taken from a .get() from an Entry for example.)
		"""

		if debug.level > 0:
			print "\nbegin .get()...."

		# set variables to instance defaults if not provided
		if not table:
			table = self.table

		if columns.lower() == 'all':
			columns == '*'

		# prep parameters for use in formatting sqlite command
		# (note alternating) and handle bad parameters
		try:
			tcol = tuple(columns.replace(",","").split())
			if search_param:
				keys_sp = []
				vals_sp = []
				l_sp = []
				for col in self.col_list:
					if col in search_param.keys():
						keys_sp.append(col)
						vals_sp.append(search_param[col])
				for i in range(len(keys_sp)):
					l_sp.append(keys_sp[i])
					l_sp.append(vals_sp[i])
				tparam = tuple(l_sp)

		except TypeError, e:
			print "TypeError in table retreival function .get():", e
			return None

		# open up connection, format sqlite statement, execute it and retrieve it
		con = lite.connect(self.db)

		with con:
			cur = con.cursor()
			statement = 'select ' + ('%s, '*(len(tcol)-1)) + '%s from ' + self.table
			if search_param:
				statement = statement + ' where ' + "%s='%s' and "*(len(search_param)-1) + "%s='%s';"
				format = tcol+tparam
				if debug.level > 0:
					print "statement =", statement
					print "format =", format
					print "statement % format =", statement % format
				cur.execute(statement % format)
			else:
				statement = statement + ';'
				cur.execute(statement % tcol)
			rows = cur.fetchall()
		con.close()
		print "\nend .get()....\n"
		return rows

	def add_row(self, table=None, data_dict={}):
		"""Adds a new row to a table based on input data.

		-accepts table and data_dict
		-data_dict should be a dictionary with 
		keys of column name and values of cell data
		-defaults to self.table initialized in __init__ (must provide data)
		-returns a tuple of the data successfully inserted into the database
		-returns None if no data dictionary is provided

		(make sure all parameters are strings or dictionaries of strings, 
		as taken from a .get() from an Entry for example)
		"""

		# set variables to instance defaults if not provided
		if not table:
			table = self.table

		# if no data_dict provided, exit
		if len(data_dict) == 0:
			return None

		# open up connection and establish cursor
		con = lite.connect(self.db)

		with con:
			cur = con.cursor()

			# format sqlite statement
			statement = ('insert into ' + self.table + '(' + '%s,'*(len(data_dict)-1) +
				'%s) values(' + "'%s',"*(len(data_dict)-1) + "'%s');")

			# sort data_dict keys and values by self.col_list 
			# (see __init__) and make them tuples
			keys = []
			vals = []
			for col in self.col_list:
				if col in data_dict.keys():
					keys.append(col)
					vals.append(data_dict[col])
			format = tuple(keys) + tuple(vals)

			# execute sqlite command
			if debug.level > 0:
				print "statement =", statement
				print "format =", format
			cur.execute(statement % format)
			con.commit()

			# return the inserted row
			lid = str(cur.lastrowid)
			result = self.get(search_param={'Id':lid})

		con.close()
		return result

	def update_row(self, table=None, data_dict={}, search_param=None):
		"""Updates a preexisting row in a table based on input data.

		-accepts table, row Id (search_param), and data_dict
			-data_dict should be a dictionary with 
			keys of column name and values of cell data (as in .add_row())
			-search_param should be dictionary of columns and values
		-defaults to self.table initialized in __init__ (must provide data and row)
		-returns a tuple of the rows specified in search_param
			-returns None if no data dictionary or search_param is provided

		(make sure all parameters are strings or dictionaries of strings, 
		as taken from a .get() from an Entry for example)
		"""

		# set variables to instance defaults if not provided
		if not table:
			table = self.table

		#if no data_dict or row ID provided, exit
		if len(data_dict) == 0 or not search_param:
			print "Error: insufficient input in function .update_row()"
			return None

		#prep search parameters for use in formatting sqlite command
		#and handle bad parameters


		#open up connection and establish cursor
		con = lite.connect(self.db)

		with con:
			cur = con.cursor()

			#format sqlite statement
			statement = ('update ' + self.table + ' set ' + "%s='%s', "*(len(data_dict)-1) +
				"%s='%s' where " + "%s='%s' and "*(len(search_param)-1) + "%s='%s';")

			#sort data_dict and search_param keys and values 
			#by self.col_list (see __init__) and make them tuples
			keys_dd = []
			vals_dd = []
			keys_sp = []
			vals_sp = []

			for col in self.col_list:
				if col in data_dict.keys():
					keys_dd.append(col)
					vals_dd.append(data_dict[col])

				if col in search_param.keys():
					keys_sp.append(col)
					vals_sp.append(search_param[col])
			l_dd = []
			l_sp = []

			for i in range(len(keys_dd)):
				l_dd.append(keys_dd[i])
				l_dd.append(vals_dd[i])
			for i in range(len(keys_sp)):
				l_sp.append(keys_sp[i])
				l_sp.append(vals_sp[i])
			format = tuple(l_dd) + tuple(l_sp)

			#execute sqlite command
			if debug.level > 0:
				print "\n\nstatement =", statement
				print "format =", format
			cur.execute(statement % format)

			#return the updated row
			result = self.get(search_param=search_param)

		con.close()
		return result


class Colony_cur(Table_cur):
	"""Object representing an entire colony, derived from Table_cur with following methods:

	--searching for mice search parameters (get)
	--adding (add_mouse) or updating (update_mouse) mice to database 
	--adding (add_sci) or updating (update_sci) scientists and permissions to system
	"""
	def __init__(self, db='test.db', table='Mice'):
		""" Sets default database for methods, and default table for .get() method """
		Table_cur.__init__(self, db, table)

	def add_mouse(self, data_dict={}):
		"""Adds a new mouse row to the table 'Mice'.

		Special version of Table_cur.add_row()

		-accepts data_dict
			-data_dict should be a dictionary with 
			keys of column name and values of cell data
		-defaults to self.table attribue initialized in __init__ (must provide data)
		-returns a tuple of the data successfully inserted into the database
		"""
		return self.add_row(data_dict=data_dict)

	def update_mouse(self, data_dict={}, search_param=None):
		"""Updates data on existing mouse/mice in database.

		-accepts data for searching (search_param) and for updating (data_dict)
			-data_dict should be a dictionary with 
			keys of column name and values of cell data (as in .add_mouse())
			-search_param should be similarly structured
		-defaults to colony table 'Mice' (must provide data and row parameters)
		-returns a tuple of the rows specified in search_param
			-returns None if no data dictionary or search_param is provided

		(make sure all parameters are strings or dictionaries of strings, 
		as taken from a .get() from an Entry for example)
		"""
		return self.update_row(data_dict=data_dict, search_param=search_param)

	def add_sci(self, table='Scientists', data_dict={}):
		"""Adds a new scientist row to the table 'Scientists'.

		Special version of Table_cur.add_row()

		-accepts data_dict
			-data_dict should be a dictionary with 
			keys of column name and values of cell data
		-defaults to 'Scientists' table (must provide data)
		-returns a tuple of the data successfully inserted into the database
		"""
		return self.add_row(table=table, data_dict=data_dict)

	def update_sci(self, table='Scientists', data_dict={}, search_param=None):
		"""Updates data on existing scientist in database.

		-accepts data for searching (search_param) and for updating (data_dict)
			-data_dict should be a dictionary with 
			keys of column name and values of cell data (as in .add_sci())
			-search_param should be similarly structured
		-defaults to 'Scientists' table (must provide data and row parameters)
		-returns a tuple of the rows specified in search_param
			-returns None if no data dictionary or search_param is provided

		(make sure all parameters are strings or dictionaries of strings, 
		as taken from a .get() from an Entry for example)
		"""
		return self.update_row(data_dict=data_dict, search_param=search_param)


if __name__ == '__main__':
	import database as acdb

	cur = acdb.Colony_cur('test.db', 'Cars')
	data = {}
	data['Name'] = 'Skoda'
	data['Price'] = 9000
	search = {}
	search['Id'] = 3
	print cur.add_row(data_dict=data)
	print cur.update_row(data_dict=data, search_param=search)
	print ''
	print cur.get()
	print ''