nominatim/db/sql_preprocessor.py

   1 """
   2 Preprocessing of SQL files.
   3 """
   4 import jinja2
   5
   6
   7 def _get_partitions(conn):
   8     """ Get the set of partitions currently in use.
   9     """
  10     with conn.cursor() as cur:
  11         cur.execute('SELECT DISTINCT partition FROM country_name')
  12         partitions = set([0])
  13         for row in cur:
  14             partitions.add(row[0])
  15
  16     return partitions
  17
  18
  19 def _get_tables(conn):
  20     """ Return the set of tables currently in use.
  21         Only includes non-partitioned
  22     """
  23     with conn.cursor() as cur:
  24         cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
  25
  26         return set((row[0] for row in list(cur)))
  27
  28
  29 def _setup_tablespace_sql(config):
  30     """ Returns a dict with tablespace expressions for the different tablespace
  31         kinds depending on whether a tablespace is configured or not.
  32     """
  33     out = {}
  34     for subset in ('ADDRESS', 'SEARCH', 'AUX'):
  35         for kind in ('DATA', 'INDEX'):
  36             tspace = getattr(config, 'TABLESPACE_{}_{}'.format(subset, kind))
  37             if tspace:
  38                 tspace = 'TABLESPACE "{}"'.format(tspace)
  39             out['{}_{}'.format(subset.lower, kind.lower())] = tspace
  40
  41     return out
  42
  43
  44 def _setup_postgres_sql(conn):
  45     """ Set up a dictionary with various Postgresql/Postgis SQL terms which
  46         are dependent on the database version in use.
  47     """
  48     out = {}
  49     pg_version = conn.server_version_tuple()
  50     # CREATE INDEX IF NOT EXISTS was introduced in PG9.5.
  51     # Note that you need to ignore failures on older versions when
  52     # using this construct.
  53     out['if_index_not_exists'] = ' IF NOT EXISTS ' if pg_version >= (9, 5, 0) else ''
  54
  55     return out
  56
  57
  58 def _setup_postgresql_features(conn):
  59     """ Set up a dictionary with various optional Postgresql/Postgis features that
  60         depend on the database version.
  61     """
  62     pg_version = conn.server_version_tuple()
  63     return {
  64         'has_index_non_key_column' : pg_version >= (11, 0, 0)
  65     }
  66
  67 class SQLPreprocessor: # pylint: disable=too-few-public-methods
  68     """ A environment for preprocessing SQL files from the
  69         lib-sql directory.
  70
  71         The preprocessor provides a number of default filters and variables.
  72         The variables may be overwritten when rendering an SQL file.
  73
  74         The preprocessing is currently based on the jinja2 templating library
  75         and follows its syntax.
  76     """
  77
  78     def __init__(self, conn, config, sqllib_dir):
  79         self.env = jinja2.Environment(autoescape=False,
  80                                       loader=jinja2.FileSystemLoader(str(sqllib_dir)))
  81
  82         db_info = {}
  83         db_info['partitions'] = _get_partitions(conn)
  84         db_info['tables'] = _get_tables(conn)
  85         db_info['reverse_only'] = 'search_name' not in db_info['tables']
  86         db_info['tablespace'] = _setup_tablespace_sql(config)
  87
  88         self.env.globals['config'] = config
  89         self.env.globals['db'] = db_info
  90         self.env.globals['sql'] = _setup_postgres_sql(conn)
  91         self.env.globals['postgres'] = _setup_postgresql_features(conn)
  92         self.env.globals['modulepath'] = config.DATABASE_MODULE_PATH or \
  93                                          str((config.project_dir / 'module').resolve())
  94
  95
  96     def run_sql_file(self, conn, name, **kwargs):
  97         """ Execute the given SQL file on the connection. The keyword arguments
  98             may supply additional parameters for preprocessing.
  99         """
 100         sql = self.env.get_template(name).render(**kwargs)
 101
 102         with conn.cursor() as cur:
 103             cur.execute(sql)
 104         conn.commit()