]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/db/sql_preprocessor.py
Merge pull request #2708 from lonvia/use-format-literals
[nominatim.git] / nominatim / db / sql_preprocessor.py
1 # SPDX-License-Identifier: GPL-2.0-only
2 #
3 # This file is part of Nominatim. (https://nominatim.org)
4 #
5 # Copyright (C) 2022 by the Nominatim developer community.
6 # For a full list of authors see the git log.
7 """
8 Preprocessing of SQL files.
9 """
10 import jinja2
11
12
13 def _get_partitions(conn):
14     """ Get the set of partitions currently in use.
15     """
16     with conn.cursor() as cur:
17         cur.execute('SELECT DISTINCT partition FROM country_name')
18         partitions = set([0])
19         for row in cur:
20             partitions.add(row[0])
21
22     return partitions
23
24
25 def _get_tables(conn):
26     """ Return the set of tables currently in use.
27         Only includes non-partitioned
28     """
29     with conn.cursor() as cur:
30         cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
31
32         return set((row[0] for row in list(cur)))
33
34
35 def _setup_tablespace_sql(config):
36     """ Returns a dict with tablespace expressions for the different tablespace
37         kinds depending on whether a tablespace is configured or not.
38     """
39     out = {}
40     for subset in ('ADDRESS', 'SEARCH', 'AUX'):
41         for kind in ('DATA', 'INDEX'):
42             tspace = getattr(config, f'TABLESPACE_{subset}_{kind}')
43             if tspace:
44                 tspace = f'TABLESPACE "{tspace}"'
45             out[f'{subset.lower()}_{kind.lower()}'] = tspace
46
47     return out
48
49
50 def _setup_postgresql_features(conn):
51     """ Set up a dictionary with various optional Postgresql/Postgis features that
52         depend on the database version.
53     """
54     pg_version = conn.server_version_tuple()
55     postgis_version = conn.postgis_version_tuple()
56     return {
57         'has_index_non_key_column': pg_version >= (11, 0, 0),
58         'spgist_geom' : 'SPGIST' if postgis_version >= (3, 0) else 'GIST'
59     }
60
61 class SQLPreprocessor:
62     """ A environment for preprocessing SQL files from the
63         lib-sql directory.
64
65         The preprocessor provides a number of default filters and variables.
66         The variables may be overwritten when rendering an SQL file.
67
68         The preprocessing is currently based on the jinja2 templating library
69         and follows its syntax.
70     """
71
72     def __init__(self, conn, config):
73         self.env = jinja2.Environment(autoescape=False,
74                                       loader=jinja2.FileSystemLoader(str(config.lib_dir.sql)))
75
76         db_info = {}
77         db_info['partitions'] = _get_partitions(conn)
78         db_info['tables'] = _get_tables(conn)
79         db_info['reverse_only'] = 'search_name' not in db_info['tables']
80         db_info['tablespace'] = _setup_tablespace_sql(config)
81
82         self.env.globals['config'] = config
83         self.env.globals['db'] = db_info
84         self.env.globals['postgres'] = _setup_postgresql_features(conn)
85
86
87     def run_sql_file(self, conn, name, **kwargs):
88         """ Execute the given SQL file on the connection. The keyword arguments
89             may supply additional parameters for preprocessing.
90         """
91         sql = self.env.get_template(name).render(**kwargs)
92
93         with conn.cursor() as cur:
94             cur.execute(sql)
95         conn.commit()