jianmu
/
openobserve
mirror of https://github.com/openobserve/openobserve
0
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: unit tests (#2542) 

- [x] change sqlite backend to a Mutex
- [x] fix unit tests
This commit is contained in:
Hengfei Yang 2024-01-24 22:42:20 +08:00 committed by GitHub
parent cc404ecb7e
commit 325aa4b3f2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
33 changed files with 710 additions and 1157 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
coverage.sh
View File

@ -3,9 +3,9 @@ set -eu -o pipefail
# set -x
export PS4='+ [${BASH_SOURCE[0]##*/}:${LINENO}${FUNCNAME[0]:+:${FUNCNAME[0]}}] '
export COVERAGE_FUNCTIONS=${COVERAGE_FUNCTIONS:-43}
export COVERAGE_LINES=${COVERAGE_LINES:-37}
export COVERAGE_REGIONS=${COVERAGE_REGIONS:-25}
export COVERAGE_FUNCTIONS=${COVERAGE_FUNCTIONS:-40}
export COVERAGE_LINES=${COVERAGE_LINES:-34}
export COVERAGE_REGIONS=${COVERAGE_REGIONS:-23}
usage() {
cat <<EOF

1
src/cli/basic/cli.rs
View File

@ -168,7 +168,6 @@ pub async fn cli() -> Result<bool, anyhow::Error> {
db::compact::stats::set_offset(0, None).await?;
// reset stream stats table data
infra::file_list::reset_stream_stats().await?;
infra::file_list::set_initialised().await?;
// load stream list
db::schema::cache().await?;
// update stats from file list

86
src/common/infra/db/mod.rs
View File

@ -18,16 +18,10 @@ use std::sync::Arc;
use ahash::HashMap;
use async_trait::async_trait;
use bytes::Bytes;
use config::{
meta::stream::{FileKey, FileMeta},
CONFIG,
};
use config::CONFIG;
use tokio::sync::{mpsc, OnceCell};
use crate::common::{
infra::errors::Result,
meta::{meta_store::MetaStore, stream::StreamStats},
};
use crate::common::{infra::errors::Result, meta::meta_store::MetaStore};
pub mod dynamo;
pub mod etcd;
@ -191,82 +185,6 @@ pub struct MetaRecord {
pub value: String,
}
#[derive(Debug)]
pub enum DbEvent {
Meta(DbEventMeta),
FileList(DbEventFileList),
FileListDeleted(DbEventFileListDeleted),
StreamStats(DbEventStreamStats),
CreateTableMeta,
CreateTableFileList,
CreateTableFileListIndex,
Shutdown,
}
pub enum DbEventMeta {
Put(String, Bytes, bool),
Delete(String, bool, bool),
}
impl std::fmt::Debug for DbEventMeta {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DbEventMeta::Put(key, ..) => write!(f, "Put({})", key),
DbEventMeta::Delete(key, ..) => write!(f, "Delete({})", key),
}
}
}
pub enum DbEventFileList {
Add(String, FileMeta),
BatchAdd(Vec<FileKey>),
BatchRemove(Vec<String>),
Initialized,
}
impl std::fmt::Debug for DbEventFileList {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DbEventFileList::Add(key, _) => write!(f, "Add({})", key),
DbEventFileList::BatchAdd(keys) => write!(f, "BatchAdd({})", keys.len()),
DbEventFileList::BatchRemove(keys) => write!(f, "BatchRemove({})", keys.len()),
DbEventFileList::Initialized => write!(f, "Initialized"),
}
}
}
pub enum DbEventFileListDeleted {
BatchAdd(String, i64, Vec<String>),
BatchRemove(Vec<String>),
}
impl std::fmt::Debug for DbEventFileListDeleted {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DbEventFileListDeleted::BatchAdd(org_id, created_at, keys) => {
write!(f, "BatchAdd({}, {}, {})", org_id, created_at, keys.len())
}
DbEventFileListDeleted::BatchRemove(keys) => write!(f, "BatchRemove({})", keys.len()),
}
}
}
pub enum DbEventStreamStats {
Set(String, Vec<(String, StreamStats)>),
ResetMinTS(String, i64),
ResetAll,
}
impl std::fmt::Debug for DbEventStreamStats {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DbEventStreamStats::Set(key, _) => write!(f, "Set({})", key),
DbEventStreamStats::ResetMinTS(key, _) => write!(f, "ResetMinTS({})", key),
DbEventStreamStats::ResetAll => write!(f, "ResetAll"),
}
}
}
#[cfg(test)]
mod tests {
use bytes::Bytes;

559
src/common/infra/db/sqlite.rs
View File

@ -13,11 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{
str::FromStr,
sync::{atomic::AtomicBool, Arc},
time::Duration,
};
use std::{str::FromStr, sync::Arc, time::Duration};
use ahash::HashMap;
use async_trait::async_trait;
@ -31,36 +27,23 @@ use sqlx::{
},
Pool, Sqlite,
};
use tokio::{
sync::{mpsc, RwLock},
time,
};
use tokio::sync::{mpsc, Mutex, RwLock};
use crate::common::infra::{
cluster,
db::{
DbEvent, DbEventFileList, DbEventFileListDeleted, DbEventMeta, DbEventStreamStats, Event,
EventData,
},
db::{Event, EventData},
errors::*,
file_list::sqlite as sqlite_file_list,
};
/// Database update retry times
const DB_RETRY_TIMES: usize = 5;
/// Database shutdown flag
static DB_SHUTDOWN: AtomicBool = AtomicBool::new(false);
pub static CLIENT_RO: Lazy<Pool<Sqlite>> = Lazy::new(connect_ro);
pub static CLIENT_RW: Lazy<Pool<Sqlite>> = Lazy::new(connect_rw);
pub static CLIENT_RW: Lazy<Arc<Mutex<Pool<Sqlite>>>> =
Lazy::new(|| Arc::new(Mutex::new(connect_rw())));
pub static CHANNEL: Lazy<SqliteDbChannel> = Lazy::new(SqliteDbChannel::new);
static WATCHERS: Lazy<RwLock<FxIndexMap<String, EventChannel>>> =
Lazy::new(|| RwLock::new(Default::default()));
type EventChannel = Arc<mpsc::Sender<Event>>;
type DbChannel = Arc<mpsc::Sender<DbEvent>>;
fn connect_rw() -> Pool<Sqlite> {
let url = format!("{}{}", CONFIG.common.data_db_dir, "metadata.sqlite");
@ -104,14 +87,12 @@ fn connect_ro() -> Pool<Sqlite> {
pub struct SqliteDbChannel {
pub watch_tx: EventChannel,
pub db_tx: DbChannel,
}
impl SqliteDbChannel {
pub fn new() -> Self {
Self {
watch_tx: SqliteDbChannel::handle_watch_channel(),
db_tx: SqliteDbChannel::handle_db_channel(),
}
}
@ -162,284 +143,6 @@ impl SqliteDbChannel {
});
Arc::new(tx)
}
fn handle_db_channel() -> DbChannel {
let (tx, mut rx) = mpsc::channel::<DbEvent>(100000);
tokio::task::spawn(async move {
loop {
let event = match rx.recv().await {
Some(v) => v,
None => {
log::info!("[SQLITE] db event channel closed");
break;
}
};
if CONFIG.common.print_key_event {
log::info!("[SQLITE] db event: {:?}", event);
}
let client = CLIENT_RW.clone();
match event {
DbEvent::Meta(DbEventMeta::Put(key, value, need_watch)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match put(&client, &key, value.clone(), need_watch).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] put meta error: {}", e);
}
}
DbEvent::Meta(DbEventMeta::Delete(key, with_prefix, need_watch)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match delete(&client, &key, with_prefix, need_watch).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] delete meta error: {}", e);
}
}
DbEvent::FileList(DbEventFileList::Add(file, meta)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::add(&client, &file, &meta).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] add file_list error: {}", e);
}
}
DbEvent::FileList(DbEventFileList::BatchAdd(files)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::batch_add(&client, &files).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] batch add file_list error: {}", e);
}
}
DbEvent::FileList(DbEventFileList::BatchRemove(files)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::batch_remove(&client, &files).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] batch remove file_list error: {}", e);
}
}
DbEvent::FileListDeleted(DbEventFileListDeleted::BatchAdd(
org_id,
created_at,
files,
)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::batch_add_deleted(
&client, &org_id, created_at, &files,
)
.await
{
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] batch add file_list_deleted error: {}", e);
}
}
DbEvent::FileListDeleted(DbEventFileListDeleted::BatchRemove(files)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::batch_remove_deleted(&client, &files).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] batch remove file_list_deleted error: {}", e);
}
}
DbEvent::FileList(DbEventFileList::Initialized) => {
sqlite_file_list::set_initialised();
}
DbEvent::StreamStats(DbEventStreamStats::Set(org_id, streams)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::set_stream_stats(&client, &org_id, &streams)
.await
{
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] set stream stats error: {}", e);
}
}
DbEvent::StreamStats(DbEventStreamStats::ResetMinTS(stream, min_ts)) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::reset_stream_stats_min_ts(
&client, &stream, min_ts,
)
.await
{
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] reset stream stats min_ts error: {}", e);
}
}
DbEvent::StreamStats(DbEventStreamStats::ResetAll) => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::reset_stream_stats(&client).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] reset stream stats error: {}", e);
}
}
DbEvent::CreateTableMeta => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match create_table(&client).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] create table meta error: {}", e);
}
}
DbEvent::CreateTableFileList => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::create_table(&client).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] create table file_list error: {}", e);
}
}
DbEvent::CreateTableFileListIndex => {
let mut err: Option<String> = None;
for _ in 0..DB_RETRY_TIMES {
match sqlite_file_list::create_table_index(&client).await {
Ok(_) => {
err = None;
break;
}
Err(e) => {
err = Some(e.to_string());
}
}
time::sleep(time::Duration::from_secs(1)).await;
}
if let Some(e) = err {
log::error!("[SQLITE] create table file_list index error: {}", e);
}
}
DbEvent::Shutdown => {
DB_SHUTDOWN.store(true, std::sync::atomic::Ordering::Release);
break;
}
}
}
log::info!("[SQLITE] db event loop exit");
});
Arc::new(tx)
}
}
impl Default for SqliteDbChannel {
@ -465,11 +168,7 @@ impl Default for SqliteDb {
#[async_trait]
impl super::Db for SqliteDb {
async fn create_table(&self) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::CreateTableMeta)
.await
.map_err(|e| Error::Message(e.to_string()))?;
Ok(())
create_table().await
}
async fn stats(&self) -> Result<super::Stats> {
@ -508,26 +207,113 @@ impl super::Db for SqliteDb {
}
async fn put(&self, key: &str, value: Bytes, need_watch: bool) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::Meta(DbEventMeta::Put(
key.to_string(),
value,
need_watch,
)))
let (module, key1, key2) = super::parse_key(key);
let client = CLIENT_RW.clone();
let client = client.lock().await;
let mut tx = client.begin().await?;
if let Err(e) = sqlx::query(
r#"INSERT OR IGNORE INTO meta (module, key1, key2, value) VALUES ($1, $2, $3, '');"#,
)
.bind(&module)
.bind(&key1)
.bind(&key2)
.execute(&mut *tx)
.await
.map_err(|e| Error::Message(e.to_string()))?;
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback put meta error: {}", e);
}
return Err(e.into());
}
if let Err(e) = sqlx::query(
r#"UPDATE meta SET value=$4 WHERE module = $1 AND key1 = $2 AND key2 = $3;"#,
)
.bind(&module)
.bind(&key1)
.bind(&key2)
.bind(String::from_utf8(value.to_vec()).unwrap_or_default())
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback put meta error: {}", e);
}
return Err(e.into());
}
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit put meta error: {}", e);
return Err(e.into());
}
// event watch
if need_watch {
if let Err(e) = CHANNEL
.watch_tx
.clone()
.send(Event::Put(EventData {
key: key.to_string(),
value: Some(value),
}))
.await
{
log::error!("[SQLITE] send event error: {}", e);
}
}
Ok(())
}
async fn delete(&self, key: &str, with_prefix: bool, need_watch: bool) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::Meta(DbEventMeta::Delete(
key.to_string(),
with_prefix,
need_watch,
)))
.await
.map_err(|e| Error::Message(e.to_string()))?;
// event watch
if need_watch {
// find all keys then send event
let items = if with_prefix {
let db = super::get_db().await;
db.list_keys(key).await?
} else {
vec![key.to_string()]
};
let tx = CHANNEL.watch_tx.clone();
tokio::task::spawn(async move {
for key in items {
if let Err(e) = tx
.send(Event::Delete(EventData {
key: key.to_string(),
value: None,
}))
.await
{
log::error!("[SQLITE] send event error: {}", e);
}
}
});
}
let (module, key1, key2) = super::parse_key(key);
let sql = if with_prefix {
if key1.is_empty() {
format!(r#"DELETE FROM meta WHERE module = '{}';"#, module)
} else if key2.is_empty() {
format!(
r#"DELETE FROM meta WHERE module = '{}' AND key1 = '{}';"#,
module, key1
)
} else {
format!(
r#"DELETE FROM meta WHERE module = '{}' AND key1 = '{}' AND key2 LIKE '{}%';"#,
module, key1, key2
)
}
} else {
format!(
r#"DELETE FROM meta WHERE module = '{}' AND key1 = '{}' AND key2 = '{}';"#,
module, key1, key2
)
};
let client = CLIENT_RW.clone();
let client = client.lock().await;
sqlx::query(&sql).execute(&*client).await?;
Ok(())
}
@ -612,132 +398,13 @@ impl super::Db for SqliteDb {
}
async fn close(&self) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::Shutdown)
.await
.map_err(|e| Error::Message(e.to_string()))?;
loop {
if DB_SHUTDOWN.load(std::sync::atomic::Ordering::Relaxed) {
break;
}
time::sleep(time::Duration::from_secs(1)).await;
}
Ok(())
}
}
async fn put(client: &Pool<Sqlite>, key: &str, value: Bytes, need_watch: bool) -> Result<()> {
let (module, key1, key2) = super::parse_key(key);
let mut tx = client.begin().await?;
if let Err(e) = sqlx::query(
r#"INSERT OR IGNORE INTO meta (module, key1, key2, value) VALUES ($1, $2, $3, '');"#,
)
.bind(&module)
.bind(&key1)
.bind(&key2)
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback put meta error: {}", e);
}
return Err(e.into());
}
if let Err(e) =
sqlx::query(r#"UPDATE meta SET value=$4 WHERE module = $1 AND key1 = $2 AND key2 = $3;"#)
.bind(&module)
.bind(&key1)
.bind(&key2)
.bind(String::from_utf8(value.to_vec()).unwrap_or_default())
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback put meta error: {}", e);
}
return Err(e.into());
}
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit put meta error: {}", e);
return Err(e.into());
}
// event watch
if need_watch {
if let Err(e) = CHANNEL
.watch_tx
.clone()
.send(Event::Put(EventData {
key: key.to_string(),
value: Some(value),
}))
.await
{
log::error!("[SQLITE] send event error: {}", e);
}
}
Ok(())
}
async fn delete(
client: &Pool<Sqlite>,
key: &str,
with_prefix: bool,
need_watch: bool,
) -> Result<()> {
// event watch
if need_watch {
// find all keys then send event
let items = if with_prefix {
let db = super::get_db().await;
db.list_keys(key).await?
} else {
vec![key.to_string()]
};
let tx = CHANNEL.watch_tx.clone();
tokio::task::spawn(async move {
for key in items {
if let Err(e) = tx
.send(Event::Delete(EventData {
key: key.to_string(),
value: None,
}))
.await
{
log::error!("[SQLITE] send event error: {}", e);
}
}
});
}
let (module, key1, key2) = super::parse_key(key);
let sql = if with_prefix {
if key1.is_empty() {
format!(r#"DELETE FROM meta WHERE module = '{}';"#, module)
} else if key2.is_empty() {
format!(
r#"DELETE FROM meta WHERE module = '{}' AND key1 = '{}';"#,
module, key1
)
} else {
format!(
r#"DELETE FROM meta WHERE module = '{}' AND key1 = '{}' AND key2 LIKE '{}%';"#,
module, key1, key2
)
}
} else {
format!(
r#"DELETE FROM meta WHERE module = '{}' AND key1 = '{}' AND key2 = '{}';"#,
module, key1, key2
)
};
sqlx::query(&sql).execute(client).await?;
Ok(())
}
async fn create_table(client: &Pool<Sqlite>) -> Result<()> {
async fn create_table() -> Result<()> {
let client = CLIENT_RW.clone();
let client = client.lock().await;
// create table
sqlx::query(
r#"
@ -751,7 +418,7 @@ CREATE TABLE IF NOT EXISTS meta
);
"#,
)
.execute(client)
.execute(&*client)
.await?;
// create table index
sqlx::query(
@ -761,7 +428,7 @@ CREATE INDEX IF NOT EXISTS meta_module_key1_idx on meta (module, key1);
CREATE UNIQUE INDEX IF NOT EXISTS meta_module_key2_idx on meta (module, key1, key2);
"#,
)
.execute(client)
.execute(&*client)
.await?;
Ok(())
}

8
src/common/infra/file_list/dynamo.rs
View File

@ -69,14 +69,6 @@ impl super::FileList for DynamoFileList {
create_table_index().await
}
async fn set_initialised(&self) -> Result<()> {
Ok(())
}
async fn get_initialised(&self) -> Result<bool> {
Ok(true)
}
async fn add(&self, file: &str, meta: &FileMeta) -> Result<()> {
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;

10
src/common/infra/file_list/mod.rs
View File

@ -53,8 +53,6 @@ pub fn connect() -> Box<dyn FileList> {
pub trait FileList: Sync + Send + 'static {
async fn create_table(&self) -> Result<()>;
async fn create_table_index(&self) -> Result<()>;
async fn set_initialised(&self) -> Result<()>;
async fn get_initialised(&self) -> Result<bool>;
async fn add(&self, file: &str, meta: &FileMeta) -> Result<()>;
async fn remove(&self, file: &str) -> Result<()>;
async fn batch_add(&self, files: &[FileKey]) -> Result<()>;
@ -120,14 +118,6 @@ pub async fn create_table_index() -> Result<()> {
CLIENT.create_table_index().await
}
pub async fn set_initialised() -> Result<()> {
CLIENT.set_initialised().await
}
pub async fn get_initialised() -> Result<bool> {
CLIENT.get_initialised().await
}
#[inline]
pub async fn add(file: &str, meta: &FileMeta) -> Result<()> {
CLIENT.add(file, meta).await

8
src/common/infra/file_list/mysql.rs
View File

@ -54,14 +54,6 @@ impl super::FileList for MysqlFileList {
create_table_index().await
}
async fn set_initialised(&self) -> Result<()> {
Ok(())
}
async fn get_initialised(&self) -> Result<bool> {
Ok(true)
}
async fn add(&self, file: &str, meta: &FileMeta) -> Result<()> {
let pool = CLIENT.clone();
let (stream_key, date_key, file_name) =

8
src/common/infra/file_list/postgres.rs
View File

@ -54,14 +54,6 @@ impl super::FileList for PostgresFileList {
create_table_index().await
}
async fn set_initialised(&self) -> Result<()> {
Ok(())
}
async fn get_initialised(&self) -> Result<bool> {
Ok(true)
}
async fn add(&self, file: &str, meta: &FileMeta) -> Result<()> {
let pool = CLIENT.clone();
let (stream_key, date_key, file_name) =

749
src/common/infra/file_list/sqlite.rs
View File

@ -13,8 +13,6 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::sync::atomic::AtomicBool;
use ahash::HashMap;
use async_trait::async_trait;
use chrono::Utc;
@ -22,22 +20,16 @@ use config::{
meta::stream::{FileKey, FileMeta, StreamType},
utils::parquet::parse_file_key_columns,
};
use sqlx::{Executor, Pool, QueryBuilder, Row, Sqlite};
use sqlx::{Executor, QueryBuilder, Row, Sqlite};
use crate::common::{
infra::{
db::{
sqlite::{CHANNEL, CLIENT_RO as CLIENT},
DbEvent, DbEventFileList, DbEventFileListDeleted, DbEventStreamStats,
},
db::sqlite::{CLIENT_RO, CLIENT_RW},
errors::{Error, Result},
},
meta::stream::{PartitionTimeLevel, StreamStats},
};
/// Table file_list inited flag
static FILE_LIST_INITED: AtomicBool = AtomicBool::new(false);
pub struct SqliteFileList {}
impl SqliteFileList {
@ -55,62 +47,113 @@ impl Default for SqliteFileList {
#[async_trait]
impl super::FileList for SqliteFileList {
async fn create_table(&self) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::CreateTableFileList)
.await
.map_err(|e| Error::Message(e.to_string()))?;
Ok(())
create_table().await
}
async fn create_table_index(&self) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::CreateTableFileListIndex)
.await
.map_err(|e| Error::Message(e.to_string()))?;
Ok(())
}
async fn set_initialised(&self) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileList(DbEventFileList::Initialized))
.await
.map_err(|e| Error::Message(e.to_string()))?;
Ok(())
}
async fn get_initialised(&self) -> Result<bool> {
Ok(FILE_LIST_INITED.load(std::sync::atomic::Ordering::Relaxed))
create_table_index().await
}
async fn add(&self, file: &str, meta: &FileMeta) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileList(DbEventFileList::Add(
file.to_string(),
meta.to_owned(),
)))
.await
.map_err(|e| Error::Message(e.to_string()))?;
Ok(())
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let org_id = stream_key[..stream_key.find('/').unwrap()].to_string();
let client = CLIENT_RW.clone();
let client = client.lock().await;
match sqlx::query(
r#"
INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10);
"#,
)
.bind(org_id)
.bind(stream_key)
.bind(date_key)
.bind(file_name)
.bind(false)
.bind(meta.min_ts)
.bind(meta.max_ts)
.bind(meta.records)
.bind(meta.original_size)
.bind(meta.compressed_size)
.execute(&*client)
.await {
Err(sqlx::Error::Database(e)) => if e.is_unique_violation() {
Ok(())
} else {
Err(Error::Message(e.to_string()))
},
Err(e) => Err(e.into()),
Ok(_) => Ok(()),
}
}
async fn remove(&self, file: &str) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileList(DbEventFileList::BatchRemove(vec![
file.to_string(),
])))
.await
.map_err(|e| Error::Message(e.to_string()))?;
Ok(())
self.batch_remove(&[file.to_string()]).await
}
async fn batch_add(&self, files: &[FileKey]) -> Result<()> {
if files.is_empty() {
return Ok(());
}
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileList(DbEventFileList::BatchAdd(files.to_vec())))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let chunks = files.chunks(100);
let client = CLIENT_RW.clone();
let client = client.lock().await;
for files in chunks {
let mut tx = client.begin().await?;
let mut query_builder: QueryBuilder<Sqlite> = QueryBuilder::new(
"INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size)",
);
query_builder.push_values(files, |mut b, item| {
let (stream_key, date_key, file_name) =
parse_file_key_columns(&item.key).expect("parse file key failed");
let org_id = stream_key[..stream_key.find('/').unwrap()].to_string();
b.push_bind(org_id)
.push_bind(stream_key)
.push_bind(date_key)
.push_bind(file_name)
.push_bind(false)
.push_bind(item.meta.min_ts)
.push_bind(item.meta.max_ts)
.push_bind(item.meta.records)
.push_bind(item.meta.original_size)
.push_bind(item.meta.compressed_size);
});
let need_single_insert = match query_builder.build().execute(&mut *tx).await {
Ok(_) => false,
Err(sqlx::Error::Database(e)) => {
if e.is_unique_violation() {
true
} else {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list batch add error: {}", e);
}
return Err(Error::Message(e.to_string()));
}
}
Err(e) => {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list batch add error: {}", e);
}
return Err(e.into());
}
};
if need_single_insert {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list batch add error: {}", e);
return Err(e.into());
}
for item in files {
if let Err(e) = self.add(&item.key, &item.meta).await {
log::error!("[SQLITE] single insert file_list add error: {}", e);
return Err(e);
}
}
} else if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit file_list batch add error: {}", e);
return Err(e.into());
}
}
Ok(())
}
@ -118,12 +161,44 @@ impl super::FileList for SqliteFileList {
if files.is_empty() {
return Ok(());
}
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileList(DbEventFileList::BatchRemove(
files.to_vec(),
)))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let chunks = files.chunks(100);
let client = CLIENT_RW.clone();
let client = client.lock().await;
for files in chunks {
// get ids of the files
let pool = client.clone();
let mut ids = Vec::with_capacity(files.len());
for file in files {
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let ret: Option<i64> = match sqlx::query_scalar(
r#"SELECT id FROM file_list WHERE stream = $1 AND date = $2 AND file = $3;"#,
)
.bind(stream_key)
.bind(date_key)
.bind(file_name)
.fetch_one(&pool)
.await
{
Ok(v) => v,
Err(sqlx::Error::RowNotFound) => continue,
Err(e) => return Err(e.into()),
};
match ret {
Some(v) => ids.push(v.to_string()),
None => {
return Err(Error::Message(
"[SQLITE] query error: id should not empty from file_list".to_string(),
));
}
}
}
// delete files by ids
if !ids.is_empty() {
let sql = format!("DELETE FROM file_list WHERE id IN({});", ids.join(","));
_ = pool.execute(sql.as_str()).await?;
}
}
Ok(())
}
@ -136,14 +211,34 @@ impl super::FileList for SqliteFileList {
if files.is_empty() {
return Ok(());
}
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileListDeleted(DbEventFileListDeleted::BatchAdd(
org_id.to_string(),
created_at,
files.to_vec(),
)))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let chunks = files.chunks(100);
let client = CLIENT_RW.clone();
let client = client.lock().await;
for files in chunks {
let mut tx = client.begin().await?;
let mut query_builder: QueryBuilder<Sqlite> = QueryBuilder::new(
"INSERT INTO file_list_deleted (org, stream, date, file, created_at)",
);
query_builder.push_values(files, |mut b, item: &String| {
let (stream_key, date_key, file_name) =
parse_file_key_columns(item).expect("parse file key failed");
b.push_bind(org_id)
.push_bind(stream_key)
.push_bind(date_key)
.push_bind(file_name)
.push_bind(created_at);
});
if let Err(e) = query_builder.build().execute(&mut *tx).await {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list_deleted batch add error: {}", e);
}
return Err(e.into());
};
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit file_list_deleted batch add error: {}", e);
return Err(e.into());
}
}
Ok(())
}
@ -151,17 +246,53 @@ impl super::FileList for SqliteFileList {
if files.is_empty() {
return Ok(());
}
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::FileListDeleted(
DbEventFileListDeleted::BatchRemove(files.to_vec()),
))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let chunks = files.chunks(100);
let client = CLIENT_RW.clone();
let client = client.lock().await;
for files in chunks {
// get ids of the files
let pool = client.clone();
let mut ids = Vec::with_capacity(files.len());
for file in files {
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let ret: Option<i64> = match sqlx::query_scalar(
r#"SELECT id FROM file_list_deleted WHERE stream = $1 AND date = $2 AND file = $3;"#,
)
.bind(stream_key)
.bind(date_key)
.bind(file_name)
.fetch_one(&pool)
.await
{
Ok(v) => v,
Err(sqlx::Error::RowNotFound) => continue,
Err(e) => return Err(e.into()),
};
match ret {
Some(v) => ids.push(v.to_string()),
None => {
return Err(Error::Message(
"[SQLITE] query error: id should not empty from file_list_deleted"
.to_string(),
));
}
}
}
// delete files by ids
if !ids.is_empty() {
let sql = format!(
"DELETE FROM file_list_deleted WHERE id IN({});",
ids.join(",")
);
_ = pool.execute(sql.as_str()).await?;
}
}
Ok(())
}
async fn get(&self, file: &str) -> Result<FileMeta> {
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let ret = sqlx::query_as::<_, super::FileRecord>(
@ -179,7 +310,7 @@ SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, comp
}
async fn contains(&self, file: &str) -> Result<bool> {
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let ret = sqlx::query(
@ -200,7 +331,7 @@ SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, comp
}
async fn list(&self) -> Result<Vec<(String, FileMeta)>> {
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret = sqlx::query_as::<_, super::FileRecord>(r#"SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size FROM file_list;"#)
.fetch_all(&pool)
.await?;
@ -230,7 +361,7 @@ SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, comp
let stream_key = format!("{org_id}/{stream_type}/{stream_name}");
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret = sqlx::query_as::<_, super::FileRecord>(
r#"
SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size
@ -258,7 +389,7 @@ SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, comp
if time_max == 0 {
return Ok(Vec::new());
}
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret = sqlx::query_as::<_, super::FileDeletedRecord>(
r#"SELECT stream, date, file FROM file_list_deleted WHERE org = $1 AND created_at < $2 LIMIT $3;"#,
)
@ -281,7 +412,7 @@ SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, comp
) -> Result<i64> {
let stream_key = format!("{org_id}/{stream_type}/{stream_name}");
let min_ts = crate::common::utils::time::BASE_TIME.timestamp_micros();
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret: Option<i64> = sqlx::query_scalar(
r#"SELECT MIN(min_ts) AS id FROM file_list WHERE stream = $1 AND min_ts > $2;"#,
)
@ -293,7 +424,7 @@ SELECT stream, date, file, deleted, min_ts, max_ts, records, original_size, comp
}
async fn get_max_pk_value(&self) -> Result<i64> {
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret: Option<i64> = sqlx::query_scalar(r#"SELECT MAX(id) AS id FROM file_list;"#)
.fetch_one(&pool)
.await?;
@ -334,7 +465,7 @@ SELECT stream, MIN(min_ts) as min_ts, MAX(max_ts) as max_ts, COUNT(*) as file_nu
format!("{} AND id > {} AND id <= {} GROUP BY stream", sql, min, max)
}
};
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret = sqlx::query_as::<_, super::StatsRecord>(&sql)
.fetch_all(&pool)
.await?;
@ -360,7 +491,7 @@ SELECT stream, MIN(min_ts) as min_ts, MAX(max_ts) as max_ts, COUNT(*) as file_nu
} else {
format!("SELECT * FROM stream_stats WHERE org = '{}';", org_id)
};
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret = sqlx::query_as::<_, super::StatsRecord>(&sql)
.fetch_all(&pool)
.await?;
@ -375,13 +506,80 @@ SELECT stream, MIN(min_ts) as min_ts, MAX(max_ts) as max_ts, COUNT(*) as file_nu
org_id: &str,
streams: &[(String, StreamStats)],
) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::StreamStats(DbEventStreamStats::Set(
org_id.to_string(),
streams.to_vec(),
)))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let old_stats = super::get_stream_stats(org_id, None, None).await?;
let old_stats = old_stats.into_iter().collect::<HashMap<_, _>>();
let mut new_streams = Vec::new();
let mut update_streams = Vec::with_capacity(streams.len());
for (stream_key, item) in streams {
let mut stats = match old_stats.get(stream_key) {
Some(s) => s.to_owned(),
None => {
new_streams.push(stream_key);
StreamStats::default()
}
};
stats.add_stream_stats(item);
update_streams.push((stream_key, stats));
}
let client = CLIENT_RW.clone();
let client = client.lock().await;
let mut tx = client.begin().await?;
for stream_key in new_streams {
let org_id = stream_key[..stream_key.find('/').unwrap()].to_string();
if let Err(e) = sqlx::query(
r#"
INSERT INTO stream_stats
(org, stream, file_num, min_ts, max_ts, records, original_size, compressed_size)
VALUES ($1, $2, 0, 0, 0, 0, 0, 0);
"#,
)
.bind(org_id)
.bind(stream_key)
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback insert stream stats error: {}", e);
}
return Err(e.into());
}
}
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit set stream stats error: {}", e);
return Err(e.into());
}
let mut tx = client.begin().await?;
for (stream_key, stats) in update_streams {
if let Err(e) = sqlx::query(
r#"
UPDATE stream_stats
SET file_num = $1, min_ts = $2, max_ts = $3, records = $4, original_size = $5, compressed_size = $6
WHERE stream = $7;
"#,
)
.bind(stats.file_num)
.bind(stats.doc_time_min)
.bind(stats.doc_time_max)
.bind(stats.doc_num)
.bind(stats.storage_size as i64)
.bind(stats.compressed_size as i64)
.bind(stream_key)
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback set stream stats error: {}", e);
}
return Err(e.into());
}
}
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit set stream stats error: {}", e);
return Err(e.into());
}
Ok(())
}
@ -391,26 +589,27 @@ SELECT stream, MIN(min_ts) as min_ts, MAX(max_ts) as max_ts, COUNT(*) as file_nu
stream: &str,
min_ts: i64,
) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::StreamStats(DbEventStreamStats::ResetMinTS(
stream.to_string(),
min_ts,
)))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let client = CLIENT_RW.clone();
let client = client.lock().await;
sqlx::query(r#"UPDATE stream_stats SET min_ts = $1 WHERE stream = $2;"#)
.bind(min_ts)
.bind(stream)
.execute(&*client)
.await?;
Ok(())
}
async fn reset_stream_stats(&self) -> Result<()> {
let tx = CHANNEL.db_tx.clone();
tx.send(DbEvent::StreamStats(DbEventStreamStats::ResetAll))
.await
.map_err(|e| Error::Message(e.to_string()))?;
let client = CLIENT_RW.clone();
let client = client.lock().await;
sqlx::query(r#"UPDATE stream_stats SET file_num = 0, min_ts = 0, max_ts = 0, records = 0, original_size = 0, compressed_size = 0;"#)
.execute(&*client)
.await?;
Ok(())
}
async fn len(&self) -> usize {
let pool = CLIENT.clone();
let pool = CLIENT_RO.clone();
let ret = match sqlx::query(r#"SELECT COUNT(*) as num FROM file_list;"#)
.fetch_one(&pool)
.await
@ -436,327 +635,9 @@ SELECT stream, MIN(min_ts) as min_ts, MAX(max_ts) as max_ts, COUNT(*) as file_nu
}
}
pub async fn add(client: &Pool<Sqlite>, file: &str, meta: &FileMeta) -> Result<()> {
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let org_id = stream_key[..stream_key.find('/').unwrap()].to_string();
match sqlx::query(
r#"
INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10);
"#,
)
.bind(org_id)
.bind(stream_key)
.bind(date_key)
.bind(file_name)
.bind(false)
.bind(meta.min_ts)
.bind(meta.max_ts)
.bind(meta.records)
.bind(meta.original_size)
.bind(meta.compressed_size)
.execute(client)
.await {
Err(sqlx::Error::Database(e)) => if e.is_unique_violation() {
Ok(())
} else {
Err(Error::Message(e.to_string()))
},
Err(e) => Err(e.into()),
Ok(_) => Ok(()),
}
}
pub async fn batch_add(client: &Pool<Sqlite>, files: &[FileKey]) -> Result<()> {
if files.is_empty() {
return Ok(());
}
let chunks = files.chunks(100);
for files in chunks {
let mut tx = client.begin().await?;
let mut query_builder: QueryBuilder<Sqlite> = QueryBuilder::new(
"INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size)",
);
query_builder.push_values(files, |mut b, item| {
let (stream_key, date_key, file_name) =
parse_file_key_columns(&item.key).expect("parse file key failed");
let org_id = stream_key[..stream_key.find('/').unwrap()].to_string();
b.push_bind(org_id)
.push_bind(stream_key)
.push_bind(date_key)
.push_bind(file_name)
.push_bind(false)
.push_bind(item.meta.min_ts)
.push_bind(item.meta.max_ts)
.push_bind(item.meta.records)
.push_bind(item.meta.original_size)
.push_bind(item.meta.compressed_size);
});
let need_single_insert = match query_builder.build().execute(&mut *tx).await {
Ok(_) => false,
Err(sqlx::Error::Database(e)) => {
if e.is_unique_violation() {
true
} else {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list batch add error: {}", e);
}
return Err(Error::Message(e.to_string()));
}
}
Err(e) => {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list batch add error: {}", e);
}
return Err(e.into());
}
};
if need_single_insert {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list batch add error: {}", e);
return Err(e.into());
}
for item in files {
if let Err(e) = add(client, &item.key, &item.meta).await {
log::error!("[SQLITE] single insert file_list add error: {}", e);
return Err(e);
}
}
} else if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit file_list batch add error: {}", e);
return Err(e.into());
}
}
Ok(())
}
pub async fn batch_remove(client: &Pool<Sqlite>, files: &[String]) -> Result<()> {
if files.is_empty() {
return Ok(());
}
let chunks = files.chunks(100);
for files in chunks {
// get ids of the files
let pool = client.clone();
let mut ids = Vec::with_capacity(files.len());
for file in files {
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let ret: Option<i64> = match sqlx::query_scalar(
r#"SELECT id FROM file_list WHERE stream = $1 AND date = $2 AND file = $3;"#,
)
.bind(stream_key)
.bind(date_key)
.bind(file_name)
.fetch_one(&pool)
.await
{
Ok(v) => v,
Err(sqlx::Error::RowNotFound) => continue,
Err(e) => return Err(e.into()),
};
match ret {
Some(v) => ids.push(v.to_string()),
None => {
return Err(Error::Message(
"[SQLITE] query error: id should not empty from file_list".to_string(),
));
}
}
}
// delete files by ids
if !ids.is_empty() {
let sql = format!("DELETE FROM file_list WHERE id IN({});", ids.join(","));
_ = pool.execute(sql.as_str()).await?;
}
}
Ok(())
}
pub async fn batch_add_deleted(
client: &Pool<Sqlite>,
org_id: &str,
created_at: i64,
files: &[String],
) -> Result<()> {
let chunks = files.chunks(100);
for files in chunks {
let mut tx = client.begin().await?;
let mut query_builder: QueryBuilder<Sqlite> = QueryBuilder::new(
"INSERT INTO file_list_deleted (org, stream, date, file, created_at)",
);
query_builder.push_values(files, |mut b, item: &String| {
let (stream_key, date_key, file_name) =
parse_file_key_columns(item).expect("parse file key failed");
b.push_bind(org_id)
.push_bind(stream_key)
.push_bind(date_key)
.push_bind(file_name)
.push_bind(created_at);
});
if let Err(e) = query_builder.build().execute(&mut *tx).await {
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback file_list_deleted batch add error: {}", e);
}
return Err(e.into());
};
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit file_list_deleted batch add error: {}", e);
return Err(e.into());
}
}
Ok(())
}
pub async fn batch_remove_deleted(client: &Pool<Sqlite>, files: &[String]) -> Result<()> {
if files.is_empty() {
return Ok(());
}
let chunks = files.chunks(100);
for files in chunks {
// get ids of the files
let pool = client.clone();
let mut ids = Vec::with_capacity(files.len());
for file in files {
let (stream_key, date_key, file_name) =
parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?;
let ret: Option<i64> = match sqlx::query_scalar(
r#"SELECT id FROM file_list_deleted WHERE stream = $1 AND date = $2 AND file = $3;"#,
)
.bind(stream_key)
.bind(date_key)
.bind(file_name)
.fetch_one(&pool)
.await
{
Ok(v) => v,
Err(sqlx::Error::RowNotFound) => continue,
Err(e) => return Err(e.into()),
};
match ret {
Some(v) => ids.push(v.to_string()),
None => {
return Err(Error::Message(
"[SQLITE] query error: id should not empty from file_list_deleted"
.to_string(),
));
}
}
}
// delete files by ids
if !ids.is_empty() {
let sql = format!(
"DELETE FROM file_list_deleted WHERE id IN({});",
ids.join(",")
);
_ = pool.execute(sql.as_str()).await?;
}
}
Ok(())
}
pub async fn set_stream_stats(
client: &Pool<Sqlite>,
org_id: &str,
streams: &[(String, StreamStats)],
) -> Result<()> {
let old_stats = super::get_stream_stats(org_id, None, None).await?;
let old_stats = old_stats.into_iter().collect::<HashMap<_, _>>();
let mut new_streams = Vec::new();
let mut update_streams = Vec::with_capacity(streams.len());
for (stream_key, item) in streams {
let mut stats = match old_stats.get(stream_key) {
Some(s) => s.to_owned(),
None => {
new_streams.push(stream_key);
StreamStats::default()
}
};
stats.add_stream_stats(item);
update_streams.push((stream_key, stats));
}
let mut tx = client.begin().await?;
for stream_key in new_streams {
let org_id = stream_key[..stream_key.find('/').unwrap()].to_string();
if let Err(e) = sqlx::query(
r#"
INSERT INTO stream_stats
(org, stream, file_num, min_ts, max_ts, records, original_size, compressed_size)
VALUES ($1, $2, 0, 0, 0, 0, 0, 0);
"#,
)
.bind(org_id)
.bind(stream_key)
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback insert stream stats error: {}", e);
}
return Err(e.into());
}
}
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit set stream stats error: {}", e);
return Err(e.into());
}
let mut tx = client.begin().await?;
for (stream_key, stats) in update_streams {
if let Err(e) = sqlx::query(
r#"
UPDATE stream_stats
SET file_num = $1, min_ts = $2, max_ts = $3, records = $4, original_size = $5, compressed_size = $6
WHERE stream = $7;
"#,
)
.bind(stats.file_num)
.bind(stats.doc_time_min)
.bind(stats.doc_time_max)
.bind(stats.doc_num)
.bind(stats.storage_size as i64)
.bind(stats.compressed_size as i64)
.bind(stream_key)
.execute(&mut *tx)
.await
{
if let Err(e) = tx.rollback().await {
log::error!("[SQLITE] rollback set stream stats error: {}", e);
}
return Err(e.into());
}
}
if let Err(e) = tx.commit().await {
log::error!("[SQLITE] commit set stream stats error: {}", e);
return Err(e.into());
}
Ok(())
}
pub async fn reset_stream_stats_min_ts(
client: &Pool<Sqlite>,
stream: &str,
min_ts: i64,
) -> Result<()> {
sqlx::query(r#"UPDATE stream_stats SET min_ts = $1 WHERE stream = $2;"#)
.bind(min_ts)
.bind(stream)
.execute(client)
.await?;
Ok(())
}
pub async fn reset_stream_stats(client: &Pool<Sqlite>) -> Result<()> {
sqlx::query(r#"UPDATE stream_stats SET file_num = 0, min_ts = 0, max_ts = 0, records = 0, original_size = 0, compressed_size = 0;"#)
.execute(client)
.await?;
Ok(())
}
pub async fn create_table(client: &Pool<Sqlite>) -> Result<()> {
pub async fn create_table() -> Result<()> {
let client = CLIENT_RW.clone();
let client = client.lock().await;
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS file_list
@ -775,7 +656,7 @@ CREATE TABLE IF NOT EXISTS file_list
);
"#,
)
.execute(client)
.execute(&*client)
.await?;
sqlx::query(
@ -791,7 +672,7 @@ CREATE TABLE IF NOT EXISTS file_list_deleted
);
"#,
)
.execute(client)
.execute(&*client)
.await?;
sqlx::query(
@ -810,13 +691,13 @@ CREATE TABLE IF NOT EXISTS stream_stats
);
"#,
)
.execute(client)
.execute(&*client)
.await?;
Ok(())
}
pub async fn create_table_index(client: &Pool<Sqlite>) -> Result<()> {
pub async fn create_table_index() -> Result<()> {
let sqls = vec![
(
"file_list",
@ -847,8 +728,11 @@ pub async fn create_table_index(client: &Pool<Sqlite>) -> Result<()> {
"CREATE UNIQUE INDEX IF NOT EXISTS stream_stats_stream_idx on stream_stats (stream);",
),
];
let client = CLIENT_RW.clone();
let client = client.lock().await;
for (table, sql) in sqls {
if let Err(e) = sqlx::query(sql).execute(client).await {
if let Err(e) = sqlx::query(sql).execute(&*client).await {
log::error!("[SQLITE] create table {} index error: {}", table, e);
return Err(e.into());
}
@ -856,7 +740,7 @@ pub async fn create_table_index(client: &Pool<Sqlite>) -> Result<()> {
// create UNIQUE index for file_list
let unique_index_sql = r#"CREATE UNIQUE INDEX IF NOT EXISTS file_list_stream_file_idx on file_list (stream, date, file);"#;
if let Err(e) = sqlx::query(unique_index_sql).execute(client).await {
if let Err(e) = sqlx::query(unique_index_sql).execute(&*client).await {
if !e.to_string().contains("UNIQUE constraint failed") {
return Err(e.into());
}
@ -864,7 +748,7 @@ pub async fn create_table_index(client: &Pool<Sqlite>) -> Result<()> {
log::warn!("[SQLITE] starting delete duplicate records");
let ret = sqlx::query(
r#"SELECT stream, date, file, min(id) as id FROM file_list GROUP BY stream, date, file HAVING COUNT(*) > 1;"#,
).fetch_all(client).await?;
).fetch_all(&*client).await?;
log::warn!("[SQLITE] total: {} duplicate records", ret.len());
for (i, r) in ret.iter().enumerate() {
let stream = r.get::<String, &str>("stream");
@ -873,7 +757,7 @@ pub async fn create_table_index(client: &Pool<Sqlite>) -> Result<()> {
let id = r.get::<i64, &str>("id");
sqlx::query(
r#"DELETE FROM file_list WHERE id != $1 AND stream = $2 AND date = $3 AND file = $4;"#,
).bind(id).bind(stream).bind(date).bind(file).execute(client).await?;
).bind(id).bind(stream).bind(date).bind(file).execute(&*client).await?;
if i / 1000 == 0 {
log::warn!("[SQLITE] delete duplicate records: {}/{}", i, ret.len());
}
@ -884,20 +768,15 @@ pub async fn create_table_index(client: &Pool<Sqlite>) -> Result<()> {
ret.len()
);
// create index again
sqlx::query(unique_index_sql).execute(client).await?;
sqlx::query(unique_index_sql).execute(&*client).await?;
log::warn!("[SQLITE] create table index(file_list_stream_file_idx) succeed");
}
// delete trigger for old version
// compitable for old version <= 0.6.4
sqlx::query(r#"DROP TRIGGER IF EXISTS update_stream_stats_delete;"#)
.execute(client)
.execute(&*client)
.await?;
Ok(())
}
/// set file list inited flag
pub fn set_initialised() {
FILE_LIST_INITED.store(true, std::sync::atomic::Ordering::Release);
}

18
src/common/meta/sql.rs
View File

@ -26,6 +26,9 @@ use sqlparser::{
parser::Parser,
};
const MAX_LIMIT: usize = 10000;
const MAX_OFFSET: usize = 100000;
/// parsed sql
#[derive(Clone, Debug, Serialize)]
pub struct Sql {
@ -34,6 +37,7 @@ pub struct Sql {
pub(crate) source: String, // table
pub(crate) order_by: Vec<(String, bool)>, // desc: true / false
pub(crate) group_by: Vec<String>, // field
pub(crate) having: bool,
pub(crate) offset: usize,
pub(crate) limit: usize,
pub(crate) time_range: Option<(i64, i64)>,
@ -109,6 +113,7 @@ impl TryFrom<&Statement> for Sql {
selection,
projection,
group_by: groups,
having,
..
} = match &q.body.as_ref() {
SetExpr::Select(statement) => statement.as_ref(),
@ -152,6 +157,7 @@ impl TryFrom<&Statement> for Sql {
source,
order_by,
group_by,
having: having.is_some(),
offset,
limit,
time_range,
@ -179,7 +185,13 @@ impl<'a> From<Offset<'a>> for usize {
SqlOffset {
value: SqlExpr::Value(Value::Number(v, _b)),
..
} => v.parse().unwrap_or(0),
} => {
let mut v: usize = v.parse().unwrap_or(0);
if v > MAX_OFFSET {
v = MAX_OFFSET;
}
v
}
_ => 0,
}
}
@ -190,8 +202,8 @@ impl<'a> From<Limit<'a>> for usize {
match l.0 {
SqlExpr::Value(Value::Number(v, _b)) => {
let mut v: usize = v.parse().unwrap_or(0);
if v > 10000 {
v = 10000;
if v > MAX_LIMIT {
v = MAX_LIMIT;
}
v
}

3
src/common/migration/file_list.rs
View File

@ -136,9 +136,6 @@ pub async fn run(prefix: &str, from: &str, to: &str) -> Result<(), anyhow::Error
}
}
infra_file_list::set_initialised()
.await
.expect("file list migration set initialised failed");
// update stream stats
update_stats_from_file_list()
.await

1
src/common/utils/json.rs
View File

@ -67,5 +67,4 @@ mod tests {
let val: Value = from_str(json).unwrap();
assert_eq!(estimate_json_bytes(&val), json.len());
}
}

61
src/handler/http/request/organization/org.rs
View File

@ -15,14 +15,17 @@
use std::{collections::HashSet, io::Error};
use actix_web::{get, post, put, web, HttpResponse, Result};
use actix_web::{get, http, post, put, web, HttpResponse, Result};
use crate::{
common::{
infra::config::{STREAM_SCHEMAS, USERS},
meta::organization::{
OrgDetails, OrgUser, Organization, OrganizationResponse, PasscodeResponse,
RumIngestionResponse, CUSTOM, DEFAULT_ORG, THRESHOLD,
meta::{
http::HttpResponse as MetaHttpResponse,
organization::{
OrgDetails, OrgUser, Organization, OrganizationResponse, PasscodeResponse,
RumIngestionResponse, CUSTOM, DEFAULT_ORG, THRESHOLD,
},
},
utils::auth::{is_root_user, UserEmail},
},
@ -156,6 +159,7 @@ async fn org_summary(org_id: web::Path<String>) -> Result<HttpResponse, Error> {
),
responses(
(status = 200, description = "Success", content_type = "application/json", body = PasscodeResponse),
(status = 404, description = "NotFound", content_type = "application/json", body = HttpResponse),
)
)]
#[get("/{org_id}/passcode")]
@ -169,8 +173,13 @@ async fn get_user_passcode(
if is_root_user(user_id) {
org_id = None;
}
let passcode = get_passcode(org_id, user_id).await;
Ok(HttpResponse::Ok().json(PasscodeResponse { data: passcode }))
match get_passcode(org_id, user_id).await {
Ok(passcode) => Ok(HttpResponse::Ok().json(PasscodeResponse { data: passcode })),
Err(e) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
http::StatusCode::NOT_FOUND.into(),
e.to_string(),
))),
}
}
/// UpdateIngestToken
@ -186,6 +195,7 @@ async fn get_user_passcode(
),
responses(
(status = 200, description = "Success", content_type = "application/json", body = PasscodeResponse),
(status = 404, description = "NotFound", content_type = "application/json", body = HttpResponse),
)
)]
#[put("/{org_id}/passcode")]
@ -199,8 +209,13 @@ async fn update_user_passcode(
if is_root_user(user_id) {
org_id = None;
}
let passcode = update_passcode(org_id, user_id).await;
Ok(HttpResponse::Ok().json(PasscodeResponse { data: passcode }))
match update_passcode(org_id, user_id).await {
Ok(passcode) => Ok(HttpResponse::Ok().json(PasscodeResponse { data: passcode })),
Err(e) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
http::StatusCode::NOT_FOUND.into(),
e.to_string(),
))),
}
}
/// GetRumIngestToken
@ -216,6 +231,7 @@ async fn update_user_passcode(
),
responses(
(status = 200, description = "Success", content_type = "application/json", body = RumIngestionResponse),
(status = 404, description = "NotFound", content_type = "application/json", body = HttpResponse),
)
)]
#[get("/{org_id}/rumtoken")]
@ -229,8 +245,13 @@ async fn get_user_rumtoken(
if is_root_user(user_id) {
org_id = None;
}
let rumtoken = get_rum_token(org_id, user_id).await;
Ok(HttpResponse::Ok().json(RumIngestionResponse { data: rumtoken }))
match get_rum_token(org_id, user_id).await {
Ok(rumtoken) => Ok(HttpResponse::Ok().json(RumIngestionResponse { data: rumtoken })),
Err(e) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
http::StatusCode::NOT_FOUND.into(),
e.to_string(),
))),
}
}
/// UpdateRumIngestToken
@ -246,6 +267,7 @@ async fn get_user_rumtoken(
),
responses(
(status = 200, description = "Success", content_type = "application/json", body = RumIngestionResponse),
(status = 404, description = "NotFound", content_type = "application/json", body = HttpResponse),
)
)]
#[put("/{org_id}/rumtoken")]
@ -259,8 +281,13 @@ async fn update_user_rumtoken(
if is_root_user(user_id) {
org_id = None;
}
let rumtoken = update_rum_token(org_id, user_id).await;
Ok(HttpResponse::Ok().json(RumIngestionResponse { data: rumtoken }))
match update_rum_token(org_id, user_id).await {
Ok(rumtoken) => Ok(HttpResponse::Ok().json(RumIngestionResponse { data: rumtoken })),
Err(e) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
http::StatusCode::NOT_FOUND.into(),
e.to_string(),
))),
}
}
/// CreateRumIngestToken
@ -276,6 +303,7 @@ async fn update_user_rumtoken(
),
responses(
(status = 200, description = "Success", content_type = "application/json", body = RumIngestionResponse),
(status = 404, description = "NotFound", content_type = "application/json", body = HttpResponse),
)
)]
#[post("/{org_id}/rumtoken")]
@ -289,8 +317,13 @@ async fn create_user_rumtoken(
if is_root_user(user_id) {
org_id = None;
}
let rumtoken = update_rum_token(org_id, user_id).await;
Ok(HttpResponse::Ok().json(RumIngestionResponse { data: rumtoken }))
match update_rum_token(org_id, user_id).await {
Ok(rumtoken) => Ok(HttpResponse::Ok().json(RumIngestionResponse { data: rumtoken })),
Err(e) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
http::StatusCode::NOT_FOUND.into(),
e.to_string(),
))),
}
}
/// CreateOrganization

