"""GenericDBServerInfo for use with pygr.sqlgraph.
Should extend the functionality of existing *ServerInfo classes.
Depends on sqlalchemy.

Author: Paul Rigor (prigor@ics.uci.edu)
Date: 04/29/2009
"""

class GenericDBServerInfo(object):
    """Picklable reference to generic database server.
    Takes advantage of sqlalchemy's connection interface.
    
    TODO: DEBUG AND SUBMIT TO PYGR
    
    """

    def __init__(self, *args, **kwargs):
        print "GenericDBServerInfo: ",args,kwargs
        if "sqlite:" in str(args) and ":memory:" in str(args):
            raise(Exception("Error: In-memory sqlite databases cannot be pickled."))
        self.kwargs = kwargs
        self.args = args

    def cursor(self):
        """Get a cursor for accessing this database."""
        try:
            return self._cursor
        except AttributeError:
            from sqlalchemy import create_engine
            self.dbengine = create_engine(*self.args, **self.kwargs)
            self._connection = self.dbengine.pool.create_connection().get_connection()
            self._cursor = self._connection.cursor()
            return self._cursor

    def close(self):
        """Close this cursor and connection. """
        self._cursor.close()
        self._connection.close()
        del self._cursor
        del self._connection
        
    def get_primary_key(self,tablename=None):
        """Returns the primary_key for hte specified tablename.
        If no table is provided, the first table is obtained.
        
        This method can be pushed as a stand-alone method.  For example,
        def get_primary_key(tablename, serverInfo):
            '''Standalone version of obtaining the primary key of a table.
            serverInfo is an instance of GenericDBServerInfo
            '''
            from sqlalchemy import MetaData
            metadata = MetaData(*serverInfo.args,**serverInfo.kwargs)
            #...
            #...
            #(and so on)
        
        """
    def get_primary_key(tablename,*args,**kwargs):
        from sqlalchemy import MetaData
        metadata = MetaData(*self.args,**self.kwargs) #
        metadata.reflect()
        try:
            table = metadata.tables[tablename]
        except KeyError:
            table = metadata.tables.values()[0]
        except IndexError:
            raise(Exception("Error: The database does not contain any tables."))
        try:
            primary_key = table.primary_key.keys()[0]
        except IndexError:
            raise(Exception("Error: The specified table has no primary key!"))
        return primary_key

    def __getstate__(self):
        """Return picklable fields."""
        return dict(args=self.args,kwargs=self.kwargs)