osm-search/Nominatim
1
1
Fork 0
mirror of https://github.com/osm-search/Nominatim.git synced 2024-12-25 05:52:32 +03:00
Code Issues Packages Projects Releases Wiki Activity

type annotations for DB utils

Browse Source
This commit is contained in:
Sarah Hoffmann 2022-07-02 10:18:10 +02:00
parent e6ee3c772c
commit 992e6f72cf
2 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nominatim/db/connection.py
View File

@ -226,7 +226,7 @@ _PG_CONNECTION_STRINGS = {
def get_pg_env(dsn: str, def get_pg_env(dsn: str,
base_env: Optional[Mapping[str, Optional[str]]] = None) -> Mapping[str, Optional[str]]: base_env: Optional[Mapping[str, str]] = None) -> Mapping[str, str]:
""" Return a copy of `base_env` with the environment variables for """ Return a copy of `base_env` with the environment variables for
PostgresSQL set up from the given database connection string. PostgresSQL set up from the given database connection string.
If `base_env` is None, then the OS environment is used as a base If `base_env` is None, then the OS environment is used as a base

12
nominatim/db/utils.py
View File

@ -7,17 +7,21 @@
""" """
Helper functions for handling DB accesses. Helper functions for handling DB accesses.
""" """
from typing import IO, Optional, Union
import subprocess import subprocess
import logging import logging
import gzip import gzip
import io import io
from pathlib import Path
from nominatim.db.connection import get_pg_env from nominatim.db.connection import get_pg_env
from nominatim.errors import UsageError from nominatim.errors import UsageError
LOG = logging.getLogger() LOG = logging.getLogger()
def _pipe_to_proc(proc, fdesc): def _pipe_to_proc(proc: subprocess.Popen[bytes],
fdesc: Union[IO[bytes], gzip.GzipFile]) -> int:
assert proc.stdin is not None
chunk = fdesc.read(2048) chunk = fdesc.read(2048)
while chunk and proc.poll() is None: while chunk and proc.poll() is None:
try: try:
@ -28,7 +32,10 @@ def _pipe_to_proc(proc, fdesc):
return len(chunk) return len(chunk)
def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None): def execute_file(dsn: str, fname: Path,
ignore_errors: bool = False,
pre_code: Optional[str] = None,
post_code: Optional[str] = None) -> None:
""" Read an SQL file and run its contents against the given database """ Read an SQL file and run its contents against the given database
using psql. Use `pre_code` and `post_code` to run extra commands using psql. Use `pre_code` and `post_code` to run extra commands
before or after executing the file. The commands are run within the before or after executing the file. The commands are run within the
@ -42,6 +49,7 @@ def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None)
cmd.append('--quiet') cmd.append('--quiet')
with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc: with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc:
assert proc.stdin is not None
try: try:
if not LOG.isEnabledFor(logging.INFO): if not LOG.isEnabledFor(logging.INFO):
proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8')) proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))