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 psycopg
  11 from psycopg import sql as pysql
  12
  13
  14 class DBManager:
  15
  16     def __init__(self, purge=False):
  17         self.purge = purge
  18
  19     def check_for_db(self, dbname):
  20         """ Check if the given DB already exists.
  21             When the purge option is set, then an existing database will
  22             be deleted and the function returns that it does not exist.
  23         """
  24         if self.purge:
  25             self.drop_db(dbname)
  26             return False
  27
  28         return self.exists_db(dbname)
  29
  30     def drop_db(self, dbname):
  31         """ Drop the given database if it exists.
  32         """
  33         with psycopg.connect(dbname='postgres') as conn:
  34             conn.autocommit = True
  35             conn.execute(pysql.SQL('DROP DATABASE IF EXISTS')
  36                          + pysql.Identifier(dbname))
  37
  38     def exists_db(self, dbname):
  39         """ Check if a database with the given name exists already.
  40         """
  41         with psycopg.connect(dbname='postgres') as conn:
  42             cur = conn.execute('select count(*) from pg_database where datname = %s',
  43                                (dbname,))
  44             return cur.fetchone()[0] == 1