Project

General

Profile

« Previous | Next » 

Revision 1965

Added by Aaron Marcuse-Kubitza over 12 years ago

csv2db: Fixed bug where extra columns were not truncated in INSERT mode. Replace empty column names with the column # to avoid errors with CSVs that have trailing ","s, etc.

View differences:

csv2db
16 16 
import sql
17 17 
import streams
18 18 
import strings
19 
import util
19 20 

  		




  20
  21
  
    
def main():


  		




  21
  22
  
    
    # Usage


  		




  ......




  50
  51
  
    
        info = csvs.stream_info(in_, parse_header=True)


  		




  51
  52
  
    
        dialect = info.dialect


  		




  52
  53
  
    
        if csvs.is_tsv(dialect): use_copy_from[0] = False


  		




  
  54
  
    
        cols = info.header


  		




  
  55
  
    
        for i, col in enumerate(cols): # replace empty column names


  		




  
  56
  
    
            if col == '': cols[i] = 'column_'+str(i)


  		




  53
  57
  
    
        


  		




  54
  58
  
    
        # Select schema and escape names


  		




  55
  59
  
    
        def esc_name(name): return sql.esc_name(db, name, preserve_case=True)


  		




  56
  60
  
    
        sql.run_query(db, 'SET search_path TO '+esc_name(schema))


  		




  57
  61
  
    
        esc_table = esc_name(table)


  		




  58
  
  
    
        esc_cols = map(esc_name, info.header)


  		




  
  62
  
    
        esc_cols = map(esc_name, cols)


  		




  59
  63
  
    
        


  		




  60
  64
  
    
        # Create CREATE TABLE statement


  		




  61
  65
  
    
        pkey = esc_name(table+'_pkey')


  		




  62
  66
  
    
        create_table = 'CREATE TABLE '+esc_table+' (\n'


  		




  63
  
  
    
        create_table += '    row_num serial NOT NULL,\n'


  		




  64
  
  
    
        for esc_col in esc_cols: create_table += '    '+esc_col+' text,\n'


  		




  65
  
  
    
        create_table += '    CONSTRAINT '+pkey+' PRIMARY KEY (row_num)\n'


  		




  
  67
  
    
        create_table += '    row_num serial NOT NULL PRIMARY KEY\n'


  		




  
  68
  
    
        for esc_col in esc_cols: create_table += '    , '+esc_col+' text\n'


  		




  66
  69
  
    
        create_table += ');\n'


  		




  67
  70
  
    
        if debug: sys.stderr.write(create_table)


  		




  68
  71
  
    
        


  		




  ......




  93
  96
  
    
                db.db.cursor().copy_expert(copy_from, line_in)


  		




  94
  97
  
    
            else:


  		




  95
  98
  
    
                sys.stderr.write('Using INSERT\n')


  		




  
  99
  
    
                cols_ct = len(cols)+1 # +1 for row_num


  		




  96
  100
  
    
                for row in csvs.make_reader(line_in, dialect):


  		




  97
  101
  
    
                    row = map(strings.to_unicode, row)


  		




  98
  102
  
    
                    row.insert(0, sql.default) # leave space for autogen row_num


  		




  
  103
  
    
                    util.list_set_length(row, cols_ct) # truncate extra cols


  		




  99
  104
  
    
                    sql.insert(db, esc_table, row, table_is_esc=True)


  		




  100
  105
  
    
        finally:


  		




  101
  106
  
    
            line_in.close() # also closes proc.stdout


  		












Also available in:  Unified diff