Geografic-Information-System/scripts/db.dropcolumn/db.dropcolumn.py

#!/usr/bin/env python3

############################################################################
#
# MODULE:       db.dropcolumn
# AUTHOR(S):   	Markus Neteler
#               Converted to Python by Glynn Clements
# PURPOSE:      Interface to db.execute to drop a column from an
#               attribute table
#               - with special trick for SQLite
# COPYRIGHT:    (C) 2007, 2012 by Markus Neteler and the GRASS Development Team
#
#               This program is free software under the GNU General
#               Public License (>=v2). Read the file COPYING that
#               comes with GRASS for details.
#
#############################################################################

#%module
#% description: Drops a column from selected attribute table.
#% keyword: database
#% keyword: attribute table
#%End

#%flag
#% key: f
#% description: Force removal (required for actual deletion of files)
#%end

#%option G_OPT_DB_TABLE
#% required : yes
#%end

#%option G_OPT_DB_COLUMN
#% required : yes
#%end

import sys
import string

from grass.exceptions import CalledModuleError
import grass.script as gscript


def main():
    table = options['table']
    column = options['column']
    force = flags['f']

    # check if DB parameters are set, and if not set them.
    gscript.run_command('db.connect', flags='c')

    kv = gscript.db_connection()
    database = kv['database']
    driver = kv['driver']
    # schema needed for PG?

    if force:
        gscript.message(_("Forcing ..."))

    if column == "cat":
        gscript.warning(_("Deleting <%s> column which may be needed to keep "
                          "table connected to a vector map") % column)

    cols = [f[0] for f in gscript.db_describe(table)['cols']]
    if column not in cols:
        gscript.fatal(_("Column <%s> not found in table") % column)

    if not force:
        gscript.message(_("Column <%s> would be deleted.") % column)
        gscript.message("")
        gscript.message(_("You must use the force flag (-f) to actually "
                          "remove it. Exiting."))
        return 0

    if driver == "sqlite":
        # echo "Using special trick for SQLite"
        # http://www.sqlite.org/faq.html#q13
        colnames = []
        coltypes = []
        for f in gscript.db_describe(table)['cols']:
            if f[0] != column:
                colnames.append(f[0])
                coltypes.append("%s %s" % (f[0], f[1]))

        colnames = ", ".join(colnames)
        coltypes = ", ".join(coltypes)

        cmds = [
            "BEGIN TRANSACTION",
            "CREATE TEMPORARY TABLE ${table}_backup(${coldef})",
            "INSERT INTO ${table}_backup SELECT ${colnames} FROM ${table}",
            "DROP TABLE ${table}",
            "CREATE TABLE ${table}(${coldef})",
            "INSERT INTO ${table} SELECT ${colnames} FROM ${table}_backup",
            "DROP TABLE ${table}_backup",
            "COMMIT"
        ]
        tmpl = string.Template(';\n'.join(cmds))
        sql = tmpl.substitute(table=table, coldef=coltypes, colnames=colnames)
    else:
        sql = "ALTER TABLE %s DROP COLUMN %s" % (table, column)

    try:
        gscript.write_command('db.execute', input='-', database=database,
                              driver=driver, stdin=sql)
    except CalledModuleError:
        gscript.fatal(_("Cannot continue (problem deleting column)"))

    return 0

if __name__ == "__main__":
    options, flags = gscript.parser()
    sys.exit(main())