<s> from psycopg2 . extensions import AsIs <EOL> from arbalest . core import PipelineStep <EOL> from arbalest . redshift . manifest import Manifest <EOL> from arbalest . redshift . runner import S3JsonStepRunner <EOL> from arbalest . s3 import normalize_path <EOL> class BulkCopyFromS3JsonStep ( PipelineStep ) : <EOL> def __init__ ( self , metadata , source , schema , aws_access_key_id , <EOL> aws_secret_access_key , bucket , table ) : <EOL> self . metadata = metadata <EOL> self . source = source <EOL> self . schema = schema <EOL> self . aws_access_key_id = aws_access_key_id <EOL> self . aws_secret_access_key = aws_secret_access_key <EOL> self . bucket = bucket <EOL> self . table = table <EOL> self . runner = S3JsonStepRunner ( metadata , schema , bucket , table ) <EOL> self . sql = SqlStep ( table . database ) <EOL> self . max_error_count = <NUM_LIT:1> <EOL> @ property <EOL> def source_key ( self ) : <EOL> return normalize_path ( self . source ) <EOL> @ property <EOL> def schema_key ( self ) : <EOL> return self . runner . schema_key <EOL> @ property <EOL> def source_url ( self ) : <EOL> return '<STR_LIT>' . format ( self . bucket . name , self . source_key ) <EOL> @ property <EOL> def schema_url ( self ) : <EOL> return self . runner . schema_url <EOL> @ property <EOL> def copy_sql ( self ) : <EOL> return "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" <EOL> @ property <EOL> def validate_sql ( self ) : <EOL> return self . copy_sql + "<STR_LIT>" <EOL> def run ( self ) : <EOL> self . runner . stage ( ) <EOL> self . __execute ( self . copy_sql ) <EOL> self . __promote ( ) <EOL> self . runner . commit ( ) <EOL> def validate ( self ) : <EOL> self . runner . stage ( ) <EOL> self . __execute ( self . validate_sql ) <EOL> self . __promote ( ) <EOL> self . runner . rollback ( ) <EOL> def __execute ( self , sql ) : <EOL> self . table . stage_update ( ) <EOL> self . sql . execute ( ( <EOL> sql , self . schema . update_table , self . source_url , <EOL> self . aws_access_key_id , <EOL> self . aws_secret_access_key , self . schema_url , self . max_error_count ) ) <EOL> def __promote ( self ) : <EOL> if self . table . exists ( ) : <EOL> self . table . drop ( ) <EOL> self . table . promote_update ( ) <EOL> class ManifestCopyFromS3JsonStep ( PipelineStep ) : <EOL> def __init__ ( self , metadata , source , schema , aws_access_key_id , <EOL> aws_secret_access_key , bucket , table ) : <EOL> self . metadata = metadata <EOL> self . source = source <EOL> self . schema = schema <EOL> self . aws_access_key_id = aws_access_key_id <EOL> self . aws_secret_access_key = aws_secret_access_key <EOL> self . bucket = bucket <EOL> self . table = table <EOL> self . runner = S3JsonStepRunner ( metadata , schema , bucket , table ) <EOL> self . manifest = Manifest ( metadata , source , schema , bucket ) <EOL> self . sql = SqlStep ( table . database ) <EOL> self . max_error_count = <NUM_LIT:1> <EOL> @ property <EOL> def schema_key ( self ) : <EOL> return self . runner . schema_key <EOL> @ property <EOL> def schema_url ( self ) : <EOL> return self . runner . schema_url <EOL> @ property <EOL> def copy_sql ( self ) : <EOL> return "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" <EOL> @ property <EOL> def validate_sql ( self ) : <EOL> return self . copy_sql + "<STR_LIT>" <EOL> def run ( self ) : <EOL> self . runner . stage ( ) <EOL> updated_journal = self . manifest . save ( ) <EOL> self . __execute ( self . copy_sql ) <EOL> self . manifest . commit ( updated_journal ) <EOL> self . runner . commit ( ) <EOL> def validate ( self ) : <EOL> self . runner . stage ( ) <EOL> self . manifest . save ( ) <EOL> self . __execute ( self . validate_sql ) <EOL> self . runner . rollback ( ) <EOL> def __execute ( self , sql ) : <EOL> if self . table . exists ( ) and not self . manifest . journal_exists ( ) : <EOL> self . table . drop ( ) <EOL> self . table . create ( ) <EOL> elif not self . table . exists ( ) : <EOL> self . table . create ( ) <EOL> self . sql . execute ( ( sql , self . schema . table , self . manifest . manifest_url , <EOL> self . aws_access_key_id , self . aws_secret_access_key , <EOL> self . schema_url , <EOL> self . max_error_count ) ) <EOL> class SqlStep ( PipelineStep ) : <EOL> def __init__ ( self , database , * args ) : <EOL> self . database = database <EOL> self . statements = list ( args ) <EOL> def run ( self ) : <EOL> self . database . open ( ) <EOL> self . execute ( self . statements ) <EOL> self . database . commit ( ) <EOL> def validate ( self ) : <EOL> self . database . open ( ) <EOL> self . execute ( self . statements ) <EOL> self . database . rollback ( ) <EOL> def execute ( self , statements ) : <EOL> if not isinstance ( statements , list ) : <EOL> self . __execute ( statements ) <EOL> else : <EOL> for statement in statements : <EOL> self . __execute ( statement ) <EOL> def __execute ( self , statement ) : <EOL> if isinstance ( statement , basestring ) : <EOL> self . database . execute ( statement ) <EOL> else : <EOL> query = statement [ <NUM_LIT:0> ] <EOL> params = statement [ <NUM_LIT:1> : ] <EOL> self . database . execute ( query , tuple ( [ AsIs ( x ) for x in params ] ) ) </s>
