test/bdd/utils/db.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) 2025 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Helper functions for managing test databases.
   9 """
  10 import asyncio
  11 import psycopg
  12 from psycopg import sql as pysql
  13
  14 from nominatim_db.tools.database_import import setup_database_skeleton, create_tables, \
  15                                                create_partition_tables, create_search_indices
  16 from nominatim_db.data.country_info import setup_country_tables, create_country_names
  17 from nominatim_db.tools.refresh import create_functions, load_address_levels_from_config
  18 from nominatim_db.tools.exec_utils import run_osm2pgsql
  19 from nominatim_db.tokenizer import factory as tokenizer_factory
  20
  21
  22 class DBManager:
  23
  24     def __init__(self, purge=False):
  25         self.purge = purge
  26
  27     def check_for_db(self, dbname):
  28         """ Check if the given DB already exists.
  29             When the purge option is set, then an existing database will
  30             be deleted and the function returns that it does not exist.
  31         """
  32         if self.purge:
  33             self.drop_db(dbname)
  34             return False
  35
  36         return self.exists_db(dbname)
  37
  38     def drop_db(self, dbname):
  39         """ Drop the given database if it exists.
  40         """
  41         with psycopg.connect(dbname='postgres') as conn:
  42             conn.autocommit = True
  43             conn.execute(pysql.SQL('DROP DATABASE IF EXISTS')
  44                          + pysql.Identifier(dbname))
  45
  46     def exists_db(self, dbname):
  47         """ Check if a database with the given name exists already.
  48         """
  49         with psycopg.connect(dbname='postgres') as conn:
  50             cur = conn.execute('select count(*) from pg_database where datname = %s',
  51                                (dbname,))
  52             return cur.fetchone()[0] == 1
  53
  54     def create_db_from_template(self, dbname, template):
  55         """ Create a new database from the given template database.
  56             Any existing database with the same name will be dropped.
  57         """
  58         with psycopg.connect(dbname='postgres') as conn:
  59             conn.autocommit = True
  60             conn.execute(pysql.SQL('DROP DATABASE IF EXISTS')
  61                          + pysql.Identifier(dbname))
  62             conn.execute(pysql.SQL('CREATE DATABASE {} WITH TEMPLATE {}')
  63                               .format(pysql.Identifier(dbname),
  64                                       pysql.Identifier(template)))
  65
  66     def setup_template_db(self, config):
  67         """ Create a template DB which contains the necessary extensions
  68             and basic static tables.
  69
  70             The template will only be created if the database does not yet
  71             exist or 'purge' is set.
  72         """
  73         dsn = config.get_libpq_dsn()
  74
  75         if self.check_for_db(config.get_database_params()['dbname']):
  76             return
  77
  78         setup_database_skeleton(dsn)
  79
  80         run_osm2pgsql(dict(osm2pgsql='osm2pgsql',
  81                            osm2pgsql_cache=1,
  82                            osm2pgsql_style=str(config.get_import_style_file()),
  83                            osm2pgsql_style_path=config.lib_dir.lua,
  84                            threads=1,
  85                            dsn=dsn,
  86                            flatnode_file='',
  87                            tablespaces=dict(slim_data='', slim_index='',
  88                                             main_data='', main_index=''),
  89                            append=False,
  90                            import_data=b'<osm version="0.6"></osm>'))
  91
  92         setup_country_tables(dsn, config.lib_dir.data)
  93
  94         with psycopg.connect(dsn) as conn:
  95             create_tables(conn, config)
  96             load_address_levels_from_config(conn, config)
  97             create_partition_tables(conn, config)
  98             create_functions(conn, config, enable_diff_updates=False)
  99             asyncio.run(create_search_indices(conn, config))
 100
 101             tokenizer = tokenizer_factory.create_tokenizer(config)
 102             create_country_names(conn, tokenizer)