use std::collections::VecDeque;
use std::{collections::HashMap, path::Path, time::Duration};
use futures::future::try_join_all;
use tokio::sync::RwLock;
use turso::{Connection, Rows};
use self::sql::*;
use crate::Error;
use crate::Result;
use crate::Topic;
use crate::message::filter::MessageFilter;
pub mod sql;
pub enum ConnectionType {
Atomic(Db3File),
Splitted(SplittedRosbag),
}
impl ConnectionType {
pub fn iter<'a>(&'a self) -> std::slice::Iter<'a, Db3File> {
match self {
ConnectionType::Atomic(db3_file) => std::slice::from_ref(db3_file).iter(),
ConnectionType::Splitted(splitted_rosbag) => splitted_rosbag.connections.iter(),
}
}
pub fn as_slice(&self) -> &[Db3File] {
match self {
ConnectionType::Atomic(db3_file) => std::slice::from_ref(db3_file),
ConnectionType::Splitted(splitted_rosbag) => &splitted_rosbag.connections,
}
}
pub async fn run_on_connections<F: AsyncFn(&Db3File) -> Result<Rows>>(
&self,
f: F,
) -> Result<VecDeque<Rows>> {
let rows = try_join_all(self.iter().map(|db3| f(db3))).await?;
Ok(rows.into())
}
}
pub struct SplittedRosbag {
pub connections: Vec<Db3File>,
pub split_after: Duration,
}
#[derive(Debug)]
pub struct Db3File {
pub connection: Connection,
pub path: std::path::PathBuf,
pub topic_ids: RwLock<HashMap<String, i64>>,
}
impl Db3File {
pub async fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
let path = path.as_ref();
let connection = Self::open_connection(path).await?;
Ok(Db3File {
path: path.to_path_buf(),
topic_ids: RwLock::new(Self::read_topics(&connection).await?),
connection,
})
}
async fn read_topics(connection: &Connection) -> Result<HashMap<String, i64>> {
let mut rows = connection.query(SQL_STATEMENT_SELECT_TOPIC_IDS, ()).await?;
let mut topics = HashMap::new();
while let Some(row) = rows.next().await? {
topics.insert(row.get(1)?, row.get(0)?);
}
Ok(topics)
}
async fn open_connection(path: &Path) -> Result<Connection> {
if !path.is_file() {
return Err(Error::Db3FileNotFound(path.to_path_buf()));
}
let path_str = path
.to_str()
.ok_or_else(|| Error::NonUtf8Path(path.to_path_buf()))?;
let db = turso::Builder::new_local(path_str).build().await?;
let connection = db.connect()?;
Ok(connection)
}
pub async fn init<P: AsRef<Path>>(path: P) -> Result<Self> {
let path = path.as_ref();
if path.exists() {
return Err(Error::Db3FileAlreadyExists(path.to_path_buf()));
}
let path_str = path
.to_str()
.ok_or_else(|| Error::NonUtf8Path(path.to_path_buf()))?;
let db = turso::Builder::new_local(path_str).build().await?;
let connection = db.connect()?;
connection.execute(SQL_STATEMENT_CREATE_TABLE, ()).await?;
connection
.execute(SQL_STATEMENT_CREATE_METADATA, ())
.await?;
connection.execute(SQL_STATEMENT_CREATE_TOPICS, ()).await?;
connection
.execute(SQL_STATEMENT_CREATE_MSG_DEFINITIONS, ())
.await?;
connection.execute(SQL_STATEMENT_CREATE_MSGS, ()).await?;
connection.execute(SQL_STATEMENT_CREATE_INDEX, ()).await?;
connection.execute(SQL_STATEMENT_INSERT_SCHEMA, ()).await?;
Ok(Db3File {
connection,
path: path.to_path_buf(),
topic_ids: RwLock::new(HashMap::new()),
})
}
pub async fn get_topics(&self) -> Result<HashMap<String, Topic>> {
let mut rows = self
.connection
.query(SQL_STATEMENT_SELECT_TOPICS, ())
.await?;
let mut topics = HashMap::new();
while let Some(row) = rows.next().await? {
let topic = Topic::try_from(row)?;
topics.insert(topic.name.clone(), topic);
}
Ok(topics)
}
pub async fn topic_id(&self, topic: &str) -> Result<i64> {
self.topic_ids
.read()
.await
.get(topic)
.copied()
.ok_or(Error::TopicNotFound(topic.to_string()))
}
pub async fn query_messages(&self, filter: &MessageFilter<'_>, delete: bool) -> Result<Rows> {
let topic_count = filter.topics.as_ref().map(|topics| topics.len());
let sql = sql_statement_message_filter(
topic_count,
filter.begin.is_some(),
filter.end.is_some(),
delete,
);
let params = sql_params_message_filter(filter, &&self.topic_ids.read().await)?;
let params = turso::params::params_from_iter(params);
let mut statement = self.connection.prepare(&sql).await?;
let rows = statement.query(params).await?;
Ok(rows)
}
pub async fn insert_message(&self, topic_id: i64, timestamp: i64, data: Vec<u8>) -> Result<()> {
self.connection
.execute(
SQL_STATEMENT_INSERT_MSG,
(turso::Value::Integer(topic_id), timestamp, data),
)
.await?;
Ok(())
}
pub async fn time_range(&self) -> Result<(Option<i64>, Option<i64>)> {
let mut rows = self
.connection
.query(SQL_STATEMENT_SELECT_TIME_RANGE, ())
.await?;
let row = rows
.next()
.await?
.ok_or(turso::Error::QueryReturnedNoRows)?;
let min = row.get_value(0)?.as_integer().copied();
let max = row.get_value(1)?.as_integer().copied();
Ok((min, max))
}
pub async fn get_index(&self) -> Result<crate::Index> {
let mut rows = self
.connection
.query(SQL_STATEMENT_SELECT_TOPIC_TIMESTAMPS, ())
.await?;
let mut index: HashMap<String, Vec<i64>> = HashMap::new();
while let Some(row) = rows.next().await? {
index.entry(row.get(0)?).or_default().push(row.get(1)?);
}
let mut index: crate::Index = index
.into_iter()
.map(|(topic_name, timestamps)| crate::IndexEntry {
topic_name,
timestamps,
})
.collect();
index.sort_by(|a, b| a.topic_name.cmp(&b.topic_name));
Ok(index)
}
}
impl SplittedRosbag {
pub fn iter(&self) -> std::slice::Iter<'_, Db3File> {
return self.connections.iter();
}
}
pub(crate) fn find_db3_files<P: AsRef<std::path::Path>>(dir: P) -> Result<Vec<std::path::PathBuf>> {
let paths = std::fs::read_dir(dir)?
.into_iter()
.filter_map(|e| e.ok())
.map(|e| e.path())
.filter(|p| p.is_file() && p.extension().is_some_and(|ext| ext == "db3"))
.collect();
Ok(paths)
}