1 #!/usr/bin/env python
   2 """
   3 Copyright (c) 2007, Muharem Hrnjadovic
   4 
   5 All rights reserved.
   6 
   7 Redistribution and use in source and binary forms, with or without
   8 modification, are permitted provided that the following conditions
   9 are met:
  10 
  11     * Redistributions of source code must retain the above copyright notice,
  12       this list of conditions and the following disclaimer.
  13     * Redistributions in binary form must reproduce the above copyright
  14       notice, this list of conditions and the following disclaimer in the
  15       documentation and/or other materials provided with the distribution.
  16     * Neither the name of Muharem Hrnjadovic nor the names of other
  17       contributors may be used to endorse or promote products derived from
  18       this software without specific prior written permission.
  19 
  20 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21 "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22 LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23 A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
  24 CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  25 EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  26 PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  27 PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  28 LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  29 NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  30 SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31 
  32 ---------------------------------------------------------------------------
  33 
  34 PostgreSQL tools
  35 """
  36 # created: Thu Oct 25 21:35:50 2007
  37 __version__ = "$Id:$"
  38 # $HeadURL $
  39 
  40 import psycopg2 as ps2
  41 
  42 class PGT(object):
  43     """Utility functions for using a PostgreSQL database with python"""
  44 
  45     def __init__(self):
  46         """initialiser"""
  47         # super(NEW_CLASS, self).__init__()
  48 
  49     @staticmethod
  50     def prepareByteaString(byteaSeq):
  51         """
  52         Given a sequence of byte arrays this function prepares a properly
  53         quoted string to be used for inserting database rows with Bytea
  54         arrays.
  55 
  56         Given e.g. a table like the following
  57             Create table baex(byteaa Bytea ARRAY[16]);
  58         the resulting string ('bas') can be used as follows:
  59             cursor.execute("INSERT INTO baex(byteaa) values('{%s}')" % bas))
  60 
  61         Parameters:
  62         - byteaSeq: a sequence of byte arrays each corresponding to a Bytea
  63                     value in the database
  64         Returns:
  65         string: containing all the byte arrays from the 'byteaSeq' properly
  66                 quoted for utilisation in an INSERT statement
  67         """
  68         # in a first step
  69         #   1. quote all the byte arrays (using the psycopg2 Binary()
  70         #      function)
  71         #   2. strip away the single quotes on the left and right side
  72         #   3. escape any double quote characters with a backslash
  73         # The last step is necessary because we will use the double quote as
  74         # the quote/delimiter for the byte arrays
  75         baSeq = [str(ps2.Binary(ba))[1:-1].replace(r'"', r'\"') for ba in byteaSeq]
  76         # join the prep'ed byte arrays into a single string
  77         bas = "\"%s\"" % '","'.join(baSeq)
  78         # double the number of backslashes (needed because we're inserting a
  79         # Bytea array as opposed to a single Bytea)
  80         bas = bas.replace('\\', '\\\\')
  81         # done!
  82         return(bas)
  83 
  84 if __name__ == '__main__':
  85     ### TEST code ********************************************************
  86     import os, sys
  87     from random import random as rand
  88 
  89     # connect to test database
  90     db = ps2.connect("dbname='test' user='postgres'")
  91     cursor = db.cursor()
  92     # create test table
  93     cursor.execute('Create Table public.baex(id Serial, byteaa Bytea ARRAY[16])')
  94 
  95     sys.stdout.flush()
  96     print "\n******************** Bytea data generated: ********************"
  97     # generate 3 byte array sequences
  98     for rowId in range(1, 4):
  99         byteaSeq = []
 100         # generate 1-7 random byte arrays
 101         for numOfSstrings in range(1, int(rand()*8)):
 102             bytea = ''.join([chr(int(rand()*256)) for x in range(int(rand()*5))])
 103             byteaSeq.append(bytea)
 104         print byteaSeq
 105         # get the INSERT string for the byte array sequence generated
 106         bas = PGT.prepareByteaString(byteaSeq)
 107         # insert the row into the table
 108         cursor.execute("INSERT INTO public.baex(id, byteaa) VALUES(%s, '{%s}')" \
 109                        % (rowId, bas))
 110     db.commit()
 111     sys.stdout.flush()
 112 
 113     print "\n******************** Bytea data inserted: ********************"
 114     sys.stdout.flush()
 115     # show the data inserted
 116     os.system("psql -d test -U postgres -c 'SELECT * FROM public.baex'")
 117     cursor.execute('DROP TABLE public.baex')
 118     db.commit()