4
src/ingester/src/immutable.rs
View File

@ -135,8 +135,8 @@ pub(crate) async fn persist() -> Result<()> {
for task in tasks {
if let Some((path, json_size, arrow_size)) = task? {
log::info!(
"[INGESTER:WAL] persist file: {:?}, json_size: {}, arrow_size: {}",
&path,
"[INGESTER:WAL] persist file: {}, json_size: {}, arrow_size: {}",
path.to_string_lossy(),
json_size,
arrow_size
);

12
src/job/files/parquet.rs
View File

@ -27,7 +27,7 @@ use config::{
meta::stream::{FileKey, FileMeta, StreamType},
metrics,
utils::parquet::{read_metadata_from_bytes, read_metadata_from_file},
CONFIG,
FxIndexMap, CONFIG,
};
use parquet::arrow::ParquetRecordBatchStreamBuilder;
use tokio::{sync::Semaphore, task::JoinHandle, time};
@ -69,9 +69,13 @@ pub async fn move_files_to_storage() -> Result<(), anyhow::Error> {
let pattern = wal_dir.join("files/");
let files = scan_files(&pattern, "parquet");
if files.is_empty() {
return Ok(());
}
log::info!("[INGESTER:JOB] move files get: {}", files.len());
// do partition by partition key
let mut partition_files_with_size: HashMap<String, Vec<FileKey>> = HashMap::default();
let mut partition_files_with_size: FxIndexMap<String, Vec<FileKey>> = FxIndexMap::default();
for file in files {
let Ok(parquet_meta) = read_metadata_from_file(&(&file).into()).await else {
continue;
@ -273,7 +277,9 @@ async fn merge_files(
let mut new_file_list = Vec::new();
let mut deleted_files = Vec::new();
for file in files_with_size.iter() {
if new_file_size + file.meta.original_size > CONFIG.compact.max_file_size as i64 {
if new_file_size > 0
&& new_file_size + file.meta.original_size > CONFIG.compact.max_file_size as i64
{
break;
}
new_file_size += file.meta.original_size;

4
src/job/mod.rs
View File

@ -159,9 +159,6 @@ pub async fn init() -> Result<(), anyhow::Error> {
infra_file_list::create_table_index()
.await
.expect("file list create table index failed");
infra_file_list::set_initialised()
.await
.expect("file list set initialised failed");
update_stats_from_file_list()
.await
.expect("file list remote calculate stats failed");
@ -169,7 +166,6 @@ pub async fn init() -> Result<(), anyhow::Error> {
}
infra_file_list::create_table_index().await?;
infra_file_list::set_initialised().await?;
db::file_list::remote::cache_stats()
.await
.expect("Load stream stats failed");

17
src/service/alerts/mod.rs
View File

@ -341,7 +341,12 @@ impl QueryCondition {
(end - start) / promql::MAX_DATA_POINTS,
),
};
let resp = promql::search::search(&alert.org_id, &req, 0).await?;
let resp = match promql::search::search(&alert.org_id, &req, 0).await {
Ok(v) => v,
Err(_) => {
return Ok(None);
}
};
let promql::value::Value::Matrix(value) = resp else {
log::warn!(
"Alert evaluate: PromQL query {} returned unexpected response: {:?}",
@ -404,8 +409,14 @@ impl QueryCondition {
timeout: 0,
};
let session_id = ider::uuid();
let resp =
SearchService::search(&session_id, &alert.org_id, alert.stream_type, &req).await?;
let resp = match SearchService::search(&session_id, &alert.org_id, alert.stream_type, &req)
.await
{
Ok(v) => v,
Err(_) => {
return Ok(None);
}
};
if resp.total < alert.trigger_condition.threshold as usize {
Ok(None)
} else {

8
src/service/compact/stats.rs
View File

@ -25,14 +25,6 @@ use crate::{
};
pub async fn update_stats_from_file_list() -> Result<(), anyhow::Error> {
// waiting for file list remote inited
loop {
if infra_file_list::get_initialised().await.unwrap_or_default() {
break;
};
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
// get last offset
let (mut offset, node) = db::compact::stats::get_offset().await;
if !node.is_empty() && LOCAL_NODE_UUID.ne(&node) && get_node_by_uuid(&node).is_some() {

2
src/service/ingestion/grpc.rs
View File

@ -217,6 +217,6 @@ mod tests {
value: Some(in_byte),
};
let resp = get_val(&Some(&byte_val));
assert!(!resp.as_array().unwrap().is_empty());
assert!(!resp.as_str().unwrap().is_empty());
}
}

3
src/service/ingestion/mod.rs
View File

@ -18,6 +18,7 @@ use std::{
sync::Arc,
};
use anyhow::{anyhow, Result};
use arrow_schema::Schema;
use chrono::{TimeZone, Utc};
use config::{meta::stream::StreamType, SIZE_IN_MB};
@ -27,8 +28,6 @@ use vrl::{
prelude::state,
};
use anyhow::{anyhow, Result};
use crate::{
common::{
infra::{

1
src/service/logs/bulk.rs
View File

@ -206,6 +206,7 @@ pub async fn ingest(
json::Value::Object(v) => v,
_ => unreachable!(),
};
// set _id
if !doc_id.is_empty() {
local_val.insert("_id".to_string(), json::Value::String(doc_id.clone()));

4
src/service/logs/ingest.rs
View File

@ -19,7 +19,7 @@ use std::{
};
use actix_web::http;
use anyhow::{anyhow, Result};
use anyhow::Result;
use chrono::{Duration, Utc};
use config::{meta::stream::StreamType, metrics, CONFIG, DISTINCT_FIELDS};
use datafusion::arrow::datatypes::Schema;
@ -201,6 +201,8 @@ pub async fn ingest(
if local_trigger.is_some() {
trigger = local_trigger;
}
// add distinct values
distinct_values.extend(to_add_distinct_values);
}

8
src/service/logs/multi.rs
View File

@ -136,11 +136,17 @@ async fn ingest_inner(
)?;
}
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
continue;
}
// End row based transform
// get json object
let mut local_val = match value.take() {
json::Value::Object(v) => v,
_ => unreachable!(),
};
// End row based transform
// handle timestamp
let timestamp = match local_val.get(&CONFIG.common.column_timestamp) {

3
src/service/logs/otlp_grpc.rs
View File

@ -26,7 +26,6 @@ use opentelemetry_proto::tonic::collector::logs::v1::{
ExportLogsServiceRequest, ExportLogsServiceResponse,
};
use prost::Message;
use anyhow::Result;
use super::StreamMeta;
use crate::{
@ -184,6 +183,8 @@ pub async fn usage_ingest(
if local_trigger.is_some() {
trigger = local_trigger;
}
// add distinct values
distinct_values.extend(to_add_distinct_values);
}

2
src/service/logs/otlp_http.rs
View File

@ -351,6 +351,7 @@ pub async fn logs_json_handler(
.unwrap();
}
// get json object
let local_val = match value.take() {
json::Value::Object(v) => v,
_ => unreachable!(),
@ -400,6 +401,7 @@ pub async fn logs_json_handler(
trigger = local_trigger;
}
// add distinct values
distinct_values.extend(to_add_distinct_values);
}
}

10
src/service/logs/syslog.rs
View File

@ -131,6 +131,16 @@ pub async fn ingest(msg: &str, addr: SocketAddr) -> Result<HttpResponse> {
)?;
}
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
return Ok(HttpResponse::Ok().json(IngestionResponse::new(
http::StatusCode::OK.into(),
vec![stream_status],
))); // just return
}
// End row based transform
// get json object
let mut local_val = match value.take() {
json::Value::Object(v) => v,
_ => unreachable!(),

71
src/service/organization.rs
View File

@ -43,38 +43,54 @@ pub async fn get_summary(org_id: &str) -> OrgSummary {
}
#[tracing::instrument]
pub async fn get_passcode(org_id: Option<&str>, user_id: &str) -> IngestionPasscode {
let user = db::user::get(org_id, user_id).await.unwrap().unwrap();
IngestionPasscode {
pub async fn get_passcode(
org_id: Option<&str>,
user_id: &str,
) -> Result<IngestionPasscode, anyhow::Error> {
let Ok(Some(user)) = db::user::get(org_id, user_id).await else {
return Err(anyhow::Error::msg("User not found"));
};
Ok(IngestionPasscode {
user: user.email,
passcode: user.token,
}
})
}
#[tracing::instrument]
pub async fn get_rum_token(org_id: Option<&str>, user_id: &str) -> RumIngestionToken {
let user = db::user::get(org_id, user_id).await.unwrap().unwrap();
RumIngestionToken {
pub async fn get_rum_token(
org_id: Option<&str>,
user_id: &str,
) -> Result<RumIngestionToken, anyhow::Error> {
let Ok(Some(user)) = db::user::get(org_id, user_id).await else {
return Err(anyhow::Error::msg("User not found"));
};
Ok(RumIngestionToken {
user: user.email,
rum_token: user.rum_token,
}
})
}
#[tracing::instrument]
pub async fn update_rum_token(org_id: Option<&str>, user_id: &str) -> RumIngestionToken {
pub async fn update_rum_token(
org_id: Option<&str>,
user_id: &str,
) -> Result<RumIngestionToken, anyhow::Error> {
let is_rum_update = true;
match update_passcode_inner(org_id, user_id, is_rum_update).await {
IngestionTokensContainer::RumToken(response) => response,
_ => panic!("This shouldn't have happened, we were expecting rum token updates"),
Ok(IngestionTokensContainer::RumToken(response)) => Ok(response),
_ => Err(anyhow::Error::msg("User not found")),
}
}
#[tracing::instrument]
pub async fn update_passcode(org_id: Option<&str>, user_id: &str) -> IngestionPasscode {
pub async fn update_passcode(
org_id: Option<&str>,
user_id: &str,
) -> Result<IngestionPasscode, anyhow::Error> {
let is_rum_update = false;
match update_passcode_inner(org_id, user_id, is_rum_update).await {
IngestionTokensContainer::Passcode(response) => response,
_ => panic!("This shouldn't have happened, we were expecting ingestion token updates"),
Ok(IngestionTokensContainer::Passcode(response)) => Ok(response),
_ => Err(anyhow::Error::msg("User not found")),
}
}
@ -83,9 +99,11 @@ async fn update_passcode_inner(
org_id: Option<&str>,
user_id: &str,
is_rum_update: bool,
) -> IngestionTokensContainer {
) -> Result<IngestionTokensContainer, anyhow::Error> {
let mut local_org_id = "dummy";
let mut db_user = db::user::get_db_user(user_id).await.unwrap();
let Ok(mut db_user) = db::user::get_db_user(user_id).await else {
return Err(anyhow::Error::msg("User not found"));
};
if org_id.is_some() {
local_org_id = org_id.unwrap();
@ -147,7 +165,7 @@ async fn update_passcode_inner(
db_user.organizations = new_orgs;
let _ = db::user::set(db_user.clone()).await;
if is_rum_update {
let ret = if is_rum_update {
IngestionTokensContainer::RumToken(RumIngestionToken {
user: db_user.email,
rum_token: Some(rum_token),
@ -157,7 +175,8 @@ async fn update_passcode_inner(
user: db_user.email,
passcode: token,
})
}
};
Ok(ret)
}
#[tracing::instrument]
@ -177,16 +196,13 @@ pub async fn create_org(org: &Organization) -> Result<Organization, Error> {
#[cfg(test)]
mod tests {
use super::*;
use crate::{
common::{infra::db as infra_db, meta::user::UserRequest},
service::users,
};
use crate::{common::meta::user::UserRequest, service::users};
#[actix_web::test]
async fn test_organization() {
infra_db::create_table().await.unwrap();
let org_id = "dummy";
let org_id = "default";
let user_id = "userone@example.com";
let init_user = "root@example.com";
// let passcode = "samplePassCode";
let resp = users::post_user(
org_id,
@ -198,16 +214,17 @@ mod tests {
last_name: "".to_owned(),
is_external: false,
},
user_id,
init_user,
)
.await;
assert!(resp.is_ok());
assert!(resp.unwrap().status().is_success());
let resp = get_passcode(Some(org_id), user_id).await;
let resp = get_passcode(Some(org_id), user_id).await.unwrap();
let passcode = resp.passcode.clone();
assert!(!resp.passcode.is_empty());
let resp = update_passcode(Some(org_id), user_id).await;
let resp = update_passcode(Some(org_id), user_id).await.unwrap();
assert_ne!(resp.passcode, passcode);
}
}

3
src/service/promql/mod.rs
View File

@ -30,14 +30,13 @@ mod engine;
mod exec;
mod functions;
pub mod name_visitor;
pub mod search;
pub mod value;
pub mod name_visitor;
pub use engine::Engine;
pub use exec::Query;
use crate::common::meta::stream::ScanStats;
pub(crate) const DEFAULT_LOOKBACK: Duration = Duration::from_secs(300); // 5m

7
src/service/schema.rs
View File

@ -846,8 +846,9 @@ mod tests {
async fn test_check_for_schema() {
let stream_name = "Sample";
let org_name = "nexus";
let record =
json::json!(r#"{"Year": 1896, "City": "Athens", "_timestamp": 1234234234234}"#);
let record: json::Value =
json::from_str(r#"{"Year": 1896, "City": "Athens", "_timestamp": 1234234234234}"#)
.unwrap();
let schema = Schema::new(vec![
Field::new("Year", DataType::Int64, false),
@ -861,7 +862,7 @@ mod tests {
stream_name,
StreamType::Logs,
&mut map,
&record.as_object().unwrap(),
record.as_object().unwrap(),
1234234234234,
)
.await

8
src/service/search/mod.rs
View File

@ -321,7 +321,9 @@ async fn search_in_cluster(mut req: cluster_rpc::SearchRequest) -> Result<search
let grpc_span = info_span!(
"service:search:cluster:grpc_search",
session_id,
org_id = req.org_id
org_id = req.org_id,
node_id = node.id,
node_addr = node_addr.as_str(),
);
let task = tokio::task::spawn(
async move {
@ -339,6 +341,8 @@ async fn search_in_cluster(mut req: cluster_rpc::SearchRequest) -> Result<search
)
});
log::info!("[session_id {session_id}] search->grpc: request node: {}, is_querier: {}", &node_addr, is_querier);
let token: MetadataValue<_> = cluster::get_internal_grpc_token()
.parse()
.map_err(|_| Error::Message("invalid token".to_string()))?;
@ -375,7 +379,7 @@ async fn search_in_cluster(mut req: cluster_rpc::SearchRequest) -> Result<search
};
log::info!(
"[session_id {session_id}] search->grpc: result node: {}, is_querier: {}, total: {}, took: {}, files: {}, scan_size: {}",
"[session_id {session_id}] search->grpc: response node: {}, is_querier: {}, total: {}, took: {}, files: {}, scan_size: {}",
&node.grpc_addr,
is_querier,
response.total,

44
src/service/search/sql.rs
View File

@ -53,7 +53,6 @@ static RE_ONLY_GROUPBY: Lazy<Regex> =
static RE_SELECT_FIELD: Lazy<Regex> =
Lazy::new(|| Regex::new(r"(?i)select (.*) from[ ]+query").unwrap());
static RE_SELECT_FROM: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i)SELECT (.*) FROM").unwrap());
static RE_TIMESTAMP_EMPTY: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i) where (.*)").unwrap());
static RE_WHERE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i) where (.*)").unwrap());
static RE_ONLY_WHERE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i) where ").unwrap());
@ -220,7 +219,7 @@ impl Sql {
};
origin_sql = origin_sql.replace(caps.get(0).unwrap().as_str(), " FROM tbl ");
// Hack _timestamp
// Hack select for _timestamp
if !sql_mode.eq(&SqlMode::Full) && meta.order_by.is_empty() && !origin_sql.contains('*') {
let caps = RE_SELECT_FROM.captures(origin_sql.as_str()).unwrap();
let cap_str = caps.get(1).unwrap().as_str();
@ -232,7 +231,7 @@ impl Sql {
}
}
// check time_range
// check time_range values
if req_time_range.0 > 0
&& req_time_range.0 < Duration::seconds(1).num_microseconds().unwrap()
{
@ -253,9 +252,8 @@ impl Sql {
)));
}
// Hack time_range
let meta_time_range_is_empty =
meta.time_range.is_none() || meta.time_range.unwrap() == (0, 0);
// Hack time_range for sql
let meta_time_range_is_empty = meta.time_range.is_none() || meta.time_range == Some((0, 0));
if meta_time_range_is_empty && (req_time_range.0 > 0 || req_time_range.1 > 0) {
if req_time_range.1 == 0 {
req_time_range.1 = chrono::Utc::now().timestamp_micros();
@ -279,22 +277,29 @@ impl Sql {
"".to_string()
};
if !time_range_sql.is_empty() && meta_time_range_is_empty {
match RE_TIMESTAMP_EMPTY.captures(origin_sql.as_str()) {
match RE_WHERE.captures(origin_sql.as_str()) {
Some(caps) => {
let mut where_str = caps.get(1).unwrap().as_str().to_string();
if !meta.order_by.is_empty() {
where_str = where_str
[0..where_str.to_lowercase().rfind(" order ").unwrap()]
.to_string();
}
if !meta.group_by.is_empty() {
where_str = where_str
[0..where_str.to_lowercase().rfind(" group ").unwrap()]
.to_string();
} else if where_str.to_lowercase().contains(" having ") {
} else if meta.having {
where_str = where_str
[0..where_str.to_lowercase().rfind(" having ").unwrap()]
.to_string();
} else if !meta.order_by.is_empty() {
where_str = where_str
[0..where_str.to_lowercase().rfind(" order ").unwrap()]
.to_string();
} else if meta.limit > 0 {
where_str = where_str
[0..where_str.to_lowercase().rfind(" limit ").unwrap()]
.to_string();
} else if meta.offset > 0 {
where_str = where_str
[0..where_str.to_lowercase().rfind(" offset ").unwrap()]
.to_string();
}
let pos_start = origin_sql.find(where_str.as_str()).unwrap();
let pos_end = pos_start + where_str.len();
@ -314,7 +319,7 @@ impl Sql {
}
}
// Hack offset limit
// Hack offset limit and sort by for sql
if meta.limit == 0 {
meta.offset = req_query.from as usize;
meta.limit = req_query.size as usize;
@ -910,6 +915,11 @@ mod tests {
false,
(0, 0),
),
(
"select abc, count(*) as cnt from table1 where match_all('abc') and str_match(log,'abc') group by abc having cnt > 1 order by _timestamp desc limit 10",
false,
(0, 0),
),
];
let org_id = "test_org";
@ -1023,6 +1033,12 @@ mod tests {
0,
(0, 0),
),
(
"SELECT trace_id, MIN(start_time) AS start_time FROM table1 WHERE service_name ='APISIX-B' GROUP BY trace_id ORDER BY start_time DESC LIMIT 10",
true,
10,
(0, 0),
),
];
let org_id = "test_org";

35
src/service/users.rs
View File

@ -15,10 +15,7 @@
use std::io::Error;
use actix_web::{
http::{self, StatusCode},
HttpResponse,
};
use actix_web::{http, HttpResponse};
use config::ider;
use crate::{
@ -44,13 +41,11 @@ pub async fn post_user(
usr_req: UserRequest,
initiator_id: &str,
) -> Result<HttpResponse, Error> {
let initiator_user = db::user::get(Some(org_id), initiator_id).await;
if is_root_user(initiator_id)
|| db::user::get(Some(org_id), initiator_id)
.await
.unwrap()
.unwrap()
.role
.eq(&UserRole::Admin)
|| (initiator_user.is_ok()
&& initiator_user.as_ref().unwrap().is_some()
&& initiator_user.unwrap().unwrap().role.eq(&UserRole::Admin))
{
let existing_user = if is_root_user(&usr_req.email) {
db::user::get(None, &usr_req.email).await
@ -83,7 +78,7 @@ pub async fn post_user(
}
} else {
Ok(HttpResponse::Unauthorized().json(MetaHttpResponse::error(
StatusCode::UNAUTHORIZED.into(),
http::StatusCode::UNAUTHORIZED.into(),
"Not Allowed".to_string(),
)))
}
@ -238,7 +233,7 @@ pub async fn update_user(
)))
}
Err(_) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
"User not found".to_string(),
))),
}
@ -253,13 +248,13 @@ pub async fn update_user(
}
}
None => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
"User not found".to_string(),
))),
}
} else {
Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
"User not found".to_string(),
)))
}
@ -315,13 +310,13 @@ pub async fn add_user_to_org(
)))
} else {
Ok(HttpResponse::Unauthorized().json(MetaHttpResponse::error(
StatusCode::UNAUTHORIZED.into(),
http::StatusCode::UNAUTHORIZED.into(),
"Not Allowed".to_string(),
)))
}
} else {
Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
"User not found".to_string(),
)))
}
@ -419,19 +414,19 @@ pub async fn remove_user_from_org(
)))
} else {
Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
"User for the organization not found".to_string(),
)))
}
}
Err(_) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
"User for the organization not found".to_string(),
))),
}
} else {
Ok(HttpResponse::Unauthorized().json(MetaHttpResponse::error(
StatusCode::UNAUTHORIZED.into(),
http::StatusCode::UNAUTHORIZED.into(),
"Not Allowed".to_string(),
)))
}
@ -445,7 +440,7 @@ pub async fn delete_user(email_id: &str) -> Result<HttpResponse, Error> {
"User deleted".to_string(),
))),
Err(e) => Ok(HttpResponse::NotFound().json(MetaHttpResponse::error(
StatusCode::NOT_FOUND.into(),
http::StatusCode::NOT_FOUND.into(),
e.to_string(),
))),
}

