Customise raster2pgsql, shp2pgsql INSERT statements

pginsert-customize.py

#! /usr/bin/env python
import sys
import os
import argparse
import re
import logging

help_text = """Get input from stdin and modify INSERT statements. Useful for customizing output from PostGIS shp2pgsql, raster2pgsql."""

def modify_statement(st, args, clist, vlist, do_delete=False):
    pat = re.compile('\s*INSERT INTO\s+(\S+)\s+\(([^\)]+)\)\s+VALUES\s+\(([^\)]+)\)\s*;\s*', re.IGNORECASE)
    match = pat.match(st)
    if match is None: return st
    otable, oclist, ovlist = match.groups()
    oclist = oclist.split(',')
    ovlist = ovlist.split(',')
    drop_indices = []
    if args.drop:
        for i, column in enumerate(oclist):
            column = column.translate(None, ''.join(["'", '"']))
            if column in args.drop:
                drop_indices.append(i)
    clist = [c for i, c in enumerate(oclist) if i not in drop_indices] + clist 
    vlist = [c for i, c in enumerate(ovlist) if i not in drop_indices] + vlist 
    sql_del = ''
    if do_delete:
        sql_del = 'DELETE FROM {} WHERE {}'.format(
            otable,
            " AND ".join(
            ['"%s" = %s'%(clist[i], vlist[i]) for i, c in enumerate(vlist)]
        ))
        if args.meta:
            sql_del += ' AND ST_Metadata(rast) = %s' % args.meta
        sql_del += ';\n'
    return '''{}INSERT INTO {} ({}) VALUES ({});'''.format(
        sql_del,
        otable,
        ','.join(clist),
        ','.join(vlist),
    )

def main(argv=None):
    if argv is None:
        argv = sys.argv

    debuglevelD = {
        'debug': logging.DEBUG,
        'info': logging.INFO,
        'warning': logging.WARNING,
        'error': logging.ERROR,
        'critical': logging.CRITICAL
    }

    defvals = {}
    parser = argparse.ArgumentParser(description=help_text)
    parser.add_argument("-A", "--add", action="append", nargs=2, metavar="COLUMN VALUE",
        help='''Column name should be quotes '"name"'. String value must be quoted "'example string value'"''')
    parser.add_argument("-f", "--field", default='rast', metavar="COLUMN",
        help="specify final raster column name")
    parser.add_argument("-d", "--delete", action="store_true",
        help="Delete before insert")
    parser.add_argument("-m", "--meta", metavar='RASTER_META',
        help="Used in WHERE-clause for delete")
    parser.add_argument("-D", "--drop", action="append", help="specify column(s) to omit, unquoted.")
    parser.add_argument("-l", "--loglevel", help="Verbosity %s"%debuglevelD.keys(), metavar='LOGLEVEL')
    parser.set_defaults(**defvals)
    args = parser.parse_args()
    if args.loglevel:
        if args.loglevel not in debuglevelD: raise AssertionError("Verbosity level must be one of: %s"%debuglevelD.keys())
        dbglvl = debuglevelD[args.loglevel]
    else:
        dbglvl = logging.WARNING
    logger = logging.getLogger(__name__)
    logger.setLevel(dbglvl)
    ch = logging.StreamHandler()
    ch.setFormatter( logging.Formatter('%(asctime)s %(lineno)d %(name)s %(funcName)s %(message)s') )
    ch.setLevel(dbglvl)
    logger.addHandler(ch)

    if args.add:
        clist, vlist = zip(*args.add)
        clist = list(clist)
        clist = list(clist)
        vlist = list(vlist)
    else:
        clist = []
        vlist = []
    #
    # INSERT INTO "test"."w_ecmwf" ( rast ) VALUES ( ('...') )
    #
    pattern = re.compile(r'\(\s*"%s"\s*\)\s+VALUES\s*\(' % args.field)
    replacement = '( %s ) VALUES ( %s,' % (', '.join('"%s"'%c for c in clist), ', '.join(vlist))
    logger.debug('pattern: %s' % pattern)
    logger.debug('replacement: %s' % replacement)
    do_delete = args.delete
    for line in sys.stdin:
        line2 = modify_statement(line, args, clist, vlist, do_delete=do_delete)
        do_delete = False
        print line2
    return 0

if __name__ == "__main__":
    sys.exit(main())