use std::str::FromStr;
use std::time::Duration;
use sqlx::mysql::{MySqlConnectOptions, MySqlPoolOptions};
use sqlx::pool::PoolOptions;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use sqlx::{Database, MySqlPool, PgPool, SqlitePool};
use tokio::sync::mpsc;
use crate::{Log, LogError, is_heartbeat};
const ASYNC_LOG_CHANNEL_CAPACITY: usize = 1024;
fn io_err<E: std::fmt::Display>(e: E) -> LogError {
LogError::Io(e.to_string())
}
#[derive(Debug, Clone, Copy)]
pub struct SqlLogPoolOptions {
pub max_connections: u32,
pub min_connections: u32,
pub acquire_timeout: Duration,
pub idle_timeout: Option<Duration>,
pub max_lifetime: Option<Duration>,
}
impl Default for SqlLogPoolOptions {
fn default() -> Self {
Self {
max_connections: 10,
min_connections: 0,
acquire_timeout: Duration::from_secs(30),
idle_timeout: None,
max_lifetime: None,
}
}
}
#[derive(Debug, Clone)]
pub struct SqlLogConfig {
pub url: String,
pub incoming_table: String,
pub outgoing_table: String,
pub event_table: String,
pub include_heartbeats: bool,
pub pool: SqlLogPoolOptions,
pub session_id: String,
}
impl SqlLogConfig {
pub fn new(url: impl Into<String>) -> Self {
Self {
url: url.into(),
incoming_table: "log_incoming".to_owned(),
outgoing_table: "log_outgoing".to_owned(),
event_table: "log_event".to_owned(),
include_heartbeats: true,
pool: SqlLogPoolOptions::default(),
session_id: "default".to_owned(),
}
}
}
fn valid_identifier(s: &str) -> Result<(), LogError> {
let ok = !s.is_empty()
&& s.len() <= 64
&& s.chars()
.next()
.is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
&& s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_');
if ok {
Ok(())
} else {
Err(LogError::Io(format!(
"{s:?} is not a valid SQL table identifier"
)))
}
}
enum Entry {
Message {
direction: &'static str,
text: String,
},
Event {
text: String,
},
}
enum Pool {
Sqlite(SqlitePool),
Postgres(PgPool),
MySql(MySqlPool),
}
fn apply_pool_options<DB: Database>(
opts: PoolOptions<DB>,
cfg: &SqlLogPoolOptions,
) -> PoolOptions<DB> {
opts.max_connections(cfg.max_connections)
.min_connections(cfg.min_connections)
.acquire_timeout(cfg.acquire_timeout)
.idle_timeout(cfg.idle_timeout)
.max_lifetime(cfg.max_lifetime)
}
async fn connect_pool(url: &str, pool_cfg: &SqlLogPoolOptions) -> Result<Pool, LogError> {
if url.starts_with("postgres://") || url.starts_with("postgresql://") {
let opts = PgConnectOptions::from_str(url).map_err(io_err)?;
let pool = apply_pool_options(PgPoolOptions::new(), pool_cfg)
.connect_with(opts)
.await
.map_err(io_err)?;
Ok(Pool::Postgres(pool))
} else if url.starts_with("mysql://") {
let opts = MySqlConnectOptions::from_str(url).map_err(io_err)?;
let pool = apply_pool_options(MySqlPoolOptions::new(), pool_cfg)
.connect_with(opts)
.await
.map_err(io_err)?;
Ok(Pool::MySql(pool))
} else {
let opts = SqliteConnectOptions::from_str(url)
.map_err(io_err)?
.create_if_missing(true);
let pool = apply_pool_options(SqlitePoolOptions::new(), pool_cfg)
.connect_with(opts)
.await
.map_err(io_err)?;
Ok(Pool::Sqlite(pool))
}
}
async fn ensure_table(pool: &Pool, table: &str) -> Result<(), LogError> {
match pool {
Pool::Sqlite(p) => {
sqlx::query(sqlx::AssertSqlSafe(format!(
"CREATE TABLE IF NOT EXISTS {table} (\
id INTEGER PRIMARY KEY AUTOINCREMENT, text TEXT NOT NULL, \
logged_at BIGINT, session_id TEXT)"
)))
.execute(p)
.await
.map_err(io_err)?;
}
Pool::Postgres(p) => {
sqlx::query(sqlx::AssertSqlSafe(format!(
"CREATE TABLE IF NOT EXISTS {table} (id BIGSERIAL PRIMARY KEY, text TEXT NOT NULL, \
logged_at BIGINT, session_id TEXT)"
)))
.execute(p)
.await
.map_err(io_err)?;
}
Pool::MySql(p) => {
sqlx::query(sqlx::AssertSqlSafe(format!(
"CREATE TABLE IF NOT EXISTS {table} (\
id BIGINT AUTO_INCREMENT PRIMARY KEY, text TEXT NOT NULL, \
logged_at BIGINT, session_id VARCHAR(255))"
)))
.execute(p)
.await
.map_err(io_err)?;
}
}
match pool {
Pool::Sqlite(p) => {
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"ALTER TABLE {table} ADD COLUMN logged_at BIGINT"
)))
.execute(p)
.await;
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"ALTER TABLE {table} ADD COLUMN session_id TEXT"
)))
.execute(p)
.await;
}
Pool::Postgres(p) => {
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS logged_at BIGINT"
)))
.execute(p)
.await;
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS session_id TEXT"
)))
.execute(p)
.await;
}
Pool::MySql(p) => {
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"ALTER TABLE {table} ADD COLUMN logged_at BIGINT"
)))
.execute(p)
.await;
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"ALTER TABLE {table} ADD COLUMN session_id VARCHAR(255)"
)))
.execute(p)
.await;
}
}
Ok(())
}
fn now_unix() -> i64 {
time::OffsetDateTime::now_utc().unix_timestamp()
}
async fn insert_text(pool: &Pool, table: &str, session_id: &str, text: &str) {
match pool {
Pool::Sqlite(p) => {
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"INSERT INTO {table} (text, logged_at, session_id) VALUES (?, ?, ?)"
)))
.bind(text)
.bind(now_unix())
.bind(session_id)
.execute(p)
.await;
}
Pool::Postgres(p) => {
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"INSERT INTO {table} (text, logged_at, session_id) VALUES ($1, $2, $3)"
)))
.bind(text)
.bind(now_unix())
.bind(session_id)
.execute(p)
.await;
}
Pool::MySql(p) => {
let _ = sqlx::query(sqlx::AssertSqlSafe(format!(
"INSERT INTO {table} (text, logged_at, session_id) VALUES (?, ?, ?)"
)))
.bind(text)
.bind(now_unix())
.bind(session_id)
.execute(p)
.await;
}
}
}
pub struct SqlLog {
tx: std::sync::Mutex<Option<mpsc::Sender<Entry>>>,
task: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
include_heartbeats: bool,
}
impl SqlLog {
pub async fn connect(url: &str) -> Result<Self, LogError> {
Self::connect_with_config(SqlLogConfig::new(url)).await
}
pub async fn connect_with_config(config: SqlLogConfig) -> Result<Self, LogError> {
valid_identifier(&config.incoming_table)?;
valid_identifier(&config.outgoing_table)?;
valid_identifier(&config.event_table)?;
let pool = connect_pool(&config.url, &config.pool).await?;
ensure_table(&pool, &config.incoming_table).await?;
ensure_table(&pool, &config.outgoing_table).await?;
ensure_table(&pool, &config.event_table).await?;
let (tx, mut rx) = mpsc::channel::<Entry>(ASYNC_LOG_CHANNEL_CAPACITY);
let incoming_table = config.incoming_table;
let outgoing_table = config.outgoing_table;
let event_table = config.event_table;
let session_id = config.session_id;
let task = tokio::spawn(async move {
while let Some(entry) = rx.recv().await {
match entry {
Entry::Message { direction, text } => {
let table = if direction == "I" {
&incoming_table
} else {
&outgoing_table
};
insert_text(&pool, table, &session_id, &text).await;
}
Entry::Event { text } => {
insert_text(&pool, &event_table, &session_id, &text).await;
}
}
}
});
Ok(Self {
tx: std::sync::Mutex::new(Some(tx)),
task: std::sync::Mutex::new(Some(task)),
include_heartbeats: config.include_heartbeats,
})
}
fn send(&self, entry: Entry) {
if let Ok(guard) = self.tx.lock()
&& let Some(tx) = guard.as_ref()
{
let _ = tx.try_send(entry);
}
}
}
#[async_trait::async_trait]
impl Log for SqlLog {
fn on_incoming(&self, message: &str) {
if is_heartbeat(message) && !self.include_heartbeats {
return;
}
self.send(Entry::Message {
direction: "I",
text: message.to_owned(),
});
}
fn on_outgoing(&self, message: &str) {
if is_heartbeat(message) && !self.include_heartbeats {
return;
}
self.send(Entry::Message {
direction: "O",
text: message.to_owned(),
});
}
fn on_event(&self, text: &str) {
self.send(Entry::Event {
text: text.to_owned(),
});
}
async fn shutdown(&self) {
let tx = self.tx.lock().ok().and_then(|mut guard| guard.take());
drop(tx);
let task = self.task.lock().ok().and_then(|mut guard| guard.take());
if let Some(task) = task {
let _ = task.await;
}
}
}