psycopg2/examples/copy_from.py

   1 # copy_from.py -- example about copy_from
   2 #
   3 # Copyright (C) 2002 Tom Jenkins <tjenkins@devis.com>
   4 # Copyright (C) 2005 Federico Di Gregorio <fog@initd.org>
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by the
   8 # Free Software Foundation; either version 2, or (at your option) any later
   9 # version.
  10 #
  11 # This program is distributed in the hope that it will be useful, but
  12 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY
  13 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  14 # for more details.
  15 #
  16
  17 ## put in DSN your DSN string
  18
  19 DSN = 'dbname=test'
  20
  21 ## don't modify anything below tis line (except for experimenting)
  22
  23 import sys
  24 import os
  25 import StringIO
  26 import psycopg2
  27
  28 if len(sys.argv) > 1:
  29     DSN = sys.argv[1]
  30
  31 print "Opening connection using dns:", DSN
  32 conn = psycopg2.connect(DSN)
  33 print "Encoding for this connection is", conn.encoding
  34
  35 curs = conn.cursor()
  36 try:
  37     curs.execute("CREATE TABLE test_copy (fld1 text, fld2 text, fld3 int4)")
  38 except:
  39     conn.rollback()
  40     curs.execute("DROP TABLE test_copy")
  41     curs.execute("CREATE TABLE test_copy (fld1 text, fld2 text, fld3 int4)")
  42 conn.commit()
  43
  44 # copy_from with default arguments, from open file
  45
  46 io = open('copy_from.txt', 'wr')
  47 data = ['Tom\tJenkins\t37\n',
  48         'Madonna\t\N\t45\n',
  49         'Federico\tDi Gregorio\t\N\n']
  50 io.writelines(data)
  51 io.close()
  52
  53 io = open('copy_from.txt', 'r')
  54 curs.copy_from(io, 'test_copy')
  55 print "1) Copy %d records from file object " % len(data) + \
  56       "using defaults (sep: \\t and null = \\N)"
  57 io.close()
  58
  59 curs.execute("SELECT * FROM test_copy")
  60 rows = curs.fetchall()
  61 print "   Select returned %d rows" % len(rows)
  62
  63 for r in rows:
  64     print "    %s %s\t%s" % (r[0], r[1], r[2])
  65 curs.execute("delete from test_copy")
  66 conn.commit()
  67
  68 # copy_from using custom separator, from open file
  69
  70 io = open('copy_from.txt', 'wr')
  71 data = ['Tom:Jenkins:37\n',
  72         'Madonna:\N:45\n',
  73         'Federico:Di Gregorio:\N\n']
  74 io.writelines(data)
  75 io.close()
  76
  77 io = open('copy_from.txt', 'r')
  78 curs.copy_from(io, 'test_copy', ':')
  79 print "2) Copy %d records from file object using sep = :" % len(data)
  80 io.close()
  81
  82 curs.execute("SELECT * FROM test_copy")
  83 rows = curs.fetchall()
  84 print "   Select returned %d rows" % len(rows)
  85
  86 for r in rows:
  87     print "    %s %s\t%s" % (r[0], r[1], r[2])
  88 curs.execute("delete from test_copy")
  89 conn.commit()
  90
  91 # copy_from using custom null identifier, from open file
  92
  93 io = open('copy_from.txt', 'wr')
  94 data = ['Tom\tJenkins\t37\n',
  95         'Madonna\tNULL\t45\n',
  96         'Federico\tDi Gregorio\tNULL\n']
  97 io.writelines(data)
  98 io.close()
  99
 100 io = open('copy_from.txt', 'r')
 101 curs.copy_from(io, 'test_copy', null='NULL')
 102 print "3) Copy %d records from file object using null = NULL" % len(data)
 103 io.close()
 104
 105 curs.execute("SELECT * FROM test_copy")
 106 rows = curs.fetchall()
 107 print "   Select using cursor returned %d rows" % len(rows)
 108
 109 for r in rows:
 110     print "    %s %s\t%s" % (r[0], r[1], r[2])
 111 curs.execute("delete from test_copy")
 112 conn.commit()
 113
 114 # copy_from using custom separator and null identifier
 115
 116 io = open('copy_from.txt', 'wr')
 117 data = ['Tom:Jenkins:37\n', 'Madonna:NULL:45\n', 'Federico:Di Gregorio:NULL\n']
 118 io.writelines(data)
 119 io.close()
 120
 121 io = open('copy_from.txt', 'r')
 122 curs.copy_from(io, 'test_copy', ':', 'NULL')
 123 print "4) Copy %d records from file object " % len(data) + \
 124       "using sep = : and null = NULL"
 125 io.close()
 126
 127 curs.execute("SELECT * FROM test_copy")
 128 rows = curs.fetchall()
 129 print "   Select using cursor returned %d rows" % len(rows)
 130
 131 for r in rows:
 132     print "    %s %s\t%s" % (r[0], r[1], r[2])
 133 curs.execute("delete from test_copy")
 134 conn.commit()
 135
 136 # anything can be used as a file if it has .read() and .readline() methods
 137
 138 data = StringIO.StringIO()
 139 data.write('\n'.join(['Tom\tJenkins\t37',
 140                       'Madonna\t\N\t45',
 141                       'Federico\tDi Gregorio\t\N']))
 142 data.seek(0)
 143
 144 curs.copy_from(data, 'test_copy')
 145 print "5) Copy 3 records from StringIO object using defaults"
 146
 147 curs.execute("SELECT * FROM test_copy")
 148 rows = curs.fetchall()
 149 print "   Select using cursor returned %d rows" % len(rows)
 150
 151 for r in rows:
 152     print "    %s %s\t%s" % (r[0], r[1], r[2])
 153 curs.execute("delete from test_copy")
 154 conn.commit()
 155
 156 # simple error test
 157
 158 print "6) About to raise an error"
 159 data = StringIO.StringIO()
 160 data.write('\n'.join(['Tom\tJenkins\t37',
 161                       'Madonna\t\N\t45',
 162                       'Federico\tDi Gregorio\taaa']))
 163 data.seek(0)
 164
 165 try:
 166     curs.copy_from(data, 'test_copy')
 167 except StandardError, err:
 168     conn.rollback()
 169     print "   Catched error (as expected):\n", err
 170
 171 conn.rollback()
 172
 173 curs.execute("DROP TABLE test_copy")
 174 os.unlink('copy_from.txt')
 175 conn.commit()
 176
 177
 178