diff --git a/src/infra/src/db/mysql.rs b/src/infra/src/db/mysql.rs
index 1c0b80808e..cfaaf05c03 100644
--- a/src/infra/src/db/mysql.rs
+++ b/src/infra/src/db/mysql.rs
@@ -166,7 +166,14 @@ impl super::Db for MysqlDb {
}
return Err(e.into());
}
+ // need commit it first to avoid the deadlock of insert and update
+ if let Err(e) = tx.commit().await {
+ log::error!("[MYSQL] commit put meta error: {}", e);
+ return Err(e.into());
+ }
+
DB_QUERY_NUMS.with_label_values(&["update", "meta"]).inc();
+ let mut tx = pool.begin().await?;
if let Err(e) = sqlx::query(
r#"UPDATE meta SET value = ? WHERE module = ? AND key1 = ? AND key2 = ? AND start_dt = ?;"#
)
diff --git a/src/infra/src/db/postgres.rs b/src/infra/src/db/postgres.rs
index f5ad6a179a..41ea5f17de 100644
--- a/src/infra/src/db/postgres.rs
+++ b/src/infra/src/db/postgres.rs
@@ -156,8 +156,14 @@ impl super::Db for PostgresDb {
}
return Err(e.into());
}
+ // need commit it first to avoid the deadlock of insert and update
+ if let Err(e) = tx.commit().await {
+ log::error!("[POSTGRES] commit put meta error: {}", e);
+ return Err(e.into());
+ }
DB_QUERY_NUMS.with_label_values(&["update", "meta"]).inc();
+ let mut tx = pool.begin().await?;
if let Err(e) = sqlx::query(
r#"UPDATE meta SET value = $1 WHERE module = $2 AND key1 = $3 AND key2 = $4 AND start_dt = $5;"#
)
diff --git a/src/infra/src/db/sqlite.rs b/src/infra/src/db/sqlite.rs
index 69c5063867..bd3add4b4b 100644
--- a/src/infra/src/db/sqlite.rs
+++ b/src/infra/src/db/sqlite.rs
@@ -272,6 +272,13 @@ impl super::Db for SqliteDb {
}
return Err(e.into());
}
+ // need commit it first to avoid the deadlock of insert and update
+ if let Err(e) = tx.commit().await {
+ log::error!("[SQLITE] commit put meta error: {}", e);
+ return Err(e.into());
+ }
+
+ let mut tx = client.begin().await?;
if let Err(e) = sqlx::query(
r#"UPDATE meta SET value = $1 WHERE module = $2 AND key1 = $3 AND key2 = $4 AND start_dt = $5;"#
)
diff --git a/src/infra/src/table/entity/README.md b/src/infra/src/table/entity/README.md
index cdf8da3269..d50f75d7e3 100644
--- a/src/infra/src/table/entity/README.md
+++ b/src/infra/src/table/entity/README.md
@@ -8,20 +8,20 @@ Be sure to replace `YOUR_DATABASE_URL` with the database connection string. Alte
```
sea-orm-cli generate entity \
---output-dir src/infra/src/table/entity \
---date-time-crate chrono \
---ignore-tables file_list_deleted \
---ignore-tables file_list_history \
---ignore-tables file_list_jobs \
---ignore-tables file_list \
---ignore-tables meta \
---ignore-tables pipeline \
---ignore-tables scheduled_jobs \
---ignore-tables schema_history \
---ignore-tables seaql_migrations \
---ignore-tables short_urls \
---ignore-tables stream_stats \
---database-url YOUR_DATABASE_URL
+ --output-dir src/infra/src/table/entity \
+ --date-time-crate chrono \
+ --ignore-tables file_list_deleted \
+ --ignore-tables file_list_history \
+ --ignore-tables file_list_jobs \
+ --ignore-tables file_list \
+ --ignore-tables meta \
+ --ignore-tables pipeline \
+ --ignore-tables scheduled_jobs \
+ --ignore-tables schema_history \
+ --ignore-tables seaql_migrations \
+ --ignore-tables short_urls \
+ --ignore-tables stream_stats \
+ --database-url YOUR_DATABASE_URL
```
We use the `--ignore-tables` parameter to omit model generation of existing tables for which automatic model generation is not needed yet. In the future we can include these tables in the automatic model generation when we want to start using `sea-orm` as the means of interacting with these tables.
diff --git a/src/infra/src/table/entity/mod.rs b/src/infra/src/table/entity/mod.rs
index 7b2a6a2931..4217965378 100644
--- a/src/infra/src/table/entity/mod.rs
+++ b/src/infra/src/table/entity/mod.rs
@@ -4,3 +4,4 @@ pub mod prelude;
pub mod dashboards;
pub mod folders;
+pub mod search_queue;
diff --git a/src/infra/src/table/entity/prelude.rs b/src/infra/src/table/entity/prelude.rs
index ca648cff5a..e390b4d063 100644
--- a/src/infra/src/table/entity/prelude.rs
+++ b/src/infra/src/table/entity/prelude.rs
@@ -1,3 +1,6 @@
//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.0
-pub use super::{dashboards::Entity as Dashboards, folders::Entity as Folders};
+pub use super::{
+ dashboards::Entity as Dashboards, folders::Entity as Folders,
+ search_queue::Entity as SearchQueue,
+};
diff --git a/src/infra/src/table/entity/search_queue.rs b/src/infra/src/table/entity/search_queue.rs
new file mode 100644
index 0000000000..e0ada0f804
--- /dev/null
+++ b/src/infra/src/table/entity/search_queue.rs
@@ -0,0 +1,19 @@
+//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.0
+
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
+#[sea_orm(table_name = "search_queue")]
+pub struct Model {
+ #[sea_orm(primary_key)]
+ pub id: i64,
+ pub work_group: String,
+ pub user_id: String,
+ pub trace_id: String,
+ pub created_at: i64,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}
diff --git a/src/infra/src/table/migration/m20241204_143100_create_table_search_queue.rs b/src/infra/src/table/migration/m20241204_143100_create_table_search_queue.rs
new file mode 100644
index 0000000000..d224fbdf08
--- /dev/null
+++ b/src/infra/src/table/migration/m20241204_143100_create_table_search_queue.rs
@@ -0,0 +1,85 @@
+use sea_orm_migration::prelude::*;
+
+#[derive(DeriveMigrationName)]
+pub struct Migration;
+
+const SEARCH_QUEUE_WORK_GROUP_IDX: &str = "search_queue_work_group_idx";
+const SEARCH_QUEUE_CREATED_AT_IDX: &str = "search_queue_created_at_idx";
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+ async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+ manager.create_table(create_table_stmt()).await?;
+ manager.create_index(create_index_work_group_stmt()).await?;
+ manager.create_index(create_index_created_at_stmt()).await?;
+ Ok(())
+ }
+
+ async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+ manager
+ .drop_index(Index::drop().name(SEARCH_QUEUE_WORK_GROUP_IDX).to_owned())
+ .await?;
+ manager
+ .drop_index(Index::drop().name(SEARCH_QUEUE_CREATED_AT_IDX).to_owned())
+ .await?;
+ manager
+ .drop_table(Table::drop().table(SearchQueue::Table).to_owned())
+ .await?;
+ Ok(())
+ }
+}
+
+/// Statement to create table.
+fn create_table_stmt() -> TableCreateStatement {
+ Table::create()
+ .table(SearchQueue::Table)
+ .if_not_exists()
+ .col(
+ ColumnDef::new(SearchQueue::Id)
+ .big_integer()
+ .not_null()
+ .auto_increment()
+ .primary_key(),
+ )
+ .col(ColumnDef::new(SearchQueue::WorkGroup).string_len(16).not_null())
+ .col(ColumnDef::new(SearchQueue::UserId).string_len(256).not_null())
+ // Identifier of the user that owns the dashboard.
+ .col(ColumnDef::new(SearchQueue::TraceId).string_len(32).not_null())
+ .col(
+ ColumnDef::new(SearchQueue::CreatedAt)
+ .big_integer()
+ .not_null(),
+ )
+ .to_owned()
+}
+
+/// Statement to create index work_group.
+fn create_index_work_group_stmt() -> IndexCreateStatement {
+ sea_query::Index::create()
+ .if_not_exists()
+ .name(SEARCH_QUEUE_WORK_GROUP_IDX)
+ .table(SearchQueue::Table)
+ .col(SearchQueue::WorkGroup)
+ .col(SearchQueue::UserId)
+ .to_owned()
+}
+
+/// Statement to create index created_at.
+fn create_index_created_at_stmt() -> IndexCreateStatement {
+ sea_query::Index::create()
+ .if_not_exists()
+ .name(SEARCH_QUEUE_CREATED_AT_IDX)
+ .table(SearchQueue::Table)
+ .col(SearchQueue::CreatedAt)
+ .to_owned()
+}
+
+#[derive(DeriveIden)]
+enum SearchQueue {
+ Table,
+ Id,
+ WorkGroup,
+ UserId,
+ TraceId,
+ CreatedAt,
+}
diff --git a/src/infra/src/table/migration/mod.rs b/src/infra/src/table/migration/mod.rs
index 4e5d51fb6a..71462e151a 100644
--- a/src/infra/src/table/migration/mod.rs
+++ b/src/infra/src/table/migration/mod.rs
@@ -22,6 +22,7 @@ mod m20241116_000002_drop_folders_created_at_column;
mod m20241119_000001_create_dashboards_table;
mod m20241119_000002_populate_dashboards_table;
mod m20241119_000003_delete_metas;
+mod m20241204_143100_create_table_search_queue;
pub struct Migrator;
@@ -36,6 +37,7 @@ impl MigratorTrait for Migrator {
Box::new(m20241119_000001_create_dashboards_table::Migration),
Box::new(m20241119_000002_populate_dashboards_table::Migration),
Box::new(m20241119_000003_delete_metas::Migration),
+ Box::new(m20241204_143100_create_table_search_queue::Migration),
]
}
}
diff --git a/src/infra/src/table/mod.rs b/src/infra/src/table/mod.rs
index d9cc4dd593..db14d26180 100644
--- a/src/infra/src/table/mod.rs
+++ b/src/infra/src/table/mod.rs
@@ -25,11 +25,12 @@ pub mod distinct_values;
mod entity;
pub mod folders;
mod migration;
+pub mod search_queue;
pub mod short_urls;
pub async fn init() -> Result<(), anyhow::Error> {
- short_urls::init().await?;
distinct_values::init().await?;
+ short_urls::init().await?;
let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
Migrator::up(client, None).await?;
Ok(())
diff --git a/src/infra/src/table/search_queue.rs b/src/infra/src/table/search_queue.rs
new file mode 100644
index 0000000000..bd00826140
--- /dev/null
+++ b/src/infra/src/table/search_queue.rs
@@ -0,0 +1,99 @@
+// Copyright 2024 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use sea_orm::{ColumnTrait, EntityTrait, PaginatorTrait, QueryFilter, Set};
+
+use super::{entity::search_queue::*, get_lock};
+use crate::{
+ db::{connect_to_orm, ORM_CLIENT},
+ errors,
+};
+
+pub async fn add(work_group: &str, user_id: &str, trace_id: &str) -> Result<(), errors::Error> {
+ let record = ActiveModel {
+ work_group: Set(work_group.to_string()),
+ user_id: Set(user_id.to_string()),
+ trace_id: Set(trace_id.to_string()),
+ created_at: Set(chrono::Utc::now().timestamp_micros()),
+ ..Default::default()
+ };
+
+ // make sure only one client is writing to the database(only for sqlite)
+ let _lock = get_lock().await;
+
+ let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
+ Entity::insert(record).exec(client).await?;
+
+ Ok(())
+}
+
+pub async fn delete(id: i64) -> Result<(), errors::Error> {
+ // make sure only one client is writing to the database(only for sqlite)
+ let _lock = get_lock().await;
+
+ let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
+ Entity::delete_many()
+ .filter(Column::Id.eq(id))
+ .exec(client)
+ .await?;
+
+ Ok(())
+}
+
+pub async fn delete_by_trace_id(trace_id: &str) -> Result<(), errors::Error> {
+ // make sure only one client is writing to the database(only for sqlite)
+ let _lock = get_lock().await;
+
+ let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
+ Entity::delete_many()
+ .filter(Column::TraceId.eq(trace_id))
+ .exec(client)
+ .await?;
+
+ Ok(())
+}
+
+pub async fn count(work_group: &str, user_id: Option<&str>) -> Result {
+ let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
+ let mut query = Entity::find().filter(Column::WorkGroup.eq(work_group));
+ if let Some(user_id) = user_id {
+ query = query.filter(Column::UserId.eq(user_id));
+ }
+ let res = query.count(client).await?;
+ Ok(res as usize)
+}
+
+pub async fn clear() -> Result<(), errors::Error> {
+ // make sure only one client is writing to the database(only for sqlite)
+ let _lock = get_lock().await;
+
+ let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
+ Entity::delete_many().exec(client).await?;
+
+ Ok(())
+}
+
+pub async fn clean_incomplete(expired: i64) -> Result<(), errors::Error> {
+ // make sure only one client is writing to the database(only for sqlite)
+ let _lock = get_lock().await;
+
+ let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
+ Entity::delete_many()
+ .filter(Column::CreatedAt.lt(expired))
+ .exec(client)
+ .await?;
+
+ Ok(())
+}