File: v.db.addcolumn.py

package info (click to toggle)
grass 8.4.2-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 277,040 kB
  • sloc: ansic: 460,798; python: 227,732; cpp: 42,026; sh: 11,262; makefile: 7,007; xml: 3,637; sql: 968; lex: 520; javascript: 484; yacc: 450; asm: 387; perl: 157; sed: 25; objc: 6; ruby: 4
file content (145 lines) | stat: -rwxr-xr-x 4,461 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/usr/bin/env python3
#
############################################################################
#
# MODULE:       v.db.addcolumnumn
# AUTHOR(S):    Moritz Lennert
#               Converted to Python by Glynn Clements
# PURPOSE:      interface to db.execute to add a column to the attribute table
#               connected to a given vector map
# COPYRIGHT:    (C) 2005 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: Adds one or more columns to the attribute table connected to a given vector map.
# % keyword: vector
# % keyword: attribute table
# % keyword: database
# %end

# %option G_OPT_V_MAP
# %end

# %option G_OPT_V_FIELD
# % label: Layer number where to add column(s)
# %end

# %option
# % key: columns
# % type: string
# % label: Name and type of the new column(s) ('name type [,name type, ...]')
# % description: Types depend on database backend, but all support VARCHAR(), INT, DOUBLE PRECISION and DATE. Example: 'label varchar(250), value integer'
# % required: yes
# % multiple: yes
# % key_desc: name type
# %end

import atexit
import os
import re

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

rm_files = []


def cleanup():
    for file in rm_files:
        if os.path.isfile(file):
            try:
                os.remove(file)
            except Exception as e:
                grass.warning(
                    _("Unable to remove file {file}: {message}").format(
                        file=file, message=e
                    )
                )


def main():
    global rm_files
    map = options["map"]
    layer = options["layer"]
    columns = options["columns"]
    columns = [col.strip() for col in columns.split(",")]

    # does map exist in CURRENT mapset?
    mapset = grass.gisenv()["MAPSET"]
    exists = bool(grass.find_file(map, element="vector", mapset=mapset)["file"])

    if not exists:
        grass.fatal(_("Vector map <{}> not found in current mapset").format(map))

    try:
        f = grass.vector_db(map)[int(layer)]
    except KeyError:
        if grass.vector_db(map):
            grass.fatal(
                _(
                    "There is no table connected to layer <{layer}> of <{name}>. "
                    "Run v.db.connect or v.db.addtable first."
                ).format(name=map, layer=layer)
            )
        grass.fatal(
            _(
                "There is no table connected to <{name}>. "
                "Run v.db.connect or v.db.addtable first."
            ).format(name=map)
        )

    table = f["table"]
    database = f["database"]
    driver = f["driver"]
    column_existing = grass.vector_columns(map, int(layer)).keys()

    add_str = "BEGIN TRANSACTION\n"
    pattern = re.compile(r"\s+")
    for col in columns:
        if not col:
            grass.fatal(_("There is an empty column. Did you leave a trailing comma?"))
        whitespace = re.search(pattern, col)
        if not whitespace:
            grass.fatal(
                _(
                    "Incorrect new column(s) format, use"
                    " <'name type [,name type, ...]'> format, please."
                )
            )
        col_name, col_type = col.split(whitespace.group(0), 1)
        if col_name in column_existing:
            grass.error(
                _("Column <{}> is already in the table. Skipping.").format(col_name)
            )
            continue
        grass.verbose(_("Adding column <{}> to the table").format(col_name))
        add_str += f'ALTER TABLE {table} ADD COLUMN "{col_name}" {col_type};\n'
    add_str += "END TRANSACTION"
    sql_file = grass.tempfile()
    rm_files.append(sql_file)
    cols_add_str = ",".join([col[0] for col in columns])
    with open(sql_file, "w") as write_file:
        write_file.write(add_str)
    try:
        grass.run_command(
            "db.execute",
            input=sql_file,
            database=database,
            driver=driver,
        )
    except CalledModuleError:
        grass.fatal(_("Error adding columns {}").format(cols_add_str))
    # write cmd history:
    grass.vector_history(map)


if __name__ == "__main__":
    options, flags = grass.parser()
    atexit.register(cleanup)
    main()