AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-04-22 14:10:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

add seed data generation (#13)

Browse Source 
* merge upstream

* add seed data generation
This commit is contained in:
KevinHuSh 2023-12-22 20:02:16 +08:00 committed by GitHub
parent 1eb186a25f
commit 8a65ad888e
8 changed files with 117 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
migration/Cargo.toml
View File

@ -10,6 +10,7 @@ path = "src/lib.rs"
[dependencies]
async-std = { version = "1", features = ["attributes", "tokio1"] }
chrono = "0.4.31"
[dependencies.sea-orm-migration]
version = "0.12.0"

66
migration/src/m20220101_000001_create_table.rs
View File

@ -1,5 +1,9 @@
use sea_orm_migration::prelude::*;
use sea_orm_migration::{prelude::*, sea_orm::Statement};
use chrono::{FixedOffset, Utc};
fn now()->chrono::DateTime<FixedOffset>{
Utc::now().with_timezone(&FixedOffset::east_opt(3600*8).unwrap())
}
#[derive(DeriveMigrationName)]
pub struct Migration;
@ -25,9 +29,9 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(UserInfo::ListStyle).string().default("list"))
.col(ColumnDef::new(UserInfo::Language).string().default("chinese"))
.col(ColumnDef::new(UserInfo::Password).string().not_null())
.col(ColumnDef::new(UserInfo::LastLoginAt).timestamp_with_time_zone())
.col(ColumnDef::new(UserInfo::CreatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(UserInfo::UpdatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(UserInfo::LastLoginAt).timestamp_with_time_zone().default(Expr::current_timestamp()))
.col(ColumnDef::new(UserInfo::CreatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(UserInfo::UpdatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(UserInfo::IsDeleted).boolean().default(false))
.to_owned(),
)
@ -51,8 +55,8 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(TagInfo::Color).tiny_unsigned().default(1))
.col(ColumnDef::new(TagInfo::Icon).tiny_unsigned().default(1))
.col(ColumnDef::new(TagInfo::FolderId).big_integer())
.col(ColumnDef::new(TagInfo::CreatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(TagInfo::UpdatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(TagInfo::CreatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(TagInfo::UpdatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(TagInfo::IsDeleted).boolean().default(false))
.to_owned(),
)
@ -92,7 +96,7 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(Kb2Doc::Did).big_integer())
.col(ColumnDef::new(Kb2Doc::KbProgress).float().default(0))
.col(ColumnDef::new(Kb2Doc::KbProgressMsg).string().default(""))
.col(ColumnDef::new(Kb2Doc::UpdatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(Kb2Doc::UpdatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(Kb2Doc::IsDeleted).boolean().default(false))
.to_owned(),
)
@ -146,8 +150,8 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(KbInfo::Uid).big_integer().not_null())
.col(ColumnDef::new(KbInfo::KbName).string().not_null())
.col(ColumnDef::new(KbInfo::Icon).tiny_unsigned().default(1))
.col(ColumnDef::new(KbInfo::CreatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(KbInfo::UpdatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(KbInfo::CreatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(KbInfo::UpdatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(KbInfo::IsDeleted).boolean().default(false))
.to_owned(),
)
@ -167,8 +171,8 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(DocInfo::Location).string().not_null())
.col(ColumnDef::new(DocInfo::Size).big_integer().not_null())
.col(ColumnDef::new(DocInfo::Type).string().not_null()).comment("doc|folder")
.col(ColumnDef::new(DocInfo::CreatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(DocInfo::UpdatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(DocInfo::CreatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(DocInfo::UpdatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(DocInfo::IsDeleted).boolean().default(false))
.to_owned(),
)
@ -188,13 +192,49 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(DialogInfo::KbId).big_integer().not_null())
.col(ColumnDef::new(DialogInfo::DialogName).string().not_null())
.col(ColumnDef::new(DialogInfo::History).string().comment("json"))
.col(ColumnDef::new(DialogInfo::CreatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(DialogInfo::UpdatedAt).timestamp_with_time_zone().not_null())
.col(ColumnDef::new(DialogInfo::CreatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(DialogInfo::UpdatedAt).timestamp_with_time_zone().default(Expr::current_timestamp()).not_null())
.col(ColumnDef::new(DialogInfo::IsDeleted).boolean().default(false))
.to_owned(),
)
.await?;
let tm = now();
let root_insert = Query::insert()
.into_table(UserInfo::Table)
.columns([UserInfo::Email, UserInfo::Nickname, UserInfo::Password])
.values_panic([
"kai.hu@infiniflow.org".into(),
"root".into(),
"123456".into()
])
.to_owned();
let doc_insert = Query::insert()
.into_table(DocInfo::Table)
.columns([DocInfo::Uid, DocInfo::DocName, DocInfo::Size, DocInfo::Type,
DocInfo::Location])
.values_panic([
1.into(),
"/".into(),
0.into(),
"folder".into(),
"".into()
])
.to_owned();
let tag_insert = Query::insert()
.into_table(TagInfo::Table)
.columns([TagInfo::Uid, TagInfo::TagName, TagInfo::Regx, TagInfo::Color, TagInfo::Icon])
.values_panic([1.into(), "视频".into(),".*\\.(mpg|mpeg|avi|rm|rmvb|mov|wmv|asf|dat|asx|wvx|mpe|mpa)".into(),1.into(),1.into()])
.values_panic([1.into(), "图片".into(),".*\\.(png|tif|gif|pcx|tga|exif|fpx|svg|psd|cdr|pcd|dxf|ufo|eps|ai|raw|WMF|webp|avif|apng)".into(),2.into(),2.into()])
.values_panic([1.into(), "音乐".into(),".*\\.(WAV|FLAC|APE|ALAC|WavPack|WV|MP3|AAC|Ogg|Vorbis|Opus)".into(),3.into(),3.into()])
.values_panic([1.into(), "文档".into(),".*\\.(pdf|doc|ppt|yml|xml|htm|json|csv|txt|ini|xsl|wps|rtf|hlp)".into(),3.into(),3.into()])
.to_owned();
manager.exec_stmt(root_insert).await?;
manager.exec_stmt(doc_insert).await?;
manager.exec_stmt(tag_insert).await?;
Ok(())
}

55
src/api/user_info.rs
View File

@ -1,15 +1,23 @@
use std::collections::HashMap;
use std::io::SeekFrom;
use std::ptr::null;
use actix_identity::Identity;
use actix_web::{HttpResponse, post, web};
use chrono::{FixedOffset, Utc};
use sea_orm::ActiveValue::NotSet;
use serde::{Deserialize, Serialize};
use crate::api::JsonResponse;
use crate::AppState;
use crate::entity::{doc_info, tag_info};
use crate::entity::user_info::Model;
use crate::errors::{AppError, UserError};
use crate::service::user_info::Mutation;
use crate::service::user_info::Query;
fn now()->chrono::DateTime<FixedOffset>{
Utc::now().with_timezone(&FixedOffset::east_opt(3600*8).unwrap())
}
pub(crate) fn create_auth_token(user: &Model) -> u64 {
use std::{
collections::hash_map::DefaultHasher,
@ -58,8 +66,51 @@ async fn login(
#[post("/v1.0/register")]
async fn register(model: web::Json<Model>, data: web::Data<AppState>) -> Result<HttpResponse, AppError> {
let mut result = HashMap::new();
let u = Query::find_user_infos(&data.conn, &model.email).await?;
if let Some(_) = u {
let json_response = JsonResponse {
code: 500,
err: "Email registered!".to_owned(),
data: (),
};
return Ok(HttpResponse::Ok()
.content_type("application/json")
.body(serde_json::to_string(&json_response)?));
}
let usr = Mutation::create_user(&data.conn, &model).await?;
result.insert("uid", usr.uid.unwrap());
result.insert("uid", usr.uid.clone().unwrap());
crate::service::doc_info::Mutation::create_doc_info(&data.conn, doc_info::Model{
did:Default::default(),
uid: usr.uid.clone().unwrap(),
doc_name: "/".into(),
size: 0,
location: "".into(),
r#type: "folder".to_string(),
created_at: now(),
updated_at: now(),
is_deleted:Default::default(),
}).await?;
let tnm = vec!["视频","图片","音乐","文档"];
let tregx = vec![
".*\\.(mpg|mpeg|avi|rm|rmvb|mov|wmv|asf|dat|asx|wvx|mpe|mpa)",
".*\\.(png|tif|gif|pcx|tga|exif|fpx|svg|psd|cdr|pcd|dxf|ufo|eps|ai|raw|WMF|webp|avif|apng)",
".*\\.(WAV|FLAC|APE|ALAC|WavPack|WV|MP3|AAC|Ogg|Vorbis|Opus)",
".*\\.(pdf|doc|ppt|yml|xml|htm|json|csv|txt|ini|xsl|wps|rtf|hlp)"
];
for i in 0..4 {
crate::service::tag_info::Mutation::create_tag(&data.conn, tag_info::Model{
tid: Default::default(),
uid: usr.uid.clone().unwrap(),
tag_name: tnm[i].to_owned(),
regx: tregx[i].to_owned(),
color: (i+1).to_owned() as i16,
icon: (i+1).to_owned() as i16,
folder_id: 0,
created_at: Default::default(),
updated_at: Default::default(),
}).await?;
}
let json_response = JsonResponse {
code: 200,
err: "".to_owned(),

3
src/entity/kb2_doc.rs
View File

@ -11,7 +11,6 @@ pub struct Model {
pub kb_id: i64,
#[sea_orm(index)]
pub did: i64,
<<<<<<< HEAD
#[serde(skip_deserializing)]
pub kb_progress: f32,
#[serde(skip_deserializing)]
@ -20,8 +19,6 @@ pub struct Model {
pub updated_at: DateTime<FixedOffset>,
#[serde(skip_deserializing)]
pub is_deleted: bool,
=======
>>>>>>> upstream/main
}
#[derive(Debug, Clone, Copy, EnumIter)]

7
src/entity/tag_info.rs
View File

@ -11,19 +11,12 @@ pub struct Model {
#[sea_orm(index)]
pub uid: i64,
pub tag_name: String,
<<<<<<< HEAD
#[serde(skip_deserializing)]
pub regx: String,
pub color: i16,
pub icon: i16,
#[serde(skip_deserializing)]
pub folder_id: i64,
=======
pub regx: Option<String>,
pub color: u16,
pub icon: u16,
pub dir: Option<String>,
>>>>>>> upstream/main
#[serde(skip_deserializing)]
pub created_at: DateTime<FixedOffset>,

3
src/main.rs
View File

@ -98,11 +98,8 @@ fn init(cfg: &mut web::ServiceConfig) {
cfg.service(api::kb_info::delete);
cfg.service(api::kb_info::list);
cfg.service(api::kb_info::add_docs_to_kb);
<<<<<<< HEAD
cfg.service(api::kb_info::anti_kb_docs);
cfg.service(api::kb_info::all_relevents);
=======
>>>>>>> upstream/main
cfg.service(api::doc_info::list);
cfg.service(api::doc_info::delete);

7
src/service/tag_info.rs
View File

@ -1,6 +1,6 @@
use chrono::{FixedOffset, Utc};
use sea_orm::{ActiveModelTrait, DbConn, DbErr, DeleteResult, EntityTrait, PaginatorTrait, QueryOrder, ColumnTrait, QueryFilter};
use sea_orm::ActiveValue::Set;
use sea_orm::ActiveValue::{Set, NotSet};
use crate::entity::tag_info;
use crate::entity::tag_info::Entity;
@ -51,7 +51,10 @@ impl Mutation {
regx: Set(form_data.regx.to_owned()),
color: Set(form_data.color.to_owned()),
icon: Set(form_data.icon.to_owned()),
folder_id: Set(form_data.folder_id.to_owned()),
folder_id: match form_data.folder_id {
0 => NotSet,
_ => Set(form_data.folder_id.to_owned())
},
created_at: Set(now()),
updated_at: Set(now()),
}

7
src/service/user_info.rs
View File

@ -23,8 +23,11 @@ impl Query {
.await
}
pub async fn find_user_infos(db: &DbConn) -> Result<Vec<user_info::Model>, DbErr> {
Entity::find().all(db).await
pub async fn find_user_infos(db: &DbConn, email:&String) -> Result<Option<user_info::Model>, DbErr> {
Entity::find()
.filter(user_info::Column::Email.eq(email))
.one(db)
.await
}
pub async fn find_user_infos_in_page(