use better SQL quoting in test cursor implementation

This commit is contained in:
Sarah Hoffmann
2026-02-12 17:30:25 +01:00
parent aa42dc8a93
commit 79682a94ce

View File

@@ -2,12 +2,13 @@
# #
# This file is part of Nominatim. (https://nominatim.org) # This file is part of Nominatim. (https://nominatim.org)
# #
# Copyright (C) 2025 by the Nominatim developer community. # Copyright (C) 2026 by the Nominatim developer community.
# For a full list of authors see the git log. # For a full list of authors see the git log.
""" """
Specialised psycopg cursor with shortcut functions useful for testing. Specialised psycopg cursor with shortcut functions useful for testing.
""" """
import psycopg import psycopg
from psycopg import sql as pysql
class CursorForTesting(psycopg.Cursor): class CursorForTesting(psycopg.Cursor):
@@ -52,7 +53,8 @@ class CursorForTesting(psycopg.Cursor):
def table_rows(self, table, where=None): def table_rows(self, table, where=None):
""" Return the number of rows in the given table. """ Return the number of rows in the given table.
""" """
if where is None: sql = pysql.SQL('SELECT count(*) FROM') + pysql.Identifier(table)
return self.scalar('SELECT count(*) FROM ' + table) if where is not None:
sql += pysql.SQL('WHERE') + pysql.SQL(where)
return self.scalar('SELECT count(*) FROM {} WHERE {}'.format(table, where)) return self.scalar(sql)