<s> from __future__ import absolute_import , division , print_function <EOL> import os <EOL> import sys <EOL> import subprocess <EOL> import uuid <EOL> import mmap <EOL> from contextlib import closing <EOL> from functools import partial <EOL> from distutils . spawn import find_executable <EOL> import sqlalchemy as sa <EOL> from sqlalchemy . ext . compiler import compiles <EOL> from sqlalchemy . sql . elements import Executable , ClauseElement <EOL> from toolz import merge <EOL> from multipledispatch import MDNotImplementedError <EOL> from . . append import append <EOL> from . . convert import convert <EOL> from . csv import CSV , infer_header <EOL> from . . temp import Temp <EOL> from . aws import S3 <EOL> from . sql import getbind <EOL> class CopyFromCSV ( Executable , ClauseElement ) : <EOL> def __init__ ( self , element , csv , delimiter = '<STR_LIT:U+002C>' , header = None , na_value = '<STR_LIT>' , <EOL> lineterminator = '<STR_LIT:\n>' , quotechar = '<STR_LIT:">' , escapechar = '<STR_LIT:\\>' , <EOL> encoding = '<STR_LIT:utf8>' , skiprows = <NUM_LIT:0> , bind = None , ** kwargs ) : <EOL> if not isinstance ( element , sa . Table ) : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> self . element = element <EOL> self . csv = csv <EOL> self . delimiter = delimiter <EOL> self . header = ( <EOL> header if header is not None else <EOL> ( csv . has_header <EOL> if csv . has_header is not None else infer_header ( csv . path ) ) <EOL> ) <EOL> self . na_value = na_value <EOL> self . lineterminator = lineterminator <EOL> self . quotechar = quotechar <EOL> self . escapechar = escapechar <EOL> self . encoding = encoding <EOL> self . skiprows = int ( skiprows or self . header ) <EOL> self . _bind = getbind ( element , bind ) <EOL> for k , v in kwargs . items ( ) : <EOL> setattr ( self , k , v ) <EOL> @ property <EOL> def bind ( self ) : <EOL> return self . _bind <EOL> @ compiles ( CopyFromCSV , '<STR_LIT>' ) <EOL> def compile_from_csv_sqlite ( element , compiler , ** kwargs ) : <EOL> if not find_executable ( '<STR_LIT>' ) : <EOL> raise MDNotImplementedError ( "<STR_LIT>" ) <EOL> t = element . element <EOL> if not element . header : <EOL> csv = element . csv <EOL> else : <EOL> csv = Temp ( CSV ) ( '<STR_LIT>' % uuid . uuid1 ( ) ) <EOL> assert csv . has_header , '<STR_LIT>' <EOL> chunksize = <NUM_LIT:1> << <NUM_LIT> <EOL> lineterminator = element . lineterminator . encode ( element . encoding ) <EOL> with open ( element . csv . path , '<STR_LIT:rb>' ) as f : <EOL> with closing ( mmap . mmap ( f . fileno ( ) , <NUM_LIT:0> , access = mmap . ACCESS_READ ) ) as mf : <EOL> index = mf . find ( lineterminator ) <EOL> if index == - <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" % lineterminator ) <EOL> mf . seek ( index + len ( lineterminator ) ) <EOL> with open ( csv . path , '<STR_LIT:wb>' ) as g : <EOL> for chunk in iter ( partial ( mf . read , chunksize ) , b'<STR_LIT>' ) : <EOL> g . write ( chunk ) <EOL> fullpath = os . path . abspath ( csv . path ) . encode ( '<STR_LIT>' ) . decode ( ) <EOL> cmd = [ '<STR_LIT>' , <EOL> '<STR_LIT>' , repr ( element . na_value ) , <EOL> '<STR_LIT>' , element . delimiter , <EOL> '<STR_LIT>' , '<STR_LIT>' % ( <EOL> fullpath , compiler . preparer . format_table ( t ) <EOL> ) , <EOL> element . bind . url . database ] <EOL> stderr = subprocess . check_output ( <EOL> cmd , <EOL> stderr = subprocess . STDOUT , <EOL> stdin = subprocess . PIPE <EOL> ) . decode ( sys . getfilesystemencoding ( ) ) <EOL> if stderr : <EOL> try : <EOL> raise OSError ( stderr ) <EOL> except OSError as e : <EOL> raise sa . exc . DatabaseError ( '<STR_LIT:U+0020>' . join ( cmd ) , [ ] , e ) <EOL> return '<STR_LIT>' <EOL> @ compiles ( CopyFromCSV , '<STR_LIT>' ) <EOL> def compile_from_csv_mysql ( element , compiler , ** kwargs ) : <EOL> if element . na_value : <EOL> raise ValueError ( <EOL> '<STR_LIT>' <EOL> ) <EOL> return compiler . process ( <EOL> sa . text ( <EOL> """<STR_LIT>""" . format ( <EOL> local = getattr ( element , '<STR_LIT>' , '<STR_LIT>' ) , <EOL> table = compiler . preparer . format_table ( element . element ) <EOL> ) <EOL> ) . bindparams ( <EOL> path = os . path . abspath ( element . csv . path ) , <EOL> encoding = element . encoding or element . bind . execute ( <EOL> '<STR_LIT>' <EOL> ) . scalar ( ) , <EOL> delimiter = element . delimiter , <EOL> quotechar = element . quotechar , <EOL> escapechar = element . escapechar , <EOL> lineterminator = element . lineterminator , <EOL> skiprows = int ( element . header ) <EOL> ) , <EOL> ** kwargs <EOL> ) <EOL> @ compiles ( CopyFromCSV , '<STR_LIT>' ) <EOL> def compile_from_csv_postgres ( element , compiler , ** kwargs ) : <EOL> if len ( element . escapechar ) != <NUM_LIT:1> : <EOL> raise ValueError ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> ) <EOL> if element . lineterminator != '<STR_LIT:\n>' : <EOL> raise ValueError ( <EOL> r'<STR_LIT>' <EOL> ) <EOL> return compiler . process ( <EOL> sa . text ( <EOL> """<STR_LIT>""" . format ( compiler . preparer . format_table ( element . element ) ) <EOL> ) . bindparams ( <EOL> path = os . path . abspath ( element . csv . path ) , <EOL> delimiter = element . delimiter , <EOL> na_value = element . na_value , <EOL> quotechar = element . quotechar , <EOL> escapechar = element . escapechar , <EOL> header = element . header , <EOL> encoding = element . encoding or element . bind ( <EOL> '<STR_LIT>' <EOL> ) . execute ( ) . scalar ( ) <EOL> ) , <EOL> ** kwargs <EOL> ) <EOL> try : <EOL> import boto <EOL> from odo . backends . aws import S3 <EOL> from redshift_sqlalchemy . dialect import CopyCommand <EOL> except ImportError : <EOL> pass <EOL> else : <EOL> @ compiles ( CopyFromCSV , '<STR_LIT>' ) <EOL> def compile_from_csv_redshift ( element , compiler , ** kwargs ) : <EOL> assert isinstance ( element . csv , S3 ( CSV ) ) <EOL> assert element . csv . path . startswith ( '<STR_LIT>' ) <EOL> cfg = boto . Config ( ) <EOL> aws_access_key_id = cfg . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> aws_secret_access_key = cfg . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> compression = getattr ( element , '<STR_LIT>' , '<STR_LIT>' ) . upper ( ) or None <EOL> cmd = CopyCommand ( <EOL> table = element . element , <EOL> data_location = element . csv . path , <EOL> access_key_id = aws_access_key_id , <EOL> secret_access_key = aws_secret_access_key , <EOL> format = '<STR_LIT>' , <EOL> delimiter = element . delimiter , <EOL> ignore_header = int ( element . header ) , <EOL> empty_as_null = True , <EOL> blanks_as_null = False , <EOL> compression = compression <EOL> ) <EOL> return compiler . process ( cmd ) <EOL> @ append . register ( sa . Table , CSV ) <EOL> def append_csv_to_sql_table ( tbl , csv , bind = None , ** kwargs ) : <EOL> bind = getbind ( tbl , bind ) <EOL> dialect = bind . dialect . name <EOL> if dialect == '<STR_LIT>' and not isinstance ( csv , S3 ( CSV ) ) : <EOL> csv = convert ( Temp ( S3 ( CSV ) ) , csv , ** kwargs ) <EOL> elif dialect != '<STR_LIT>' and isinstance ( csv , S3 ( CSV ) ) : <EOL> csv = convert ( Temp ( CSV ) , csv , has_header = csv . has_header , ** kwargs ) <EOL> elif dialect == '<STR_LIT>' : <EOL> from . ssh import SSH <EOL> return append ( tbl , convert ( Temp ( SSH ( CSV ) ) , csv , ** kwargs ) , ** kwargs ) <EOL> kwargs = merge ( csv . dialect , kwargs ) <EOL> stmt = CopyFromCSV ( tbl , csv , bind = bind , ** kwargs ) <EOL> with bind . begin ( ) as conn : <EOL> conn . execute ( stmt ) <EOL> return tbl </s>
