diff --git a/docker/db.py b/docker/database.py similarity index 83% rename from docker/db.py rename to docker/database.py index 7361a0a0..d5002baa 100644 --- a/docker/db.py +++ b/docker/database.py @@ -11,8 +11,12 @@ import sys import subprocess import time +from typing import cast import psycopg +from psycopg.abc import Query +from psycopg import sql import sh + from urllib import parse import qgis_styles, helpers @@ -67,6 +71,13 @@ def connection_string(admin: bool=False) -> str: return conn_string +def set_db_env_vars(): + """Set the Postgres connection string environment variables""" + # PGOSM_CONN is required to be set by the Lua styles used by osm2pgsql + os.environ['PGOSM_CONN'] = connection_string() + # Connection to DB for admin purposes, e.g. drop/create main database + os.environ['PGOSM_CONN_PG'] = connection_string(admin=True) + def pg_conn_parts() -> dict: """Returns dictionary of connection parts based on environment variables @@ -203,10 +214,10 @@ def prepare_pgosm_db( , import_mode: helpers.ImportMode , schema_name: str ): - """Runs through steps to prepare the target database for PgOSM Flex. + """Prepares the target database for PgOSM Flex. - Includes additional preparation for using --replication and --updated=append - modes. + Includes additional preparation for using `--replication` + and `--updated=append` modes. """ if pg_conn_parts()['pg_host'] == 'localhost': drop_it = True @@ -222,15 +233,21 @@ def prepare_pgosm_db( LOGGER.debug('Dropping local database if exists') drop_pgosm_db() else: - LOGGER.debug('Not dropping local DB. This is expected with subsequent import via --replication OR --update=append.') + msg = 'Not dropping local DB. ' + msg += 'This is expected with subsequent import via --replication OR --update=append.' + LOGGER.debug(msg) create_pgosm_db() else: - LOGGER.info('Using external database. Ensure the target database is setup properly with proper permissions.') - - prepare_osm_schema(db_path=db_path, skip_qgis_style=skip_qgis_style, - schema_name=schema_name) + msg = 'Using external database. ' + msg += 'Ensure the target database is setup properly with proper permissions.' + LOGGER.info(msg) + + prepare_osm_schema( + db_path=db_path + , skip_qgis_style=skip_qgis_style + , schema_name=schema_name) run_insert_pgosm_road(db_path=db_path, schema_name=schema_name) if import_mode.replication_update or import_mode.update == 'append': @@ -249,12 +266,7 @@ def start_import( , schema_name: str , input_file: str ) -> int: - """Creates record in osm.pgosm_flex table. - - Returns - ---------------------------- - import_id : int - Value from the `id` column in `osm.pgosm_flex`. + """Creates record in `osm.pgosm_flex` table and returns `id` from `osm.pgosm_flex`. """ params = {'pgosm_region': pgosm_region , 'pgosm_date': pgosm_date @@ -281,7 +293,8 @@ def start_import( ; """ sql_raw = sql_raw.format(schema_name=schema_name) - with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn: + # FIXME: Why os environ here instead of get conn string??? + with get_db_conn(conn_string=connection_string()) as conn: cur = conn.cursor() cur.execute(sql_raw, params=params) import_id = cur.fetchone()[0] @@ -318,7 +331,7 @@ def pg_version_check() -> int: def drop_pgosm_db() -> bool: """Drops the pgosm database if it exists. - Intentionally hard coded to `pgosm` database for in-Docker use only. + This is for in-Docker use only. Intentionally hard coded to `pgosm` database. """ if not pg_conn_parts()['pg_host'] == 'localhost': LOGGER.error('Attempted to drop database external from Docker. Not doing that') @@ -327,7 +340,7 @@ def drop_pgosm_db() -> bool: sql_raw = 'DROP DATABASE IF EXISTS pgosm;' conn = get_db_conn(conn_string=os.environ['PGOSM_CONN_PG']) - LOGGER.debug('Setting Pg conn to enable autocommit - required for drop/create DB') + # Setting Pg conn to enable autocommit - required for drop/create DB' conn.autocommit = True conn.execute(sql_raw) conn.close() @@ -396,11 +409,15 @@ def prepare_osm_schema(db_path: str, skip_qgis_style: bool, schema_name: str): def run_insert_pgosm_road(db_path: str, schema_name: str): - """Runs script to load data to pgosm.road table. + """Runs script to load data to `pgosm.road` table. """ sql_filename = 'roads-us.sql' - run_deploy_file(db_path=db_path, sql_filename=sql_filename, - schema_name=schema_name, subfolder='data') + run_deploy_file( + db_path=db_path + , sql_filename=sql_filename + , schema_name=schema_name + , subfolder='data' + ) def run_deploy_file( @@ -410,34 +427,24 @@ def run_deploy_file( , subfolder: str='deploy' ): """Run a SQL script under the deploy path. Used to setup PgOSM Flex DB. - - Parameters - --------------------------- - db_path : str - Path to folder with SQL scripts. - sql_filename : sql_filename - subfolder : str - Set subfolder under `db_path`. - Default: deploy - schema_name : str - Schema name for OpenStreetMap data """ full_path = os.path.join(db_path, subfolder, sql_filename) LOGGER.info(f'Deploying {full_path}') with open(full_path) as f: - deploy_sql = f.read() + deploy_sql_raw = f.read() - deploy_sql = deploy_sql.format(schema_name=schema_name) + deploy_sql = deploy_sql_raw.format(schema_name=schema_name) + query = cast(Query, deploy_sql) - with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn: + with get_db_conn(conn_string=connection_string()) as conn: cur = conn.cursor() - cur.execute(deploy_sql) + cur.execute(query) LOGGER.debug(f'Ran SQL in {sql_filename}') -def get_db_conn(conn_string: str) -> psycopg.Connection | bool: - """Establishes psycopg database connection. +def get_db_conn(conn_string: str) -> psycopg.Connection: + """Return a `psycopg` database connection. """ try: conn = psycopg.connect(conn_string) @@ -445,7 +452,7 @@ def get_db_conn(conn_string: str) -> psycopg.Connection | bool: except psycopg.OperationalError as err: err_msg = 'Database connection error. Error: {}'.format(err) LOGGER.error(err_msg) - return False + raise # previously returned False, re-raising seems more appropriate return conn @@ -476,47 +483,36 @@ def pgosm_after_import(flex_path: str) -> bool: return True -def pgosm_nested_admin_polygons(flex_path: str, schema_name: str): +def pgosm_nested_admin_polygons(schema_name: str): """Runs two stored procedures to calculate nested admin polygons via psql. """ - # Populate the table - sql_raw_1 = f'CALL {schema_name}.populate_place_polygon_nested();' - - conn_string = os.environ['PGOSM_CONN'] - cmds = ['psql', '-d', conn_string, '-c', sql_raw_1] + # Prepare the base data LOGGER.info('Populating place_polygon_nested table (osm.populate_place_polygon_nested() )') - output = subprocess.run(cmds, - text=True, - cwd=flex_path, - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT) - LOGGER.info(f'Nested polygon output: \n {output.stdout}') - - if output.returncode != 0: - err_msg = f'Failed to populate nested polygon data. Return code: {output.returncode}' - LOGGER.error(err_msg) - sys.exit(f'{err_msg} - Check the log output for details.') + sql_raw_1 = 'CALL {schema_name}.populate_place_polygon_nested();' + query_1 = sql.SQL(sql_raw_1).format( + schema_name = sql.Identifier(schema_name) + ) + with get_db_conn(conn_string=connection_string()) as conn: + # Setting autocommit to avoid issues with the explicit transaction control + # inside the SQL procedures. + conn.autocommit = True + cur = conn.cursor() + cur.execute(query_1) - # Build the data - sql_raw_2 = f' CALL {schema_name}.build_nested_admin_polygons();' - - conn_string = os.environ['PGOSM_CONN'] - cmds = ['psql', '-d', conn_string, '-c', sql_raw_2] + # Build the nested data LOGGER.info('Building nested polygons... (this can take a while)') - output = subprocess.run(cmds, - text=True, - cwd=flex_path, - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT) - LOGGER.info(f'Nested polygon output: \n {output.stdout}') + sql_raw_2 = ' CALL {schema_name}.build_nested_admin_polygons();' + query_2 = sql.SQL(sql_raw_2).format( + schema_name = sql.Identifier(schema_name) + ) - if output.returncode != 0: - err_msg = f'Failed to build nested polygons. Return code: {output.returncode}' - LOGGER.error(err_msg) - sys.exit(f'{err_msg} - Check the log output for details.') + with get_db_conn(conn_string=connection_string()) as conn: + # Setting autocommit to avoid issues with the explicit transaction control + # inside the SQL procedures. + conn.autocommit = True + cur = conn.cursor() + cur.execute(query_2) def osm2pgsql_replication_start(): @@ -559,13 +555,7 @@ def osm2pgsql_replication_finish(skip_nested: bool): def run_pg_dump(export_path: str, skip_qgis_style: bool): - """Runs pg_dump to save processed data to load into other PostGIS DBs. - - Parameters - --------------------------- - export_path : str - Absolute path to output .sql file - skip_qgis_style : bool + """Runs `pg_dump` to save processed data to load into other PostGIS DBs. """ logger = logging.getLogger('pgosm-flex') conn_string = os.environ['PGOSM_CONN'] @@ -607,8 +597,6 @@ def fix_pg_dump_create_public(export_path: str): def log_import_message(import_id: int, msg: str, schema_name: str): """Logs msg to database in `osm.pgosm_flex` for `import_id`. - - Overwrites `osm_date` if `pbf_timestamp` is set. """ sql_raw = """ UPDATE {schema_name}.pgosm_flex pf @@ -651,5 +639,3 @@ def get_prior_import(schema_name: str) -> dict: results = {} return results - - \ No newline at end of file diff --git a/docker/geofabrik.py b/docker/geofabrik.py index 685451a1..85f00ad7 100644 --- a/docker/geofabrik.py +++ b/docker/geofabrik.py @@ -4,7 +4,9 @@ import json import os import shutil -import subprocess +from pathlib import Path +import requests +from tqdm import tqdm import helpers @@ -36,9 +38,9 @@ def prepare_data(out_path: str) -> str: pbf_file : str Full path to PBF file """ - region = os.environ.get('PGOSM_REGION') - subregion = os.environ.get('PGOSM_SUBREGION') - pgosm_date = os.environ.get('PGOSM_DATE') + region = os.environ.get('PGOSM_REGION', '') + subregion = os.environ.get('PGOSM_SUBREGION', '') + pgosm_date = os.environ.get('PGOSM_DATE', '') pbf_filename = get_region_filename() @@ -54,12 +56,13 @@ def prepare_data(out_path: str) -> str: archive_data(pbf_file, md5_file, pbf_file_with_date, md5_file_with_date) else: logging.getLogger('pgosm-flex').info('Copying Archived files') - unarchive_data(pbf_file, - md5_file, - pbf_file_with_date, - md5_file_with_date) + unarchive_data(pbf_file=pbf_file + , md5_file=md5_file + , pbf_file_with_date=pbf_file_with_date + , md5_file_with_date=md5_file_with_date + ) - helpers.verify_checksum(md5_file, out_path) + helpers.verify_checksum(md5_file=md5_file, path=out_path) set_date_from_metadata(pbf_file=pbf_file) return pbf_file @@ -93,7 +96,7 @@ def set_date_from_metadata(pbf_file: str): try: meta_timestamp = meta_options['osmosis_replication_timestamp'] except KeyError: - meta_timestamp = None + meta_timestamp = '' logger.info(f'PBF Meta timestamp: {meta_timestamp}') os.environ['PBF_TIMESTAMP'] = meta_timestamp @@ -155,28 +158,53 @@ def download_data(region: str, subregion: str, pbf_file: str, md5_file: str): logger.info(f'Downloading PBF data to {pbf_file}') pbf_url = get_pbf_url(region, subregion) - subprocess.run( - ['/usr/bin/wget', pbf_url, - "-O", pbf_file , "--quiet" - ], - capture_output=True, - text=True, - check=True - ) - - logger.info(f'Downloading MD5 checksum to {md5_file}') - subprocess.run( - ['/usr/bin/wget', f'{pbf_url}.md5', - "-O", md5_file , "--quiet" - ], - capture_output=True, - text=True, - check=True - ) - - -def archive_data(pbf_file: str, md5_file: str, pbf_file_with_date: str, - md5_file_with_date: str): + downloads = [ + (pbf_url, pbf_file, 'PBF'), + (f'{pbf_url}.md5', md5_file, 'MD5 checksum'), + ] + + for url, outfile, label in downloads: + logger.info("Downloading %s to %s", label, outfile) + download_file(url, Path(outfile)) + + +def download_file(url: str, dest: Path): + """Downloads file from `url` with progress bar. + """ + with requests.get(url, stream=True, timeout=60) as request: + request.raise_for_status() + + total_size = int(request.headers.get('Content-Length', 0)) + chunk_size = 8192 + #chunk_size = 1024 * 1024 + + with ( + open(dest, "wb") as f, + tqdm( + total=total_size, + unit="B", + unit_scale=True, + unit_divisor=1024, + desc=dest.name, + disable=total_size == 0, # fallback if server doesn't send size + mininterval=0.5 + ) as bar, + ): + for chunk in request.iter_content(chunk_size=chunk_size): + if chunk: + f.write(chunk) + bar.update(len(chunk)) + + bar.update(bar.total - bar.n) + bar.close() + + +def archive_data( + pbf_file: str + , md5_file: str + , pbf_file_with_date: str + , md5_file_with_date: str + ): """Copies `pbf_file` and `md5_file` to `pbf_file_with_date` and `md5_file_with_date`. diff --git a/docker/helpers.py b/docker/helpers.py index a9c9c776..dac6807d 100644 --- a/docker/helpers.py +++ b/docker/helpers.py @@ -10,8 +10,6 @@ from time import sleep import git -import db - DEFAULT_SRID = '3857' @@ -23,11 +21,12 @@ def get_today() -> str: return today -def run_command_via_subprocess(cmd: list, - cwd: str, - output_lines: list=[], - print_to_log: bool=False - ) -> int: +def run_command_via_subprocess( + cmd: list + , cwd: str + , output_lines: list=[] + , print_to_log: bool=False + ) -> int: """Wraps around subprocess.Popen() to run commands outside of Python. Prints output as it goes, returns the status code from the command. @@ -48,9 +47,12 @@ def run_command_via_subprocess(cmd: list, Return code from command """ logger = logging.getLogger('pgosm-flex') - with subprocess.Popen(cmd, cwd=cwd, stdout=subprocess.PIPE, - stderr=subprocess.STDOUT - ) as process: + with subprocess.Popen( + cmd + , cwd=cwd + , stdout=subprocess.PIPE + , stderr=subprocess.STDOUT + ) as process: while True: output = process.stdout.readline() if process.poll() is not None and output == b'': @@ -132,10 +134,7 @@ def set_env_vars( os.environ['PGOSM_LAYERSET'] = layerset os.environ['SCHEMA_NAME'] = schema_name - # PGOSM_CONN is required to be set by the Lua styles used by osm2pgsql - os.environ['PGOSM_CONN'] = db.connection_string() - # Connection to DB for admin purposes, e.g. drop/create main database - os.environ['PGOSM_CONN_PG'] = db.connection_string(admin=True) + # Moved DB Conn string details to database.py pgosm_region = get_region_combined(region, subregion) logger.debug(f'PGOSM_REGION_COMBINED: {pgosm_region}') @@ -215,7 +214,7 @@ class ImportMode(): """Determines logical variables used to control program flow. WARNING: The values for `append_first_run` and `replication_update` - are used to determine when to drop the local DB. Be careful with any + are used to determine when to drop the local database. Be careful with any changes to these values. """ def __init__( @@ -225,7 +224,7 @@ def __init__( , update: str , force: bool ): - """Computes two variables, slim_no_drop and append_first_run + """Computes two variables, `slim_no_drop` and `append_first_run` based on inputs. Parameters @@ -273,9 +272,9 @@ def okay_to_run(self, prior_import: dict) -> bool: Parameters ------------------- prior_import : dict - Details about the latest import from osm.pgosm_flex table. + Details about the latest import from `osm.pgosm_flex` table. - An empty dictionary (len==0) indicates no prior import. + An empty dictionary (`len==0`) indicates no prior import. Only the replication key is specifically used """ self.logger.debug(f'Checking if it is okay to run...') diff --git a/docker/osm2pgsql_recommendation.py b/docker/osm2pgsql_recommendation.py index d3c89f7c..17952aac 100644 --- a/docker/osm2pgsql_recommendation.py +++ b/docker/osm2pgsql_recommendation.py @@ -5,7 +5,7 @@ import os import osm2pgsql_tuner as tuner -import db, helpers +import database, helpers LOGGER = logging.getLogger('pgosm-flex') @@ -81,7 +81,7 @@ def get_recommended_script(system_ram_gb: float, LOGGER.debug(f'Generic command to run: {osm2pgsql_cmd}') # Replace generic connection string with specific conn string - conn_string = db.connection_string() + conn_string = database.connection_string() osm2pgsql_cmd = osm2pgsql_cmd.replace('-d $PGOSM_CONN', f'-d {conn_string}') # Warning: Do not print() this string any more! Includes password return osm2pgsql_cmd diff --git a/docker/pgosm_flex.py b/docker/pgosm_flex.py index f4130174..a4814395 100644 --- a/docker/pgosm_flex.py +++ b/docker/pgosm_flex.py @@ -15,7 +15,7 @@ import click import osm2pgsql_recommendation as rec -import db, geofabrik, helpers +import database, geofabrik, helpers @click.command() @@ -95,8 +95,10 @@ def run_pgosm_flex(ram, region, subregion, debug, force, helpers.set_env_vars(region, subregion, srid, language, pgosm_date, layerset, layerset_path, schema_name, skip_nested) - db.wait_for_postgres() - if force and db.pg_conn_parts()['pg_host'] == 'localhost': + database.set_db_env_vars() + + database.wait_for_postgres() + if force and database.pg_conn_parts()['pg_host'] == 'localhost': msg = 'Using --force with the built-in database is unnecessary.' msg += ' The pgosm database is always dropped and recreated when' msg += ' running on localhost (in Docker).' @@ -127,12 +129,12 @@ def run_pgosm_flex(ram, region, subregion, debug, force, update=update, force=force) - db.prepare_pgosm_db(skip_qgis_style=skip_qgis_style, + database.prepare_pgosm_db(skip_qgis_style=skip_qgis_style, db_path=paths['db_path'], import_mode=import_mode, schema_name=schema_name) - prior_import = db.get_prior_import(schema_name=schema_name) + prior_import = database.get_prior_import(schema_name=schema_name) if not import_mode.okay_to_run(prior_import): msg = 'Not okay to run PgOSM Flex. Exiting' @@ -147,7 +149,7 @@ def run_pgosm_flex(ram, region, subregion, debug, force, cwd='/usr/bin/', output_lines=vers_lines) - import_id = db.start_import(pgosm_region=helpers.get_region_combined(region, subregion), + import_id = database.start_import(pgosm_region=helpers.get_region_combined(region, subregion), pgosm_date=pgosm_date, srid=srid, language=language, @@ -179,12 +181,12 @@ def run_pgosm_flex(ram, region, subregion, debug, force, if not success: msg = 'PgOSM Flex completed with errors. Details in output' - db.log_import_message(import_id=import_id, msg='Failed', + database.log_import_message(import_id=import_id, msg='Failed', schema_name=schema_name) logger.warning(msg) sys.exit(msg) - db.log_import_message(import_id=import_id, msg='Completed', + database.log_import_message(import_id=import_id, msg='Completed', schema_name=schema_name) dump_database(input_file=input_file, @@ -208,11 +210,6 @@ def run_osm2pgsql_standard( ) -> bool: """Runs standard osm2pgsql command and optionally inits for replication (osm2pgsql-replication) mode. - - Returns - --------------------------- - post_processing : boolean - Indicates overall success/failure of the steps within this function. """ logger = logging.getLogger('pgosm-flex') @@ -227,8 +224,7 @@ def run_osm2pgsql_standard( out_path=out_path, import_mode=import_mode) - run_osm2pgsql(osm2pgsql_command=osm2pgsql_command, flex_path=flex_path, - debug=debug) + run_osm2pgsql(osm2pgsql_command=osm2pgsql_command, flex_path=flex_path) if not skip_nested: # Don't expect user to use --skip-nested when place isn't included @@ -237,8 +233,7 @@ def run_osm2pgsql_standard( post_processing = run_post_processing(flex_path=flex_path, skip_nested=skip_nested, import_mode=import_mode, - schema_name=schema_name, - import_id=import_id) + schema_name=schema_name) if import_mode.replication: run_osm2pgsql_replication_init(pbf_path=out_path, @@ -256,7 +251,7 @@ def run_replication_update(skip_nested: bool, flex_path: str) -> bool: """Runs osm2pgsql-replication between the DB start/finish steps. """ logger = logging.getLogger('pgosm-flex') - conn_string = db.connection_string() + conn_string = database.connection_string() update_cmd = """ osm2pgsql-replication update -d $PGOSM_CONN \ @@ -266,16 +261,18 @@ def run_replication_update(skip_nested: bool, flex_path: str) -> bool: """ update_cmd = update_cmd.replace('-d $PGOSM_CONN', f'-d {conn_string}') - returncode = helpers.run_command_via_subprocess(cmd=update_cmd.split(), - cwd=flex_path, - print_to_log=True) + returncode = helpers.run_command_via_subprocess( + cmd=update_cmd.split() + , cwd=flex_path + , print_to_log=True + ) if returncode != 0: err_msg = f'Failure. Return code: {returncode}' logger.warning(err_msg) return False - db.osm2pgsql_replication_finish(skip_nested=skip_nested) + database.osm2pgsql_replication_finish(skip_nested=skip_nested) logger.info('osm2pgsql-replication update complete') return True @@ -300,12 +297,7 @@ def validate_region_inputs(region: str, subregion: str, input_file: str): def setup_logger(debug: bool): - """Prepares logging. - - Parameters - ------------------------------ - debug : bool - Enables debug mode when True. INFO when False. + """Prepares logging. Enables debug mode when True. INFO when False. """ if debug: log_level = logging.DEBUG @@ -378,7 +370,7 @@ def get_export_full_path(out_path: str, export_filename: str) -> str: return export_path -def run_osm2pgsql(osm2pgsql_command: str, flex_path: str, debug: bool): +def run_osm2pgsql(osm2pgsql_command: str, flex_path: str): """Runs the provided osm2pgsql command. """ logger = logging.getLogger('pgosm-flex') @@ -469,7 +461,6 @@ def run_post_processing( , skip_nested: bool , import_mode: helpers.ImportMode , schema_name: str - , import_id: int ) -> bool: """Runs steps following osm2pgsql import. @@ -481,16 +472,18 @@ def run_post_processing( msg = 'Running with --update append: Skipping post-processing SQL.' msg += ' Running osm2pgsql_replication_finish() instead.' logger.info(msg) - db.osm2pgsql_replication_finish(skip_nested=skip_nested, import_id=import_id) + database.osm2pgsql_replication_finish( + skip_nested=skip_nested + ) return True - post_processing_sql = db.pgosm_after_import(flex_path=flex_path) + post_processing_sql = database.pgosm_after_import(flex_path=flex_path) if skip_nested: logger.info('Skipping calculating nested polygons') else: logger.info('Calculating nested polygons') - db.pgosm_nested_admin_polygons(flex_path, schema_name) + database.pgosm_nested_admin_polygons(schema_name=schema_name) if not post_processing_sql: return False @@ -504,7 +497,7 @@ def dump_database(input_file: str, out_path: str, pg_dump: bool, skip_qgis_style export_filename = get_export_filename(input_file) export_path = get_export_full_path(out_path, export_filename) - db.run_pg_dump(export_path=export_path, + database.run_pg_dump(export_path=export_path, skip_qgis_style=skip_qgis_style) else: logging.getLogger('pgosm-flex').info('Skipping pg_dump') @@ -516,7 +509,7 @@ def check_replication_exists() -> bool: logger = logging.getLogger('pgosm-flex') check_cmd = "osm2pgsql-replication status -d $PGOSM_CONN " logger.debug(f'Command to check DB for replication status:\n{check_cmd}') - conn_string = db.connection_string() + conn_string = database.connection_string() check_cmd = check_cmd.replace('-d $PGOSM_CONN', f'-d {conn_string}') returncode = helpers.run_command_via_subprocess(cmd=check_cmd.split(), @@ -539,7 +532,7 @@ def run_osm2pgsql_replication_init(pbf_path: str, pbf_filename: str): init_cmd = 'osm2pgsql-replication init -d $PGOSM_CONN ' init_cmd += f'--osm-file {pbf_path}' logger.debug(f'Initializing DB for replication with command:\n{init_cmd}') - conn_string = db.connection_string() + conn_string = database.connection_string() init_cmd = init_cmd.replace('-d $PGOSM_CONN', f'-d {conn_string}') returncode = helpers.run_command_via_subprocess(cmd=init_cmd.split(), diff --git a/docker/qgis_styles.py b/docker/qgis_styles.py index dda40ada..a3bfe6cb 100644 --- a/docker/qgis_styles.py +++ b/docker/qgis_styles.py @@ -4,7 +4,7 @@ import os import subprocess -import db +import database LOGGER = logging.getLogger('pgosm-flex') @@ -38,7 +38,7 @@ def create_layer_style_table(db_path: str, conn_string: str): with open(create_path, 'r') as file_in: create_sql = file_in.read() - with db.get_db_conn(conn_string=conn_string) as conn: + with database.get_db_conn(conn_string=conn_string) as conn: cur = conn.cursor() cur.execute(create_sql) LOGGER.debug('QGIS Style table created') @@ -81,13 +81,13 @@ def load_staging_to_prod(db_path, conn_string): with open(load_path, 'r') as file_in: load_sql = file_in.read() - with db.get_db_conn(conn_string=conn_string) as conn: + with database.get_db_conn(conn_string=conn_string) as conn: cur = conn.cursor() cur.execute(load_sql) LOGGER.info('QGIS Style table populated') - with db.get_db_conn(conn_string=conn_string) as conn: + with database.get_db_conn(conn_string=conn_string) as conn: sql_clean = 'DELETE FROM public.layer_styles_staging;' cur = conn.cursor() cur.execute(sql_clean) @@ -112,10 +112,9 @@ def update_styles_db_name(db_name, conn_string): ; """ params = {'db_name': db_name} - with db.get_db_conn(conn_string=conn_string) as conn: + with database.get_db_conn(conn_string=conn_string) as conn: cur = conn.cursor() cur.execute(sql_raw, params=params) conn.commit() LOGGER.info(f'Updated QGIS layer styles for database {db_name}') - diff --git a/docker/tests/test_db.py b/docker/tests/test_database.py similarity index 84% rename from docker/tests/test_db.py rename to docker/tests/test_database.py index ce1b0378..01a6f89b 100644 --- a/docker/tests/test_db.py +++ b/docker/tests/test_database.py @@ -4,7 +4,7 @@ from urllib import parse from unittest import mock -import db +import database POSTGRES_USER = 'my_pg_user' POSTGRES_PASSWORD = 'here_for_fun!@#$%^&*()' @@ -14,8 +14,8 @@ 'POSTGRES_PASSWORD': ''} PG_USER_AND_PW = {'POSTGRES_USER': POSTGRES_USER, 'POSTGRES_PASSWORD': POSTGRES_PASSWORD, - 'PGOSM_CONN_PG': db.connection_string(admin=True), - 'PGOSM_CONN': db.connection_string()} + 'PGOSM_CONN_PG': database.connection_string(admin=True), + 'PGOSM_CONN': database.connection_string()} POSTGRES_HOST_NON_LOCAL = {'POSTGRES_HOST': POSTGRES_HOST_EXTERNAL, 'POSTGRES_USER': POSTGRES_USER, 'POSTGRES_PASSWORD': POSTGRES_PASSWORD} @@ -27,7 +27,7 @@ class DBTests(unittest.TestCase): def test_pg_conn_parts_user_only_returns_expected_values(self): expected_user = POSTGRES_USER expected_pw = None - results = db.pg_conn_parts() + results = database.pg_conn_parts() self.assertEqual(expected_user, results['pg_user']) self.assertEqual(expected_pw, results['pg_pass']) @@ -36,7 +36,7 @@ def test_pg_conn_parts_user_only_returns_expected_values(self): def test_pg_conn_parts_user_w_pw_returns_expected_values(self): expected_user = POSTGRES_USER expected_pw = POSTGRES_PASSWORD - results = db.pg_conn_parts() + results = database.pg_conn_parts() self.assertEqual(expected_user, results['pg_user']) self.assertEqual(expected_pw, results['pg_pass']) @@ -44,7 +44,7 @@ def test_pg_conn_parts_user_w_pw_returns_expected_values(self): @mock.patch.dict(os.environ, PG_USER_ONLY) def test_connection_string_user_only_returns_expected_string(self): expected = f'postgresql://{POSTGRES_USER}@localhost:5432/pgosm?application_name=pgosm-flex' - result = db.connection_string() + result = database.connection_string() self.assertEqual(expected, result) @@ -52,7 +52,7 @@ def test_connection_string_user_only_returns_expected_string(self): def test_connection_string_user_w_pw_returns_expected_string(self): password_safe = parse.quote(POSTGRES_PASSWORD) expected = f'postgresql://{POSTGRES_USER}:{password_safe}@localhost:5432/pgosm?application_name=pgosm-flex' - result = db.connection_string() + result = database.connection_string() self.assertEqual(expected, result) @@ -64,8 +64,8 @@ def test_admin_connection_string_external_returns_expected_string(self): """ password_safe = parse.quote(POSTGRES_PASSWORD) expected = f'postgresql://{POSTGRES_USER}:{password_safe}@{POSTGRES_HOST_EXTERNAL}:5432/pgosm?application_name=pgosm-flex' - result_standard = db.connection_string() - result_admin = db.connection_string(admin=True) + result_standard = database.connection_string() + result_admin = database.connection_string(admin=True) self.assertEqual(expected, result_standard) self.assertEqual(expected, result_admin) @@ -75,7 +75,7 @@ def test_drop_pgosm_db_with_non_localhost_returns_False(self): """Tests the function returns False instead of attempting to drop the DB """ expected = False - result = db.drop_pgosm_db() + result = database.drop_pgosm_db() self.assertEqual(expected, result) @@ -84,20 +84,20 @@ def test_create_pgosm_db_with_non_localhost_returns_False(self): """Tests the function returns False instead of attempting to create the DB """ expected = False - result = db.create_pgosm_db() + result = database.create_pgosm_db() self.assertEqual(expected, result) @mock.patch.dict(os.environ, PG_USER_AND_PW) def test_pg_version_check_returns_int(self): expected = int - result = db.pg_version_check() + result = database.pg_version_check() self.assertEqual(expected, type(result)) @mock.patch.dict(os.environ, PG_USER_AND_PW) def test_get_prior_import_returns_expected_type(self): - result = db.get_prior_import(schema_name='osm') + result = database.get_prior_import(schema_name='osm') actual = type(result) expected = dict self.assertEqual(expected, actual) diff --git a/requirements.txt b/requirements.txt index 01b6803a..b1d313d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ packaging>=23.0 psycopg>=3.1 psycopg-binary>=3.1 sh>=1.14.2 +tqdm==4.64.1