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 from typing import Set, Dict, Any
  11 import jinja2
  12
  13 from nominatim.db.connection import Connection
  14 from nominatim.db.async_connection import WorkerPool
  15 from nominatim.config import Configuration
  16
  17 def _get_partitions(conn: Connection) -> Set[int]:
  18     """ Get the set of partitions currently in use.
  19     """
  20     with conn.cursor() as cur:
  21         cur.execute('SELECT DISTINCT partition FROM country_name')
  22         partitions = set([0])
  23         for row in cur:
  24             partitions.add(row[0])
  25
  26     return partitions
  27
  28
  29 def _get_tables(conn: Connection) -> Set[str]:
  30     """ Return the set of tables currently in use.
  31         Only includes non-partitioned
  32     """
  33     with conn.cursor() as cur:
  34         cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
  35
  36         return set((row[0] for row in list(cur)))
  37
  38
  39 def _setup_tablespace_sql(config: Configuration) -> Dict[str, str]:
  40     """ Returns a dict with tablespace expressions for the different tablespace
  41         kinds depending on whether a tablespace is configured or not.
  42     """
  43     out = {}
  44     for subset in ('ADDRESS', 'SEARCH', 'AUX'):
  45         for kind in ('DATA', 'INDEX'):
  46             tspace = getattr(config, f'TABLESPACE_{subset}_{kind}')
  47             if tspace:
  48                 tspace = f'TABLESPACE "{tspace}"'
  49             out[f'{subset.lower()}_{kind.lower()}'] = tspace
  50
  51     return out
  52
  53
  54 def _setup_postgresql_features(conn: Connection) -> Dict[str, Any]:
  55     """ Set up a dictionary with various optional Postgresql/Postgis features that
  56         depend on the database version.
  57     """
  58     pg_version = conn.server_version_tuple()
  59     postgis_version = conn.postgis_version_tuple()
  60     return {
  61         'has_index_non_key_column': pg_version >= (11, 0, 0),
  62         'spgist_geom' : 'SPGIST' if postgis_version >= (3, 0) else 'GIST'
  63     }
  64
  65 class SQLPreprocessor:
  66     """ A environment for preprocessing SQL files from the
  67         lib-sql directory.
  68
  69         The preprocessor provides a number of default filters and variables.
  70         The variables may be overwritten when rendering an SQL file.
  71
  72         The preprocessing is currently based on the jinja2 templating library
  73         and follows its syntax.
  74     """
  75
  76     def __init__(self, conn: Connection, config: Configuration) -> None:
  77         self.env = jinja2.Environment(autoescape=False,
  78                                       loader=jinja2.FileSystemLoader(str(config.lib_dir.sql)))
  79
  80         db_info: Dict[str, Any] = {}
  81         db_info['partitions'] = _get_partitions(conn)
  82         db_info['tables'] = _get_tables(conn)
  83         db_info['reverse_only'] = 'search_name' not in db_info['tables']
  84         db_info['tablespace'] = _setup_tablespace_sql(config)
  85
  86         self.env.globals['config'] = config
  87         self.env.globals['db'] = db_info
  88         self.env.globals['postgres'] = _setup_postgresql_features(conn)
  89
  90
  91     def run_sql_file(self, conn: Connection, name: str, **kwargs: Any) -> None:
  92         """ Execute the given SQL file on the connection. The keyword arguments
  93             may supply additional parameters for preprocessing.
  94         """
  95         sql = self.env.get_template(name).render(**kwargs)
  96
  97         with conn.cursor() as cur:
  98             cur.execute(sql)
  99         conn.commit()
 100
 101
 102     def run_parallel_sql_file(self, dsn: str, name: str, num_threads: int = 1,
 103                               **kwargs: Any) -> None:
 104         """ Execure the given SQL files using parallel asynchronous connections.
 105             The keyword arguments may supply additional parameters for
 106             preprocessing.
 107
 108             After preprocessing the SQL code is cut at lines containing only
 109             '---'. Each chunk is sent to one of the `num_threads` workers.
 110         """
 111         sql = self.env.get_template(name).render(**kwargs)
 112
 113         parts = sql.split('\n---\n')
 114
 115         with WorkerPool(dsn, num_threads) as pool:
 116             for part in parts:
 117                 pool.next_free_worker().perform(part)