Merge pull request #303 from rustprooflabs/track-import-failure

Add ability to track import errors in osm.pgosm_flex table

docker/db.py

@@ -616,3 +616,24 @@ def fix_pg_dump_create_public(export_path):
            export_path)
     LOGGER.debug('Completed replacement to not fail when public schema exists')
     LOGGER.debug(result)
+
+
+def log_import_message(import_uuid, msg):
+    """Logs msg to database in osm.pgosm_flex for import_uuid.
+
+    Parameters
+    -------------------------------
+    import_uuid : uuid
+    msg : str
+    """
+    sql_raw = """
+UPDATE osm.pgosm_flex
+    SET import_status = %(msg)s
+    WHERE import_uuid = %(import_uuid)s
+;
+"""
+    with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn:
+        params = {'import_uuid': str(import_uuid), 'msg': msg}
+        cur = conn.cursor()
+        cur.execute(sql_raw, params=params)
+

docker/helpers.py

@@ -82,7 +82,7 @@ def verify_checksum(md5_file, path):
 
 
 def set_env_vars(region, subregion, srid, language, pgosm_date, layerset,
-                 layerset_path, sp_gist, replication):
+                 layerset_path, sp_gist, replication, import_uuid):
     """Sets environment variables needed by PgOSM Flex.
 
     See /docs/MANUAL-STEPS-RUN.md for usage examples of environment variables.
@@ -101,13 +101,16 @@ def set_env_vars(region, subregion, srid, language, pgosm_date, layerset,
         When `True` uses SP-GIST index instead of GIST for spatial indexes.
     replication : bool
         Indicates when osm2pgsql-replication is used
+    import_uuid : uuid
+        Required to track import failures in Postgres between Python and Lua
     """
     logger = logging.getLogger('pgosm-flex')
     logger.debug('Ensuring env vars are not set from prior run')
     unset_env_vars()
     logger.debug('Setting environment variables')
 
     os.environ['PGOSM_REGION'] = region
+    os.environ['PGOSM_IMPORT_UUID'] = str(import_uuid)
 
     if subregion is None:
         pgosm_region = f'{region}'
