geolite 0.0.13

use rusqlite::{Connection, OpenFlags};

// bumped whenever the on-disk schema changes in a way that makes builds incompatible;
// stamped into every writable database via PRAGMA user_version. open_write_main refuses a
// database stamped with another version, and `geolite merge` refuses to combine one.
pub const SCHEMA_VERSION: u32 = 3;

pub mod jsonb;
pub mod merge;

#[cfg(test)]
#[path = "jsonb_fixtures.test.rs"]
pub(crate) mod jsonb_fixtures;

pub(crate) trait table {
  const CREATE: &str;
  const INDEXES: &str;

  fn create_table(conn: &rusqlite::Connection) {
    conn
      .execute_batch(Self::CREATE)
      .expect("failed to create table");
  }

  fn create_indexes(conn: &rusqlite::Connection) {
    conn
      .execute_batch(Self::INDEXES)
      .expect("failed to create indexes");
  }
}

pub(crate) fn placeholders_for(count: usize) -> String {
  vec!["?"; count].join(",")
}

pub(crate) fn query_by_ids<T>(
  conn: &Connection,
  sql: &str,
  ids: impl IntoIterator<Item = i64>,
  row_of: impl FnMut(&rusqlite::Row<'_>) -> rusqlite::Result<T>,
) -> Vec<T> {
  let params: Vec<rusqlite::types::Value> = ids
    .into_iter()
    .map(rusqlite::types::Value::Integer)
    .collect();
  let mut stmt = conn
    .prepare(sql)
    .expect("failed to prepare query by ids");
  stmt
    .query_map(rusqlite::params_from_iter(params.iter()), row_of)
    .expect("failed to query by ids")
    .map(|r| r.expect("failed to read row"))
    .collect()
}

const INSERT_CHUNK_SIZE: usize = 10_000;

pub(crate) fn insert_in_chunks<const N: usize>(
  conn: &Connection,
  sql_head: &str,
  rows: &[[&dyn rusqlite::ToSql; N]],
) {
  // SQLITE_MAX_VARIABLE_NUMBER is 32766 on the bundled build: a wider row would overflow it
  const { assert!(N * INSERT_CHUNK_SIZE <= 32_766) };
  for chunk in rows.chunks(INSERT_CHUNK_SIZE) {
    let sql = multi_insert_sql(sql_head, N, chunk.len());
    let mut stmt = conn
      .prepare_cached(&sql)
      .expect("failed to prepare chunked insert");
    stmt
      .execute(rusqlite::params_from_iter(chunk.iter().flatten()))
      .expect("failed to insert chunk");
  }
}

fn multi_insert_sql(sql_head: &str, columns: usize, rows: usize) -> String {
  use std::fmt::Write;
  let mut sql = String::with_capacity(sql_head.len() + rows * (columns * 8 + 4));
  sql.push_str(sql_head);
  for row in 0..rows {
    if row > 0 {
      sql.push_str(",\n");
    }
    sql.push_str("  (");
    for column in 0..columns {
      if column > 0 {
        sql.push_str(", ");
      }
      write!(sql, "?{}", row * columns + column + 1).unwrap();
    }
    sql.push(')');
  }
  sql
}

pub fn osm_data_path(main_path: &str) -> String {
  if main_path == ":memory:" {
    return ":memory:".to_string();
  }
  let p = std::path::Path::new(main_path);
  let stem = p.file_stem().and_then(|s| s.to_str()).unwrap_or("database");
  let parent = p.parent().unwrap_or(std::path::Path::new(""));
  parent
    .join(format!("{stem}.osm_data.sqlite3"))
    .to_string_lossy()
    .into_owned()
}

pub fn remove_osm_data_files(main_path: &str) -> u64 {
  if main_path == ":memory:" {
    return 0;
  }
  let sibling = osm_data_path(main_path);
  let mut bytes = 0u64;
  for suffix in ["", "-wal", "-shm"] {
    let p = format!("{sibling}{suffix}");
    if std::path::Path::new(&p).exists() {
      bytes += std::fs::metadata(&p).map(|m| m.len()).unwrap_or(0);
      std::fs::remove_file(&p).expect("failed to remove osm_data sibling");
    }
  }
  bytes
}

pub fn compact(conn: &Connection) -> (u64, u64) {
  const SQL_PAGE_SIZE: &str = "PRAGMA page_size";
  const SQL_PAGE_COUNT: &str = "PRAGMA page_count";
  const SQL_COMPACT: &str = "
    ANALYZE;
    PRAGMA optimize;
    PRAGMA wal_checkpoint(TRUNCATE);
    VACUUM;
  ";

  let page_size: u64 = conn
    .query_row(SQL_PAGE_SIZE, [], |r| r.get(0))
    .unwrap_or(4096);
  let pages_before: u64 = conn
    .query_row(SQL_PAGE_COUNT, [], |r| r.get(0))
    .unwrap_or(0);
  conn
    .execute_batch(SQL_COMPACT)
    .expect("failed to compact sqlite file");
  let pages_after: u64 = conn
    .query_row(SQL_PAGE_COUNT, [], |r| r.get(0))
    .unwrap_or(0);
  (page_size * pages_before, page_size * pages_after)
}

pub fn destroy_data(
  path: &str,
  osm_pbf_blob_chunks: bool,
  osm_data: bool,
  admin_levels: bool,
  house_numbers: bool,
) {
  if osm_pbf_blob_chunks {
    remove_osm_data_files(path);
  }
  let conn = if osm_data && !osm_pbf_blob_chunks {
    let conn = open_write(path);
    crate::osm_node::repository::drop_table(&conn);
    crate::osm_way::repository::drop_table(&conn);
    crate::osm_relation::repository::drop_table(&conn);
    conn
  } else {
    open_write_main(path)
  };
  if house_numbers {
    crate::house_number::repository::drop_table(&conn);
  }
  if admin_levels {
    crate::admin_level_hierarchy::repository::drop_table(&conn);
    crate::admin_level::spatial_index::drop_table(&conn);
    crate::admin_level::repository::drop_table(&conn);
    crate::house_number::repository::drop_table(&conn);
  }
  conn.execute_batch("VACUUM;").expect("failed to vacuum");
}

pub(crate) fn has_column(conn: &Connection, schema: &str, table: &str, column: &str) -> bool {
  const SQL_HAS_COLUMN: &str = "
    SELECT COUNT(*)
    FROM PRAGMA_TABLE_INFO(?1, ?2)
    WHERE name = ?3
  ";

  conn
    .query_row(
      SQL_HAS_COLUMN,
      rusqlite::params![table, schema, column],
      |row| row.get::<_, i64>(0),
    )
    .expect("failed to read table info")
    > 0
}

pub fn open_write_main(path: &str) -> Connection {
  // create_dir_all races between concurrent geolite processes sharing an ancestor directory:
  // std can return AlreadyExists spuriously, so a parent that already exists is success.
  if let Some(parent) = std::path::Path::new(path).parent()
    && !parent.as_os_str().is_empty()
    && let Err(error) = std::fs::create_dir_all(parent)
    && !parent.is_dir()
  {
    panic!("failed to create sqlite parent dir: {error}");
  }
  let conn = Connection::open(path).expect("failed to open sqlite");
  require_schema_version(&conn, path);
  conn
    .execute_batch(
      "PRAGMA journal_mode=WAL;
       PRAGMA synchronous=NORMAL;
       PRAGMA cache_size=-32768;
       PRAGMA temp_store=MEMORY;",
    )
    .expect("failed to set pragmas");
  conn
    .pragma_update(None, "user_version", SCHEMA_VERSION)
    .expect("failed to set user_version");
  crate::osm_pbf_file::osm_pbf_files::create_table(&conn);
  crate::osm_pbf_file::repository::add_origin_wkt(&conn);
  crate::admin_level::admin_levels::create_table(&conn);
  crate::admin_level::repository::add_merged_way_ids(&conn);
  crate::admin_level_hierarchy::admin_levels_hierarchy::create_table(&conn);
  crate::admin_level::spatial_index::create_table(&conn);
  crate::house_number::house_numbers::create_table(&conn);
  conn
}

pub fn open_write(path: &str) -> Connection {
  let conn = open_write_main(path);
  let osm_data = osm_data_path(path);
  conn
    .execute("ATTACH DATABASE ?1 AS osm_data", [&osm_data])
    .expect("failed to attach osm_data");
  conn
    .execute_batch(
      "PRAGMA osm_data.journal_mode=WAL;
       PRAGMA osm_data.synchronous=NORMAL;
       PRAGMA osm_data.cache_size=-32768;",
    )
    .expect("failed to set osm_data pragmas");
  crate::osm_pbf_file::osm_pbf_blob_chunks::create_table(&conn);
  crate::osm_node::osm_nodes::create_table(&conn);
  crate::osm_way::osm_ways::create_table(&conn);
  crate::osm_relation::osm_relations::create_table(&conn);
  conn
}

fn require_schema_version(conn: &Connection, path: &str) {
  let stamped: u32 = conn
    .pragma_query_value(None, "user_version", |row| row.get::<_, i64>(0))
    .map(|v| v as u32)
    .expect("failed to read user_version");
  if stamped != 0 && stamped != SCHEMA_VERSION {
    panic!(
      "incompatible schema version on {path}: found {stamped}, expected {SCHEMA_VERSION} \u{2014} rebuild required"
    );
  }
}

pub fn open_readonly(path: &str) -> Connection {
  let conn = Connection::open_with_flags(
    path,
    OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
  )
  .expect("failed to open sqlite readonly");
  conn
    .execute_batch(
      "PRAGMA cache_size=-262144;
       PRAGMA mmap_size=2147483648;
       PRAGMA temp_store=MEMORY;
       PRAGMA query_only=1;",
    )
    .expect("failed to set pragmas");
  let osm_data = osm_data_path(path);
  if std::path::Path::new(&osm_data).exists() {
    let canonical = std::fs::canonicalize(&osm_data)
      .map(|p| p.to_string_lossy().into_owned())
      .unwrap_or(osm_data);
    let uri = format!("file:{canonical}?mode=ro");
    conn
      .execute("ATTACH DATABASE ?1 AS osm_data", [&uri])
      .expect("failed to attach osm_data readonly");
  }
  conn
}

// reads the schema version stamped into a database (PRAGMA user_version). databases built before
// versioning existed return 0. used by `geolite merge` to reject incompatible builds before mutating.
pub fn read_user_version(path: &str) -> u32 {
  let conn = open_readonly(path);
  conn
    .pragma_query_value(None, "user_version", |row| row.get::<_, i64>(0))
    .map(|v| v as u32)
    .unwrap_or(0)
}