test/python/mock_legacy_word_table.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) 2024 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Legacy word table for testing with functions to prefil and test contents
   9 of the table.
  10 """
  11 from nominatim_db.db.connection import execute_scalar
  12
  13 class MockLegacyWordTable:
  14     """ A word table for testing using legacy word table structure.
  15     """
  16     def __init__(self, conn):
  17         self.conn = conn
  18         with conn.cursor() as cur:
  19             cur.execute("""CREATE TABLE word (word_id INTEGER,
  20                                               word_token text,
  21                                               word text,
  22                                               class text,
  23                                               type text,
  24                                               country_code varchar(2),
  25                                               search_name_count INTEGER,
  26                                               operator TEXT)""")
  27
  28         conn.commit()
  29
  30     def add_full_word(self, word_id, word, word_token=None):
  31         with self.conn.cursor() as cur:
  32             cur.execute("""INSERT INTO word (word_id, word_token, word)
  33                            VALUES (%s, %s, %s)
  34                         """, (word_id, ' ' + (word_token or word), word))
  35         self.conn.commit()
  36
  37
  38     def add_special(self, word_token, word, cls, typ, oper):
  39         with self.conn.cursor() as cur:
  40             cur.execute("""INSERT INTO word (word_token, word, class, type, operator)
  41                               VALUES (%s, %s, %s, %s, %s)
  42                         """, (word_token, word, cls, typ, oper))
  43         self.conn.commit()
  44
  45
  46     def add_country(self, country_code, word_token):
  47         with self.conn.cursor() as cur:
  48             cur.execute("INSERT INTO word (word_token, country_code) VALUES(%s, %s)",
  49                         (word_token, country_code))
  50         self.conn.commit()
  51
  52
  53     def add_postcode(self, word_token, postcode):
  54         with self.conn.cursor() as cur:
  55             cur.execute("""INSERT INTO word (word_token, word, class, type)
  56                               VALUES (%s, %s, 'place', 'postcode')
  57                         """, (word_token, postcode))
  58         self.conn.commit()
  59
  60
  61     def count(self):
  62         return execute_scalar(self.conn, "SELECT count(*) FROM word")
  63
  64
  65     def count_special(self):
  66         return execute_scalar(self.conn, "SELECT count(*) FROM word WHERE class != 'place'")
  67
  68
  69     def get_special(self):
  70         with self.conn.cursor() as cur:
  71             cur.execute("""SELECT word_token, word, class, type, operator
  72                            FROM word WHERE class != 'place'""")
  73             result = set((tuple(row) for row in cur))
  74             assert len(result) == cur.rowcount, "Word table has duplicates."
  75             return result
  76
  77
  78     def get_country(self):
  79         with self.conn.cursor() as cur:
  80             cur.execute("""SELECT country_code, word_token
  81                            FROM word WHERE country_code is not null""")
  82             result = set((tuple(row) for row in cur))
  83             assert len(result) == cur.rowcount, "Word table has duplicates."
  84             return result
  85
  86
  87     def get_postcodes(self):
  88         with self.conn.cursor() as cur:
  89             cur.execute("""SELECT word FROM word
  90                            WHERE class = 'place' and type = 'postcode'""")
  91             return set((row[0] for row in cur))
  92
  93     def get_partial_words(self):
  94         with self.conn.cursor() as cur:
  95             cur.execute("""SELECT word_token, search_name_count FROM word
  96                            WHERE class is null and country_code is null
  97                                  and not word_token like ' %'""")
  98             return set((tuple(row) for row in cur))
  99