@@ -149,6 +152,7 @@ def set_env_vars(region, subregion, srid, language, pgosm_date, layerset,
         os.environ['PGOSM_REPLICATION'] = 'false'
 
 
+
 def unset_env_vars():
     """Unsets environment variables used by PgOSM Flex.
 
@@ -166,3 +170,4 @@ def unset_env_vars():
     os.environ.pop('PGOSM_CONN_PG', None)
     os.environ.pop('PGOSM_GIST_TYPE', None)
     os.environ.pop('PGOSM_REPLICATION', None)
+    os.environ.pop('PGOSM_IMPORT_UUID', None)

docker/pgosm_flex.py

@@ -14,6 +14,7 @@
 import sys
 
 import click
+import uuid
 
 import osm2pgsql_recommendation as rec
 import db
@@ -102,8 +103,10 @@ def run_pgosm_flex(ram, region, subregion, data_only, debug,
     if region is None and input_file:
         region = input_file
 
+    import_uuid = uuid.uuid4()
     helpers.set_env_vars(region, subregion, srid, language, pgosm_date,
-                         layerset, layerset_path, sp_gist, replication)
+                         layerset, layerset_path, sp_gist, replication,
+                         import_uuid)
     db.wait_for_postgres()
 
     if replication:
@@ -140,6 +143,12 @@ def run_pgosm_flex(ram, region, subregion, data_only, debug,
                                          skip_nested=skip_nested,
                                          import_mode=import_mode)
 
+    if not success:
+        msg = 'PgOSM Flex completed with errors. Details in output'
+        db.log_import_message(import_uuid=import_uuid, msg='Import failed')
+        logger.warning(msg)
+        sys.exit(msg)
+
     if schema_name != 'osm':
         db.rename_schema(schema_name)
 
@@ -149,10 +158,7 @@ def run_pgosm_flex(ram, region, subregion, data_only, debug,
                   data_only=data_only,
                   schema_name=schema_name)
 
-    if success:
-        logger.info('PgOSM Flex complete!')
-    else:
-        logger.warning('PgOSM Flex completed with errors. Details in output')
+    logger.info('PgOSM Flex complete!')
 
 
 def run_osm2pgsql_standard(input_file, out_path, flex_path, ram, skip_nested,

flex-config/sql/pgosm-meta.sql

@@ -12,7 +12,8 @@ COMMENT ON COLUMN osm.pgosm_flex.region IS 'Region specified at run time via env
 COMMENT ON COLUMN osm.pgosm_flex.language IS 'Preferred language specified at run time via env var PGOSM_LANGUAGE.  Empty string when not defined.';
 COMMENT ON COLUMN osm.pgosm_flex.osm2pgsql_mode IS 'Indicates which osm2pgsql mode was used, create or append.';
 COMMENT ON COLUMN osm.pgosm_flex.osm2pgsql_replication IS 'True indicates when osm2pgsql-replication was used for the import.';
-
+COMMENT ON COLUMN osm.pgosm_flex.import_uuid IS 'Used during import to set import_status on failure.';
+COMMENT ON COLUMN osm.pgosm_flex.import_status IS 'Set when import failure is detected.';
 
 -- Helper Procedures for allowing updates via osm2pgsql-replication or similar
 

flex-config/style/pgosm-meta.lua

@@ -20,6 +20,7 @@ else
     error('ENV VAR PGOSM_REPLICATION must be set')
 end
 
+local import_uuid = os.getenv("PGOSM_IMPORT_UUID")
 
 local tables = {}
 
@@ -37,6 +38,8 @@ CREATE TABLE IF NOT EXISTS osm.pgosm_flex (
     "language" text NOT NULL,
     osm2pgsql_mode TEXT NOT NULL DEFAULT 'create',
     osm2pgsql_replication BOOLEAN NOT NULL DEFAULT False,
+    import_uuid TEXT NULL,
+    import_status TEXT NULL,
     CONSTRAINT pk_osm_pgosm_flex PRIMARY KEY (id)
 );
 ]=]
@@ -50,6 +53,12 @@ ALTER TABLE osm.pgosm_flex
 ALTER TABLE osm.pgosm_flex
     ADD COLUMN IF NOT EXISTS osm2pgsql_replication
     BOOLEAN NOT NULL DEFAULT False;
+
+ALTER TABLE osm.pgosm_flex
+    ADD COLUMN IF NOT EXISTS import_uuid TEXT NULL;
+
+ALTER TABLE osm.pgosm_flex
+    ADD COLUMN IF NOT EXISTS import_status TEXT;
 ]=]
 
 
@@ -100,7 +109,7 @@ end
 
 local sql_insert = [[ INSERT INTO osm.pgosm_flex (osm_date, default_date, region,
         pgosm_flex_version, srid, project_url, osm2pgsql_version, "language",
-        osm2pgsql_mode, osm2pgsql_replication) ]] ..
+        osm2pgsql_mode, osm2pgsql_replication, import_uuid) ]] ..
  [[ VALUES (']] ..
  con:escape(pgosm_date) .. [[', ]] ..
  default_date_str .. [[ , ']] .. -- special handling for boolean
@@ -111,7 +120,8 @@ local sql_insert = [[ INSERT INTO osm.pgosm_flex (osm_date, default_date, region
  con:escape(osm2pgsql_version) .. [[', ']] ..
  con:escape(pgosm_language) .. [[', ']] ..
  con:escape(osm2pgsql_mode) .. [[', ]] ..
- pgosm_replication .. [[ );]]
+ pgosm_replication .. [[ , ']] ..
+ con:escape(import_uuid) .. [[' );]]
 
 
 -- simple query to verify connection