cynici · May 31, 2016 11:25 · May 31, 2016
diff --git a/pginsert-customize.py b/pginsert-customize.py
@@ -0,0 +1,103 @@
+#! /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())