Skip to content

Instantly share code, notes, and snippets.

@cynici
Created May 31, 2016 11:25
Show Gist options
  • Save cynici/a87d145ecf4bdb34c8d2c9f64a775125 to your computer and use it in GitHub Desktop.
Save cynici/a87d145ecf4bdb34c8d2c9f64a775125 to your computer and use it in GitHub Desktop.
Customise raster2pgsql, shp2pgsql INSERT statements
#! /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())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment