type annotations for DB utils

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

View File

@ -226,7 +226,7 @@ _PG_CONNECTION_STRINGS = {
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
PostgresSQL set up from the given database connection string.
If `base_env` is None, then the OS environment is used as a base

View File

@ -7,17 +7,21 @@
"""
Helper functions for handling DB accesses.
"""
from typing import IO, Optional, Union
import subprocess
import logging
import gzip
import io
from pathlib import Path
from nominatim.db.connection import get_pg_env
from nominatim.errors import UsageError
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)
while chunk and proc.poll() is None:
try:
@ -28,7 +32,10 @@ def _pipe_to_proc(proc, fdesc):
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
using psql. Use `pre_code` and `post_code` to run extra commands
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')
with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc:
assert proc.stdin is not None
try:
if not LOG.isEnabledFor(logging.INFO):
proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))