add type annotations for SQL preprocessor

This commit is contained in:
Sarah Hoffmann
2022-07-05 11:24:53 +02:00
parent 26f30bff28
commit 0dff71a410
2 changed files with 15 additions and 8 deletions

View File

@@ -66,7 +66,11 @@ class Configuration:
self._config.update(dotenv_values(str((project_dir / '.env').resolve()))) self._config.update(dotenv_values(str((project_dir / '.env').resolve())))
class _LibDirs: class _LibDirs:
pass module: Path
osm2pgsql: Path
php: Path
sql: Path
data: Path
self.lib_dir = _LibDirs() self.lib_dir = _LibDirs()

View File

@@ -7,10 +7,13 @@
""" """
Preprocessing of SQL files. Preprocessing of SQL files.
""" """
from typing import Set, Dict, Any
import jinja2 import jinja2
from nominatim.db.connection import Connection
from nominatim.config import Configuration
def _get_partitions(conn): def _get_partitions(conn: Connection) -> Set[int]:
""" Get the set of partitions currently in use. """ Get the set of partitions currently in use.
""" """
with conn.cursor() as cur: with conn.cursor() as cur:
@@ -22,7 +25,7 @@ def _get_partitions(conn):
return partitions return partitions
def _get_tables(conn): def _get_tables(conn: Connection) -> Set[str]:
""" Return the set of tables currently in use. """ Return the set of tables currently in use.
Only includes non-partitioned Only includes non-partitioned
""" """
@@ -32,7 +35,7 @@ def _get_tables(conn):
return set((row[0] for row in list(cur))) return set((row[0] for row in list(cur)))
def _setup_tablespace_sql(config): def _setup_tablespace_sql(config: Configuration) -> Dict[str, str]:
""" Returns a dict with tablespace expressions for the different tablespace """ Returns a dict with tablespace expressions for the different tablespace
kinds depending on whether a tablespace is configured or not. kinds depending on whether a tablespace is configured or not.
""" """
@@ -47,7 +50,7 @@ def _setup_tablespace_sql(config):
return out return out
def _setup_postgresql_features(conn): def _setup_postgresql_features(conn: Connection) -> Dict[str, Any]:
""" Set up a dictionary with various optional Postgresql/Postgis features that """ Set up a dictionary with various optional Postgresql/Postgis features that
depend on the database version. depend on the database version.
""" """
@@ -69,11 +72,11 @@ class SQLPreprocessor:
and follows its syntax. and follows its syntax.
""" """
def __init__(self, conn, config): def __init__(self, conn: Connection, config: Configuration) -> None:
self.env = jinja2.Environment(autoescape=False, self.env = jinja2.Environment(autoescape=False,
loader=jinja2.FileSystemLoader(str(config.lib_dir.sql))) loader=jinja2.FileSystemLoader(str(config.lib_dir.sql)))
db_info = {} db_info: Dict[str, Any] = {}
db_info['partitions'] = _get_partitions(conn) db_info['partitions'] = _get_partitions(conn)
db_info['tables'] = _get_tables(conn) db_info['tables'] = _get_tables(conn)
db_info['reverse_only'] = 'search_name' not in db_info['tables'] db_info['reverse_only'] = 'search_name' not in db_info['tables']
@@ -84,7 +87,7 @@ class SQLPreprocessor:
self.env.globals['postgres'] = _setup_postgresql_features(conn) self.env.globals['postgres'] = _setup_postgresql_features(conn)
def run_sql_file(self, conn, name, **kwargs): def run_sql_file(self, conn: Connection, name: str, **kwargs: Any) -> None:
""" Execute the given SQL file on the connection. The keyword arguments """ Execute the given SQL file on the connection. The keyword arguments
may supply additional parameters for preprocessing. may supply additional parameters for preprocessing.
""" """