103
tests/integration_test.rs
View File

@ -70,8 +70,19 @@ mod tests {
.unwrap_or_else(|e| log::info!("Error deleting local dir: {}", e));
setup();
let _ = openobserve::common::infra::init().await;
let _ = openobserve::job::init().await;
// register node
_ = openobserve::common::infra::cluster::register_and_keepalive()
.await
.unwrap();
// init config
_ = config::init().await.unwrap();
// init infra
_ = openobserve::common::infra::init().await.unwrap();
// ingester init
_ = ingester::init().await.unwrap();
// init job
_ = openobserve::job::init().await.unwrap();
for _i in 0..3 {
e2e_1_post_bulk().await;
@ -89,7 +100,6 @@ mod tests {
e2e_get_stream_schema().await;
e2e_get_org_summary().await;
e2e_post_stream_settings().await;
e2e_delete_stream().await;
e2e_get_org().await;
// functions
@ -163,13 +173,16 @@ mod tests {
e2e_health_check().await;
e2e_config().await;
e2e_100_tear_down().await;
// clear
e2e_delete_stream().await;
}
async fn e2e_1_post_bulk() {
let auth = setup();
let path = "./tests/input.json";
let body_str = fs::read_to_string(path).expect("Read file failed");
let thread_id: usize = 1;
let thread_id: usize = 0;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -192,7 +205,7 @@ mod tests {
async fn e2e_post_json() {
let auth = setup();
let body_str = "[{\"Year\": 1896, \"City\": \"Athens\", \"Sport\": \"Aquatics\", \"Discipline\": \"Swimming\", \"Athlete\": \"HERSCHMANN, Otto\", \"Country\": \"AUT\", \"Gender\": \"Men\", \"Event\": \"100M Freestyle\", \"Medal\": \"Silver\", \"Season\": \"summer\",\"_timestamp\":1665136888163792}]";
let thread_id: usize = 1;
let thread_id: usize = 0;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -215,7 +228,7 @@ mod tests {
async fn e2e_post_multi() {
let auth = setup();
let body_str = "{\"Year\": 1896, \"City\": \"Athens\", \"Sport\": \"Aquatics\", \"Discipline\": \"Swimming\", \"Athlete\": \"HERSCHMANN, Otto\", \"Country\": \"AUT\", \"Gender\": \"Men\", \"Event\": \"100M Freestyle\", \"Medal\": \"Silver\", \"Season\": \"summer\",\"_timestamp\":1665136888163792}";
let thread_id: usize = 1;
let thread_id: usize = 0;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -267,7 +280,10 @@ mod tests {
)
.await;
let req = test::TestRequest::get()
.uri(&format!("/api/{}/{}/schema", "e2e", "olympics_schema"))
.uri(&format!(
"/api/{}/streams/{}/schema",
"e2e", "olympics_schema"
))
.insert_header(ContentType::json())
.append_header(auth)
.to_request();
@ -277,9 +293,9 @@ mod tests {
async fn e2e_post_stream_settings() {
let auth = setup();
let body_str = r#"{ "partition_keys": ["test_key"], "full_text_search_keys": ["log"]}"#;
let body_str = r#"{"partition_keys": ["test_key"], "full_text_search_keys": ["log"]}"#;
// app
let thread_id: usize = 1;
let thread_id: usize = 0;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -289,8 +305,11 @@ mod tests {
.configure(get_basic_routes),
)
.await;
let req = test::TestRequest::post()
.uri(&format!("/api/{}/{}/settings", "e2e", "olympics"))
let req = test::TestRequest::put()
.uri(&format!(
"/api/{}/streams/{}/settings",
"e2e", "olympics_schema"
))
.insert_header(ContentType::json())
.append_header(auth)
.set_payload(body_str)
@ -310,7 +329,7 @@ mod tests {
)
.await;
let req = test::TestRequest::delete()
.uri(&format!("/api/{}/{}", "e2e", "olympics"))
.uri(&format!("/api/{}/streams/{}", "e2e", "olympics_schema"))
.insert_header(ContentType::json())
.append_header(auth)
.to_request();
@ -375,9 +394,9 @@ mod tests {
.configure(get_basic_routes),
)
.await;
let req = test::TestRequest::post()
let req = test::TestRequest::put()
.uri(&format!(
"/api/{}/{}/functions/{}",
"/api/{}/streams/{}/functions/{}",
"e2e", "olympics_schema", "e2etestfn"
))
.insert_header(ContentType::json())
@ -418,7 +437,10 @@ mod tests {
)
.await;
let req = test::TestRequest::get()
.uri(&format!("/api/{}/{}/functions", "e2e", "olympics_schema"))
.uri(&format!(
"/api/{}/streams/{}/functions",
"e2e", "olympics_schema"
))
.insert_header(ContentType::json())
.append_header(auth)
.to_request();
@ -457,7 +479,7 @@ mod tests {
.await;
let req = test::TestRequest::delete()
.uri(&format!(
"/api/{}/{}/functions/{}",
"/api/{}/streams/{}/functions/{}",
"e2e", "olympics_schema", "e2etestfn"
))
.insert_header(ContentType::json())
@ -549,7 +571,7 @@ mod tests {
)
.await;
let req = test::TestRequest::get()
.uri(&format!("/api/{}/organizations", "e2e",))
.uri("/api/organizations")
.insert_header(ContentType::json())
.append_header(auth)
.to_request();
@ -568,7 +590,7 @@ mod tests {
)
.await;
let req = test::TestRequest::get()
.uri(&format!("/api/{}/organizations/passcode", "e2e"))
.uri(&format!("/api/{}/passcode", "e2e"))
.insert_header(ContentType::json())
.append_header(auth)
.to_request();
@ -588,7 +610,7 @@ mod tests {
)
.await;
let req = test::TestRequest::put()
.uri(&format!("/api/{}/organizations/passcode", "e2e"))
.uri(&format!("/api/{}/passcode", "e2e"))
.insert_header(ContentType::json())
.append_header(auth)
.set_payload(body_str)
@ -867,7 +889,7 @@ mod tests {
let body_str = fs::read_to_string(path).expect("Read file failed");
// app
let thread_id: usize = 1;
let thread_id: usize = 0;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -949,7 +971,7 @@ mod tests {
.expect("Out of memory");
// app
let thread_id: usize = 1;
let thread_id: usize = 0;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -993,7 +1015,7 @@ mod tests {
async fn e2e_post_alert_template() {
let auth = setup();
let body_str = r#"{"body": {"text": "For stream {stream_name} of organization {org_name} alert {alert_name} of type {alert_type} is active app_name {app_name}" }}"#;
let body_str = r#"{"name":"slackTemplate","body":"{\"text\":\"For stream {stream_name} of organization {org_name} alert {alert_name} of type {alert_type} is active app_name {app_name}\"}"}"#;
let app = test::init_service(
App::new()
.app_data(web::JsonConfig::default().limit(CONFIG.limit.req_json_limit))
@ -1003,10 +1025,7 @@ mod tests {
)
.await;
let req = test::TestRequest::post()
.uri(&format!(
"/api/{}/alerts/templates/{}",
"e2e", "slackTemplate"
))
.uri(&format!("/api/{}/alerts/templates", "e2e"))
.insert_header(ContentType::json())
.append_header(auth)
.set_payload(body_str)
@ -1081,6 +1100,7 @@ mod tests {
async fn e2e_post_alert_destination() {
let auth = setup();
let body_str = r#"{
"name": "slack",
"url": "https://dummy/alert",
"method": "post",
"template": "slackTemplate",
@ -1097,7 +1117,7 @@ mod tests {
)
.await;
let req = test::TestRequest::post()
.uri(&format!("/api/{}/alerts/destinations/{}", "e2e", "slack"))
.uri(&format!("/api/{}/alerts/destinations", "e2e"))
.insert_header(ContentType::json())
.append_header(auth)
.set_payload(body_str)
@ -1167,15 +1187,23 @@ mod tests {
async fn e2e_post_alert() {
let auth = setup();
let body_str = r#"{
"condition": {
"column": "country",
"operator": "=",
"value": "USA"
"name": "alertChk",
"stream_type": "logs",
"stream_name": "olympics_schema",
"is_real_time": false,
"query_condition": {
"conditions": [{
"column": "country",
"operator": "=",
"value": "USA"
}]
},
"duration": 5,
"frequency": 1,
"time_between_alerts": 10,
"destination":"slack",
"trigger_condition": {
"period": 5,
"threshold": 1,
"silence": 10
},
"destinations": ["slack"],
"context_attributes":{
"app_name":"App1"
}
@ -1189,10 +1217,7 @@ mod tests {
)
.await;
let req = test::TestRequest::post()
.uri(&format!(
"/api/{}/{}/alerts/{}",
"e2e", "olympics_schema", "alertChk"
))
.uri(&format!("/api/{}/{}/alerts", "e2e", "olympics_schema"))
.insert_header(ContentType::json())
.append_header(auth)
.set_payload(body_str)