use rusqlite::Connection;
use super::entity::admin_level;
use geozero::ToWkt;
use super::geometry::{admin_geometry, bounding_box, mbr_center, mbr_of, merged_way_ids};
use super::id::osm_element_kind;
use super::scale::level;
use crate::database::table;
const SQL_CREATE: &str = "
CREATE TABLE IF NOT EXISTS admin_levels (
id INTEGER PRIMARY KEY,
relation_id INTEGER,
way_id INTEGER,
admin_level INTEGER NOT NULL,
wkb BLOB NOT NULL,
name VARCHAR(128) NOT NULL,
-- country_iso_code: ISO 3166-1 alpha-2 (2 chars, e.g. 'BR', 'US')
country_iso_code VARCHAR(3),
post_code VARCHAR(12),
-- merged_way_ids: jsonb array where element i is the osm way of line i of wkb (NULL unless folded)
merged_way_ids BLOB
);
";
const SQL_CREATE_INDEXES: &str = "
CREATE INDEX IF NOT EXISTS admin_levels_search_by_level
ON admin_levels (admin_level);
";
const SQL_DROP: &str = "DROP TABLE IF EXISTS admin_levels;";
pub struct admin_levels;
impl table for admin_levels {
const CREATE: &str = SQL_CREATE;
const INDEXES: &str = SQL_CREATE_INDEXES;
}
fn level_of(id: i64, raw: u8) -> Option<level> {
let level = level::new(raw);
if level.is_none() {
eprintln!("warn: admin_levels row {id} carries level {raw}, outside the scale; skipped");
}
level
}
pub(crate) fn add_merged_way_ids(conn: &Connection) {
const SQL_ADD_MERGED_WAY_IDS: &str = "ALTER TABLE admin_levels ADD COLUMN merged_way_ids BLOB";
if crate::database::has_column(conn, "main", "admin_levels", "merged_way_ids") {
return;
}
conn
.execute_batch(SQL_ADD_MERGED_WAY_IDS)
.expect("failed to add merged_way_ids to admin_levels");
}
pub(crate) fn drop_table(conn: &Connection) {
conn
.execute_batch(SQL_DROP)
.expect("failed to drop admin_levels");
}
pub struct admin_level_geom_row {
pub id: i64,
pub admin_level: level,
pub wkb: Option<admin_geometry>,
}
fn map_geom_row(row: &rusqlite::Row) -> rusqlite::Result<Option<admin_level_geom_row>> {
let id: i64 = row.get(0)?;
let Some(admin_level) = level_of(id, row.get(1)?) else {
return Ok(None);
};
Ok(Some(admin_level_geom_row {
id,
admin_level,
wkb: row.get(2)?,
}))
}
fn by_ids<T>(
conn: &Connection,
sql_prefix: &str,
ids: &[i64],
row_of: impl FnMut(&rusqlite::Row<'_>) -> rusqlite::Result<T>,
) -> Vec<T> {
if ids.is_empty() {
return vec![];
}
let sql = format!(
"{} ({})",
sql_prefix.trim(),
crate::database::placeholders_for(ids.len())
);
crate::database::query_by_ids(conn, &sql, ids.iter().copied(), row_of)
}
pub fn load_all_below_street(conn: &Connection) -> Vec<admin_level_geom_row> {
const SQL_LOAD_ALL_BELOW_STREET: &str = "
SELECT
id,
admin_level,
wkb
FROM admin_levels
WHERE admin_level < ?1
ORDER BY admin_level ASC, wkb IS NULL ASC, id ASC
";
let mut stmt = conn
.prepare(SQL_LOAD_ALL_BELOW_STREET)
.expect("failed to prepare load ancestors");
stmt
.query_map([level::street.value()], map_geom_row)
.expect("failed to query ancestors")
.filter_map(|r| r.expect("failed to read ancestor row"))
.collect()
}
pub fn load_by_ids(conn: &Connection, ids: &[i64]) -> Vec<admin_level_geom_row> {
const SQL_LOAD_BY_IDS: &str = "
SELECT
id,
admin_level,
wkb
FROM admin_levels
WHERE id IN
";
by_ids(conn, SQL_LOAD_BY_IDS, ids, map_geom_row)
.into_iter()
.flatten()
.collect()
}
pub struct admin_area_row {
pub id: i64,
pub name: String,
pub admin_level: level,
pub relation_id: Option<u64>,
pub way_id: Option<u64>,
pub wkb: Option<admin_geometry>,
}
pub fn load_full_by_ids(conn: &Connection, ids: &[i64]) -> Vec<admin_area_row> {
const SQL_LOAD_FULL_BY_IDS: &str = "
SELECT id, name, admin_level, relation_id, way_id, wkb
FROM admin_levels
WHERE id IN
";
by_ids(conn, SQL_LOAD_FULL_BY_IDS, ids, |row| {
let id: i64 = row.get(0)?;
Ok(level_of(id, row.get(2)?).map(|admin_level| admin_area_row {
id,
name: row.get(1).unwrap_or_default(),
admin_level,
relation_id: row.get(3).ok().flatten(),
way_id: row.get(4).ok().flatten(),
wkb: row.get(5).ok().flatten(),
}))
})
.into_iter()
.flatten()
.collect()
}
pub struct admin_meta_row {
pub id: i64,
pub name: String,
pub admin_level: level,
pub relation_id: Option<u64>,
pub way_id: Option<u64>,
pub country_iso_code: Option<String>,
pub post_code: Option<String>,
}
pub fn load_metadata_by_ids(
conn: &Connection,
ids: &[i64],
) -> std::collections::HashMap<i64, admin_meta_row> {
const SQL_LOAD_METADATA_BY_IDS: &str = "
SELECT id, name, admin_level, relation_id, way_id, country_iso_code, post_code
FROM admin_levels
WHERE id IN
";
by_ids(conn, SQL_LOAD_METADATA_BY_IDS, ids, |row| {
let id: i64 = row.get(0)?;
let Some(admin_level) = level_of(id, row.get(2)?) else {
return Ok(None);
};
Ok(Some(admin_meta_row {
id,
name: row.get(1)?,
admin_level,
relation_id: row.get(3)?,
way_id: row.get(4)?,
country_iso_code: row.get(5)?,
post_code: row.get(6)?,
}))
})
.into_iter()
.flatten()
.map(|r| (r.id, r))
.collect()
}
pub struct street_centre_row {
pub id: i64,
pub name: String,
pub country_iso_code: Option<String>,
pub cx: f64,
pub cy: f64,
}
pub fn streets_with_centroid(conn: &Connection) -> Vec<street_centre_row> {
const SQL_STREETS_WITH_GEOMETRY: &str = "
SELECT id, name, country_iso_code, wkb
FROM admin_levels
WHERE admin_level = ?1
AND wkb IS NOT NULL
ORDER BY id
";
let mut stmt = conn
.prepare(SQL_STREETS_WITH_GEOMETRY)
.expect("failed to prepare streets with geometry");
stmt
.query_map([level::street.value()], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, String>(1)?,
row.get::<_, Option<String>>(2)?,
row.get::<_, Vec<u8>>(3)?,
))
})
.expect("failed to query streets")
.filter_map(|r| {
let (id, name, country_iso_code, blob) = r.expect("failed to read street row");
let (cx, cy) = mbr_center(&blob)?;
Some(street_centre_row {
id,
name,
country_iso_code,
cx,
cy,
})
})
.collect()
}
pub fn geometry_by_ids(conn: &Connection, ids: &[i64]) -> Vec<(i64, admin_geometry)> {
const SQL_GEOMETRY_BY_IDS: &str = "
SELECT id, wkb
FROM admin_levels
WHERE wkb IS NOT NULL
AND id IN
";
by_ids(conn, SQL_GEOMETRY_BY_IDS, ids, |row| Ok((row.get(0)?, row.get(1)?)))
}
pub fn merged_way_ids_by_ids(
conn: &Connection,
ids: &[i64],
) -> std::collections::HashMap<i64, merged_way_ids> {
const SQL_MERGED_WAY_IDS_BY_IDS: &str = "
SELECT id, JSON(merged_way_ids)
FROM admin_levels
WHERE merged_way_ids IS NOT NULL
AND id IN
";
if !crate::database::has_column(conn, "main", "admin_levels", "merged_way_ids") {
return std::collections::HashMap::new();
}
by_ids(conn, SQL_MERGED_WAY_IDS_BY_IDS, ids, |row| {
Ok((row.get(0)?, row.get(1)?))
})
.into_iter()
.collect()
}
pub fn boxes_by_ids(
conn: &Connection,
ids: &[i64],
) -> std::collections::HashMap<i64, bounding_box> {
const SQL_HEADERS_BY_IDS: &str = "
SELECT id, SUBSTR(wkb, 1, 38)
FROM admin_levels
WHERE wkb IS NOT NULL
AND id IN
";
by_ids(conn, SQL_HEADERS_BY_IDS, ids, |row| {
let id: i64 = row.get(0)?;
let header: Vec<u8> = row.get(1)?;
Ok(mbr_of(&header).map(|mbr| (id, mbr)))
})
.into_iter()
.flatten()
.collect()
}
pub fn wkt_by_ids(conn: &Connection, ids: &[i64]) -> std::collections::HashMap<i64, String> {
geometry_by_ids(conn, ids)
.into_iter()
.filter_map(|(id, wkb)| Some((id, wkb.geometry().to_wkt().ok()?)))
.collect()
}
pub struct name_row {
pub id: i64,
pub name: String,
pub admin_level: level,
pub country_iso_code: Option<String>,
pub post_code: Option<String>,
}
pub fn load_all_names(conn: &Connection) -> Vec<name_row> {
const SQL_LOAD_ALL_NAMES: &str = "
SELECT id, name, post_code, admin_level, country_iso_code
FROM admin_levels
";
let mut stmt = conn
.prepare(SQL_LOAD_ALL_NAMES)
.expect("failed to prepare load all names");
stmt
.query_map([], |row| {
let id: i64 = row.get(0)?;
let name: String = row.get(1)?;
let post_code: Option<String> = row.get(2)?;
let country_iso_code: Option<String> = row.get(4)?;
Ok(level_of(id, row.get(3)?).map(|admin_level| name_row {
id,
name,
admin_level,
country_iso_code,
post_code,
}))
})
.expect("failed to query all names")
.filter_map(|r| r.expect("failed to read name row"))
.collect()
}
pub fn relation_ids_at_level(conn: &Connection, level: level) -> Vec<u64> {
const SQL_RELATION_IDS_AT_LEVEL: &str = "
SELECT relation_id
FROM admin_levels
WHERE admin_level = ?1
AND relation_id IS NOT NULL
";
conn
.prepare(SQL_RELATION_IDS_AT_LEVEL)
.expect("failed to prepare the relation ids of a level")
.query_map([level.value()], |row| row.get::<_, u64>(0))
.expect("failed to query the relation ids of a level")
.map(|r| r.expect("failed to read a relation id"))
.collect()
}
pub fn way_ids_at_level(conn: &Connection, level: level) -> Vec<u64> {
const SQL_WAY_IDS_AT_LEVEL: &str = "
SELECT way_id
FROM admin_levels
WHERE admin_level = ?1
AND way_id IS NOT NULL
";
conn
.prepare(SQL_WAY_IDS_AT_LEVEL)
.expect("failed to prepare the way ids of a level")
.query_map([level.value()], |row| row.get::<_, u64>(0))
.expect("failed to query the way ids of a level")
.map(|r| r.expect("failed to read a way id"))
.collect()
}
pub fn levels_present(conn: &Connection) -> Vec<level> {
const SQL_LEVELS_PRESENT: &str = "
SELECT DISTINCT admin_level
FROM admin_levels
ORDER BY admin_level
";
conn
.prepare(SQL_LEVELS_PRESENT)
.expect("failed to prepare the levels present")
.query_map([], |row| row.get::<_, u8>(0))
.expect("failed to query the levels present")
.filter_map(|raw| level::new(raw.expect("failed to read a level")))
.collect()
}
pub fn count(conn: &Connection) -> i64 {
const SQL_COUNT: &str = "
SELECT COUNT(*)
FROM admin_levels
";
conn
.query_row(SQL_COUNT, [], |row| row.get::<_, i64>(0))
.expect("failed to count admin_levels")
}
pub fn count_with_geometry(conn: &Connection) -> i64 {
const SQL_COUNT_WITH_GEOMETRY: &str = "
SELECT COUNT(*)
FROM admin_levels
WHERE wkb IS NOT NULL
";
conn
.query_row(SQL_COUNT_WITH_GEOMETRY, [], |row| row.get::<_, i64>(0))
.expect("failed to count admin_levels with geometry")
}
pub fn id_range_with_geometry(conn: &Connection) -> (i64, i64) {
const SQL_ID_RANGE_WITH_GEOMETRY: &str = "
SELECT MIN(id), MAX(id)
FROM admin_levels
WHERE wkb IS NOT NULL
";
conn
.query_row(SQL_ID_RANGE_WITH_GEOMETRY, [], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
})
.expect("failed to query id range with geometry")
}
pub fn load_wkb_page(
conn: &Connection,
last_id: i64,
max_id: i64,
limit: usize,
) -> Vec<(i64, admin_geometry)> {
const SQL_LOAD_WKB_PAGE: &str = "
SELECT id, wkb
FROM admin_levels
WHERE wkb IS NOT NULL
AND id > ?1
AND id <= ?2
ORDER BY id ASC
LIMIT ?3
";
let mut stmt = conn
.prepare_cached(SQL_LOAD_WKB_PAGE)
.expect("failed to prepare load_wkb_page");
stmt
.query_map(rusqlite::params![last_id, max_id, limit as i64], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, admin_geometry>(1)?))
})
.expect("failed to query load_wkb_page")
.map(|r| r.expect("failed to read wkb page row"))
.collect()
}
pub fn delete_by_ids(conn: &Connection, ids: &[i64]) -> usize {
const SQL_DELETE_BY_IDS: &str = "
DELETE FROM admin_levels
WHERE id IN
";
if ids.is_empty() {
return 0;
}
let sql = format!(
"{} ({})",
SQL_DELETE_BY_IDS.trim(),
crate::database::placeholders_for(ids.len())
);
conn
.execute(&sql, rusqlite::params_from_iter(ids))
.expect("failed to delete admin_levels by ids")
}
pub fn batch_upsert(conn: &Connection, rows: &[admin_level]) -> i64 {
upsert(conn, rows.iter().map(|row| (row, None)))
}
pub fn batch_upsert_folded(conn: &Connection, rows: &[(admin_level, merged_way_ids)]) -> i64 {
upsert(conn, rows.iter().map(|(row, way_ids)| (row, Some(way_ids))))
}
fn upsert<'a>(
conn: &Connection,
rows: impl Iterator<Item = (&'a admin_level, Option<&'a merged_way_ids>)>,
) -> i64 {
const SQL_UPSERT: &str = "
INSERT INTO admin_levels (
id,
relation_id,
way_id,
admin_level,
name,
country_iso_code,
post_code,
wkb,
merged_way_ids
) VALUES (
?1,
?2,
?3,
?4,
?5,
?6,
?7,
?8,
JSONB(?9)
)
ON CONFLICT (id)
DO UPDATE SET
name = excluded.name,
country_iso_code = excluded.country_iso_code,
post_code = excluded.post_code,
wkb = excluded.wkb,
merged_way_ids = excluded.merged_way_ids
";
let tx = conn
.unchecked_transaction()
.expect("failed to begin transaction");
let mut total_changes: i64 = 0;
{
let mut stmt = tx.prepare(SQL_UPSERT).expect("failed to prepare upsert");
for (row, way_ids) in rows {
let osm_id = Some(row.id.osm_id());
let (relation_id, way_id) = match row.id.kind() {
osm_element_kind::relation => (osm_id, None),
osm_element_kind::way => (None, osm_id),
};
let changes = stmt
.execute(rusqlite::params![
row.id.raw(),
relation_id,
way_id,
row.level.value(),
row.name,
row.country_iso_code,
row.post_code,
row.wkb,
way_ids,
])
.expect("failed to upsert admin_level row");
total_changes += changes as i64;
}
}
tx.commit().expect("failed to commit transaction");
total_changes
}