import-borders

#!/usr/bin/env bash
set -o errexit
set -o pipefail
set -o nounset

: "${PBF_DATA_DIR:=/import}"
: "${BORDERS_CLEANUP:=false}"
: "${BORDERS_CLEANUP_FILE:="$PBF_DATA_DIR/borders/cleanup.pbf"}"
: "${BORDERS_PBF_FILE:="$PBF_DATA_DIR/borders/filtered.pbf"}"
: "${BORDERS_CSV_FILE:="$PBF_DATA_DIR/borders/lines.csv"}"
: "${BORDERS_TABLE_NAME:=osm_border_linestring}"

function usage() {
  >&2 echo "Usage"
  >&2 echo ""
  >&2 echo "import-borders                      - Parse and import first PBF file in PBF_DATA_DIR"
  >&2 echo "import-borders [import] planet.pbf  - Parse and import planet.pbf"
  >&2 echo "import-borders parse planet.pbf     - Parse planet.pbf into a CSV file, but do not import"
  >&2 echo "import-borders load borders.csv     - Load borders.csv into a table"
  >&2 echo ""
  >&2 echo "Expected variables:   PGHOST, PGDATABASE, PGUSER, PGPASSWORD"
  >&2 echo "Supported variables:  PBF_DATA_DIR, BORDERS_PBF_FILE, BORDERS_CSV_FILE, BORDERS_TABLE_NAME, PGPORT,"
  >&2 echo "                      BORDERS_CLEANUP, BORDERS_CLEANUP_FILE"
  >&2 echo ""
  >&2 echo "If first parameter is a file, treat it as 'import'"
  >&2 echo ""
  exit 1
}

# First argument is the PBF file to import.
# If there are no arguments, searches for the first *.pbf file in the $PBF_DATA_DIR
if [[ $# -eq 0 ]]; then
  if [ "$(ls -A "$PBF_DATA_DIR"/*.pbf 2> /dev/null)" ]; then
    for pbf_file in "$PBF_DATA_DIR"/*.pbf; do
      set -- import "$pbf_file"
      break
    done
  else
    echo "No PBF files found in the $PBF_DATA_DIR dir."
    usage
  fi
elif [[ "$1" != "import" && "$1" != "load" && "$1" != "parse" ]]; then
  if [[ -f "$1" ]]; then
    # first parameter is a filename, treat it as "import"
    set -- import "$@"
  else
    echo "Unexpected first parameter '$1'"
    usage
  fi
fi

if [[ $# != 2 ]]; then
  echo "Unexpected parameters"
  usage
fi

task="$1"

if [[ "$task" == "import" || "$task" == "parse" ]]; then
  # Parse PBF file to generate borders CSV file, and import it into PG
  IMPORT_PBF_FILE="$2"

  if [[ "$BORDERS_CLEANUP" == "true" ]]; then
    # Some OSM extracts may include references to elements that are not included in the extract
    echo "Cleaning-up broken references from $IMPORT_PBF_FILE into $BORDERS_CLEANUP_FILE"
    echo "Set BORDERS_CLEANUP=false to skip the cleanup -- it is not needed for the full planet and Geofabrik downloads"
    mkdir -p "$(dirname "${BORDERS_CLEANUP_FILE:?}")"
    rm -rf "$BORDERS_CLEANUP_FILE"
    osmconvert --drop-broken-refs -b=-180,-90,180,90 --out-pbf -o="$BORDERS_CLEANUP_FILE" "$IMPORT_PBF_FILE"
    IMPORT_PBF_FILE="$BORDERS_CLEANUP_FILE"
  else
    echo "Skipping broken references cleanup because BORDERS_CLEANUP is not set to 'true'"
    echo "Enable borders cleanup if you see 'location for one or more nodes not found in node location index' error"
  fi

  echo "Filtering $IMPORT_PBF_FILE into $BORDERS_PBF_FILE"
  mkdir -p "$(dirname "${BORDERS_PBF_FILE:?}")"
  rm -rf "$BORDERS_PBF_FILE"
  osmborder_filter -o "$BORDERS_PBF_FILE" "$IMPORT_PBF_FILE"

  echo "Creating a CSV borders file $BORDERS_CSV_FILE"
  mkdir -p "$(dirname "${BORDERS_CSV_FILE:?}")"
  rm -rf "$BORDERS_CSV_FILE"
  osmborder -o "$BORDERS_CSV_FILE" "$BORDERS_PBF_FILE"
elif [[ "$task" == "load" ]]; then
  # In "load" mode, the second parameter is the CSV file to be loaded
  BORDERS_CSV_FILE="$2"
else
  echo "Unexpected first parameter '$1'"
  usage
fi


if [[ "$task" != "parse" ]]; then
  # For backward compatibility, allow both PG* and POSTGRES_* forms,
  # with the non-standard POSTGRES_* form taking precedence.
  # An error will be raised if neither form is given, except for the PGPORT
  export PGHOST="${POSTGRES_HOST:-${PGHOST?}}"
  export PGDATABASE="${POSTGRES_DB:-${PGDATABASE?}}"
  export PGUSER="${POSTGRES_USER:-${PGUSER?}}"
  export PGPASSWORD="${POSTGRES_PASSWORD:-${PGPASSWORD?}}"
  export PGPORT="${POSTGRES_PORT:-${PGPORT:-5432}}"

  echo "Importing $BORDERS_CSV_FILE into $PGHOST:$PGPORT/$PGDATABASE as table $BORDERS_TABLE_NAME..."
  psql -c "DROP TABLE IF EXISTS $BORDERS_TABLE_NAME CASCADE;" \
       -c "CREATE TABLE $BORDERS_TABLE_NAME (osm_id bigint, admin_level int, dividing_line bool, disputed bool, maritime bool, geometry Geometry(LineString, 3857));" \
       -c "CREATE INDEX ON $BORDERS_TABLE_NAME USING gist (geometry);" \
       -c "\copy $BORDERS_TABLE_NAME FROM '$BORDERS_CSV_FILE' DELIMITER E'\t' CSV;"
fi