]> git.openstreetmap.org Git - nominatim.git/blob - src/nominatim_db/tools/exec_utils.py
contributions: some additional rules for AI use
[nominatim.git] / src / nominatim_db / tools / exec_utils.py
1 # SPDX-License-Identifier: GPL-3.0-or-later
2 #
3 # This file is part of Nominatim. (https://nominatim.org)
4 #
5 # Copyright (C) 2024 by the Nominatim developer community.
6 # For a full list of authors see the git log.
7 """
8 Helper functions for executing external programs.
9 """
10 from typing import Any, Mapping, List, Optional
11 import logging
12 import os
13 import re
14 import subprocess
15 import shutil
16
17 from ..db.connection import get_pg_env
18 from ..errors import UsageError
19 from ..version import OSM2PGSQL_REQUIRED_VERSION
20
21 LOG = logging.getLogger()
22
23
24 def run_osm2pgsql(options: Mapping[str, Any]) -> None:
25     """ Run osm2pgsql with the given options.
26     """
27     _check_osm2pgsql_version(options['osm2pgsql'])
28
29     env = get_pg_env(options['dsn'])
30
31     cmd = [_find_osm2pgsql_cmd(options['osm2pgsql']),
32            '--append' if options['append'] else '--create',
33            '--slim',
34            '--log-progress', 'true',
35            '--number-processes', '1' if options['append'] else str(options['threads']),
36            '--cache', str(options['osm2pgsql_cache']),
37            '--style', str(options['osm2pgsql_style'])
38           ]
39
40     if str(options['osm2pgsql_style']).endswith('.lua'):
41         env['LUA_PATH'] = ';'.join((str(options['osm2pgsql_style_path'] / '?.lua'),
42                                     os.environ.get('LUAPATH', ';')))
43         cmd.extend(('--output', 'flex'))
44
45         for flavour in ('data', 'index'):
46             if options['tablespaces'][f"main_{flavour}"]:
47                 env[f"NOMINATIM_TABLESPACE_PLACE_{flavour.upper()}"] = \
48                     options['tablespaces'][f"main_{flavour}"]
49     else:
50         cmd.extend(('--output', 'gazetteer', '--hstore', '--latlon'))
51         cmd.extend(_mk_tablespace_options('main', options))
52
53
54     if options['flatnode_file']:
55         cmd.extend(('--flat-nodes', options['flatnode_file']))
56
57     cmd.extend(_mk_tablespace_options('slim', options))
58
59     if options.get('disable_jit', False):
60         env['PGOPTIONS'] = '-c jit=off -c max_parallel_workers_per_gather=0'
61
62     if 'import_data' in options:
63         cmd.extend(('-r', 'xml', '-'))
64     elif isinstance(options['import_file'], list):
65         for fname in options['import_file']:
66             cmd.append(str(fname))
67     else:
68         cmd.append(str(options['import_file']))
69
70     subprocess.run(cmd, cwd=options.get('cwd', '.'),
71                    input=options.get('import_data'),
72                    env=env, check=True)
73
74
75 def _mk_tablespace_options(ttype: str, options: Mapping[str, Any]) -> List[str]:
76     cmds: List[str] = []
77     for flavour in ('data', 'index'):
78         if options['tablespaces'][f"{ttype}_{flavour}"]:
79             cmds.extend((f"--tablespace-{ttype}-{flavour}",
80                          options['tablespaces'][f"{ttype}_{flavour}"]))
81
82     return cmds
83
84
85 def _find_osm2pgsql_cmd(cmdline: Optional[str]) -> str:
86     if cmdline is not None:
87         return cmdline
88
89     in_path = shutil.which('osm2pgsql')
90     if in_path is None:
91         raise UsageError('osm2pgsql executable not found. Please install osm2pgsql first.')
92
93     return str(in_path)
94
95
96 def _check_osm2pgsql_version(cmdline: Optional[str]) -> None:
97     cmd = [_find_osm2pgsql_cmd(cmdline), '--version']
98
99     result = subprocess.run(cmd, capture_output=True, check=True)
100
101     if not result.stderr:
102         raise UsageError("osm2pgsql does not print version information.")
103
104     verinfo = result.stderr.decode('UTF-8')
105
106     match = re.search(r'osm2pgsql version (\d+)\.(\d+)', verinfo)
107     if match is None:
108         raise UsageError(f"No version information found in output: {verinfo}")
109
110     if (int(match[1]), int(match[2])) < OSM2PGSQL_REQUIRED_VERSION:
111         raise UsageError(f"osm2pgsql is too old. Found version {match[1]}.{match[2]}. "
112                          f"Need at least version {'.'.join(map(str, OSM2PGSQL_REQUIRED_VERSION))}.")