Merge pull request #355 from rustprooflabs/flexible-schema-name

Add --schema-name option to handle dynamic name

Author: rustprooflabs

db/deploy/osm.sql

@@ -7,5 +7,5 @@
         https://pgosm-flex.com/postgres-permissions.html
 */
 CREATE EXTENSION IF NOT EXISTS postgis;
-CREATE SCHEMA IF NOT EXISTS osm;
-COMMENT ON SCHEMA osm IS 'Schema populated by PgOSM Flex.  SELECT * FROM osm.pgosm_flex; for details.';
+CREATE SCHEMA IF NOT EXISTS {schema_name};
+COMMENT ON SCHEMA {schema_name} IS 'Schema populated by PgOSM Flex.  SELECT * FROM {schema_name}.pgosm_flex; for details.';

db/deploy/osm_pgosm_flex.sql

@@ -2,7 +2,7 @@
 BEGIN;
 
 
-CREATE TABLE IF NOT EXISTS osm.pgosm_flex (
+CREATE TABLE IF NOT EXISTS {schema_name}.pgosm_flex (
     id BIGINT NOT NULL GENERATED BY DEFAULT AS IDENTITY,
     imported TIMESTAMPTZ NOT NULL DEFAULT NOW(),
     osm_date date NOT NULL,
@@ -17,39 +17,39 @@ CREATE TABLE IF NOT EXISTS osm.pgosm_flex (
     CONSTRAINT pk_osm_pgosm_flex PRIMARY KEY (id)
 );
 
-ALTER TABLE osm.pgosm_flex
+ALTER TABLE {schema_name}.pgosm_flex
     DROP COLUMN IF EXISTS osm2pgsql_mode
 ;
 
-ALTER TABLE osm.pgosm_flex
+ALTER TABLE {schema_name}.pgosm_flex
     DROP COLUMN IF EXISTS osm2pgsql_replication;
 
-ALTER TABLE osm.pgosm_flex
+ALTER TABLE {schema_name}.pgosm_flex
     ADD COLUMN IF NOT EXISTS import_mode JSONB NULL;
 
-ALTER TABLE osm.pgosm_flex
+ALTER TABLE {schema_name}.pgosm_flex
     DROP COLUMN IF EXISTS import_uuid;
 
-ALTER TABLE osm.pgosm_flex
+ALTER TABLE {schema_name}.pgosm_flex
     ADD COLUMN IF NOT EXISTS import_status TEXT NOT NULL DEFAULT 'Initializing';
 
-ALTER TABLE osm.pgosm_flex
+ALTER TABLE {schema_name}.pgosm_flex
     ADD COLUMN IF NOT EXISTS layerset TEXT NULL;
 
-ALTER TABLE osm.pgosm_flex DROP COLUMN IF EXISTS project_url;
-ALTER TABLE osm.pgosm_flex DROP COLUMN IF EXISTS default_date;
+ALTER TABLE {schema_name}.pgosm_flex DROP COLUMN IF EXISTS project_url;
+ALTER TABLE {schema_name}.pgosm_flex DROP COLUMN IF EXISTS default_date;
 
-COMMENT ON TABLE osm.pgosm_flex IS 'Provides meta information on the PgOSM-Flex project including version and SRID used during the import. One row per import.';
+COMMENT ON TABLE {schema_name}.pgosm_flex IS 'Provides meta information on the PgOSM-Flex project including version and SRID used during the import. One row per import.';
 
-COMMENT ON COLUMN osm.pgosm_flex.imported IS 'Indicates when the import was ran.';
-COMMENT ON COLUMN osm.pgosm_flex.osm_date IS 'Indicates the date of the OpenStreetMap data loaded.  Recommended to set PGOSM_DATE env var at runtime, otherwise defaults to the date PgOSM-Flex was run.';
-COMMENT ON COLUMN osm.pgosm_flex.srid IS 'SRID of imported data.';
-COMMENT ON COLUMN osm.pgosm_flex.pgosm_flex_version IS 'Version of PgOSM-Flex used to generate schema.';
-COMMENT ON COLUMN osm.pgosm_flex.osm2pgsql_version IS 'Version of osm2pgsql used to load data.';
-COMMENT ON COLUMN osm.pgosm_flex.region IS 'Region specified at run time via env var PGOSM_REGION.';
-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.layerset IS 'PgOSM Flex layerset used for the import style.';
-COMMENT ON COLUMN osm.pgosm_flex.import_status IS 'Status of the import. Starts as initialized, tracks status during imports and final success/failure.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.imported IS 'Indicates when the import was ran.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.osm_date IS 'Indicates the date of the OpenStreetMap data loaded.  Recommended to set PGOSM_DATE env var at runtime, otherwise defaults to the date PgOSM-Flex was run.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.srid IS 'SRID of imported data.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.pgosm_flex_version IS 'Version of PgOSM-Flex used to generate schema.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.osm2pgsql_version IS 'Version of osm2pgsql used to load data.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.region IS 'Region specified at run time via env var PGOSM_REGION.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.language IS 'Preferred language specified at run time via env var PGOSM_LANGUAGE.  Empty string when not defined.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.layerset IS 'PgOSM Flex layerset used for the import style.';
+COMMENT ON COLUMN {schema_name}.pgosm_flex.import_status IS 'Status of the import. Starts as initialized, tracks status during imports and final success/failure.';
 
 
 COMMIT;

db/deploy/replication_functions.sql

@@ -2,39 +2,39 @@
 BEGIN;
 
 
-CREATE OR REPLACE PROCEDURE osm.append_data_start()
+CREATE OR REPLACE PROCEDURE {schema_name}.append_data_start()
  LANGUAGE plpgsql
  AS $$
 
  BEGIN
 
-    RAISE NOTICE 'Truncating table osm.place_polygon_nested;';
-    TRUNCATE TABLE osm.place_polygon_nested;
+    RAISE NOTICE 'Truncating table {schema_name}.place_polygon_nested;';
+    TRUNCATE TABLE {schema_name}.place_polygon_nested;
 
 END $$;
 
 
-CREATE OR REPLACE PROCEDURE osm.append_data_finish(skip_nested BOOLEAN = False)
+CREATE OR REPLACE PROCEDURE {schema_name}.append_data_finish(skip_nested BOOLEAN = False)
  LANGUAGE plpgsql
  AS $$
  BEGIN
 
-    REFRESH MATERIALIZED VIEW osm.vplace_polygon_subdivide;
+    REFRESH MATERIALIZED VIEW {schema_name}.vplace_polygon_subdivide;
 
     IF $1 = False THEN
         RAISE NOTICE 'Populating nested place table';
-        CALL osm.populate_place_polygon_nested();
+        CALL {schema_name}.populate_place_polygon_nested();
         RAISE NOTICE 'Calculating nesting of place polygons';
-        CALL osm.build_nested_admin_polygons();
+        CALL {schema_name}.build_nested_admin_polygons();
 
     END IF;
 
 
 END $$;
 
 
-COMMENT ON PROCEDURE osm.append_data_start() IS 'Prepares PgOSM Flex database for running osm2pgsql in append mode.  Removes records from place_polygon_nested if they existed.';
-COMMENT ON PROCEDURE osm.append_data_finish(BOOLEAN) IS 'Finalizes PgOSM Flex after osm2pgsql-replication.  Refreshes materialized view and (optionally) processes the place_polygon_nested data.';
+COMMENT ON PROCEDURE {schema_name}.append_data_start() IS 'Prepares PgOSM Flex database for running osm2pgsql in append mode.  Removes records from place_polygon_nested if they existed.';
+COMMENT ON PROCEDURE {schema_name}.append_data_finish(BOOLEAN) IS 'Finalizes PgOSM Flex after osm2pgsql-replication.  Refreshes materialized view and (optionally) processes the place_polygon_nested data.';
 
 
 

docker/db.py

@@ -197,14 +197,16 @@ def log_pg_details():
     LOGGER.info(msg)
 
 
-def prepare_pgosm_db(skip_qgis_style, db_path, import_mode):
+def prepare_pgosm_db(skip_qgis_style, db_path, import_mode, schema_name):
     """Runs through series of steps to prepare database for PgOSM.
 
     Parameters
     --------------------------
     skip_qgis_style : bool
     db_path : str
     import_mode : import_mode.ImportMode
+    schema_name : str
+        Schema name for OpenStreetMap data
     """
     if pg_conn_parts()['pg_host'] == 'localhost':
         drop_it = True
@@ -227,12 +229,13 @@ def prepare_pgosm_db(skip_qgis_style, db_path, import_mode):
     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)
-    run_insert_pgosm_road(db_path=db_path)
+    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)
 
 
 def start_import(pgosm_region, pgosm_date, srid, language, layerset, git_info,
-                 osm2pgsql_version, import_mode):
+                 osm2pgsql_version, import_mode, schema_name):
     """Creates record in osm.pgosm_flex table.
 
     Parameters
@@ -245,6 +248,7 @@ def start_import(pgosm_region, pgosm_date, srid, language, layerset, git_info,
     git_info : str
     osm2pgsql_version : str
     import_mode : import_mode.ImportMode
+    schema_name : str
 
     Returns
     ----------------------------
@@ -257,7 +261,7 @@ def start_import(pgosm_region, pgosm_date, srid, language, layerset, git_info,
               'import_mode': import_mode.as_json()}
 
     sql_raw = """
