Geografic-Information-System/scripts/v.db.reconnect.all/v.db.reconnect.all.py

#!/usr/bin/env python3
############################################################################
#
# MODULE:       v.db.reconnect.all
# AUTHOR(S):    Radim Blazek
#               Converted to Python by Glynn Clements
#               Update for GRASS 7 by Markus Metz
# PURPOSE:      Reconnect all vector maps from the current mapset
# COPYRIGHT:    (C) 2004, 2012 by 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: Reconnects attribute tables for all vector maps from the current mapset to a new database.
# % keyword: vector
# % keyword: attribute table
# % keyword: database
# %end
# %flag
# % key: c
# % description: Copy attribute tables to the target database if not exist
# %end
# %flag
# % key: d
# % description: Delete attribute tables from the source database
# %end
# %option G_OPT_DB_DATABASE
# % key: old_database
# % description: Name of source database
# %end
# %option G_OPT_DB_SCHEMA
# % key: old_schema
# % label: Name of source database schema
# %end
# %option
# % key: new_driver
# % description: Name for target driver
# %end
# %option G_OPT_DB_DATABASE
# % key: new_database
# % description: Name for target database
# %end
# %option G_OPT_DB_SCHEMA
# % key: new_schema
# % label: Name for target database schema
# %end

import sys
import os
import string

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

# substitute variables (gisdbase, location_name, mapset)


def substitute_db(database):
    gisenv = gscript.gisenv()
    tmpl = string.Template(database)

    return tmpl.substitute(
        GISDBASE=gisenv["GISDBASE"],
        LOCATION_NAME=gisenv["LOCATION_NAME"],
        MAPSET=gisenv["MAPSET"],
    )


# create database if doesn't exist


def create_db(driver, database):
    subst_database = substitute_db(database)
    if driver == "dbf":
        path = subst_database
        # check if destination directory exists
        if not os.path.isdir(path):
            # create dbf database
            os.makedirs(path)
            return True
        return False

    if driver == "sqlite":
        path = os.path.dirname(subst_database)
        # check if destination directory exists
        if not os.path.isdir(path):
            os.makedirs(path)

    if (
        subst_database
        in gscript.read_command("db.databases", quiet=True, driver=driver).splitlines()
    ):
        return False

    gscript.info(
        _(
            "Target database doesn't exist, "
            "creating a new database using <%s> driver..."
        )
        % driver
    )
    try:
        gscript.run_command("db.createdb", driver=driver, database=subst_database)
    except CalledModuleError:
        gscript.fatal(
            _("Unable to create database <%s> by driver <%s>")
            % (subst_database, driver)
        )

    return False


# copy tables if required (-c)


def copy_tab(from_driver, from_database, from_table, to_driver, to_database, to_table):
    if (
        to_table
        in gscript.read_command(
            "db.tables",
            quiet=True,
            driver=to_driver,
            database=to_database,
            stderr=nuldev,
        ).splitlines()
    ):
        return False

    gscript.info("Copying table <%s> to target database..." % to_table)
    try:
        gscript.run_command(
            "db.copy",
            from_driver=from_driver,
            from_database=from_database,
            from_table=from_table,
            to_driver=to_driver,
            to_database=to_database,
            to_table=to_table,
        )
    except CalledModuleError:
        gscript.fatal(_("Unable to copy table <%s>") % from_table)

    return True


# drop tables if required (-d)


def drop_tab(vector, layer, table, driver, database):
    # disconnect
    try:
        gscript.run_command(
            "v.db.connect", flags="d", quiet=True, map=vector, layer=layer, table=table
        )
    except CalledModuleError:
        gscript.warning(
            _("Unable to disconnect table <%s> from vector <%s>") % (table, vector)
        )
    # drop table
    try:
        gscript.run_command(
            "db.droptable",
            quiet=True,
            flags="f",
            driver=driver,
            database=database,
            table=table,
        )
    except CalledModuleError:
        gscript.fatal(_("Unable to drop table <%s>") % table)


# create index on key column


