use rusqlite::Connection;
use crate::database::table;
#[derive(Clone, Copy)]
pub enum chunk_type {
header = 0,
data = 1,
}
impl rusqlite::types::FromSql for chunk_type {
fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
match value.as_i64()? {
0 => Ok(chunk_type::header),
1 => Ok(chunk_type::data),
other => Err(rusqlite::types::FromSqlError::OutOfRange(other)),
}
}
}
impl rusqlite::types::ToSql for chunk_type {
fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
Ok((*self as i8).into())
}
}
pub struct osm_pbf_blob_chunk {
pub id: u32,
pub file_id: u32,
pub first_byte: u64,
pub chunk_size: u64,
pub data_first_byte: u64,
pub data_size: u64,
pub chunk_type: chunk_type,
}
const SQL_CREATE: &str = "
CREATE TABLE IF NOT EXISTS osm_pbf_blob_chunks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_id INTEGER NOT NULL,
first_byte INTEGER NOT NULL,
chunk_size INTEGER NOT NULL,
data_first_byte INTEGER NOT NULL,
data_size INTEGER NOT NULL,
chunk_type INTEGER NOT NULL,
UNIQUE (file_id, first_byte)
);
";
const SQL_CREATE_INDEXES: &str = "
CREATE INDEX IF NOT EXISTS osm_pbf_blob_chunks_search_by_file_and_type
ON osm_pbf_blob_chunks(file_id, chunk_type);
";
pub struct osm_pbf_blob_chunks;
impl table for osm_pbf_blob_chunks {
const CREATE: &str = SQL_CREATE;
const INDEXES: &str = SQL_CREATE_INDEXES;
}
pub fn count_by_file_id(conn: &Connection, file_id: u32) -> i64 {
const SQL_COUNT_BY_FILE_ID: &str = "
SELECT COUNT(*)
FROM osm_pbf_blob_chunks
WHERE file_id = ?1
";
conn
.query_row(SQL_COUNT_BY_FILE_ID, rusqlite::params![file_id], |row| {
row.get(0)
})
.expect("failed to count blob chunks")
}
pub fn file_ids(conn: &Connection) -> Vec<u32> {
const SQL_FILE_IDS: &str = "
SELECT DISTINCT file_id
FROM osm_pbf_blob_chunks
ORDER BY file_id
";
let mut stmt = conn
.prepare(SQL_FILE_IDS)
.expect("failed to prepare the files of the osm_data");
stmt
.query_map([], |row| row.get(0))
.expect("failed to list the files of the osm_data")
.map(|row| row.expect("failed to read a file id"))
.collect()
}
pub fn delete_by_file_id(conn: &Connection, file_id: u32) {
const SQL_DELETE_BY_FILE_ID: &str = "
DELETE FROM osm_pbf_blob_chunks
WHERE file_id = ?1
";
conn
.execute(SQL_DELETE_BY_FILE_ID, rusqlite::params![file_id])
.expect("failed to delete blob chunks");
}
pub fn get_header_chunk(conn: &Connection, file_id: u32) -> Option<osm_pbf_blob_chunk> {
const SQL_HEADER_CHUNK: &str = "
SELECT
id,
first_byte,
chunk_size,
data_first_byte,
data_size,
chunk_type
FROM osm_pbf_blob_chunks
WHERE chunk_type = 0
AND file_id = ?1
LIMIT 1
";
let mut stmt = conn
.prepare(SQL_HEADER_CHUNK)
.expect("failed to prepare header chunk query");
stmt
.query_row(rusqlite::params![file_id], |row| {
Ok(osm_pbf_blob_chunk {
id: row.get(0).unwrap(),
file_id,
first_byte: row.get(1).unwrap(),
chunk_size: row.get(2).unwrap(),
data_first_byte: row.get(3).unwrap(),
data_size: row.get(4).unwrap(),
chunk_type: row.get(5).unwrap(),
})
})
.ok()
}
pub fn get_data_chunks(conn: &Connection, file_id: u32) -> Vec<osm_pbf_blob_chunk> {
const SQL_DATA_CHUNKS: &str = "
SELECT
id,
first_byte,
chunk_size,
data_first_byte,
data_size,
chunk_type
FROM osm_pbf_blob_chunks
WHERE file_id = ?1
ORDER BY first_byte ASC
";
let mut statement = conn
.prepare(SQL_DATA_CHUNKS)
.expect("failed to prepare data chunks query");
statement
.query_map(rusqlite::params![file_id], |row| {
Ok(osm_pbf_blob_chunk {
id: row.get(0).unwrap(),
file_id,
first_byte: row.get(1).unwrap(),
chunk_size: row.get(2).unwrap(),
data_first_byte: row.get(3).unwrap(),
data_size: row.get(4).unwrap(),
chunk_type: row.get(5).unwrap(),
})
})
.expect("failed to query blob chunks")
.collect::<Result<Vec<_>, _>>()
.expect("failed to collect blob chunks")
}
pub fn batch_insert(conn: &Connection, chunks: &[osm_pbf_blob_chunk]) {
const SQL_INSERT: &str = "
INSERT INTO osm_pbf_blob_chunks (
file_id,
first_byte,
chunk_size,
data_first_byte,
data_size,
chunk_type
) VALUES (
?1,
?2,
?3,
?4,
?5,
?6
)
ON CONFLICT (file_id, first_byte) DO UPDATE SET
chunk_size = excluded.chunk_size,
data_first_byte = excluded.data_first_byte,
data_size = excluded.data_size,
chunk_type = excluded.chunk_type
";
let tx = conn
.unchecked_transaction()
.expect("failed to begin transaction");
{
let mut stmt = tx
.prepare(SQL_INSERT)
.expect("failed to prepare osm_pbf_blob_chunks insert");
for chunk in chunks {
stmt
.execute(rusqlite::params![
chunk.file_id,
chunk.first_byte,
chunk.chunk_size,
chunk.data_first_byte,
chunk.data_size,
chunk.chunk_type,
])
.expect("failed to insert blob chunk");
}
}
tx.commit().expect("failed to commit");
}