<s> import sys <EOL> from conary import dbstore <EOL> from schema import getTables , getIndexes <EOL> from tablelist import TableList <EOL> class Loader : <EOL> binaryCols = [ <EOL> "<STR_LIT>" , "<STR_LIT:password>" , "<STR_LIT:data>" , "<STR_LIT>" , "<STR_LIT>" , <EOL> "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> def __init__ ( self , db , table , fields ) : <EOL> self . db = db <EOL> self . table = table <EOL> self . sql = "<STR_LIT>" % ( <EOL> table , "<STR_LIT:U+002CU+0020>" . join ( fields ) , "<STR_LIT:U+002C>" . join ( [ "<STR_LIT:?>" ] * len ( fields ) ) ) <EOL> self . cu = self . db . cursor ( ) <EOL> self . cu . execute ( "<STR_LIT>" % ( table , ) ) <EOL> set1 = set ( [ x . lower ( ) for x in fields ] ) <EOL> set2 = set ( [ x . lower ( ) for x in self . cu . fields ( ) ] ) <EOL> assert ( set1 . difference ( set2 ) == set ( ) ) , "<STR_LIT>" % ( set1 , set2 ) <EOL> self . fields = fields <EOL> self . funcs = [ self . __getfunc ( x ) for x in fields ] <EOL> def announce ( self ) : <EOL> sys . stdout . write ( "<STR_LIT>" % ( self . table , ) ) <EOL> sys . stdout . flush ( ) <EOL> def __getfunc ( self , field ) : <EOL> if field . lower ( ) in self . binaryCols : <EOL> return self . cu . binary <EOL> return lambda a : a <EOL> def rowvals ( self , row ) : <EOL> return tuple ( map ( lambda ( f , x ) : f ( x ) , zip ( self . funcs , row ) ) ) <EOL> def insertRows ( self , rows , callback = None ) : <EOL> self . cu . executemany ( self . sql , [ self . rowvals ( row ) for row in rows ] ) <EOL> if callback : <EOL> callback . increment ( len ( rows ) ) <EOL> return len ( rows ) <EOL> class PgSQLLoader ( Loader ) : <EOL> def __init__ ( self , db , table , fields ) : <EOL> Loader . __init__ ( self , db , table , fields ) <EOL> self . __usebulk = False <EOL> if hasattr ( self . db . dbh , "<STR_LIT>" ) : <EOL> self . __usebulk = True <EOL> if not self . __usebulk : <EOL> print "<STR_LIT>" <EOL> def bulkInsert ( self , rows , callback ) : <EOL> self . db . dbh . bulkload ( self . table , ( self . rowvals ( row ) for row in rows ) , self . fields ) <EOL> if callback : <EOL> callback . increment ( len ( rows ) ) <EOL> return len ( rows ) <EOL> def insertRows ( self , rows , callback = None ) : <EOL> if self . __usebulk : <EOL> return self . bulkInsert ( rows , callback ) <EOL> else : <EOL> return Loader . insertRows ( self , rows , callback ) <EOL> class Database : <EOL> def __init__ ( self , driver , db , verbose = True ) : <EOL> self . db = dbstore . connect ( db , driver ) <EOL> self . driver = driver <EOL> self . db . loadSchema ( ) <EOL> self . verbose = verbose <EOL> self . _hint = '<STR_LIT>' <EOL> def createSchema ( self ) : <EOL> cu = self . db . cursor ( ) <EOL> for stmt in getTables ( self . driver ) : <EOL> if self . verbose : <EOL> print stmt <EOL> cu . execute ( stmt ) <EOL> self . db . loadSchema ( ) <EOL> def createIndexes ( self ) : <EOL> cu = self . db . cursor ( ) <EOL> for stmt in getIndexes ( self . driver ) : <EOL> if self . verbose : <EOL> print stmt <EOL> cu . execute ( stmt ) <EOL> self . db . loadSchema ( ) <EOL> def checkTablesList ( self , isSrc = True ) : <EOL> self . db . loadSchema ( ) <EOL> skip = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> knowns = [ x . lower ( ) for x in TableList ] <EOL> haves = [ x . lower ( ) for x in self . db . tables ] <EOL> if isSrc : <EOL> which = "<STR_LIT>" <EOL> else : <EOL> which = "<STR_LIT>" <EOL> onlyKnowns = set ( knowns ) . difference ( set ( haves ) ) . difference ( set ( skip ) ) <EOL> if onlyKnowns : <EOL> raise RuntimeError ( "<STR_LIT>" % ( <EOL> which , self . driver , onlyKnowns ) ) <EOL> onlyHaves = set ( haves ) . difference ( set ( knowns ) ) . difference ( set ( skip ) ) <EOL> if onlyHaves and isSrc : <EOL> raise RuntimeError ( "<STR_LIT>" , onlyHaves ) <EOL> return True <EOL> def getCount ( self , table ) : <EOL> cu = self . db . cursor ( ) <EOL> cu . execute ( "<STR_LIT>" % ( table , ) ) <EOL> return cu . fetchall ( ) [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] <EOL> def getFields ( self , table ) : <EOL> cu = self . db . cursor ( ) <EOL> cu . execute ( "<STR_LIT>" % ( table , ) ) <EOL> return [ x . lower ( ) for x in cu . fields ( ) ] <EOL> def getTables ( self ) : <EOL> return [ x . lower ( ) for x in self . db . tables ] <EOL> def iterRows ( self , table , fields = "<STR_LIT:*>" ) : <EOL> cu = self . db . itercursor ( ) <EOL> cu . execute ( "<STR_LIT>" % ( self . _hint , fields , table , <EOL> '<STR_LIT>' if fields == '<STR_LIT:*>' else ( '<STR_LIT>' + fields ) ) ) <EOL> return cu <EOL> def prepareInsert ( self , table , fields ) : <EOL> return Loader ( self . db , table , fields ) <EOL> def finalize ( self , version ) : <EOL> self . db . setVersion ( version ) <EOL> self . db . commit ( ) <EOL> def commit ( self ) : <EOL> return self . db . commit ( ) <EOL> def close ( self ) : <EOL> self . db . close ( ) <EOL> class PgSQLDatabase ( Database ) : <EOL> def __init__ ( self , db , verbose = True ) : <EOL> Database . __init__ ( self , "<STR_LIT>" , db , verbose ) <EOL> def createIndexes ( self ) : <EOL> Database . createIndexes ( self ) <EOL> self . fix_primary_keys ( ) <EOL> def fix_primary_keys ( self ) : <EOL> cu = self . db . cursor ( ) <EOL> cu . execute ( """<STR_LIT>""" ) <EOL> for ( t , col ) in cu . fetchall ( ) : <EOL> table = t . lower ( ) <EOL> cu . execute ( "<STR_LIT>" , ( table , col ) ) <EOL> seqname = cu . fetchone ( ) [ <NUM_LIT:0> ] <EOL> if seqname is None : <EOL> continue <EOL> cu . execute ( "<STR_LIT>" % ( col , table ) ) <EOL> seqval = cu . fetchone ( ) [ <NUM_LIT:0> ] <EOL> if not seqval : <EOL> seqval = <NUM_LIT:1> <EOL> else : <EOL> seqval += <NUM_LIT:1> <EOL> cu . execute ( "<STR_LIT>" , ( seqname , seqval ) ) <EOL> ret = cu . fetchone ( ) [ <NUM_LIT:0> ] <EOL> assert ( ret == seqval ) <EOL> if self . verbose : <EOL> print "<STR_LIT>" % ( seqname , ret , table , col ) <EOL> def prepareInsert ( self , table , fields ) : <EOL> return PgSQLLoader ( self . db , table , fields ) <EOL> def finalize ( self , version ) : <EOL> Database . finalize ( self , version ) <EOL> cu = self . db . cursor ( ) <EOL> if self . verbose : <EOL> print "<STR_LIT>" <EOL> cu . execute ( "<STR_LIT>" ) <EOL> class MySQLDatabase ( Database ) : <EOL> def __init__ ( self , db , verbose = True ) : <EOL> Database . __init__ ( self , "<STR_LIT>" , db , verbose ) <EOL> self . _hint = "<STR_LIT>" <EOL> def finalize ( self , version ) : <EOL> Database . finalize ( self , version ) <EOL> cu = self . db . cursor ( ) <EOL> if self . verbose : <EOL> print "<STR_LIT>" <EOL> for t in TableList : <EOL> cu . execute ( "<STR_LIT>" % ( t , ) ) <EOL> def getdb ( driver , db , verbose = True ) : <EOL> if driver == "<STR_LIT>" : <EOL> return PgSQLDatabase ( db , verbose ) <EOL> elif driver == "<STR_LIT>" : <EOL> return MySQLDatabase ( db , verbose ) <EOL> return Database ( driver , db , verbose ) </s>
