<s> import abc <EOL> import json <EOL> import logging <EOL> import time <EOL> import os <EOL> import luigi <EOL> from luigi import postgres <EOL> from luigi . contrib import rdbms <EOL> from luigi . s3 import S3PathTask , S3Target <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> try : <EOL> import psycopg2 <EOL> import psycopg2 . errorcodes <EOL> except ImportError : <EOL> logger . warning ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> class RedshiftTarget ( postgres . PostgresTarget ) : <EOL> """<STR_LIT>""" <EOL> marker_table = luigi . configuration . get_config ( ) . get ( <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ) <EOL> use_db_timestamps = False <EOL> class S3CopyToTable ( rdbms . CopyToTable ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def s3_load_path ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ abc . abstractproperty <EOL> def aws_access_key_id ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ abc . abstractproperty <EOL> def aws_secret_access_key ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ abc . abstractproperty <EOL> def copy_options ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' <EOL> @ property <EOL> def prune_table ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ property <EOL> def prune_column ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ property <EOL> def prune_date ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ property <EOL> def table_attributes ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' <EOL> @ property <EOL> def do_truncate_table ( self ) : <EOL> """<STR_LIT>""" <EOL> return False <EOL> def do_prune ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . prune_table and self . prune_column and self . prune_date : <EOL> return True <EOL> elif self . prune_table or self . prune_column or self . prune_date : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> else : <EOL> return False <EOL> @ property <EOL> def table_type ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' <EOL> @ property <EOL> def queries ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ ] <EOL> def truncate_table ( self , connection ) : <EOL> query = "<STR_LIT>" % self . table <EOL> cursor = connection . cursor ( ) <EOL> try : <EOL> cursor . execute ( query ) <EOL> finally : <EOL> cursor . close ( ) <EOL> def prune ( self , connection ) : <EOL> query = "<STR_LIT>" % ( self . prune_table , self . prune_column , self . prune_date ) <EOL> cursor = connection . cursor ( ) <EOL> try : <EOL> cursor . execute ( query ) <EOL> finally : <EOL> cursor . close ( ) <EOL> def create_table ( self , connection ) : <EOL> """<STR_LIT>""" <EOL> if len ( self . columns [ <NUM_LIT:0> ] ) == <NUM_LIT:1> : <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % self . table ) <EOL> elif len ( self . columns [ <NUM_LIT:0> ] ) == <NUM_LIT:2> : <EOL> coldefs = '<STR_LIT:U+002C>' . join ( <EOL> '<STR_LIT>' . format ( <EOL> name = name , <EOL> type = type ) for name , type in self . columns <EOL> ) <EOL> query = ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) . format ( <EOL> type = self . table_type , <EOL> table = self . table , <EOL> coldefs = coldefs , <EOL> table_attributes = self . table_attributes ) <EOL> connection . cursor ( ) . execute ( query ) <EOL> def run ( self ) : <EOL> """<STR_LIT>""" <EOL> if not ( self . table ) : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> path = self . s3_load_path ( ) <EOL> output = self . output ( ) <EOL> connection = output . connect ( ) <EOL> cursor = connection . cursor ( ) <EOL> self . init_copy ( connection ) <EOL> self . copy ( cursor , path ) <EOL> self . post_copy ( cursor ) <EOL> output . touch ( connection ) <EOL> connection . commit ( ) <EOL> connection . close ( ) <EOL> def copy ( self , cursor , f ) : <EOL> """<STR_LIT>""" <EOL> logger . info ( "<STR_LIT>" , f ) <EOL> cursor . execute ( """<STR_LIT>""" % ( self . table , f , self . aws_access_key_id , <EOL> self . aws_secret_access_key , self . copy_options ) ) <EOL> def output ( self ) : <EOL> """<STR_LIT>""" <EOL> return RedshiftTarget ( <EOL> host = self . host , <EOL> database = self . database , <EOL> user = self . user , <EOL> password = self . password , <EOL> table = self . table , <EOL> update_id = self . update_id ) <EOL> def does_table_exist ( self , connection ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT:.>' in self . table : <EOL> query = ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> else : <EOL> query = ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> cursor = connection . cursor ( ) <EOL> try : <EOL> cursor . execute ( query , tuple ( self . table . split ( '<STR_LIT:.>' ) ) ) <EOL> result = cursor . fetchone ( ) <EOL> return bool ( result ) <EOL> finally : <EOL> cursor . close ( ) <EOL> def init_copy ( self , connection ) : <EOL> """<STR_LIT>""" <EOL> if not self . does_table_exist ( connection ) : <EOL> logger . info ( "<STR_LIT>" , self . table ) <EOL> connection . reset ( ) <EOL> self . create_table ( connection ) <EOL> if self . do_truncate_table : <EOL> logger . info ( "<STR_LIT>" , self . table ) <EOL> self . truncate_table ( connection ) <EOL> if self . do_prune ( ) : <EOL> logger . info ( "<STR_LIT>" , self . prune_column , self . prune_date , self . prune_table ) <EOL> self . prune ( connection ) <EOL> def post_copy ( self , cursor ) : <EOL> """<STR_LIT>""" <EOL> logger . info ( '<STR_LIT>' ) <EOL> for query in self . queries : <EOL> cursor . execute ( query ) <EOL> class S3CopyJSONToTable ( S3CopyToTable ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractproperty <EOL> def jsonpath ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' <EOL> @ abc . abstractproperty <EOL> def copy_json_options ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' <EOL> def copy ( self , cursor , f ) : <EOL> """<STR_LIT>""" <EOL> cursor . execute ( """<STR_LIT>""" % ( self . table , f , self . aws_access_key_id , <EOL> self . aws_secret_access_key , self . jsonpath , <EOL> self . copy_json_options , self . copy_options ) ) <EOL> class RedshiftManifestTask ( S3PathTask ) : <EOL> """<STR_LIT>""" <EOL> folder_paths = luigi . Parameter ( ) <EOL> text_target = True <EOL> def run ( self ) : <EOL> entries = [ ] <EOL> for folder_path in self . folder_paths : <EOL> s3 = S3Target ( folder_path ) <EOL> client = s3 . fs <EOL> for file_name in client . list ( s3 . path ) : <EOL> entries . append ( { <EOL> '<STR_LIT:url>' : '<STR_LIT>' % ( folder_path , file_name ) , <EOL> '<STR_LIT>' : True <EOL> } ) <EOL> manifest = { '<STR_LIT>' : entries } <EOL> target = self . output ( ) . open ( '<STR_LIT:w>' ) <EOL> dump = json . dumps ( manifest ) <EOL> if not self . text_target : <EOL> dump = dump . encode ( '<STR_LIT:utf8>' ) <EOL> target . write ( dump ) <EOL> target . close ( ) <EOL> class KillOpenRedshiftSessions ( luigi . Task ) : <EOL> """<STR_LIT>""" <EOL> connection_reset_wait_seconds = luigi . IntParameter ( default = <NUM_LIT> ) <EOL> @ abc . abstractproperty <EOL> def host ( self ) : <EOL> return None <EOL> @ abc . abstractproperty <EOL> def database ( self ) : <EOL> return None <EOL> @ abc . abstractproperty <EOL> def user ( self ) : <EOL> return None <EOL> @ abc . abstractproperty <EOL> def password ( self ) : <EOL> return None <EOL> @ property <EOL> def update_id ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . task_id <EOL> def output ( self ) : <EOL> """<STR_LIT>""" <EOL> return RedshiftTarget ( <EOL> host = self . host , <EOL> database = self . database , <EOL> user = self . user , <EOL> password = self . password , <EOL> table = self . __class__ . __name__ , <EOL> update_id = self . update_id ) <EOL> def run ( self ) : <EOL> """<STR_LIT>""" <EOL> connection = self . output ( ) . connect ( ) <EOL> query = ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> cursor = connection . cursor ( ) <EOL> logger . info ( '<STR_LIT>' , self . database ) <EOL> try : <EOL> cursor . execute ( query , ( self . database , ) ) <EOL> cursor . close ( ) <EOL> connection . commit ( ) <EOL> except psycopg2 . DatabaseError as e : <EOL> if e . message and '<STR_LIT>' in e . message : <EOL> connection . close ( ) <EOL> logger . info ( '<STR_LIT>' , self . connection_reset_wait_seconds ) <EOL> time . sleep ( self . connection_reset_wait_seconds ) <EOL> logger . info ( '<STR_LIT>' ) <EOL> connection = self . output ( ) . connect ( ) <EOL> else : <EOL> raise <EOL> try : <EOL> self . output ( ) . touch ( connection ) <EOL> connection . commit ( ) <EOL> finally : <EOL> connection . close ( ) <EOL> logger . info ( '<STR_LIT>' , self . database ) <EOL> class RedshiftQuery ( postgres . PostgresQuery ) : <EOL> """<STR_LIT>""" <EOL> def output ( self ) : <EOL> """<STR_LIT>""" <EOL> return RedshiftTarget ( <EOL> host = self . host , <EOL> database = self . database , <EOL> user = self . user , <EOL> password = self . password , <EOL> table = self . table , <EOL> update_id = self . update_id <EOL> ) <EOL> class RedshiftUnloadTask ( postgres . PostgresQuery ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractproperty <EOL> def aws_access_key_id ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ abc . abstractproperty <EOL> def aws_secret_access_key ( self ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> @ property <EOL> def s3_unload_path ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' <EOL> @ property <EOL> def unload_options ( self ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" <EOL> @ property <EOL> def unload_query ( self ) : <EOL> """<STR_LIT>""" <EOL> return ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> def run ( self ) : <EOL> connection = self . output ( ) . connect ( ) <EOL> cursor = connection . cursor ( ) <EOL> config = luigi . configuration . get_config ( ) <EOL> if self . aws_access_key_id is None or self . aws_secret_access_key is None : <EOL> self . aws_access_key_id = config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . aws_secret_access_key = config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> if self . aws_access_key_id is None or self . aws_access_key_id . strip ( ) == '<STR_LIT>' or self . aws_secret_access_key is None or self . aws_secret_access_key . strip ( ) == '<STR_LIT>' : <EOL> self . aws_access_key_id = os . environ [ '<STR_LIT>' ] <EOL> self . aws_secret_access_key = os . environ [ '<STR_LIT>' ] <EOL> unload_query = self . unload_query . format ( <EOL> query = self . query ( ) . replace ( "<STR_LIT:'>" , "<STR_LIT>" ) , <EOL> s3_unload_path = self . s3_unload_path , <EOL> unload_options = self . unload_options , <EOL> s3_access_key = self . aws_access_key_id , <EOL> s3_security_key = self . aws_secret_access_key ) <EOL> logger . info ( '<STR_LIT>' . format ( name = self . __class__ ) ) <EOL> try : <EOL> cursor = connection . cursor ( ) <EOL> cursor . execute ( unload_query ) <EOL> logger . info ( cursor . statusmessage ) <EOL> except : <EOL> raise <EOL> self . output ( ) . touch ( connection ) <EOL> connection . commit ( ) <EOL> connection . close ( ) <EOL> def output ( self ) : <EOL> """<STR_LIT>""" <EOL> return RedshiftTarget ( <EOL> host = self . host , <EOL> database = self . database , <EOL> user = self . user , <EOL> password = self . password , <EOL> table = self . table , <EOL> update_id = self . update_id <EOL> ) </s>