-INSERT INTO osm.pgosm_flex
+INSERT INTO {schema_name}.pgosm_flex
     (osm_date, region, pgosm_flex_version, srid,
         osm2pgsql_version, "language", import_mode,
         layerset)
@@ -268,6 +272,7 @@ def start_import(pgosm_region, pgosm_date, srid, language, layerset, git_info,
     RETURNING id
 ;
 """
+    sql_raw = sql_raw.format(schema_name=schema_name)
     with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn:
         cur = conn.cursor()
         cur.execute(sql_raw, params=params)
@@ -360,7 +365,7 @@ def create_pgosm_db():
     return True
 
 
-def prepare_osm_schema(db_path: str, skip_qgis_style: bool):
+def prepare_osm_schema(db_path: str, skip_qgis_style: bool, schema_name: str):
     """Runs deploy scripts to prepare the PgOSM Flex database.
 
     This function's code could be simplified, but currently I like the verbosity
@@ -371,17 +376,18 @@ def prepare_osm_schema(db_path: str, skip_qgis_style: bool):
     db_path : str
         Path to folder with SQL scripts.
     skip_qgis_style : bool
+    scheme_name : str
     """
-    LOGGER.info('Preparing database schema')
+    LOGGER.info(f'Preparing database schema: {schema_name}')
     create_osm_file = 'osm.sql'
     create_osm_pgosm_flex_file = 'osm_pgosm_flex.sql'
     create_pgosm_road_file = 'pgosm_road.sql'
     create_replication_functions = 'replication_functions.sql'
 
-    run_deploy_file(db_path=db_path, sql_filename=create_osm_file)
-    run_deploy_file(db_path=db_path, sql_filename=create_osm_pgosm_flex_file)
-    run_deploy_file(db_path=db_path, sql_filename=create_pgosm_road_file)
-    run_deploy_file(db_path=db_path, sql_filename=create_replication_functions)
+    run_deploy_file(db_path=db_path, sql_filename=create_osm_file, schema_name=schema_name)
+    run_deploy_file(db_path=db_path, sql_filename=create_osm_pgosm_flex_file, schema_name=schema_name)
+    run_deploy_file(db_path=db_path, sql_filename=create_pgosm_road_file, schema_name=schema_name)
+    run_deploy_file(db_path=db_path, sql_filename=create_replication_functions, schema_name=schema_name)
 
     if skip_qgis_style:
         LOGGER.info('Skipping QGIS styles')
@@ -391,19 +397,22 @@ def prepare_osm_schema(db_path: str, skip_qgis_style: bool):
                                      db_name=pg_conn_parts()['pg_db'])
 
 
-def run_insert_pgosm_road(db_path: str):
+def run_insert_pgosm_road(db_path: str, schema_name: str):
     """Runs script to load data to pgosm.road table.
 
     Parameters
     ------------------------
     db_path : str
+    schema_name : str
+        Schema name for OpenStreetMap data
     """
     sql_filename = 'roads-us.sql'
     run_deploy_file(db_path=db_path, sql_filename=sql_filename,
-                    subfolder='data')
+                    schema_name=schema_name, subfolder='data')
 
 
-def run_deploy_file(db_path: str, sql_filename: str, subfolder: str='deploy'):
+def run_deploy_file(db_path: str, sql_filename: str, schema_name: str,
+                    subfolder: str='deploy'):
     """Run a SQL script under the deploy path.  Used to setup PgOSM Flex DB.
 
     Parameters
@@ -414,13 +423,17 @@ def run_deploy_file(db_path: str, sql_filename: str, subfolder: str='deploy'):
     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 = deploy_sql.format(schema_name=schema_name)
+
     with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn:
         cur = conn.cursor()
         cur.execute(deploy_sql)
@@ -479,14 +492,15 @@ def pgosm_after_import(flex_path):
     return True
 
 
-def pgosm_nested_admin_polygons(flex_path):
+def pgosm_nested_admin_polygons(flex_path: str, schema_name: str):
     """Runs stored procedure to calculate nested admin polygons via psql.
 
     Parameters
     ----------------------
     flex_path : str
+    schema_name : str
     """
-    sql_raw = 'CALL osm.build_nested_admin_polygons();'
+    sql_raw = f'CALL {schema_name}.build_nested_admin_polygons();'
 
     conn_string = os.environ['PGOSM_CONN']
     cmds = ['psql', '-d', conn_string, '-c', sql_raw]
@@ -600,29 +614,35 @@ def fix_pg_dump_create_public(export_path):
     LOGGER.debug(result)
 
 
-def log_import_message(import_id, msg):
+def log_import_message(import_id, msg, schema_name):
     """Logs msg to database in osm.pgosm_flex for import_uuid.
 
     Parameters
     -------------------------------
     import_id : int
     msg : str
+    schema_name: str
     """
     sql_raw = """
-UPDATE osm.pgosm_flex
+UPDATE {schema_name}.pgosm_flex
     SET import_status = %(msg)s
     WHERE id = %(import_id)s
 ;
 """
+    sql_raw = sql_raw.format(schema_name=schema_name)
     with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn:
         params = {'import_id': import_id, 'msg': msg}
         cur = conn.cursor()
         cur.execute(sql_raw, params=params)
 
 
-def get_prior_import() -> dict:
+def get_prior_import(schema_name: str) -> dict:
     """Gets the latest import details from osm.pgosm_flex.
 
+    Parameters
+    --------------------
+    schema_name : str
+
     Returns
     --------------------
     results : dict
@@ -632,11 +652,12 @@ def get_prior_import() -> dict:
         import_mode ->> 'replication' AS replication,
         import_mode ->> 'update' AS use_update,
         import_mode
-    FROM osm.pgosm_flex
+    FROM {schema_name}.pgosm_flex
     ORDER BY imported DESC
     LIMIT 1
 ;
 """
+    sql_raw = sql_raw.format(schema_name=schema_name)
     with get_db_conn(conn_string=os.environ['PGOSM_CONN']) as conn:
         cur = conn.cursor(row_factory=psycopg.rows.dict_row)
         results = cur.execute(sql_raw).fetchone()

docker/helpers.py

@@ -90,7 +90,7 @@ def verify_checksum(md5_file, path):
 
 
 def set_env_vars(region, subregion, srid, language, pgosm_date, layerset,
-                 layerset_path, replication):
+                 layerset_path, replication, schema_name):
     """Sets environment variables needed by PgOSM Flex. Also creates DB
     record in `osm.pgosm_flex` table.
 
@@ -106,6 +106,7 @@ def set_env_vars(region, subregion, srid, language, pgosm_date, layerset,
         str when set, or None
     replication : bool
         Indicates when osm2pgsql-replication is used
+    schema_name : str
     """
     logger = logging.getLogger('pgosm-flex')
     logger.debug('Ensuring env vars are not set from prior run')
@@ -128,6 +129,7 @@ def set_env_vars(region, subregion, srid, language, pgosm_date, layerset,
 
     os.environ['PGOSM_DATE'] = pgosm_date
     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()
@@ -204,3 +206,4 @@ def unset_env_vars():
     os.environ.pop('PGOSM_CONN', None)
     os.environ.pop('PGOSM_CONN_PG', None)
     os.environ.pop('PGOSM_REPLICATION', None)
+    os.environ.pop('SCHEMA_NAME', None)