def create_index(driver, database, table, index_name, key):
    if driver == "dbf":
        return False

    gscript.info(_("Creating index <%s>...") % index_name)
    try:
        gscript.run_command(
            "db.execute",
            quiet=True,
            driver=driver,
            database=database,
            sql="create unique index %s on %s(%s)" % (index_name, table, key),
        )
    except CalledModuleError:
        gscript.warning(_("Unable to create index <%s>") % index_name)


def main():
    # old connection
    old_database = options["old_database"]
    old_schema = options["old_schema"]
    # new connection
    default_connection = gscript.db_connection()
    if options["new_driver"]:
        new_driver = options["new_driver"]
    else:
        new_driver = default_connection["driver"]
    if options["new_database"]:
        new_database = options["new_database"]
    else:
        new_database = default_connection["database"]
    if options["new_schema"]:
        new_schema = options["new_schema"]
    else:
        new_schema = default_connection["schema"]

    if old_database == "":
        old_database = None
    old_database_subst = None
    if old_database is not None:
        old_database_subst = substitute_db(old_database)

    new_database_subst = substitute_db(new_database)

    if old_database_subst == new_database_subst and old_schema == new_schema:
        gscript.fatal(
            _("Old and new database connection is identical. " "Nothing to do.")
        )

    mapset = gscript.gisenv()["MAPSET"]

    vectors = gscript.list_grouped("vect")[mapset]
    num_vectors = len(vectors)

    if flags["c"]:
        # create new database if not existing
        create_db(new_driver, new_database)

    i = 0
    for vect in vectors:
        vect = "%s@%s" % (vect, mapset)
        i += 1
        gscript.message(
            _("%s\nReconnecting vector map <%s> " "(%d of %d)...\n%s")
            % ("-" * 80, vect, i, num_vectors, "-" * 80)
        )
        for f in gscript.vector_db(vect, stderr=nuldev).values():
            layer = f["layer"]
            schema_table = f["table"]
            key = f["key"]
            database = f["database"]
            driver = f["driver"]

            # split schema.table
            if "." in schema_table:
                schema, table = schema_table.split(".", 1)
            else:
                schema = ""
                table = schema_table

            if new_schema:
                new_schema_table = "%s.%s" % (new_schema, table)
            else:
                new_schema_table = table

            gscript.debug(
                "DATABASE = '%s' SCHEMA = '%s' TABLE = '%s' ->\n"
                "      NEW_DATABASE = '%s' NEW_SCHEMA_TABLE = '%s'"
                % (old_database, schema, table, new_database, new_schema_table)
            )

            do_reconnect = True
            if old_database_subst is not None:
                if database != old_database_subst:
                    do_reconnect = False
            if database == new_database_subst:
                do_reconnect = False
            if schema != old_schema:
                do_reconnect = False

            if do_reconnect:
                gscript.verbose(_("Reconnecting layer %d...") % layer)

                if flags["c"]:
                    # check if table exists in new database
                    copy_tab(
                        driver,
                        database,
                        schema_table,
                        new_driver,
                        new_database,
                        new_schema_table,
                    )

                # drop original table if required
                if flags["d"]:
                    drop_tab(vect, layer, schema_table, driver, substitute_db(database))

                # reconnect tables (don't use substituted new_database)
                # NOTE: v.db.connect creates an index on the key column
                try:
                    gscript.run_command(
                        "v.db.connect",
                        flags="o",
                        quiet=True,
                        map=vect,
                        layer=layer,
                        driver=new_driver,
                        database=new_database,
                        table=new_schema_table,
                        key=key,
                    )
                except CalledModuleError:
                    gscript.warning(
                        _(
                            "Unable to connect table <%s> to vector "
                            "<%s> on layer <%s>"
                        )
                        % (table, vect, str(layer))
                    )

            else:
                if database != new_database_subst:
                    gscript.warning(
                        _(
                            "Layer <%d> will not be reconnected "
                            "because database or schema do not "
                            "match."
                        )
                        % layer
                    )
    return 0


if __name__ == "__main__":
    options, flags = gscript.parser()
    nuldev = open(os.devnull, "w")
    sys.exit(main())