From da350813aa80c30600eb1606f925a8d56754e0b6 Mon Sep 17 00:00:00 2001 From: hongwangliu-nju <1649209217@qq.com> Date: Thu, 11 Sep 2025 19:08:18 +0800 Subject: [PATCH] Add the cratespro crate and import the tugraph module Signed-off-by: hongwangliu-nju <1649209217@qq.com> --- extensions/cratespro/Cargo.toml | 26 + .../cratespro/common/data_transporter/BUCK | 38 + .../common/data_transporter/Cargo.toml | 33 + .../data_transporter/src/data_packer.rs | 28 + .../data_transporter/src/data_reader.rs | 1151 ++++++++++++ .../common/data_transporter/src/db.rs | 1537 +++++++++++++++++ .../common/data_transporter/src/handler.rs | 965 +++++++++++ .../common/data_transporter/src/lib.rs | 303 ++++ .../data_transporter/src/redis_store.rs | 155 ++ .../data_transporter/src/transporter.rs | 36 + extensions/cratespro/common/database/.env | 1 + extensions/cratespro/common/database/BUCK | 19 + .../cratespro/common/database/Cargo.toml | 21 + .../cratespro/common/database/entity/BUCK | 38 + .../common/database/entity/Cargo.toml | 13 + .../entity/src/contributor_location.rs | 21 + .../database/entity/src/crate_owners.rs | 54 + .../common/database/entity/src/crate_users.rs | 35 + .../common/database/entity/src/crates.rs | 35 + .../common/database/entity/src/crates_info.rs | 34 + .../common/database/entity/src/cves.rs | 21 + .../database/entity/src/dependency_cache.rs | 20 + .../database/entity/src/dependent_cache.rs | 20 + .../database/entity/src/github_sync_status.rs | 19 + .../common/database/entity/src/github_user.rs | 35 + .../common/database/entity/src/graph_info.rs | 18 + .../common/database/entity/src/license.rs | 22 + .../common/database/entity/src/mod.rs | 27 + .../common/database/entity/src/prelude.rs | 24 + .../entity/src/program_dependencies.rs | 20 + .../database/entity/src/program_versions.rs | 23 + .../common/database/entity/src/programs.rs | 34 + .../database/entity/src/repo_sync_result.rs | 32 + .../entity/src/repository_contributor.rs | 21 + .../database/entity/src/rustsec_info.rs | 38 + .../common/database/entity/src/rustsecs.rs | 28 + .../entity/src/sea_orm_active_enums.rs | 27 + .../database/entity/src/senseleak_res.rs | 18 + .../common/database/entity/src/userloginfo.rs | 22 + .../database/entity/src/version_info.rs | 18 + .../cratespro/common/database/src/lib.rs | 1 + extensions/cratespro/common/model/BUCK | 21 + extensions/cratespro/common/model/Cargo.toml | 16 + .../common/model/src/general_model.rs | 50 + .../cratespro/common/model/src/github.rs | 192 ++ extensions/cratespro/common/model/src/lib.rs | 4 + .../common/model/src/repo_sync_model.rs | 58 + .../common/model/src/tugraph_model.rs | 233 +++ extensions/cratespro/common/repo_import/BUCK | 42 + .../cratespro/common/repo_import/Cargo.toml | 44 + .../common/repo_import/src/crate_info.rs | 215 +++ .../cratespro/common/repo_import/src/git.rs | 181 ++ .../common/repo_import/src/kafka_handler.rs | 176 ++ .../cratespro/common/repo_import/src/lib.rs | 1165 +++++++++++++ .../cratespro/common/repo_import/src/utils.rs | 197 +++ .../common/repo_import/src/version_info.rs | 461 +++++ extensions/cratespro/common/search/BUCK | 26 + extensions/cratespro/common/search/Cargo.toml | 13 + extensions/cratespro/common/search/src/ai.rs | 144 ++ .../common/search/src/crates_search.rs | 156 ++ .../cratespro/common/search/src/embedding.rs | 134 ++ extensions/cratespro/common/search/src/lib.rs | 4 + .../common/search/src/search_prepare.rs | 201 +++ extensions/cratespro/common/tudriver/BUCK | 17 + .../cratespro/common/tudriver/Cargo.toml | 12 + .../cratespro/common/tudriver/src/lib.rs | 81 + .../common/tudriver/src/tugraph_client.rs | 380 ++++ extensions/cratespro/import_tugraph/.env | 44 + extensions/cratespro/import_tugraph/BUCK | 30 + .../cratespro/import_tugraph/Cargo.toml | 29 + .../cratespro/import_tugraph/Dockerfile | 33 + .../import_tugraph/deploy-import-tugraph.sh | 57 + .../cratespro/import_tugraph/deployment.yaml | 54 + .../cratespro/import_tugraph/import.config | 201 +++ .../cratespro/import_tugraph/src/cli.rs | 19 + .../import_tugraph/src/core_controller.rs | 115 ++ .../cratespro/import_tugraph/src/main.rs | 41 + 77 files changed, 9877 insertions(+) create mode 100644 extensions/cratespro/Cargo.toml create mode 100644 extensions/cratespro/common/data_transporter/BUCK create mode 100644 extensions/cratespro/common/data_transporter/Cargo.toml create mode 100644 extensions/cratespro/common/data_transporter/src/data_packer.rs create mode 100644 extensions/cratespro/common/data_transporter/src/data_reader.rs create mode 100644 extensions/cratespro/common/data_transporter/src/db.rs create mode 100644 extensions/cratespro/common/data_transporter/src/handler.rs create mode 100644 extensions/cratespro/common/data_transporter/src/lib.rs create mode 100644 extensions/cratespro/common/data_transporter/src/redis_store.rs create mode 100644 extensions/cratespro/common/data_transporter/src/transporter.rs create mode 100644 extensions/cratespro/common/database/.env create mode 100644 extensions/cratespro/common/database/BUCK create mode 100644 extensions/cratespro/common/database/Cargo.toml create mode 100644 extensions/cratespro/common/database/entity/BUCK create mode 100644 extensions/cratespro/common/database/entity/Cargo.toml create mode 100644 extensions/cratespro/common/database/entity/src/contributor_location.rs create mode 100644 extensions/cratespro/common/database/entity/src/crate_owners.rs create mode 100644 extensions/cratespro/common/database/entity/src/crate_users.rs create mode 100644 extensions/cratespro/common/database/entity/src/crates.rs create mode 100644 extensions/cratespro/common/database/entity/src/crates_info.rs create mode 100644 extensions/cratespro/common/database/entity/src/cves.rs create mode 100644 extensions/cratespro/common/database/entity/src/dependency_cache.rs create mode 100644 extensions/cratespro/common/database/entity/src/dependent_cache.rs create mode 100644 extensions/cratespro/common/database/entity/src/github_sync_status.rs create mode 100644 extensions/cratespro/common/database/entity/src/github_user.rs create mode 100644 extensions/cratespro/common/database/entity/src/graph_info.rs create mode 100644 extensions/cratespro/common/database/entity/src/license.rs create mode 100644 extensions/cratespro/common/database/entity/src/mod.rs create mode 100644 extensions/cratespro/common/database/entity/src/prelude.rs create mode 100644 extensions/cratespro/common/database/entity/src/program_dependencies.rs create mode 100644 extensions/cratespro/common/database/entity/src/program_versions.rs create mode 100644 extensions/cratespro/common/database/entity/src/programs.rs create mode 100644 extensions/cratespro/common/database/entity/src/repo_sync_result.rs create mode 100644 extensions/cratespro/common/database/entity/src/repository_contributor.rs create mode 100644 extensions/cratespro/common/database/entity/src/rustsec_info.rs create mode 100644 extensions/cratespro/common/database/entity/src/rustsecs.rs create mode 100644 extensions/cratespro/common/database/entity/src/sea_orm_active_enums.rs create mode 100644 extensions/cratespro/common/database/entity/src/senseleak_res.rs create mode 100644 extensions/cratespro/common/database/entity/src/userloginfo.rs create mode 100644 extensions/cratespro/common/database/entity/src/version_info.rs create mode 100644 extensions/cratespro/common/database/src/lib.rs create mode 100644 extensions/cratespro/common/model/BUCK create mode 100644 extensions/cratespro/common/model/Cargo.toml create mode 100644 extensions/cratespro/common/model/src/general_model.rs create mode 100644 extensions/cratespro/common/model/src/github.rs create mode 100644 extensions/cratespro/common/model/src/lib.rs create mode 100644 extensions/cratespro/common/model/src/repo_sync_model.rs create mode 100644 extensions/cratespro/common/model/src/tugraph_model.rs create mode 100644 extensions/cratespro/common/repo_import/BUCK create mode 100644 extensions/cratespro/common/repo_import/Cargo.toml create mode 100644 extensions/cratespro/common/repo_import/src/crate_info.rs create mode 100644 extensions/cratespro/common/repo_import/src/git.rs create mode 100644 extensions/cratespro/common/repo_import/src/kafka_handler.rs create mode 100644 extensions/cratespro/common/repo_import/src/lib.rs create mode 100644 extensions/cratespro/common/repo_import/src/utils.rs create mode 100644 extensions/cratespro/common/repo_import/src/version_info.rs create mode 100644 extensions/cratespro/common/search/BUCK create mode 100644 extensions/cratespro/common/search/Cargo.toml create mode 100644 extensions/cratespro/common/search/src/ai.rs create mode 100644 extensions/cratespro/common/search/src/crates_search.rs create mode 100644 extensions/cratespro/common/search/src/embedding.rs create mode 100644 extensions/cratespro/common/search/src/lib.rs create mode 100644 extensions/cratespro/common/search/src/search_prepare.rs create mode 100644 extensions/cratespro/common/tudriver/BUCK create mode 100644 extensions/cratespro/common/tudriver/Cargo.toml create mode 100644 extensions/cratespro/common/tudriver/src/lib.rs create mode 100644 extensions/cratespro/common/tudriver/src/tugraph_client.rs create mode 100644 extensions/cratespro/import_tugraph/.env create mode 100644 extensions/cratespro/import_tugraph/BUCK create mode 100644 extensions/cratespro/import_tugraph/Cargo.toml create mode 100644 extensions/cratespro/import_tugraph/Dockerfile create mode 100755 extensions/cratespro/import_tugraph/deploy-import-tugraph.sh create mode 100644 extensions/cratespro/import_tugraph/deployment.yaml create mode 100644 extensions/cratespro/import_tugraph/import.config create mode 100644 extensions/cratespro/import_tugraph/src/cli.rs create mode 100644 extensions/cratespro/import_tugraph/src/core_controller.rs create mode 100644 extensions/cratespro/import_tugraph/src/main.rs diff --git a/extensions/cratespro/Cargo.toml b/extensions/cratespro/Cargo.toml new file mode 100644 index 000000000..f28389da2 --- /dev/null +++ b/extensions/cratespro/Cargo.toml @@ -0,0 +1,26 @@ +[workspace] +members = [ + "common/data_transporter", + "common/database", + "common/database/entity", + "common/model", # internal info struct,such as `crate info` `library info` `app info` + "common/repo_import", # extract repo and import data into tugraph + "common/search", + "common/tudriver", # tugraph client, deal with the quary from router + "import_tugraph", +] + +resolver = "2" # use resolver of version 2 + +[workspace.dependencies] +# internal (项目内部依赖) +data_transporter = { path = "common/data_transporter" } +model = { path = "common/model" } +repo_import = { path = "common/repo_import" } +search = { path = "common/search" } +tudriver = { path = "common/tudriver" } +database = { path = "common/database" } +entity = { path = "common/database/entity" } +import_tugraph = {path = "import_tugraph"} + +# third-party (所有第三方依赖, 按字母表排序, 均保留1个小数点) \ No newline at end of file diff --git a/extensions/cratespro/common/data_transporter/BUCK b/extensions/cratespro/common/data_transporter/BUCK new file mode 100644 index 000000000..598bd989b --- /dev/null +++ b/extensions/cratespro/common/data_transporter/BUCK @@ -0,0 +1,38 @@ +rust_library( + name = "data_transporter", + srcs = [ + "src/data_packer.rs", + "src/data_reader.rs", + "src/db.rs", + "src/handler.rs", + "src/lib.rs", + "src/transporter.rs", + "src/redis_store.rs", + ], + crate_root = "src/lib.rs", + edition = "2021", + deps = [ + "//project/crates-pro:model", + "//project/crates-pro:repo_import", + "//project/crates-pro:search", + "//project/crates-pro:tudriver", + "//third-party:actix-multipart", + "//third-party:actix-web", + "//third-party:async-trait", + "//third-party:chrono", + "//third-party:futures-util", + "//third-party:redis", + "//third-party:sanitize-filename", + "//third-party:semver", + "//third-party:serde", + "//third-party:serde_json", + "//third-party:tokio", + "//third-party:tokio-postgres", + "//third-party:tracing", + "//third-party:utoipa", + "//third-party:utoipa-swagger-ui", + "//third-party:uuid", + "//third-party:zip", + ], + visibility = ["PUBLIC"], +) diff --git a/extensions/cratespro/common/data_transporter/Cargo.toml b/extensions/cratespro/common/data_transporter/Cargo.toml new file mode 100644 index 000000000..5ddf04662 --- /dev/null +++ b/extensions/cratespro/common/data_transporter/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "data_transporter" +version = "0.1.0" +edition = "2021" + +[dependencies] +model = { workspace = true } +repo_import = { workspace = true } +search = { workspace = true } +tudriver = { workspace = true } + +# third-party (第三方依赖, 不写具体版本号, 具体版本只在根目录 Cargo.toml 中出现) +actix-multipart = "0.7" +actix-web = "4.10" +async-trait = "0.1" +chrono = { version = "0.4", features = ["clock"] } +futures-util = "0.3" +redis = "0.32" +sanitize-filename = "0.6" +semver = "1.0" +serde = "1.0" +serde_json = "1.0" +tokio = { version = "1.44", features = ["full"] } +tokio-postgres = { version = "0.7", features = ["with-chrono-0_4"] } +tracing = "0.1" +utoipa = { version = "5.3", features = ["actix_extras"] } +utoipa-swagger-ui = { version = "9.0", features = ["actix-web"] } +uuid = { version = "1.16", features = [ + "v4", # Lets you generate random UUIDs + "fast-rng", # Use a faster (but still sufficiently random) RNG + "macro-diagnostics", # Enable better diagnostics for compile-time UUIDs +] } +zip = "=2.4.2" diff --git a/extensions/cratespro/common/data_transporter/src/data_packer.rs b/extensions/cratespro/common/data_transporter/src/data_packer.rs new file mode 100644 index 000000000..f60e617db --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/data_packer.rs @@ -0,0 +1,28 @@ +use model::tugraph_model::{Program, UProgram}; + +use crate::db::DBHandler; + +pub struct DataPacker { + db: DBHandler, +} + +impl DataPacker { + pub async fn new() -> Self { + let db = DBHandler::connect().await.unwrap(); + db.clear_database().await.unwrap(); + db.create_tables().await.unwrap(); + Self { db } + } + + pub async fn pack_into_db( + &self, + program: Program, + uprogram: UProgram, + versions: Vec, + ) { + self.db + .insert_program_data(program, uprogram, versions) + .await + .unwrap(); + } +} diff --git a/extensions/cratespro/common/data_transporter/src/data_reader.rs b/extensions/cratespro/common/data_transporter/src/data_reader.rs new file mode 100644 index 000000000..b62485e82 --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/data_reader.rs @@ -0,0 +1,1151 @@ +use model::tugraph_model::{ + Application, ApplicationVersion, Library, LibraryVersion, Program, UProgram, UVersion, +}; +use semver::Version; +use serde_json::Value; +use std::{ + cmp::Ordering, + collections::{HashSet, VecDeque}, + error::Error, + time::Instant, +}; +use tokio_postgres::NoTls; +use tudriver::tugraph_client::TuGraphClient; + +use crate::{ + db::{db_connection_config_from_env, db_cratesio_connection_config_from_env, DBHandler}, + handler::{ + Crateinfo, DependencyCount, DependencyCrateInfo, DependencyInfo, DependentCount, + DependentData, DependentInfo, Deptree, Versionpage, + }, + NameVersion, +}; + +pub trait DataReaderTrait { + async fn get_all_programs_id(&self) -> Vec; + async fn get_program(&self, program_id: &str) -> Result>; + async fn get_type(&self, program_id: &str) -> Result<(UProgram, bool), Box>; + async fn get_versions( + &self, + program_id: &str, + is_lib: bool, + ) -> Result, Box>; + async fn get_direct_dependency_nodes( + &self, + name_and_version: &str, + ) -> Result, Box>; + async fn new_get_direct_dependency_nodes( + &self, + namespace: &str, + nameversion: &str, + ) -> Result, Box>; + #[allow(dead_code)] + async fn get_program_by_name(&self, program_name: &str) + -> Result, Box>; + #[allow(dead_code)] + async fn get_indirect_dependency_nodes( + &self, + nameversion: NameVersion, + ) -> Result, Box>; + #[allow(dead_code)] + async fn count_dependencies(&self, nameversion: NameVersion) -> Result>; + #[allow(dead_code)] + async fn get_direct_dependent_nodes( + &self, + name_and_version: &str, + ) -> Result, Box>; + async fn new_get_direct_dependent_nodes( + &self, + namespace: &str, + nameversion: &str, + ) -> Result, Box>; + #[allow(dead_code)] + async fn get_indirect_dependent_nodes( + &self, + nameversion: NameVersion, + ) -> Result, Box>; + + //async fn get_max_version(&self, name: String) -> Result>; + #[allow(dead_code)] + async fn get_lib_version(&self, name: String) -> Result, Box>; + async fn new_get_lib_version( + &self, + namespace: String, + name: String, + ) -> Result, Box>; + #[allow(dead_code)] + async fn get_app_version(&self, name: String) -> Result, Box>; + #[allow(dead_code)] + async fn new_get_app_version( + &self, + namespace: String, + name: String, + ) -> Result, Box>; + + #[allow(dead_code)] + async fn get_all_dependencies( + &self, + nameversion: NameVersion, + ) -> Result, Box>; + async fn new_get_all_dependencies( + &self, + nodes: Vec, + ) -> Result, Box>; + + #[allow(dead_code)] + async fn new_get_all_dependents( + &self, + namespace: String, + nameversion: String, + ) -> Result, Box>; + async fn get_github_url( + &self, + namespace: String, + name: String, + ) -> Result>; + async fn get_doc_url(&self, namespace: String, name: String) -> Result>; + async fn build_graph( + &self, + rootnode: &mut Deptree, + visited: &mut HashSet, + ) -> Result<(), Box>; + async fn get_version_page_from_tg( + &self, + nsfront: String, + nsbehind: String, + nname: String, + ) -> Result, Box>; + async fn get_crates_front_info_from_tg( + &self, + nname: String, + nversion: String, + nsfront: String, + nsbehind: String, + ) -> Result>; + async fn get_dependency_from_tg( + &self, + name: String, + version: String, + nsfront: String, + nsbehind: String, + ) -> Result>; + async fn get_dependent_from_tg( + &self, + name: String, + version: String, + nsfront: String, + nsbehind: String, + ) -> Result>; +} + +#[derive(Clone)] +pub struct DataReader { + pub client: TuGraphClient, +} +impl DataReader { + pub async fn new( + uri: &str, + user: &str, + password: &str, + db: &str, + ) -> Result> { + let client = TuGraphClient::new(uri, user, password, db).await?; + Ok(DataReader { client }) + } +} + +impl DataReaderTrait for DataReader { + async fn get_dependent_from_tg( + &self, + name: String, + version: String, + nsfront: String, + nsbehind: String, + ) -> Result> { + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let nameversion = name.clone() + "/" + &version.clone(); + let mut direct_nodes = vec![]; + if version.clone() == *"all" { + let lib_versions = self + .new_get_lib_version(namespace.clone(), name.clone()) + .await + .unwrap(); + let mut getversions = vec![]; + for version in lib_versions { + getversions.push(version); + } + getversions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), // 从高到低排序 + (Ok(_), Err(_)) => Ordering::Less, // 无法解析的版本号认为更小 + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + let mut visited = HashSet::new(); + for nversion in &getversions { + let tmp_name_and_version = name.clone() + "/" + nversion; + let tmp_direct_nodes = self + .new_get_direct_dependent_nodes(&namespace, &tmp_name_and_version) + .await + .unwrap(); + for node in tmp_direct_nodes { + if visited.insert(node.clone()) { + direct_nodes.push(node.clone()); + } + } + } + } else { + direct_nodes = self + .new_get_direct_dependent_nodes(&namespace, &nameversion) + .await + .unwrap(); + } + let getdirect_count = direct_nodes.len(); + + let mut deps = vec![]; + let mut count1 = 0; + for item in direct_nodes { + let dep = DependentData { + crate_name: item.clone().name, + version: item.clone().version, + relation: "Direct".to_string(), + }; + deps.push(dep); + count1 += 1; + if count1 == 50 { + break; + } + } + + + let res_deps = DependentInfo { + direct_count: getdirect_count, + indirect_count: 0, + data: deps, + }; + Ok(res_deps) + } + async fn get_dependency_from_tg( + &self, + name: String, + version: String, + nsfront: String, + nsbehind: String, + ) -> Result> { + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let nameversion = name.clone() + "/" + &version.clone(); + tracing::info!("{} {}", namespace.clone(), nameversion.clone()); + let mut direct_nodes = vec![]; + if version.clone() == *"all" { + let lib_versions = self + .new_get_lib_version(namespace.clone(), name.clone()) + .await + .unwrap(); + let mut getversions = vec![]; + for version in lib_versions { + getversions.push(version); + } + getversions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), // 从高到低排序 + (Ok(_), Err(_)) => Ordering::Less, // 无法解析的版本号认为更小 + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + let mut visited = HashSet::new(); + for nversion in &getversions { + let tmp_name_and_version = name.clone() + "/" + nversion; + let tmp_direct_nodes = self + .new_get_direct_dependency_nodes(&namespace, &tmp_name_and_version) + .await + .unwrap(); + for node in tmp_direct_nodes { + if visited.insert(node.clone()) { + direct_nodes.push(node.clone()); + } + } + } + } else { + direct_nodes = self + .new_get_direct_dependency_nodes(&namespace, &nameversion) + .await + .unwrap(); + } + let getdirect_count = direct_nodes.len(); + let all_dependency_nodes = self + .new_get_all_dependencies(direct_nodes.clone()) + .await + .unwrap(); + let mut indirect_dependency = vec![]; + for node in all_dependency_nodes { + let mut dr = false; + for node2 in direct_nodes.clone() { + let nv = node2.name.clone() + "/" + &node2.version.clone(); + if node == nv { + dr = true; + break; + } + } + if !dr { + indirect_dependency.push(node); + } + } + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + let indirect_dependency_count = indirect_dependency.len(); + let mut deps = vec![]; + for item in direct_nodes { + let dep_count = dbhandler + .get_all_dependency_counts_from_pg(item.clone().name, item.clone().version) + .await + .unwrap(); + let dep = DependencyCrateInfo { + crate_name: item.clone().name, + version: item.clone().version, + relation: "Direct".to_string(), + license: "".to_string(), + dependencies: dep_count, + }; + deps.push(dep); + } + for item in indirect_dependency { + let parts: Vec<&str> = item.split('/').collect(); + let newitem = NameVersion { + name: parts[0].to_string(), + version: parts[1].to_string(), + }; + let dep_count = dbhandler + .get_all_dependency_counts_from_pg(newitem.clone().name, newitem.clone().version) + .await + .unwrap(); + + let dep = DependencyCrateInfo { + crate_name: parts[0].to_string(), + version: parts[1].to_string(), + relation: "Indirect".to_string(), + license: "".to_string(), + dependencies: dep_count, + }; + deps.push(dep); + } + + let res_deps = DependencyInfo { + direct_count: getdirect_count, + indirect_count: indirect_dependency_count, + data: deps, + }; + Ok(res_deps) + } + #[allow(unused_assignments)] + async fn get_crates_front_info_from_tg( + &self, + nname: String, + nversion: String, + nsfront: String, + nsbehind: String, + ) -> Result> { + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let name_and_version = nname.clone() + "/" + &nversion.clone(); + let mut githuburl = self + .get_github_url(namespace.clone(), nname.clone()) + .await + .unwrap(); + if githuburl == *"null" || githuburl == *"None" { + githuburl = "".to_string(); + } + let mut docurl = self + .get_doc_url(namespace.clone(), nname.clone()) + .await + .unwrap(); + if docurl == *"null" || docurl == *"None" { + docurl = "".to_string(); + } + let lib_versions = self + .new_get_lib_version(namespace.clone(), nname.clone()) + .await + .unwrap(); + let mut getversions = vec![]; + for version in lib_versions { + getversions.push(version); + } + getversions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), // 从高到低排序 + (Ok(_), Err(_)) => Ordering::Less, // 无法解析的版本号认为更小 + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + let mut direct_dependency_nodes = vec![]; + let mut direct_dependent_nodes = vec![]; + let mut getcves = vec![]; + let mut get_dependency_cves = vec![]; + let mut direct_dependency_count = 0; + let mut indirect_dependency_count = 0; + let mut direct_dependent_count = 0; + if nversion.clone() == *"all" { + let (direct_depends, indirect_depends, direct_depended_by, depends_cves_id) = dbhandler + .get_depends_of_all_version(nname.clone()) + .await + .unwrap(); + direct_dependency_count = direct_depends; + indirect_dependency_count = indirect_depends; + direct_dependent_count = direct_depended_by; + let mut visited = HashSet::new(); + for version in &getversions { + let tmp_direct_nodes = dbhandler.get_direct_rustsec(&nname, version).await.unwrap(); + for node in tmp_direct_nodes { + if visited.insert(node.clone()) { + getcves.push(node.clone()); + } + } + } + get_dependency_cves = dbhandler + .get_rustsec_by_depends_cveid(depends_cves_id.clone()) + .await + .unwrap(); + } else { + direct_dependency_nodes = self + .new_get_direct_dependency_nodes(&namespace, &name_and_version) + .await + .unwrap(); + direct_dependency_count = direct_dependency_nodes.len(); + let all_dependency_nodes = self + .new_get_all_dependencies(direct_dependency_nodes.clone()) + .await + .unwrap(); + let mut indirect_dependency = vec![]; + for node in all_dependency_nodes.clone() { + let mut dr = false; + for node2 in direct_dependency_nodes.clone() { + let nv = node2.name.clone() + "/" + &node2.version.clone(); + if node == nv { + dr = true; + break; + } + } + if !dr { + indirect_dependency.push(node); + } + } + indirect_dependency_count = indirect_dependency.len(); + direct_dependent_nodes = self + .new_get_direct_dependent_nodes(&namespace, &name_and_version) + .await + .unwrap(); + direct_dependent_count = direct_dependent_nodes.len(); + getcves = dbhandler + .get_direct_rustsec(&nname, &nversion) + .await + .unwrap(); + get_dependency_cves = dbhandler + .get_dependency_rustsec(all_dependency_nodes.clone()) + .await + .unwrap(); + } + + let getlicense = dbhandler + .get_license_by_name(&namespace, &nname) + .await + .unwrap(); + let dcy_count = DependencyCount { + direct: direct_dependency_count, + indirect: indirect_dependency_count, + }; + let dt_count = DependentCount { + direct: direct_dependent_count, + indirect: 0, + }; + let res = Crateinfo { + crate_name: nname.clone(), + description: "".to_string(), + dependencies: dcy_count, + dependents: dt_count, + cves: getcves, + versions: getversions, + license: getlicense[0].clone(), + github_url: githuburl, + doc_url: docurl, + dep_cves: get_dependency_cves, + }; + Ok(res) + } + #[allow(unused_assignments)] + async fn get_version_page_from_tg( + &self, + nsfront: String, + nsbehind: String, + nname: String, + ) -> Result, Box> { + let namespace = nsfront.clone() + "/" + &nsbehind; + let time_get_version = Instant::now(); + let all_versions = self + .new_get_lib_version(namespace.clone(), nname.clone()) + .await + .unwrap(); + tracing::info!("finish get all versions"); + let mut getversions = vec![]; + for version in all_versions { + getversions.push(version); + } + getversions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + tracing::info!("get_version_time:{:?}", time_get_version.elapsed()); + let mut every_version = vec![]; + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + tracing::info!("finish connect pg"); + let dbhandler = DBHandler { client }; + let _db_cratesio_connection_config = db_cratesio_connection_config_from_env(); + #[allow(unused_variables)] + let (client2, connection2) = + tokio_postgres::connect(&_db_cratesio_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection2.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler2 = DBHandler { client: client2 }; + for version in getversions { + + let all_dts = dbhandler + .get_all_dependent_counts_from_pg(nname.clone(), version.clone()) + .await + .unwrap(); + let res = dbhandler2 + .get_dump_from_cratesio_pg(nname.clone(), version.clone()) + .await + .unwrap(); + tracing::info!("finish get dump from pg"); + if !res.is_empty() { + let parts: Vec<&str> = res.split("/").collect(); + if parts.len() == 2 { + let versionpage = Versionpage { + version, + updated_at: parts[0].to_string(), + downloads: parts[1].to_string(), + dependents: all_dts, + }; + every_version.push(versionpage); + } + } + } + Ok(every_version) + } + async fn build_graph( + &self, + rootnode: &mut Deptree, + visited: &mut HashSet, + ) -> Result<(), Box> { + let name_and_version = &rootnode.name_and_version; + let res = self + .get_direct_dependency_nodes(name_and_version) + .await + .unwrap(); + tracing::info!("direct dep count:{}", res.len()); + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + for node in res { + let name = node.name.clone(); + let version = node.version.clone(); + let dep_nv = name.clone() + "/" + &version; + let rustcve = dbhandler.get_direct_rustsec(&name, &version).await.unwrap(); + let mut dn = Deptree { + name_and_version: dep_nv.clone(), + cve_count: rustcve.len(), + direct_dependency: Vec::new(), + }; + if visited.insert(dep_nv.clone()) { + Box::pin(self.build_graph(&mut dn, visited)).await?; + rootnode.direct_dependency.push(dn); + } + } + Ok(()) + } + async fn get_github_url( + &self, + namespace: String, + name: String, + ) -> Result> { + tracing::info!("{}|{}", namespace.clone(), name.clone()); + let query = format!( + " + MATCH (n:program {{namespace:'{namespace}'}}) WHERE n.name='{name}' +RETURN n.github_url + " + ); + let results = self.client.exec_query(&query).await?; + tracing::info!("finish get github_url"); + let mut res = vec![]; + for node in results { + res.push(node); + } + let unique_items: HashSet = res.clone().into_iter().collect(); + let mut nodes = vec![]; + for res in unique_items { + let parsed: Value = serde_json::from_str(&res).unwrap(); + if let Some(url) = parsed.get("n.github_url").and_then(|v| v.as_str()) { + nodes.push(url.to_string()); + } + } + nodes.push("None".to_string()); + Ok(nodes[0].clone()) + } + async fn get_doc_url(&self, namespace: String, name: String) -> Result> { + let query = format!( + " + MATCH (n:program {{namespace:'{namespace}'}}) WHERE n.name='{name}' +RETURN n.doc_url + " + ); + let results = self.client.exec_query(&query).await?; + let mut res = vec![]; + for node in results { + res.push(node); + } + let unique_items: HashSet = res.clone().into_iter().collect(); + let mut nodes = vec![]; + for res in unique_items { + let parsed: Value = serde_json::from_str(&res).unwrap(); + if let Some(url) = parsed.get("n.doc_url").and_then(|v| v.as_str()) { + nodes.push(url.to_string()); + } + } + nodes.push("None".to_string()); + Ok(nodes[0].clone()) + } + async fn get_all_dependencies( + &self, + nameversion: NameVersion, + ) -> Result, Box> { + let mut queue = VecDeque::new(); + let mut visited = HashSet::new(); + let name_and_version = nameversion.name.clone() + "/" + &nameversion.version.clone(); + queue.push_back(name_and_version.to_string()); + let mut count = 0; + while let Some(current) = queue.pop_front() { + count += 1; + if count == 2000 { + break; + } + if visited.insert(current.clone()) { + for dep in self.get_direct_dependency_nodes(¤t).await.unwrap() { + let tmp = dep.name.clone() + "/" + &dep.version.clone(); + queue.push_back(tmp); + } + } + } + visited.remove(&name_and_version); + + Ok(visited) + } + async fn new_get_all_dependencies( + &self, + nodes: Vec, + ) -> Result, Box> { + let mut queue = VecDeque::new(); + let mut visited = HashSet::new(); + for node in nodes { + let nameandversion = node.clone().name + "/" + &node.clone().version; + queue.push_back(nameandversion.clone()); + } + let mut count = 0; + while let Some(current) = queue.pop_front() { + count += 1; + if count == 2000 { + break; + } + if visited.insert(current.clone()) { + for dep in self.get_direct_dependency_nodes(¤t).await.unwrap() { + let tmp = dep.name.clone() + "/" + &dep.version.clone(); + queue.push_back(tmp); + } + } + } + Ok(visited) + } + + async fn new_get_all_dependents( + &self, + namespace: String, + nameversion: String, + ) -> Result, Box> { + let mut queue = VecDeque::new(); + let mut visited = HashSet::new(); + for node in self + .new_get_direct_dependent_nodes(&namespace, &nameversion) + .await + .unwrap() + { + let nameandversion = node.clone().name + "/" + &node.clone().version; + queue.push_back(nameandversion.clone()); + } + let len = queue.len(); + if len < 500 { + while let Some(current) = queue.pop_front() { + if visited.insert(current.clone()) { + for dep in self.get_direct_dependent_nodes(¤t).await.unwrap() { + let tmp = dep.name.clone() + "/" + &dep.version.clone(); + queue.push_back(tmp); + } + } + } + } + Ok(visited) + } + async fn get_all_programs_id(&self) -> Vec { + let query = " + MATCH (p: program) + RETURN p + "; + + let results = self.client.exec_query(query).await.unwrap(); + + let mut programs = vec![]; + for result in results { + let programs_json: Value = serde_json::from_str(&result).unwrap(); + + let pro = programs_json["p"].clone(); + + let program: Program = serde_json::from_value(pro).unwrap(); + + programs.push(program.id); + } + + programs + } + + async fn get_program(&self, program_id: &str) -> Result> { + let query = format!( + " + MATCH (p: program {{id: '{program_id}'}}) + RETURN p + " + ); + let results = self.client.exec_query(&query).await?; + let programs_json: Value = serde_json::from_str(&results[0]).unwrap(); + let pro = programs_json["p"].clone(); + let program: Program = serde_json::from_value(pro).unwrap(); + Ok(program) + } + + async fn get_type(&self, program_id: &str) -> Result<(UProgram, bool), Box> { + let mut islib = false; + + let query = format!( + " + MATCH (p: program {{id: '{program_id}'}})-[:has_type]->(o) + RETURN o, label(o) as o_label + " + ); + + let results = self.client.exec_query(&query).await?; + let mut uprograms = vec![]; + for result in results { + let result_json: Value = serde_json::from_str(&result).unwrap(); + + let label: String = serde_json::from_value(result_json["o_label"].clone()).unwrap(); + + let o = result_json["o"].clone(); + if label.eq(&"library".to_string()) { + islib = true; + let library: Library = serde_json::from_value(o).unwrap(); + uprograms.push(UProgram::Library(library)); + } else if label.eq(&"application".to_string()) { + let application: Application = serde_json::from_value(o).unwrap(); + uprograms.push(UProgram::Application(application)); + } + } + Ok((uprograms[0].clone(), islib)) + } + + async fn get_versions( + &self, + program_id: &str, + is_lib: bool, + ) -> Result, Box> { + let query = if is_lib { + format!( + " + MATCH (l: library {{id: '{program_id}'}})-[:has_version]->(o) + RETURN o + " + ) + } else { + format!( + " + MATCH (l: application {{id: '{program_id}'}})-[:has_version]->(o) + RETURN o + " + ) + }; + + let results = self.client.exec_query(&query).await?; + + let mut versions: Vec = vec![]; + for result in results { + let result_json: Value = serde_json::from_str(&result).unwrap(); + + let o = result_json["o"].clone(); + + let (version_base, name_version) = if is_lib { + let library_version: LibraryVersion = serde_json::from_value(o).unwrap(); + ( + UVersion::LibraryVersion(library_version.clone()), + library_version.name_and_version.clone(), + ) + } else { + let application_version: ApplicationVersion = serde_json::from_value(o).unwrap(); + ( + UVersion::ApplicationVersion(application_version.clone()), + application_version.name_and_version.clone(), + ) + }; + tracing::debug!("Read version for id {program_id}: {:?}", version_base); + + // get dependencies + let dependencies = self + .get_direct_dependency_nodes(&name_version) + .await + .unwrap(); + + versions.push(crate::VersionInfo { + version_base, + dependencies, + }) + } + Ok(versions) + } + + async fn get_direct_dependency_nodes( + &self, + name_and_version: &str, + ) -> Result, Box> { + let query = format!( + " + MATCH (n:version {{name_and_version: '{name_and_version}'}})-[:depends_on]->(m:version) + RETURN m.name_and_version as name_and_version + " + ); + + let results = self.client.exec_query(&query).await?; + let unique_items: HashSet = results.clone().into_iter().collect(); + let mut nodes = vec![]; + + for result in unique_items { + let result_json: Value = serde_json::from_str(&result).unwrap(); + let name_version_str: String = + serde_json::from_value(result_json["name_and_version"].clone()).unwrap(); + + if let Some(name_version) = crate::NameVersion::from_string(&name_version_str) { + nodes.push(name_version); + } + } + + Ok(nodes) + } + async fn new_get_direct_dependency_nodes( + &self, + namespace: &str, + nameversion: &str, + ) -> Result, Box> { + tracing::info!("enter get_direct_dependency_nodes"); + let query1 = format!( + " + MATCH (p:program {{namespace: '{namespace}'}})-[:has_type]->(l)-[:has_version]->(lv {{name_and_version: '{nameversion}'}})-[:has_dep_version]->(vs:version)-[:depends_on]->(m:version) +RETURN m.name_and_version as name_and_version + " + ); + let results1 = self.client.exec_query(&query1).await?; + tracing::info!("finish get_direct_dep"); + let mut res = vec![]; + for node in results1 { + res.push(node); + } + let unique_items: HashSet = res.clone().into_iter().collect(); + let mut nodes = vec![]; + + for result in unique_items { + let result_json: Value = serde_json::from_str(&result).unwrap(); + let name_version_str: String = + serde_json::from_value(result_json["name_and_version"].clone()).unwrap(); + + if let Some(name_version) = crate::NameVersion::from_string(&name_version_str) { + nodes.push(name_version); + } + } + + Ok(nodes) + } + async fn get_indirect_dependency_nodes( + &self, + nameversion: NameVersion, + ) -> Result, Box> { + let name_and_version = nameversion.name + "/" + &nameversion.version; + let mut nodes = self + .get_direct_dependency_nodes(&name_and_version) + .await + .unwrap(); + for node in nodes.clone() { + tracing::info!("{} {}", node.clone().name, node.clone().version); + let new_nodes = Box::pin(self.get_indirect_dependency_nodes(node)) + .await + .unwrap(); + for new_node in new_nodes { + tracing::info!("{} {}", new_node.clone().name, new_node.clone().version); + nodes.push(new_node); + } + } + Ok(nodes) + } + async fn get_program_by_name( + &self, + program_name: &str, + ) -> Result, Box> { + let query = format!( + " + MATCH (p:program) + WHERE p.name CONTAINS '{program_name}' + RETURN p + " + ); + let results = self.client.exec_query(&query).await?; + let mut programs = vec![]; + for result in results { + let programs_json: Value = serde_json::from_str(&result).unwrap(); + let pro = programs_json["p"].clone(); + let program: Program = serde_json::from_value(pro).unwrap(); + programs.push(program); + } + Ok(programs) + } + async fn count_dependencies(&self, nameversion: NameVersion) -> Result> { + let all_nodes = self.get_all_dependencies(nameversion).await.unwrap(); + Ok(all_nodes.len()) + } + async fn get_direct_dependent_nodes( + &self, + name_and_version: &str, + ) -> Result, Box> { + let query = format!( + " + MATCH (n:version {{name_and_version: '{name_and_version}'}})<-[:depends_on]-(m:version) + RETURN m.name_and_version as name_and_version + " + ); + + let results = self.client.exec_query(&query).await?; + let unique_items: HashSet = results.clone().into_iter().collect(); + let mut nodes = vec![]; + + for result in unique_items { + let result_json: Value = serde_json::from_str(&result).unwrap(); + let name_version_str: String = + serde_json::from_value(result_json["name_and_version"].clone()).unwrap(); + + if let Some(name_version) = crate::NameVersion::from_string(&name_version_str) { + nodes.push(name_version); + } + } + + Ok(nodes) + } + async fn new_get_direct_dependent_nodes( + &self, + namespace: &str, + nameversion: &str, + ) -> Result, Box> { + let query1 = format!( + " + MATCH (p:program {{namespace: '{namespace}'}})-[:has_type]->(l)-[:has_version]->(lv {{name_and_version:'{nameversion}'}})-[:has_dep_version]->(vs:version)<-[:depends_on]-(m:version) +RETURN m.name_and_version as name_and_version + " + ); + let results1 = self.client.exec_query(&query1).await?; + let mut res = vec![]; + for node in results1 { + res.push(node); + } + let unique_items: HashSet = res.clone().into_iter().collect(); + + let mut nodes = vec![]; + + for result in unique_items { + let result_json: Value = serde_json::from_str(&result).unwrap(); + let name_version_str: String = + serde_json::from_value(result_json["name_and_version"].clone()).unwrap(); + + if let Some(name_version) = crate::NameVersion::from_string(&name_version_str) { + nodes.push(name_version); + } + } + + Ok(nodes) + } + async fn get_indirect_dependent_nodes( + &self, + nameversion: NameVersion, + ) -> Result, Box> { + let name_and_version = nameversion.name + "/" + &nameversion.version; + let mut nodes = self + .get_direct_dependent_nodes(&name_and_version) + .await + .unwrap(); + for node in nodes.clone() { + let new_nodes = Box::pin(self.get_indirect_dependent_nodes(node)) + .await + .unwrap(); + for new_node in new_nodes { + nodes.push(new_node); + } + } + Ok(nodes) + } + + async fn get_lib_version(&self, name: String) -> Result, Box> { + let query = format!( + " + MATCH (n:library_version {{name: '{name}'}}) RETURN n.version LIMIT 100" + ); + + let results = self.client.exec_query(&query).await.unwrap(); + + let mut realres = vec![]; + + for res in results { + let parsed: Value = serde_json::from_str(&res).unwrap(); + if let Some(version) = parsed.get("n.version").and_then(|v| v.as_str()) { + realres.push(version.to_string()); + } + } + + Ok(realres) + } + async fn new_get_lib_version( + &self, + namespace: String, + name: String, + ) -> Result, Box> { + let query = format!( + " + MATCH (p:program {{namespace: '{namespace}'}})-[:has_type]->(l)-[:has_version]->(lv {{name:'{name}'}}) +RETURN lv.version" + ); + let time1 = Instant::now(); + let results = self.client.exec_query(&query).await.unwrap(); + let query_time = time1.elapsed(); + tracing::info!("query_statement_need_time:{:?}", query_time); + let unique_items: HashSet = results.clone().into_iter().collect(); + + let mut realres = vec![]; + + for res in unique_items { + let parsed: Value = serde_json::from_str(&res).unwrap(); + if let Some(version) = parsed.get("lv.version").and_then(|v| v.as_str()) { + realres.push(version.to_string()); + } + } + + Ok(realres) + } + async fn get_app_version(&self, name: String) -> Result, Box> { + let query = format!( + " + MATCH (n:application_version {{name: '{name}'}}) RETURN n.version LIMIT 100" + ); + let results = self.client.exec_query(&query).await.unwrap(); + let mut realres = vec![]; + for res in results { + let parsed: Value = serde_json::from_str(&res).unwrap(); + if let Some(version) = parsed.get("n.version").and_then(|v| v.as_str()) { + realres.push(version.to_string()); + } + } + Ok(realres) + } + async fn new_get_app_version( + &self, + namespace: String, + name: String, + ) -> Result, Box> { + let query = format!( + " + MATCH (p:program {{namespace: '{namespace}'}})-[:has_type]->(a:application)-[:has_version]->(av:application_version {{name:'{name}'}}) +RETURN av.version" + + ); + + let results = self.client.exec_query(&query).await.unwrap(); + let unique_items: HashSet = results.clone().into_iter().collect(); + + let mut realres = vec![]; + + for res in unique_items { + let parsed: Value = serde_json::from_str(&res).unwrap(); + if let Some(version) = parsed.get("av.version").and_then(|v| v.as_str()) { + realres.push(version.to_string()); + } + } + + Ok(realres) + } +} diff --git a/extensions/cratespro/common/data_transporter/src/db.rs b/extensions/cratespro/common/data_transporter/src/db.rs new file mode 100644 index 000000000..fbe24d8e1 --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/db.rs @@ -0,0 +1,1537 @@ +use std::{cmp::Ordering, collections::HashSet, env}; + +use crate::{ + handler::{ + Crateinfo, DependencyCount, DependencyCrateInfo, DependencyInfo, DependentCount, + DependentData, DependentInfo, NewRustsec, RustSec, Versionpage, + }, + UploadedCrate, Userinfo, +}; +use chrono::{NaiveDateTime, SubsecRound}; +use model::tugraph_model::{Program, UProgram}; +use semver::Version; +use serde::{Deserialize, Serialize}; +use tokio_postgres::{Error, NoTls}; +use utoipa::ToSchema; +pub struct DBHandler { + pub client: tokio_postgres::Client, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct CveInfo { + cve_id: String, + url: String, + description: String, + crate_name: String, + start_version: String, + end_version: String, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct Allcve { + cves: Vec, +} + +pub fn db_connection_config_from_env() -> String { + format!( + "host={} port={} user={} password={} dbname={}", + env::var("POSTGRES_HOST_IP").unwrap(), + env::var("POSTGRES_HOST_PORT").unwrap(), + env::var("POSTGRES_USER_NAME").unwrap(), + env::var("POSTGRES_USER_PASSWORD").unwrap(), + env::var("POSTGRES_CRATESPRO_DB").unwrap() + ) +} +pub fn db_cratesio_connection_config_from_env() -> String { + format!( + "host={} port={} user={} password={} dbname={}", + env::var("POSTGRES_HOST_IP").unwrap(), + env::var("POSTGRES_HOST_PORT").unwrap(), + env::var("POSTGRES_USER_NAME").unwrap(), + env::var("POSTGRES_USER_PASSWORD").unwrap(), + env::var("POSTGRES_CRATESIO_DB").unwrap() + ) +} + +impl DBHandler { + pub async fn connect() -> Result { + let db_connection_config = db_connection_config_from_env(); + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls).await?; + + // Spawn the connection on a separate task + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + + // 创建 cratespro 数据库 + client + .execute("CREATE DATABASE cratespro", &[]) + .await + .or_else(|err| { + if let Some(db_err) = err.as_db_error() { + if db_err.code() == &tokio_postgres::error::SqlState::DUPLICATE_DATABASE { + return Ok(0); + } + } + Err(err) + })?; + + // 重新连接到 cratespro 数据库 + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls).await?; + + // Spawn the connection on a separate task + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + + Ok(DBHandler { client }) + } + + pub async fn clear_database(&self) -> Result<(), Error> { + self.client + .batch_execute( + " + DO $$ + BEGIN + IF EXISTS (SELECT 1 FROM pg_tables WHERE tablename = 'programs') THEN + DROP TABLE programs CASCADE; + END IF; + + + IF EXISTS (SELECT 1 FROM pg_tables WHERE tablename = 'program_versions') THEN + DROP TABLE program_versions CASCADE; + END IF; + + IF EXISTS (SELECT 1 FROM pg_tables WHERE tablename = 'program_dependencies') THEN + DROP TABLE program_dependencies CASCADE; + END IF; + + + END $$; + ", + ) + .await + } + + pub async fn create_tables(&self) -> Result<(), Error> { + let create_programs_table = " + CREATE TABLE IF NOT EXISTS programs ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + namespace TEXT, + max_version TEXT, + github_url TEXT, + mega_url TEXT, + doc_url TEXT, + program_type TEXT NOT NULL, + downloads BIGINT, + cratesio TEXT + ); + "; + + let create_program_versions_table = " + CREATE TABLE IF NOT EXISTS program_versions ( + name_and_version TEXT PRIMARY KEY, + id TEXT NOT NULL, + name TEXT NOT NULL, + version TEXT NOT NULL, + documentation TEXT, + version_type TEXT NOT NULL, + created_at TIMESTAMPTZ DEFAULT NOW() + ); + "; + + let create_program_dependencies_table = " + CREATE TABLE IF NOT EXISTS program_dependencies ( + name_and_version TEXT NOT NULL, + dependency_name TEXT NOT NULL, + dependency_version TEXT NOT NULL, + PRIMARY KEY (name_and_version, dependency_name, dependency_version) + ); + "; + + // 执行创建表的 SQL 语句 + let result = self + .client + .batch_execute(&format!( + "{create_programs_table}{create_program_versions_table}{create_program_dependencies_table}" + )) + .await; + + match result { + Ok(_) => { + tracing::info!("Tables created successfully."); + Ok(()) + } + Err(e) => { + tracing::error!("Error creating tables: {:?}", e); + Err(e) + } + } + } + pub async fn insert_program_data( + &self, + program: Program, + uprogram: UProgram, + _versions: Vec, + ) -> Result<(), Error> { + let (program_type, downloads, cratesio) = match &uprogram { + UProgram::Library(lib) => ("Library", Some(lib.downloads), lib.cratesio.clone()), + UProgram::Application(_) => ("Application", None, None), + }; + + self.client + .execute( + " + INSERT INTO programs ( + id, name, description, namespace, + max_version, github_url, mega_url, doc_url, + program_type, downloads, cratesio + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + ", + &[ + &program.id, + &program.name, + &program.description.unwrap_or_default(), + &program.namespace.unwrap_or_default(), + &program.max_version.unwrap_or_default(), + &program.github_url.unwrap_or_default(), + &program.mega_url.unwrap_or_default(), + &program.doc_url.unwrap_or_default(), + &program_type, + &downloads.unwrap_or_default(), + &cratesio.unwrap_or_default(), + ], + ) + .await + .map_err(|e| { + eprintln!("Error inserting program: {e}"); + e + }) + .unwrap(); + + tracing::info!("finish to insert program."); + + // 插入 UVersion 数据 + /*for version in versions { + let name_and_version = version.version_base.get_name_and_version(); + + match version.version_base { + UVersion::LibraryVersion(lib_ver) => { + self.client + .execute( + " + INSERT INTO program_versions ( + name_and_version, id, name, version, + documentation, version_type, created_at + ) VALUES ($1, $2, $3, $4, $5, $6, NOW()) + ", + &[ + &lib_ver.name_and_version, + &lib_ver.id, + &lib_ver.name, + &lib_ver.version, + &Some(lib_ver.documentation), + &"LibraryVersion", + ], + ) + .await + .unwrap(); + } + UVersion::ApplicationVersion(app_ver) => { + self.client + .execute( + " + INSERT INTO program_versions ( + name_and_version, id, name, version, + documentation, version_type, created_at + ) VALUES ($1, $2, $3, $4, $5, $6, NOW()) + ", + &[ + &app_ver.name_and_version, + &app_ver.id, + &app_ver.name, + &app_ver.version, + &None::, // ApplicationVersion 没有 documentation 字段 + &"ApplicationVersion", + ], + ) + .await + .unwrap(); + } + } + + // 插入该版本的所有依赖项 + for dep in version.dependencies { + self.client + .execute( + " + INSERT INTO program_dependencies ( + name_and_version, dependency_name, dependency_version + ) VALUES ($1, $2, $3) + ", + &[&name_and_version, &dep.name, &dep.version], + ) + .await?; + } + }*/ + //tracing::info!("Finish to insert all versions."); + + Ok(()) + } + pub async fn get_all_cvelist(&self) -> Result { + + let raws = self + .client + .query( + "SELECT cve_id, name, start_version, end_version,description FROM cves;", + &[], + ) + .await?; + let mut getcves = vec![]; + for raw in raws { + let front = "https://www.cve.org/CVERecord?id="; + let cve_id: String = raw.get(0); + let cve_url = front.to_string() + &cve_id; + let cve_info = CveInfo { + cve_id: raw.get(0), + url: cve_url, + description: raw.get(4), + crate_name: raw.get(1), + start_version: raw.get(2), + end_version: raw.get(3), + }; + getcves.push(cve_info); + } + let res = Allcve { cves: getcves }; + + Ok(res) + } + pub async fn process_closed_interval_of_match_version( + &self, + oneline_patched: String, + version: String, + ) -> Result { + let mut matched = false; + let mut two_versions = vec![]; + let newparts: Vec<&str> = oneline_patched.split(',').collect(); + for part in newparts { + let one_version = part.to_string(); + let res_one_version = one_version.trim(); + two_versions.push(res_one_version.to_string()); + } + let mut left = "".to_string(); + let mut right = "".to_string(); + if two_versions.len() == 2 { + if two_versions[0].clone().starts_with(">") || two_versions[0].clone().starts_with(">=") + { + left = two_versions[0].clone(); + right = two_versions[1].clone(); + } else if two_versions[0].clone().starts_with("<") + || two_versions[0].clone().starts_with("<=") + { + left = two_versions[1].clone(); + right = two_versions[0].clone(); + } + } + if (left.starts_with(">") && !left.starts_with(">=")) + && (right.starts_with("<") && !right.starts_with("<=")) + { + //> < + let mut versions = vec![]; + let tmp_left = &left[1..]; + let left_version = tmp_left.to_string(); + let tmp_right = &right[1..]; + let right_version = tmp_right.to_string(); + versions.push(version.clone()); + versions.push(left_version.clone()); + versions.push(right_version.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if version.clone() == versions[1].clone() + && (versions[0].clone() != version.clone() + && versions[2].clone() != version.clone()) + { + matched = true; + } + } else if (left.starts_with(">") && !left.starts_with(">=")) && right.starts_with("<=") { + //> <= + let mut versions = vec![]; + let tmp_left = &left[1..]; + let left_version = tmp_left.to_string(); + let tmp_right = &right[2..]; + let right_version = tmp_right.to_string(); + versions.push(version.clone()); + versions.push(left_version.clone()); + versions.push(right_version.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if version.clone() == versions[1].clone() && (versions[2].clone() != version.clone()) { + matched = true; + } + } else if left.starts_with(">=") && (right.starts_with("<") && !right.starts_with("<=")) { + //>= < + let mut versions = vec![]; + let tmp_left = &left[2..]; + let left_version = tmp_left.to_string(); + let tmp_right = &right[1..]; + let right_version = tmp_right.to_string(); + versions.push(version.clone()); + versions.push(left_version.clone()); + versions.push(right_version.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if version.clone() == versions[1].clone() && (versions[0].clone() != version.clone()) { + matched = true; + } + } else if left.starts_with(">=") && right.starts_with("<=") { + //>= <= + let mut versions = vec![]; + let tmp_left = &left[2..]; + let left_version = tmp_left.to_string(); + let tmp_right = &right[2..]; + let right_version = tmp_right.to_string(); + versions.push(version.clone()); + versions.push(left_version.clone()); + versions.push(right_version.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if version.clone() == versions[1].clone() { + matched = true; + } + } + Ok(matched) + } + pub async fn process_open_interval_of_match_version( + &self, + oneline_patched: String, + version: String, + ) -> Result { + let mut matched = false; + if oneline_patched.starts_with(">") && !oneline_patched.starts_with(">=") { + let mut versions = vec![]; + let trimmed = &oneline_patched[1..]; + let res = trimmed.to_string(); + versions.push(version.clone()); + versions.push(res.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if versions[0].clone() == version.clone() && res.clone() != version.clone() { + matched = true; + } + } else if let Some(trimmed) = oneline_patched.strip_prefix(">=") { + let mut versions = vec![]; + let res = trimmed.to_string(); + versions.push(version.clone()); + versions.push(res.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if versions[0].clone() == version.clone() { + matched = true; + } + } else if oneline_patched.starts_with("<") && !oneline_patched.starts_with("<=") { + let mut versions = vec![]; + let trimmed = &oneline_patched[1..]; + let res = trimmed.to_string(); + versions.push(version.clone()); + versions.push(res.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if versions[1].clone() == version.clone() && res.clone() != version.clone() { + matched = true; + } + } else if let Some(trimmed) = oneline_patched.strip_prefix("<=") { + let mut versions = vec![]; + let res = trimmed.to_string(); + versions.push(version.clone()); + versions.push(res.clone()); + versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + if versions[1].clone() == version.clone() { + matched = true; + } + } + Ok(matched) + } + pub async fn match_version(&self, patched: String, version: String) -> Result { + let mut matched = false; + let mut part_petched = vec![]; + let parts: Vec<&str> = patched.split('|').collect(); + for part in parts { + part_petched.push(part); + } + for np in part_petched { + let oneline_patched = np.to_string(); + if oneline_patched.clone().contains(",") { + //closed interval + matched = matched + || self + .process_closed_interval_of_match_version( + oneline_patched.clone(), + version.clone(), + ) + .await + .unwrap(); + } else if oneline_patched.clone().contains("^") { + //specific version + if let Some(trimmed) = oneline_patched.strip_prefix("^") { + let res = trimmed.to_string(); + if version == res { + matched = true; + } + } + } else { + //open interval + matched = matched + || self + .process_open_interval_of_match_version( + oneline_patched.clone(), + version.clone(), + ) + .await?; + } + } + Ok(matched) + } + pub async fn get_direct_rustsec( + &self, + cname: &str, + version: &str, + ) -> Result, Error> { + tracing::info!("enter get direct_rustsec"); + let rows = self + .client + .query("SELECT * FROM rustsecs WHERE cratename=$1;", &[&cname]) + .await + .unwrap(); + let mut get_direct_rust_sec = vec![]; + for row in rows { + let t_aliases: String = row.get("aliases"); + let parts: Vec<&str> = t_aliases.split(';').collect(); + let mut real_aliases = vec![]; + for part in parts { + real_aliases.push(part.to_string()); + } + let rs = RustSec { + id: row.get("id"), + cratename: row.get("cratename"), + patched: row.get("patched"), + aliases: real_aliases.clone(), + small_desc: row.get("small_desc"), + }; + get_direct_rust_sec.push(rs.clone()); + } + let mut getres = vec![]; + for rc in get_direct_rust_sec { + if rc.cratename.clone() == *cname { + let matched = self + .match_version(rc.clone().patched, version.to_string()) + .await + .unwrap(); + if !matched { + let rows2 = self + .client + .query("SELECT * FROM rustsec_info WHERE id=$1;", &[&rc.clone().id]) + .await + .unwrap(); + for row in rows2 { + let tmp_id: String = row.get("id"); + let rs_url = + "https://rustsec.org/advisories/".to_string() + &tmp_id + ".html"; + let nrs = NewRustsec { + id: row.get("id"), + subtitle: row.get("subtitle"), + reported: row.get("reported"), + issued: row.get("issued"), + package: row.get("package"), + ttype: row.get("type"), + keywords: row.get("keywords"), + aliases: row.get("aliases"), + reference: row.get("reference"), + patched: row.get("patched"), + unaffected: row.get("unaffected"), + description: row.get("description"), + url: rs_url.clone(), + }; + getres.push(nrs.clone()); + } + } + } + } + tracing::info!("finish get direct_rustsec"); + Ok(getres) + } + pub async fn get_dependency_rustsec( + &self, + nameversion: HashSet, + ) -> Result, Error> { + let mut getres = vec![]; + for nv in nameversion { + let parts: Vec<&str> = nv.split('/').collect(); + let cname = parts[0].to_string(); + let version = parts[1].to_string(); + let rows = self + .client + .query("SELECT * FROM rustsecs where cratename=$1;", &[&cname]) + .await + .unwrap(); + let mut get_all_rust_sec = vec![]; + for row in rows { + let t_aliases: String = row.get("aliases"); + let parts: Vec<&str> = t_aliases.split(';').collect(); + let mut real_aliases = vec![]; + for part in parts { + real_aliases.push(part.to_string()); + } + let rs = RustSec { + id: row.get("id"), + cratename: row.get("cratename"), + patched: row.get("patched"), + aliases: real_aliases.clone(), + small_desc: row.get("small_desc"), + }; + get_all_rust_sec.push(rs.clone()); + } + for rc in get_all_rust_sec.clone() { + if rc.cratename.clone() == cname.clone() { + let matched = self + .match_version(rc.clone().patched, version.to_string()) + .await + .unwrap(); + if !matched { + let rows2 = self + .client + .query("SELECT * FROM rustsec_info WHERE id=$1;", &[&rc.clone().id]) + .await + .unwrap(); + for row in rows2 { + let tmp_id: String = row.get("id"); + let rs_url = + "https://rustsec.org/advisories/".to_string() + &tmp_id + ".html"; + let nrs = NewRustsec { + id: row.get("id"), + subtitle: row.get("subtitle"), + reported: row.get("reported"), + issued: row.get("issued"), + package: row.get("package"), + ttype: row.get("type"), + keywords: row.get("keywords"), + aliases: row.get("aliases"), + reference: row.get("reference"), + patched: row.get("patched"), + unaffected: row.get("unaffected"), + description: row.get("description"), + url: rs_url.clone(), + }; + getres.push(nrs.clone()); + } + //getres.push(rc.clone()); + } + } + } + } + let unique: Vec = getres + .into_iter() + .collect::>() + .into_iter() + .collect(); + Ok(unique) + } + + pub async fn get_rustsec_by_depends_cveid( + &self, + depends_cves_id: String, + ) -> Result, Error> { + let parts: Vec = depends_cves_id.split('/').map(|x| x.to_string()).collect(); + let mut res = vec![]; + for part in parts { + let rows = self + .client + .query("SELECT * FROM rustsec_info WHERE id=$1;", &[&part]) + .await + .unwrap(); + for row in rows { + let rs_url = + "https://rustsec.org/advisories/".to_string() + &part.clone() + ".html"; + let nrs = NewRustsec { + id: part.clone(), + subtitle: row.get("subtitle"), + reported: row.get("reported"), + issued: row.get("issued"), + package: row.get("package"), + ttype: row.get("type"), + keywords: row.get("keywords"), + aliases: row.get("aliases"), + reference: row.get("reference"), + patched: row.get("patched"), + unaffected: row.get("unaffected"), + description: row.get("description"), + url: rs_url.clone(), + }; + res.push(nrs.clone()); + } + } + Ok(res) + } + + pub async fn get_license_by_name( + &self, + namespace: &str, + name: &str, + ) -> Result, Error> { + let rows = self + .client + .query( + "SELECT license FROM license WHERE program_namespace = $1 and program_name = $2;", + &[&namespace.to_string(), &name.to_string()], + ) + .await + .unwrap(); + let mut licenses = vec![]; + for row in rows { + let new_license: String = row.get(0); + licenses.push(new_license); + } + licenses.push("None".to_string()); + Ok(licenses) + } + pub async fn query_crates_info_from_pg( + &self, + id: &str, + name: String, + ) -> Result, Box> { + tracing::info!("start query crates from pg"); + let rows = self + .client + .query( + "SELECT * FROM crates_info WHERE id = $1;", + &[&id.to_string()], + ) + .await + .unwrap(); + + let mut cf = vec![]; + for row in rows { + let desc: String = row.get("description"); + let dcyct: i32 = row.get("direct_dependency"); + let indcyct: i32 = row.get("indirect_dependency"); + let dtct: i32 = row.get("direct_dependent"); + let indtct: i32 = row.get("indirect_dependent"); + let cs: String = row.get("cves"); + let vs: String = row.get("versions"); + let lcs: String = row.get("license"); + let gu: String = row.get("github_url"); + let du: String = row.get("doc_url"); + let dep_cs: String = row.get("dep_cves"); + let mut getcves = vec![]; + let everypartscs: Vec<&str> = cs.split("||||||").collect(); + for part in everypartscs { + let new_part = part.to_string(); + let parts2: Vec<&str> = new_part.split("------").collect(); + if parts2.len() == 13 { + let onecve = NewRustsec { + id: parts2[0].to_string(), + subtitle: parts2[1].to_string(), + reported: parts2[2].to_string(), + issued: parts2[3].to_string(), + package: parts2[4].to_string(), + ttype: parts2[5].to_string(), + keywords: parts2[6].to_string(), + aliases: parts2[7].to_string(), + reference: parts2[8].to_string(), + patched: parts2[9].to_string(), + unaffected: parts2[10].to_string(), + description: parts2[12].to_string(), + url: parts2[11].to_string(), + }; + getcves.push(onecve); + } + } + + let mut getdepcs = vec![]; + let everypartsdepcs: Vec<&str> = dep_cs.split("||||||").collect(); + for part in everypartsdepcs { + let new_part = part.to_string(); + let parts2: Vec<&str> = new_part.split("------").collect(); + if parts2.len() == 13 { + let onecve = NewRustsec { + id: parts2[0].to_string(), + subtitle: parts2[1].to_string(), + reported: parts2[2].to_string(), + issued: parts2[3].to_string(), + package: parts2[4].to_string(), + ttype: parts2[5].to_string(), + keywords: parts2[6].to_string(), + aliases: parts2[7].to_string(), + reference: parts2[8].to_string(), + patched: parts2[9].to_string(), + unaffected: parts2[10].to_string(), + description: parts2[12].to_string(), + url: parts2[11].to_string(), + }; + getdepcs.push(onecve); + } + } + + let mut getversions = vec![]; + let partsvs: Vec<&str> = vs.split('/').collect(); + for part in partsvs { + getversions.push(part.to_string()); + } + let res_crates_info = Crateinfo { + crate_name: name.clone(), + description: desc.clone(), + dependencies: DependencyCount { + direct: dcyct as usize, + indirect: indcyct as usize, + }, + dependents: DependentCount { + direct: dtct as usize, + indirect: indtct as usize, + }, + cves: getcves, + license: lcs.clone(), + github_url: gu.clone(), + doc_url: du.clone(), + versions: getversions, + dep_cves: getdepcs, + }; + cf.push(res_crates_info); + } + Ok(cf) + } + pub async fn process_cves( + &self, + cves: Vec, + ) -> Result> { + let mut every_cs = vec![]; + for rs in cves { + let t_id = rs.clone().id; + let mut t_subtitle = rs.clone().subtitle; + if t_subtitle.is_empty() { + t_subtitle = "Null".to_string(); + } + let mut t_reported = rs.clone().reported; + if t_reported.is_empty() { + t_reported = "Null".to_string(); + } + let mut t_issued = rs.clone().issued; + if t_issued.is_empty() { + t_issued = "Null".to_string(); + } + let mut t_package = rs.clone().package; + if t_package.is_empty() { + t_package = "Null".to_string(); + } + let mut t_type = rs.clone().ttype; + if t_type.is_empty() { + t_type = "Null".to_string(); + } + let mut t_keywords = rs.clone().keywords; + if t_keywords.is_empty() { + t_keywords = "Null".to_string(); + } + let mut t_aliases = rs.clone().aliases; + if t_aliases.is_empty() { + t_aliases = "Null".to_string(); + } + let mut t_reference = rs.clone().reference; + if t_reference.is_empty() { + t_reference = "Null".to_string(); + } + let mut t_patched = rs.clone().patched; + if t_patched.is_empty() { + t_patched = "Null".to_string(); + } + let mut t_unaffected = rs.clone().unaffected; + if t_unaffected.is_empty() { + t_unaffected = "Null".to_string(); + } + let mut t_desc = rs.clone().description; + if t_desc.is_empty() { + t_desc = "Null".to_string(); + } + let t_url = rs.clone().url; + let tmp_strings = [ + t_id, + t_subtitle, + t_reported, + t_issued, + t_package, + t_type, + t_keywords, + t_aliases, + t_reference, + t_patched, + t_unaffected, + t_url, + t_desc, + ]; + let result: String = tmp_strings + .iter() + .filter(|&s| !s.is_empty()) + .cloned() // 复制引用的字符串 + .collect::>() + .join("------"); + every_cs.push(result); + } + let cs = every_cs.clone().join("||||||"); + Ok(cs) + } + pub async fn insert_crates_info_into_pg( + &self, + crateinfo: Crateinfo, + namespace: String, + name: String, + version: String, + ) -> Result<(), Box> { + let id = namespace.clone() + "/" + &name + "/" + &version; + let dcyct = crateinfo.dependencies.direct as i32; + let indcyct = crateinfo.dependencies.indirect as i32; + let dtct = crateinfo.dependents.direct as i32; + let indtct = crateinfo.dependents.indirect as i32; + let vs = crateinfo.versions.clone().join("/"); + let cs = self.process_cves(crateinfo.clone().cves).await.unwrap(); + let depcs = self.process_cves(crateinfo.clone().dep_cves).await.unwrap(); + self.client + .execute( + " + INSERT INTO crates_info ( + id,description,direct_dependency,indirect_dependency, + direct_dependent,indirect_dependent,cves,dep_cves,versions, + license,github_url,doc_url + ) VALUES ($1, $2, $3, $4, $5, $6, $7,$8,$9,$10,$11,$12); + ", + &[ + &id, + &crateinfo.description, + &dcyct, + &indcyct, + &dtct, + &indtct, + &cs, + &depcs, + &vs, + &crateinfo.license, + &crateinfo.github_url, + &crateinfo.doc_url, + ], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn get_graph_from_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + version: String, + ) -> Result, Box> { + let id = nsfront + "/" + &nsbehind + "/" + &name + "/" + &version; + let rows = self + .client + .query("SELECT * FROM graph_info WHERE id = $1;", &[&id]) + .await + .unwrap(); + let mut res = vec![]; + for row in rows { + let graph: String = row.get("graph"); + res.push(graph); + } + Ok(res) + } + pub async fn insert_graph_into_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + version: String, + graph: String, + ) -> Result<(), Box> { + let id = nsfront + "/" + &nsbehind + "/" + &name + "/" + &version; + self.client + .execute( + " + INSERT INTO graph_info ( + id,graph + ) VALUES ($1, $2); + ", + &[&id, &graph], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn get_version_from_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + ) -> Result, Box> { + let id = nsfront + "/" + &nsbehind + "/" + &name; + let rows = self + .client + .query("SELECT * FROM version_info WHERE id = $1;", &[&id]) + .await + .unwrap(); + let mut res = vec![]; + for row in rows { + let newversion: String = row.get("versions"); + res.push(newversion); + } + Ok(res) + } + pub async fn insert_version_into_pg( + &self, + nsbehind: String, + nsfront: String, + name: String, + versionpg: Vec, + ) -> Result<(), Box> { + let id = nsfront + "/" + &nsbehind + "/" + &name; + let mut every_version = vec![]; + for vp in versionpg { + let dts_count = vp.dependents.to_string(); + let one_version = vp.version.clone() + + "|" + + &vp.updated_at.clone() + + "|" + + &vp.downloads.clone() + + "|" + + &dts_count; + every_version.push(one_version); + } + let versions = every_version.join("/"); + self.client + .execute( + " + INSERT INTO version_info ( + id,versions + ) VALUES ($1, $2); + ", + &[&id, &versions], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn get_dump_from_cratesio_pg( + &self, + name: String, + version: String, + ) -> Result> { + tracing::info!("enter get dump"); + let rows1 = self + .client + .query("SELECT * FROM crates WHERE name=$1 LIMIT 1", &[&name]) + .await + .unwrap(); + tracing::info!("finish get id"); + let mut res = "".to_string(); + for row in rows1 { + tracing::info!("enter rows1"); + let crate_id: i32 = row.get("id"); + tracing::info!("id:{}", crate_id.clone()); + tracing::info!("start get num,up,dl"); + let rows = self + .client + .query("SELECT * FROM versions WHERE crate_id=$1;", &[&crate_id]) + .await + .unwrap(); + tracing::info!("finish get num,up,dl"); + for row in rows { + let num: String = row.get("num"); + let updated_at: NaiveDateTime = row.get("updated_at"); + let updated_at_without_ms = updated_at.trunc_subsecs(0); + let downloads: i32 = row.get("downloads"); + let downloads_string = downloads.to_string(); + let updated_at_string = updated_at_without_ms.to_string(); + if num == version { + res = updated_at_string + "/" + &downloads_string; + } + } + tracing::info!("finish get dump"); + } + Ok(res) + } + pub async fn get_dependency_from_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + version: String, + ) -> Result, Box> { + let id = nsfront.clone() + "/" + &nsbehind + "/" + &name + "/" + &version; + let rows = self + .client + .query("SELECT * FROM dependency_cache WHERE id = $1;", &[&id]) + .await + .unwrap(); + let mut res = vec![]; + for row in rows { + let all_dependency: String = row.get("dependency"); + let direct: i32 = row.get("direct_count"); + let indirect: i32 = row.get("indirect_count"); + let mut deps = vec![]; + let parts1: Vec<&str> = all_dependency.split("|").collect(); + for part in parts1 { + let one_dep = part.to_string(); + let parts2: Vec<&str> = one_dep.split("/").collect(); + if parts2.len() == 5 { + let dcs = parts2[4].to_string(); + let dcc = dcs.parse::().unwrap(); + let one_res = DependencyCrateInfo { + crate_name: parts2[0].to_string(), + version: parts2[1].to_string(), + relation: parts2[2].to_string(), + license: parts2[3].to_string(), + dependencies: dcc, + }; + deps.push(one_res); + } + } + let real_res = DependencyInfo { + direct_count: direct as usize, + indirect_count: indirect as usize, + data: deps, + }; + res.push(real_res); + } + Ok(res) + } + pub async fn get_dependent_from_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + version: String, + ) -> Result, Box> { + let id = nsfront.clone() + "/" + &nsbehind + "/" + &name + "/" + &version; + let rows = self + .client + .query("SELECT * FROM dependent_cache WHERE id = $1;", &[&id]) + .await + .unwrap(); + let mut res = vec![]; + for row in rows { + let all_dependent: String = row.get("dependent"); + let direct: i32 = row.get("direct_count"); + let indirect: i32 = row.get("indirect_count"); + let mut deps = vec![]; + let parts1: Vec<&str> = all_dependent.split("|").collect(); + for part in parts1 { + let one_dep = part.to_string(); + let parts2: Vec<&str> = one_dep.split("/").collect(); + if parts2.len() == 3 { + let one_res = DependentData { + crate_name: parts2[0].to_string(), + version: parts2[1].to_string(), + relation: parts2[2].to_string(), + }; + deps.push(one_res); + } + } + let real_res = DependentInfo { + direct_count: direct as usize, + indirect_count: indirect as usize, + data: deps, + }; + res.push(real_res); + } + Ok(res) + } + pub async fn insert_dependency_into_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + version: String, + dep_info: DependencyInfo, + ) -> Result<(), Box> { + let id = nsfront.clone() + "/" + &nsbehind + "/" + &name + "/" + &version; + let mut every_dep = vec![]; + for one_dep in dep_info.data { + let dcs = one_dep.dependencies.to_string(); + let one_res = one_dep.crate_name.clone() + + "/" + + &one_dep.version + + "/" + + &one_dep.relation + + "/" + + &one_dep.license + + "/" + + &dcs; + every_dep.push(one_res); + } + let real_dep = every_dep.join("|"); + self.client + .execute( + " + INSERT INTO dependency_cache ( + id,direct_count,indirect_count,dependency + ) VALUES ($1, $2,$3,$4); + ", + &[ + &id, + &(dep_info.direct_count as i32), + &(dep_info.indirect_count as i32), + &real_dep, + ], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn insert_dependent_into_pg( + &self, + nsfront: String, + nsbehind: String, + name: String, + version: String, + dep_info: DependentInfo, + ) -> Result<(), Box> { + let id = nsfront.clone() + "/" + &nsbehind + "/" + &name + "/" + &version; + let mut every_dep = vec![]; + for one_dep in dep_info.data { + let one_res = + one_dep.crate_name.clone() + "/" + &one_dep.version + "/" + &one_dep.relation; + every_dep.push(one_res); + } + let real_dep = every_dep.join("|"); + self.client + .execute( + " + INSERT INTO dependent_cache ( + id,direct_count,indirect_count,dependent + ) VALUES ($1, $2,$3,$4); + ", + &[ + &id, + &(dep_info.direct_count as i32), + &(dep_info.indirect_count as i32), + &real_dep, + ], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn insert_userinfo_into_pg( + &self, + info: Userinfo, + ) -> Result<(), Box> { + tracing::info!("enter insert userinfo into pg"); + self.client + .execute( + "INSERT INTO userloginfo( + id,image,name,expires) VALUES ($1, $2,$3,$4) + ON CONFLICT (id) + DO UPDATE SET image = $2, name = $3, expires = $4;", + &[ + &info.user.email, + &info.user.image, + &info.user.name, + &info.expires, + ], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn query_uploaded_crates_from_pg( + &self, + email: String, + ) -> Result, Box> { + let rows = self + .client + .query("SELECT * FROM uploadedcrate WHERE email=$1", &[&email]) + .await + .unwrap(); + let mut res = vec![]; + for row in rows { + let name: String = row.get("filename"); + let time: String = row.get("uploadtime"); + let tmp_res = UploadedCrate { name, time }; + res.push(tmp_res); + } + Ok(res) + } + pub async fn query_uploaded_url_from_pg( + &self, + email: String, + ) -> Result, Box> { + let rows = self + .client + .query("SELECT * FROM uploadedurl WHERE email=$1", &[&email]) + .await + .unwrap(); + let mut res = vec![]; + for row in rows { + let name: String = row.get("githuburl"); + let time: String = row.get("uploadtime"); + let tmp_res = UploadedCrate { name, time }; + res.push(tmp_res); + } + Ok(res) + } + pub async fn insert_sensleak_result_into_pg( + &self, + id: String, + result: String, + ) -> Result<(), Box> { + self.client + .execute( + "INSERT INTO senseleak_res( + id,res) VALUES ($1, $2) + ON CONFLICT (id) + DO UPDATE SET res=$2;", + &[&id, &result], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn insert_mirchecker_result_into_pg( + &self, + id: String, + result: String, + ) -> Result<(), Box> { + self.client + .execute( + "INSERT INTO mirchecker_res( + id,res) VALUES ($1, $2) + ON CONFLICT (id) + DO UPDATE SET res=$2;", + &[&id, &result], + ) + .await + .unwrap(); + Ok(()) + } + pub async fn insert_mirchecker_failed_into_pg( + &self, + id: String, + ) -> Result<(), Box> { + self.client + .execute( + "INSERT INTO mirchecker_run_failed( + id) VALUES ($1) + ON CONFLICT (id) + DO NOTHING;", + &[&id], + ) + .await + .unwrap(); + Ok(()) + } + #[allow(clippy::len_zero)] + pub async fn get_senseleak_from_pg( + &self, + id: String, + ) -> Result> { + let rows = self + .client + .query("SELECT * FROM senseleak_res WHERE id=$1", &[&id]) + .await + .unwrap(); + let mut tmp_res = vec![]; + for row in rows { + let s_res: String = row.get("res"); + tmp_res.push(s_res); + } + let mut real_res = "[]".to_string(); + if tmp_res.len() != 0 { + real_res = tmp_res[0].clone(); + } + Ok(real_res) + } + #[allow(clippy::len_zero)] + pub async fn get_mirchecker_from_pg( + &self, + id: String, + ) -> Result> { + let rows = self + .client + .query("SELECT * FROM mirchecker_res WHERE id=$1", &[&id]) + .await + .unwrap(); + let mut tmp_res = vec![]; + for row in rows { + let s_res: String = row.get("res"); + tmp_res.push(s_res); + } + let mut real_res = "".to_string(); + if tmp_res.len() != 0 { + real_res = tmp_res[0].clone(); + } + Ok(real_res) + } + #[allow(clippy::len_zero)] + pub async fn get_mirchecker_run_state_from_pg( + &self, + id: String, + ) -> Result> { + let rows = self + .client + .query("SELECT * FROM mirchecker_run_failed WHERE id=$1", &[&id]) + .await + .unwrap(); + let mut tmp_res = vec![]; + for row in rows { + let s_res: String = row.get("id"); + tmp_res.push(s_res); + } + let mut real_res = true; + if tmp_res.len() != 0 { + real_res = false; + } + Ok(real_res) + } + pub async fn get_all_dependency_counts_from_pg( + &self, + crate_name: String, + version: String, + ) -> Result> { + let rows = self + .client + .query( + "SELECT count FROM all_dependency_counts WHERE crate_name=$1 AND version=$2", + &[&crate_name, &version], + ) + .await + .unwrap(); + let mut tmp_res = vec![]; + for row in rows { + let count: i32 = row.get("count"); + tmp_res.push(count as usize); + } + let mut real_res = 0; + if !tmp_res.is_empty() { + real_res = tmp_res[0]; + } + Ok(real_res) + } + pub async fn get_all_dependent_counts_from_pg( + &self, + crate_name: String, + version: String, + ) -> Result> { + let rows = self + .client + .query( + "SELECT count FROM direct_dependent_counts WHERE crate_name=$1 AND version=$2", + &[&crate_name, &version], + ) + .await + .unwrap(); + let mut tmp_res = vec![]; + for row in rows { + let count: i32 = row.get("count"); + tmp_res.push(count as usize); + } + let mut real_res = 0; + if !tmp_res.is_empty() { + real_res = tmp_res[0]; + } + Ok(real_res) + } + pub async fn get_depends_of_all_version( + &self, + crate_name: String, + ) -> Result<(usize, usize, usize, String), Box> { + let rows = self + .client + .query( + "SELECT * FROM crate_noversion_depends WHERE crate_name=$1", + &[&crate_name], + ) + .await + .unwrap(); + let mut tmp_res = vec![]; + for row in rows { + let direct_depends: i32 = row.get("direct_depends"); + let indirect_depends: i32 = row.get("indirect_depends"); + let direct_depended_by: i32 = row.get("direct_depended_by"); + let depends_cves_id: String = row.get("depends_cves_id"); + tmp_res.push(( + direct_depends as usize, + indirect_depends as usize, + direct_depended_by as usize, + depends_cves_id, + )); + } + let mut real_res = (0, 0, 0, "".to_string()); + if !tmp_res.is_empty() { + real_res = tmp_res[0].clone(); + } + Ok(real_res) + } +} diff --git a/extensions/cratespro/common/data_transporter/src/handler.rs b/extensions/cratespro/common/data_transporter/src/handler.rs new file mode 100644 index 000000000..7692ea10c --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/handler.rs @@ -0,0 +1,965 @@ +use std::collections::HashSet; +#[allow(unused_imports)] +use std::env; +use std::error::Error; +//use std::error::Error; +use std::time::Instant; + +use crate::data_reader::{DataReader, DataReaderTrait}; +use crate::db::{db_connection_config_from_env, DBHandler}; +use crate::redis_store::{get_redis_connection, RedisHandler}; +use crate::{get_tugraph_api_handler, NameVersion, Userinfo}; +use crate::{Query, VersionInfo}; +use actix_multipart::{Field, Multipart}; +use actix_web::http::header::ContentDisposition; +use actix_web::{web, HttpResponse, Responder}; +use futures_util::StreamExt; +//use model::repo_sync_model; +//use model::repo_sync_model::CrateType; +use model::tugraph_model::{Program, UProgram}; +//use repo_import::ImportDriver; +use sanitize_filename::sanitize; +use search::crates_search::RecommendCrate; +//use search::crates_search::RecommendCrate; +use search::crates_search::SearchModule; +use search::crates_search::SearchSortCriteria; +use semver::Version; +use serde::Deserialize; +use serde::Serialize; +use std::cmp::Ordering; +use std::io::Cursor; +use std::io::Read; +//use std::time::Instant; +//use semver::Version; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio_postgres::NoTls; +use utoipa::ToSchema; +use zip::ZipArchive; +pub struct ApiHandler { + reader: DataReader, +} +impl ApiHandler { + pub async fn new(reader: DataReader) -> Self { + Self { reader } + } +} + +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct QueryCratesInfo { + code: u32, + message: String, + data: QueryData, +} +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct QueryData { + total_page: usize, + items: Vec, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct QueryItem { + name: String, + version: String, + date: String, + nsfront: String, + nsbehind: String, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct DependencyCrateInfo { + pub crate_name: String, + pub version: String, + pub relation: String, + pub license: String, + pub dependencies: usize, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct DependencyInfo { + pub direct_count: usize, + pub indirect_count: usize, + pub data: Vec, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct DependentInfo { + pub direct_count: usize, + pub indirect_count: usize, + pub data: Vec, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct DependentData { + pub crate_name: String, + pub version: String, + pub relation: String, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct Crateinfo { + pub crate_name: String, + pub description: String, + pub dependencies: DependencyCount, + pub dependents: DependentCount, + pub cves: Vec, + pub dep_cves: Vec, + pub license: String, + pub github_url: String, + pub doc_url: String, + pub versions: Vec, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct DependencyCount { + pub direct: usize, + pub indirect: usize, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct DependentCount { + pub direct: usize, + pub indirect: usize, +} +#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash, ToSchema)] +pub struct RustSec { + pub id: String, + pub cratename: String, + pub patched: String, + pub aliases: Vec, + pub small_desc: String, +} +#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash, ToSchema)] +pub struct NewRustsec { + pub id: String, + pub subtitle: String, + pub reported: String, + pub issued: String, + pub package: String, + pub ttype: String, + pub keywords: String, + pub aliases: String, + pub reference: String, + pub patched: String, + pub unaffected: String, + pub description: String, + pub url: String, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct Deptree { + pub name_and_version: String, + pub cve_count: usize, + #[schema(value_type = Vec)] + pub direct_dependency: Vec, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct Versionpage { + pub version: String, + pub updated_at: String, + pub downloads: String, + pub dependents: usize, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct SenseleakRes { + pub exist: bool, + pub res: String, +} +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub struct MircheckerRes { + pub run_state: bool, + pub exist: bool, + pub res: Vec, +} +impl MircheckerRes { + pub fn new() -> Self { + Self { + run_state: true, + exist: false, + res: Vec::new(), + } + } +} +/// 获取cve信息 +#[utoipa::path( + get, + path = "/api/cvelist", + responses( + (status = 200, description = "成功获取crate信息", body = crate::db::Allcve), + (status = 404, description = "未找到crate信息") + ), + tag = "security" +)] +pub async fn get_cves() -> impl Responder { + let _handler = get_tugraph_api_handler().await; + + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhd = DBHandler { client }; + let cves = dbhd.get_all_cvelist().await.unwrap(); + + HttpResponse::Ok().json(cves) +} + +/// 获取所有crates +#[utoipa::path( + get, + path = "/api/crates", + responses( + (status = 200, description = "成功获取所有crate的id", body = Vec), + (status = 404, description = "未找到crate id") + ), + tag = "crates" +)] +pub async fn get_all_crates() -> impl Responder { + tracing::info!("get all crates func run"); + let handler = get_tugraph_api_handler().await; + let ids = handler.reader.get_all_programs_id().await; + + let mut programs = vec![]; + for id in &ids { + let program = handler.reader.get_program(id).await.unwrap(); + programs.push(program); + } + + tracing::info!("finish get all crates func"); + + HttpResponse::Ok().json(programs) // 返回 JSON 格式 +} + +/// 获取crate详细信息,ok +#[utoipa::path( + get, + path = "/api/crates/{cratename}", + params( + ("cratename" = String, Path, description = "crate 名称") + ), + responses( + (status = 200, description = "成功获取crate详细信息", body= (Program, UProgram, Vec)), + (status = 404, description = "未找到crate"), + (status = 500, description = "服务器内部错误") + ), + tag = "crates" +)] +pub async fn get_crate_details(crate_name: web::Path) -> impl Responder { + let handler = get_tugraph_api_handler().await; + match handler.reader.get_program(&crate_name).await { + Ok(program) => { + match handler.reader.get_type(&crate_name).await { + Ok((uprogram, islib)) => { + match handler.reader.get_versions(&crate_name, islib).await { + Ok(versions) => { + HttpResponse::Ok().json((program, uprogram, versions)) + // 返回 JSON 格式 + } + Err(_) => { + HttpResponse::InternalServerError().body("Failed to get versions.") + } + } + } + Err(_) => HttpResponse::InternalServerError().body("Failed to get type."), + } + } + Err(_) => HttpResponse::NotFound().body("Crate not found."), + } +} + +/// 获取直接依赖关系图,ok +#[utoipa::path( + get, + path = "/api/graph/{cratename}/{version}/direct", + params( + ("cratename" = String, Path, description = "crate 名称"), + ("version" = String, Path, description = "版本号") + ), + responses( + (status = 200, description = "成功获取依赖关系图", body = Vec), + (status = 404, description = "未找到依赖关系图") + ), + tag = "dependencies" +)] +pub async fn get_direct_dep_for_graph(nname: String, nversion: String) -> impl Responder { + let handler = get_tugraph_api_handler().await; + let name_and_version = nname + "/" + &nversion; + let res = handler + .reader + .get_direct_dependency_nodes(&name_and_version) + .await + .unwrap(); + HttpResponse::Ok().json(res) +} + + + +/// 查询 crates +#[utoipa::path( + post, + path = "/api/search", + request_body = Query, + responses( + (status = 200, description = "查询成功", body = QueryCratesInfo), + (status = 400, description = "无效的查询参数") + ), + tag = "search" +)] +pub async fn query_crates(q: Query) -> impl Responder { + let _handler = get_tugraph_api_handler().await; + //add yj's search module + let name = q.query; + let page = q.pagination.page; + let per_page = q.pagination.per_page; + tracing::info!("name:{},page:{},per_page:{}", name, page, per_page); + let db_connection_config = db_connection_config_from_env(); + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let start_time2 = Instant::now(); + let question = name.clone(); + let search_module = SearchModule::new(&client).await; + let res = search_module + .search_crate(&question, SearchSortCriteria::Relavance) + .await + .unwrap(); + tracing::trace!("search need time:{:?}", start_time2.elapsed()); + let mut seen = HashSet::new(); + let uniq_res: Vec = res + .into_iter() + .filter(|x| seen.insert((x.name.clone(), x.namespace.clone()))) + .collect(); + tracing::trace!("total programs: {}", uniq_res.len()); + let mut gettotal_page = uniq_res.len() / per_page; + if uniq_res.is_empty() || uniq_res.len() % per_page != 0 { + gettotal_page += 1; + } + let mut getitems = vec![]; + for i in (page - 1) * 20..(page - 1) * 20 + 20 { + if i >= uniq_res.len() { + break; + } + let mut mv = vec![]; + let program_name = uniq_res[i].clone().name; + let getnamespace = uniq_res[i].clone().namespace; + let parts: Vec<&str> = getnamespace.split('/').collect(); + let nsf = parts[0].to_string(); + let nsb = parts[1].to_string(); + + mv.push(uniq_res[i].clone().max_version); + + if mv[0].clone() == *"null" { + mv[0] = "0.0.0".to_string(); + } + let query_item = QueryItem { + name: program_name.clone(), + version: mv[0].clone(), + date: "".to_string(), + nsfront: nsf, + nsbehind: nsb, + }; + getitems.push(query_item); + } + let response = QueryCratesInfo { + code: 200, + message: "成功".to_string(), + data: QueryData { + total_page: gettotal_page, + items: getitems, + }, + }; + + HttpResponse::Ok().json(response) +} +//post of upload +#[allow(clippy::let_unit_value)] +pub async fn upload_crate(mut payload: Multipart) -> impl Responder { + tracing::info!("enter upload crate"); + use futures_util::StreamExt as _; + let mut upload_time: Option = None; + let mut user_email: Option = None; + let mut github_link: Option = None; + let mut file_name: Option = None; + while let Some(Ok(mut field)) = payload.next().await { + tracing::info!("enter while"); + if let Some(content_disposition) = field.content_disposition().cloned() { + tracing::info!("enter first if"); + if let Some(name) = content_disposition.get_name() { + tracing::info!("enter second if"); + match name { + "file" => { + tracing::info!("enter match file"); + file_name = process_file_of_upload_crate(&content_disposition, &mut field) + .await + .unwrap(); + // analyze + } + "githubLink" => { + github_link = process_githublink_of_upload_crate(&mut field) + .await + .unwrap(); + } + "uploadTime" => { + tracing::info!("enter match uploadtime"); + upload_time = process_uploadtime_of_upload_crate(&mut field) + .await + .unwrap(); + } + "user_email" => { + tracing::info!("enter match user_email"); + user_email = process_useremail_of_upload_crate(&mut field).await.unwrap(); + } + _ => { + tracing::info!("enter match nothing"); + } + } + } + } + } + let _ = process_insertintopg_of_upload_crate(file_name, upload_time, github_link, user_email) + .await + .unwrap(); + HttpResponse::Ok().json(()) +} +pub async fn process_insertintopg_of_upload_crate( + file_name: Option, + upload_time: Option, + github_link: Option, + user_email: Option, +) -> Result<(), Box> { + if let Some(filename) = file_name { + tracing::info!("enter 1/2 if let"); + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + if let Some(uploadtime) = upload_time.clone() { + tracing::info!("enter upload time:{}", uploadtime.clone()); + if let Some(useremail) = user_email.clone() { + tracing::info!("enter user email:{}", useremail.clone()); + dbhandler + .client + .execute( + "INSERT INTO uploadedcrate(email,filename,uploadtime) VALUES ($1, $2,$3);", + &[&useremail.clone(), &filename.clone(), &uploadtime.clone()], + ) + .await + .unwrap(); + } + } + }; + if let Some(githublink) = github_link { + tracing::info!("enter 2/2 if let"); + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + if let Some(uploadtime) = upload_time.clone() { + if let Some(useremail) = user_email.clone() { + dbhandler + .client + .execute( + "INSERT INTO uploadedurl(email,githuburl,uploadtime) VALUES ($1, $2,$3);", + &[&useremail.clone(), &githublink.clone(), &uploadtime.clone()], + ) + .await + .unwrap(); + } + } + } + Ok(()) +} +#[allow(unused_assignments)] +pub async fn process_useremail_of_upload_crate( + field: &mut Field, +) -> Result, Box> { + let mut user_email: Option = None; + let mut email_data = Vec::new(); + + while let Some(chunk) = field.next().await { + let data = chunk.unwrap(); + email_data.extend_from_slice(&data); + } + user_email = Some(String::from_utf8(email_data).unwrap_or_default()); + tracing::info!("user_email:{:?}", user_email); + Ok(user_email) +} +#[allow(unused_assignments)] +pub async fn process_uploadtime_of_upload_crate( + field: &mut Field, +) -> Result, Box> { + let mut upload_time: Option = None; + let mut time_data = Vec::new(); + + while let Some(chunk) = field.next().await { + let data = chunk.unwrap(); + time_data.extend_from_slice(&data); + } + upload_time = Some(String::from_utf8(time_data).unwrap_or_default()); + tracing::info!("uploadtime:{:?}", upload_time); + Ok(upload_time) +} +#[allow(unused_assignments)] +pub async fn process_githublink_of_upload_crate( + field: &mut Field, +) -> Result, Box> { + let mut github_link: Option = None; + let mut url_data = Vec::new(); + while let Some(chunk) = field.next().await { + let data = chunk.unwrap(); + url_data.extend_from_slice(&data); + } + github_link = Some(String::from_utf8(url_data).unwrap_or_default()); + Ok(github_link) +} +#[allow(unused_assignments)] +pub async fn process_file_of_upload_crate( + content_disposition: &ContentDisposition, + field: &mut Field, +) -> Result, Box> { + tracing::info!("enter match file"); + let mut file_name: Option = None; + let filename = if let Some(file_name) = content_disposition.get_filename() { + file_name.to_string() + } else { + "default.zip".to_string() + }; + tracing::info!("filename:{}", filename.clone()); + let sanitized_filename = sanitize(filename.clone()); + file_name = Some(filename.clone()); + tracing::info!("file_name:{:?}", file_name.clone()); + if sanitized_filename.ends_with(".zip") { + tracing::info!("enter file zip"); + let zip_filepath = format!("target/zip/upload/{sanitized_filename}"); + let _ = tokio::fs::create_dir_all("target/zip/upload/").await; + let mut f = tokio::fs::File::create(&zip_filepath).await.unwrap(); + while let Some(chunk) = field.next().await { + let data = chunk.unwrap(); + f.write_all(&data).await.unwrap(); + } + let parts: Vec<&str> = sanitized_filename.split('.').collect(); + let mut filename = "".to_string(); + if parts.len() >= 2 { + filename = parts[0].to_string(); + tracing::info!("filename without zip: {}", filename); + } + let mut zip_file = tokio::fs::File::open(&zip_filepath).await.unwrap(); + let mut buffer = Vec::new(); + zip_file.read_to_end(&mut buffer).await.unwrap(); + let reader = Cursor::new(buffer.clone()); + let mut archive = ZipArchive::new(reader).unwrap(); + for i in 0..archive.len() { + let mut file = archive.by_index(i).unwrap(); + let outpath = match file.enclosed_name() { + Some(path) => { + format!("target/www/uploads/{}/{}", filename, path.display()) + } + None => continue, + }; + + if file.name().ends_with('/') { + // This is a directory, create it + tokio::fs::create_dir_all(&outpath).await.unwrap(); + } else { + // Ensure the parent directory exists + if let Some(parent) = std::path::Path::new(&outpath).parent() { + if !parent.exists() { + tokio::fs::create_dir_all(&parent).await.unwrap(); + } + } + + // Write the file + let mut outfile = tokio::fs::File::create(&outpath).await.unwrap(); + while let Ok(bytes_read) = file.read(&mut buffer) { + if bytes_read == 0 { + break; + } + outfile.write_all(&buffer[..bytes_read]).await.unwrap(); + } + } + } + tracing::info!("finish match file"); + //send message + } else { + tracing::info!("enter else"); + let filepath = format!("/home/rust/output/www/uploads/{sanitized_filename}"); + let mut f = tokio::fs::File::create(&filepath).await.unwrap(); + + while let Some(chunk) = field.next().await { + let data = chunk.unwrap(); + f.write_all(&data).await.unwrap(); + } + } + Ok(file_name) +} +//post of log in +pub async fn submituserinfo(info: Userinfo) -> impl Responder { + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + tracing::info!("enter submituserinfo and set db client"); + #[allow(clippy::let_unit_value)] + let _ = dbhandler + .insert_userinfo_into_pg(info.clone()) + .await + .unwrap(); + HttpResponse::Ok().json(()) +} +pub async fn query_upload_crate(email: String) -> impl Responder { + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + let mut real_res = vec![]; + let res = dbhandler + .query_uploaded_crates_from_pg(email.clone()) + .await + .unwrap(); + for row in res { + real_res.push(row); + } + let res2 = dbhandler + .query_uploaded_url_from_pg(email.clone()) + .await + .unwrap(); + for row in res2 { + real_res.push(row); + } + HttpResponse::Ok().json(real_res) +} +pub async fn get_senseleak(nsfront: String, nsbehind: String) -> impl Responder { + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + let id = nsfront.clone() + "/" + &nsbehind; + let res = dbhandler.get_senseleak_from_pg(id).await.unwrap(); + let mut exist = true; + if res.clone() == *"[]" { + exist = false; + } + let return_val = SenseleakRes { exist, res }; + HttpResponse::Ok().json(return_val) +} +#[allow(unused_assignments)] +pub async fn get_mirchecker( + nsfront: String, + nsbehind: String, + name: String, + version: String, +) -> impl Responder { + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + let mut return_val = MircheckerRes::new(); + if version.clone() == *"all" { + let handler = get_tugraph_api_handler().await; + let namespace = nsfront.clone() + "/" + &nsbehind; + let mut all_versions = handler + .reader + .new_get_lib_version(namespace.clone(), name.clone()) + .await + .unwrap(); + all_versions.sort_by(|a, b| { + let version_a = Version::parse(a); + let version_b = Version::parse(b); + match (version_a, version_b) { + (Ok(v_a), Ok(v_b)) => v_b.cmp(&v_a), + (Ok(_), Err(_)) => Ordering::Less, + (Err(_), Ok(_)) => Ordering::Greater, + (Err(_), Err(_)) => Ordering::Equal, + } + }); + let run_state = true; + let mut exist = false; + let mut tmp_res = vec![]; + for single_version in all_versions.clone() { + let id = nsfront.clone() + "/" + &nsbehind + "/" + &name + "/" + &single_version; + let single_res = dbhandler.get_mirchecker_from_pg(id.clone()).await.unwrap(); + if single_res.contains("warning: [MirChecker]") { + exist = true; + tmp_res.push("Version-".to_string() + &single_version + ":\n" + &single_res); + } + } + return_val = MircheckerRes { + run_state, + exist, + res: tmp_res, + }; + } else { + let id = nsfront.clone() + "/" + &nsbehind + "/" + &name + "/" + &version; + let run_state = dbhandler + .get_mirchecker_run_state_from_pg(id.clone()) + .await + .unwrap(); + let res = dbhandler.get_mirchecker_from_pg(id.clone()).await.unwrap(); + let mut real_res = vec![]; + let mut exist = false; + if res.contains("warning: [MirChecker]") { + exist = true; + real_res.push(res.clone()); + } + return_val = MircheckerRes { + run_state, + exist, + res: real_res, + }; + } + HttpResponse::Ok().json(return_val) +} + +pub async fn new_get_crates_front_info_from_redis( + nname: String, + nversion: String, + nsfront: String, + nsbehind: String, +) -> impl Responder { + let handler = get_tugraph_api_handler().await; + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + + let conn = get_redis_connection().await.unwrap(); + let mut redisconn = RedisHandler { connection: conn }; + let qid = format!("crates_info:{namespace}:{nname}:{nversion}"); + let qres = redisconn.query_from_redis(qid).await.unwrap(); + println!("finish query crates from reids"); + if qres.is_empty() { + println!("qres is empty"); + let res = handler + .reader + .get_crates_front_info_from_tg( + nname.clone(), + nversion.clone(), + nsfront.clone(), + nsbehind.clone(), + ) + .await + .unwrap(); + println!("finish get crates_info from tugraph"); + let val = serde_json::to_string(&res).unwrap(); + redisconn + .insert_crates_info_into_redis( + namespace.clone(), + nname.clone(), + nversion.clone(), + val.clone(), + ) + .await + .unwrap(); + HttpResponse::Ok().json(res) + } else { + let res: Crateinfo = serde_json::from_str(&qres).unwrap(); + HttpResponse::Ok().json(res.clone()) + } +} +pub async fn dependency_redis_cache( + name: String, + version: String, + nsfront: String, + nsbehind: String, +) -> impl Responder { + let handler = get_tugraph_api_handler().await; + let conn = get_redis_connection().await.unwrap(); + let mut redisconn = RedisHandler { connection: conn }; + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let qid = format!("dependency:{namespace}:{name}:{version}"); + let res = redisconn.query_from_redis(qid.clone()).await.unwrap(); + if res.is_empty() { + let res_deps = handler + .reader + .get_dependency_from_tg( + name.clone(), + version.clone(), + nsfront.clone(), + nsbehind.clone(), + ) + .await + .unwrap(); + let val = serde_json::to_string(&res_deps).unwrap(); + redisconn + .insert_dependency_into_redis( + namespace.clone(), + name.clone(), + version.clone(), + val.clone(), + ) + .await + .unwrap(); + HttpResponse::Ok().json(res_deps.clone()) + } else { + let res_deps: DependencyInfo = serde_json::from_str(&res).unwrap(); + HttpResponse::Ok().json(res_deps.clone()) + } +} +pub async fn new_get_graph( + nsfront: String, + nsbehind: String, + nname: String, + nversion: String, +) -> impl Responder { + let handler = get_tugraph_api_handler().await; + let db_connection_config = db_connection_config_from_env(); + #[allow(unused_variables)] + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let dbhandler = DBHandler { client }; + let conn = get_redis_connection().await.unwrap(); + let mut redisconn = RedisHandler { connection: conn }; + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let qid = format!("dependencygraph:{namespace}:{nname}:{nversion}"); + let qres = redisconn.query_from_redis(qid.clone()).await.unwrap(); + if qres.is_empty() { + tracing::info!("first time"); + let nav = nname.clone() + "/" + &nversion; + let rustcve = dbhandler + .get_direct_rustsec(&nname, &nversion) + .await + .unwrap(); + let mut res = Deptree { + name_and_version: nav.clone(), + cve_count: rustcve.len(), + direct_dependency: Vec::new(), + }; + let mut visited = HashSet::new(); + visited.insert(nav.clone()); + handler + .reader + .build_graph(&mut res, &mut visited) + .await + .unwrap(); + let graph = serde_json::to_string(&res).unwrap(); + redisconn + .insert_dependency_graph_into_redis( + namespace.clone(), + nname.clone(), + nversion.clone(), + graph.clone(), + ) + .await + .unwrap(); + HttpResponse::Ok().json(res) + } else { + tracing::info!("second time"); + let res_tree: Deptree = serde_json::from_str(&qres).unwrap(); + HttpResponse::Ok().json(res_tree) + } +} +pub async fn dependent_redis_cache( + name: String, + version: String, + nsfront: String, + nsbehind: String, +) -> impl Responder { + let handler = get_tugraph_api_handler().await; + let conn = get_redis_connection().await.unwrap(); + let mut redisconn = RedisHandler { connection: conn }; + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let qid = format!("dependent:{namespace}:{name}:{version}"); + let qres = redisconn.query_from_redis(qid.clone()).await.unwrap(); + if qres.is_empty() { + let res_deps = handler + .reader + .get_dependent_from_tg( + name.clone(), + version.clone(), + nsfront.clone(), + nsbehind.clone(), + ) + .await + .unwrap(); + let val = serde_json::to_string(&res_deps).unwrap(); + redisconn + .insert_dependent_into_redis( + namespace.clone(), + name.clone(), + version.clone(), + val.clone(), + ) + .await + .unwrap(); + HttpResponse::Ok().json(res_deps) + } else { + let res: DependentInfo = serde_json::from_str(&qres).unwrap(); + HttpResponse::Ok().json(res.clone()) + } +} +pub async fn new_get_version_page( + nsfront: String, + nsbehind: String, + nname: String, + _nversion: String, +) -> impl Responder { + let handler = get_tugraph_api_handler().await; + let conn = get_redis_connection().await.unwrap(); + let mut redisconn = RedisHandler { connection: conn }; + let namespace = nsfront.clone() + "/" + &nsbehind.clone(); + let qid = format!("versionpage:{namespace}:{nname}"); + let res = redisconn.query_from_redis(qid.clone()).await.unwrap(); + if res.is_empty() { + let every_version = handler + .reader + .get_version_page_from_tg(nsfront.clone(), nsbehind.clone(), nname.clone()) + .await + .unwrap(); + let val = serde_json::to_string(&every_version).unwrap(); + redisconn + .insert_versionpage_into_redis(namespace, nname.clone(), val.clone()) + .await + .unwrap(); + HttpResponse::Ok().json(every_version) + } else { + let every_version: Vec = serde_json::from_str(&res).unwrap(); + HttpResponse::Ok().json(every_version) + } +} diff --git a/extensions/cratespro/common/data_transporter/src/lib.rs b/extensions/cratespro/common/data_transporter/src/lib.rs new file mode 100644 index 000000000..598f3f017 --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/lib.rs @@ -0,0 +1,303 @@ +mod data_packer; +mod data_reader; +pub mod db; +mod handler; +mod transporter; +mod redis_store; + +use model::tugraph_model::UVersion; +use search::search_prepare; +use serde::{Deserialize, Serialize}; +use tokio_postgres::NoTls; +pub use transporter::Transporter; + +use crate::data_reader::DataReader; // 确保导入你的 DataReader +use crate::db::db_connection_config_from_env; +use crate::handler::ApiHandler; + +use actix_multipart::Multipart; +use actix_web::{web, App, HttpResponse, HttpServer}; +use utoipa::{OpenApi, ToSchema}; +use utoipa_swagger_ui::SwaggerUi; + +#[derive(Deserialize, Debug, ToSchema)] +pub struct Query { + query: String, + pagination: Pagination, +} +#[derive(Deserialize, Debug, ToSchema)] +pub struct Pagination { + page: usize, + per_page: usize, +} +#[derive(Deserialize, Debug, ToSchema,Serialize,Clone)] +pub struct Loginfo { + email: String, + image:String, + name: String, +} +#[derive(Deserialize, Debug, ToSchema,Serialize,Clone)] +pub struct Userinfo { + user:Loginfo, + expires:String, +} +#[derive(Deserialize, Debug, ToSchema,Serialize,Clone)] +pub struct UploadedCrate { + name:String, + time:String, +} +#[allow(non_snake_case)] +#[derive(Debug, Deserialize)] +struct Root { + requestBody: RequestBody, +} + +#[derive(Debug, Deserialize)] +struct RequestBody { + session: Userinfo, +} +#[allow(non_snake_case)] +#[derive(Debug, Deserialize)] +struct RequestBody2 { + requestBody: String, +} +async fn get_tugraph_api_handler() -> ApiHandler { + let tugraph_bolt_url = &std::env::var("TUGRAPH_BOLT_URL").unwrap(); + let tugraph_user_name = &std::env::var("TUGRAPH_USER_NAME").unwrap(); + let tugraph_user_password = &std::env::var("TUGRAPH_USER_PASSWORD").unwrap(); + let tugraph_cratespro_db = &std::env::var("TUGRAPH_CRATESPRO_DB").unwrap(); + let reader = DataReader::new( + tugraph_bolt_url, + tugraph_user_name, + tugraph_user_password, + tugraph_cratespro_db, + ) + .await + .unwrap(); + ApiHandler::new(reader).await +} + +#[derive(OpenApi)] +#[openapi( + paths( + handler::get_cves, + handler::get_all_crates, + //handler::get_graph, + handler::get_crate_details, + handler::query_crates, + //handler::get_graph, + //route::get_version_page, + // route::get_graph, + // route::get_direct_dep_for_graph, + // route::new_get_crates_front_info, + // route::new_get_dependency, + // route::dependency_cache, + // route::new_get_dependent, + // route::dependent_cache, + // route::query_crates, + + // route::get_crate_details, + ), + components( + schemas( + model::tugraph_model::Program, + db::Allcve, + handler::Versionpage, + //handler::Deptree, + //handler::Crateinfo, + handler::DependencyInfo, + handler::DependentInfo, + handler::Crateinfo, + model::tugraph_model::UProgram, + VersionInfo, + Query, + handler::QueryCratesInfo, + //handler::Deptree, + // Query, + // Pagination, + // route::QueryCratesInfo, + // route::QueryData, + // route::QueryItem, + // route::DependencyInfo, + // route::DependencyCrateInfo, + // route::DependentInfo, + // route::DependentData, + // route::Crateinfo, + + // route::Versionpage, + // route::NewRustsec, + // NameVersion, + // route::Deptree, + ) + ), + tags( + (name = "crates", description = "Crates API"), + (name = "dependencies", description = "Dependencies API"), + (name = "search", description = "Search API"), + (name = "security", description = "Security API"), + (name = "versions", description = "Version API"), + //(name = "upload", description = "Upload API"), + ) +)] +struct ApiDoc; + +pub async fn run_api_server() -> std::io::Result<()> { + tracing::info!("Start run_api_server"); + let db_connection_config = db_connection_config_from_env(); + let (client, connection) = tokio_postgres::connect(&db_connection_config, NoTls) + .await + .unwrap(); + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + let pre_search = search_prepare::SearchPrepare::new(&client).await; + pre_search.prepare_tsv().await.unwrap(); + HttpServer::new(move || { + tracing::info!("start route"); + App::new() + .service( + SwaggerUi::new("/swagger-ui/{_:.*}") + .url("/api-docs/openapi.json", ApiDoc::openapi()) + ) + .route( + "/api/cvelist", + web::get().to( + || async move { + handler::get_cves().await + }, + ), + ) + .route( + "/api/crates", + web::get().to(|| async move { + handler::get_all_crates().await + }), + ) + .route( + "/api/crates/{cratename}", + web::get().to( + |name: web::Path| async move { + handler::get_crate_details(name.into_inner().into()).await + }, + ), + ) + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/versions", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,cratename, version) = path.into_inner(); + handler::new_get_version_page(nsfront,nsbehind,cratename,version).await + })) + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/dependencies/graphpage", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,cratename, version) = path.into_inner(); + handler::new_get_graph(nsfront,nsbehind,cratename,version).await + })) + + .route( + "/api/submit", + web::post().to( + | payload: Multipart| async move { + handler::upload_crate(payload).await + }, + ), + ) + .route( + "/api/submitCrate", + web::post().to( + | payload: Multipart| async move { + tracing::info!("enter submitcrate"); + handler::upload_crate(payload).await + }, + ), + ) + .route("/api/submitUserinfo", web::post().to( + |payload: String| async move{ + //web::Json + tracing::info!("enter submitUserinfo"); + tracing::info!("payload:{}",payload.clone()); + let query:Root = serde_json::from_str(&payload).unwrap(); + tracing::info!("userinfo {:?}",query); + handler::submituserinfo(query.requestBody.session).await + },),) + .route("/api/profile", web::post().to( + |payload: String| async move{ + tracing::info!("enter profile"); + tracing::info!("payload:{}",payload.clone()); + let query:RequestBody2 = serde_json::from_str(&payload).unwrap(); + tracing::info!("profile email:{}",query.requestBody.clone()); + handler::query_upload_crate(query.requestBody).await + },),) + .route("/api/search", web::post().to( + |payload: web::Json| async move{ + let query = payload.into_inner(); + handler::query_crates(query).await + },),) + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/dependencies", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,cratename, version) = path.into_inner(); + handler::dependency_redis_cache(cratename,version,nsfront,nsbehind).await + })) + + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/dependencies/graph", + web::get().to(|_path: web::Path<(String, String,String,String)>|async move{ + HttpResponse::Ok().json(()) + })) + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/dependents", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,cratename, version) = path.into_inner(); + handler::dependent_redis_cache(cratename,version,nsfront,nsbehind).await + })) + + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,cratename, version) = path.into_inner(); + handler::new_get_crates_front_info_from_redis(cratename,version,nsfront,nsbehind).await + })) + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/senseleak", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,_cratename,_versionn) = path.into_inner(); + handler::get_senseleak(nsfront, nsbehind).await + })) + .route("/api/graph/{cratename}/{version}/direct", + web::get().to(|path: web::Path<(String,String)>|async move{ + let (cratename, version) = path.into_inner(); + handler::get_direct_dep_for_graph(cratename,version).await + })) + .route("/api/crates/{nsfront}/{nsbehind}/{cratename}/{version}/unsafechecker", + web::get().to(|path: web::Path<(String, String,String,String)>|async move{ + let (nsfront,nsbehind,cratename,version) = path.into_inner(); + handler::get_mirchecker(nsfront, nsbehind,cratename,version).await + })) + }) + .bind("0.0.0.0:6888")? + .run() + .await +} + +#[derive(Debug, Serialize, Deserialize, Clone, ToSchema,Hash,PartialEq,Eq)] +pub struct NameVersion { + pub name: String, + pub version: String, +} + +impl NameVersion { + // 解析 "name/version" 格式的字符串 + pub fn from_string(name_version: &str) -> Option { + let parts: Vec<&str> = name_version.split('/').collect(); + if parts.len() == 2 { + Some(NameVersion { + name: parts[0].to_string(), + version: parts[1].to_string(), + }) + } else { + None + } + } +} + +#[derive(Debug, Serialize, Deserialize, ToSchema)] +pub struct VersionInfo { + pub version_base: UVersion, + pub dependencies: Vec, +} diff --git a/extensions/cratespro/common/data_transporter/src/redis_store.rs b/extensions/cratespro/common/data_transporter/src/redis_store.rs new file mode 100644 index 000000000..b3feeff63 --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/redis_store.rs @@ -0,0 +1,155 @@ +use std::env; + +use redis::{Commands, Connection}; + +pub struct RedisHandler { + pub connection: Connection, +} + +pub async fn get_redis_connection() -> Result> { + let host = env::var("REDIS_HOST").unwrap_or_else(|_| "172.17.0.1".to_string()); + let password = env::var("REDIS_PASSWORD").unwrap_or_else(|_| "".to_string()); + println!("host:{host}"); + // 构建连接字符串 + let conn_string = format!("redis://:{password}@{host}:6379/"); + println!("conn_string:{conn_string}"); + println!("尝试连接 Redis: {conn_string}"); + + // 创建客户端 + let client = redis::Client::open(&*conn_string).unwrap(); + println!("finish get client"); + // 获取连接 + let conn = client.get_connection().unwrap(); + Ok(conn) +} + +impl RedisHandler { + pub fn get_connection_mut(&mut self) -> &mut Connection { + &mut self.connection + } + pub async fn query_from_redis( + &mut self, + qid: String, + ) -> Result> { + match self.get_connection_mut().get::<_, String>(&qid) { + Ok(value) => Ok(value), + Err(e) => { + if e.kind() == redis::ErrorKind::TypeError { + Ok("".to_string()) + } else { + Err(Box::new(e)) + } + } + } + } + + pub async fn insert_crates_info_into_redis( + &mut self, + namespace: String, + name: String, + version: String, + value: String, + ) -> Result<(), Box> { + println!("start insert crates_info"); + let key = format!("crates_info:{namespace}:{name}:{version}"); + + // 使用管道设置值和过期时间 + let _: () = redis::pipe() + .cmd("SET") + .arg(&key) + .arg(&value) + .cmd("EXPIRE") + .arg(&key) + .arg(7 * 24 * 60 * 60) // 一周的秒数 + .query(&mut self.connection)?; + println!("finish insert crates_info"); + Ok(()) + } + pub async fn insert_dependency_into_redis( + &mut self, + namespace: String, + name: String, + version: String, + value: String, + ) -> Result<(), Box> { + println!("start insert dependency"); + let key = format!("dependency:{namespace}:{name}:{version}"); + + // 使用管道设置值和过期时间 + let _: () = redis::pipe() + .cmd("SET") + .arg(&key) + .arg(&value) + .cmd("EXPIRE") + .arg(&key) + .arg(7 * 24 * 60 * 60) // 一周的秒数 + .query(&mut self.connection)?; + println!("finish insert dependency"); + Ok(()) + } + pub async fn insert_dependency_graph_into_redis( + &mut self, + namespace: String, + name: String, + version: String, + value: String, + ) -> Result<(), Box> { + println!("start insert dependencygraph"); + let key = format!("dependencygraph:{namespace}:{name}:{version}"); + + // 使用管道设置值和过期时间 + let _: () = redis::pipe() + .cmd("SET") + .arg(&key) + .arg(&value) + .cmd("EXPIRE") + .arg(&key) + .arg(7 * 24 * 60 * 60) // 一周的秒数 + .query(&mut self.connection)?; + println!("finish insert dependencygraph"); + Ok(()) + } + pub async fn insert_dependent_into_redis( + &mut self, + namespace: String, + name: String, + version: String, + value: String, + ) -> Result<(), Box> { + println!("start insert dependent"); + let key = format!("dependent:{namespace}:{name}:{version}"); + + // 使用管道设置值和过期时间 + let _: () = redis::pipe() + .cmd("SET") + .arg(&key) + .arg(&value) + .cmd("EXPIRE") + .arg(&key) + .arg(7 * 24 * 60 * 60) // 一周的秒数 + .query(&mut self.connection)?; + println!("finish insert dependent"); + Ok(()) + } + pub async fn insert_versionpage_into_redis( + &mut self, + namespace: String, + name: String, + value: String, + ) -> Result<(), Box> { + println!("start insert versionpage"); + let key = format!("versionpage:{namespace}:{name}"); + + // 使用管道设置值和过期时间 + let _: () = redis::pipe() + .cmd("SET") + .arg(&key) + .arg(&value) + .cmd("EXPIRE") + .arg(&key) + .arg(7 * 24 * 60 * 60) // 一周的秒数 + .query(&mut self.connection)?; + println!("finish insert versionpage"); + Ok(()) + } +} diff --git a/extensions/cratespro/common/data_transporter/src/transporter.rs b/extensions/cratespro/common/data_transporter/src/transporter.rs new file mode 100644 index 000000000..497fcb25c --- /dev/null +++ b/extensions/cratespro/common/data_transporter/src/transporter.rs @@ -0,0 +1,36 @@ +use crate::{ + data_packer::DataPacker, + data_reader::{DataReader, DataReaderTrait}, +}; + +pub struct Transporter { + pub reader: DataReader, + pub packer: DataPacker, +} + +impl Transporter { + pub async fn new(uri: &str, user: &str, password: &str, db: &str) -> Self { + Self { + reader: DataReader::new(uri, user, password, db).await.unwrap(), + packer: DataPacker::new().await, + } + } + + pub async fn transport_data(&mut self) -> Result<(), ()> { + tracing::info!("Start to pack the data"); + let ids = self.reader.get_all_programs_id().await; + for id in ids { + tracing::info!("id:{}", id); + let program: model::tugraph_model::Program = + self.reader.get_program(&id).await.unwrap(); + let (uprogram, islib): (model::tugraph_model::UProgram, bool) = + self.reader.get_type(&id).await.unwrap(); + let versions: Vec = + self.reader.get_versions(&id, islib).await.unwrap(); + + self.packer.pack_into_db(program, uprogram, versions).await; + } + tracing::info!("finish to pack the data"); + Ok(()) + } +} diff --git a/extensions/cratespro/common/database/.env b/extensions/cratespro/common/database/.env new file mode 100644 index 000000000..1a29dec05 --- /dev/null +++ b/extensions/cratespro/common/database/.env @@ -0,0 +1 @@ +DATABASE_URL="postgres://mega:mega@localhost:5432/cratesio2" \ No newline at end of file diff --git a/extensions/cratespro/common/database/BUCK b/extensions/cratespro/common/database/BUCK new file mode 100644 index 000000000..1e42be850 --- /dev/null +++ b/extensions/cratespro/common/database/BUCK @@ -0,0 +1,19 @@ +rust_library( + name = "database", + srcs = [ + "src/lib.rs", + ], + crate_root = "src/lib.rs", + edition = "2021", + deps = [ + "//project/crates-pro:entity", + "//project/crates-pro:model", + "//third-party:tracing", + "//third-party:serde", + "//third-party:serde_json", + "//third-party:sea-orm", + "//third-party:chrono", + "//third-party:futures", + ], + visibility = ["PUBLIC"], +) \ No newline at end of file diff --git a/extensions/cratespro/common/database/Cargo.toml b/extensions/cratespro/common/database/Cargo.toml new file mode 100644 index 000000000..664efdedc --- /dev/null +++ b/extensions/cratespro/common/database/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "database" +version = "0.1.0" +edition = "2021" + + +[lib] +name = "database" +path = "src/lib.rs" + + +[dependencies] +entity = { workspace = true } +model = { workspace = true } + +tracing = "0.1" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sea-orm = { version = "1.1", features = ["sqlx-postgres"] } +chrono = "0.4" +futures = "0.3" diff --git a/extensions/cratespro/common/database/entity/BUCK b/extensions/cratespro/common/database/entity/BUCK new file mode 100644 index 000000000..922c8ac14 --- /dev/null +++ b/extensions/cratespro/common/database/entity/BUCK @@ -0,0 +1,38 @@ +rust_library( + name = "entity", + srcs = [ + "src/prelude.rs", + "src/mod.rs", + "src/contributor_location.rs", + "src/crate_owners.rs", + "src/crate_users.rs", + "src/crates.rs", + "src/crates_info.rs", + "src/cves.rs", + "src/dependency_cache.rs", + "src/dependent_cache.rs", + "src/github_sync_status.rs", + "src/github_user.rs", + "src/graph_info.rs", + "src/license.rs", + "src/program_dependencies.rs", + "src/program_versions.rs", + "src/programs.rs", + "src/repo_sync_result.rs", + "src/repository_contributor.rs", + "src/rustsec_info.rs", + "src/rustsecs.rs", + "src/sea_orm_active_enums.rs", + "src/senseleak_res.rs", + "src/userloginfo.rs", + "src/version_info.rs", + ], + crate_root = "src/mod.rs", + edition = "2021", + deps = [ + "//third-party:serde", + "//third-party:serde_json", + "//third-party:sea-orm", + ], + visibility = ["PUBLIC"], +) diff --git a/extensions/cratespro/common/database/entity/Cargo.toml b/extensions/cratespro/common/database/entity/Cargo.toml new file mode 100644 index 000000000..5ce4b80cc --- /dev/null +++ b/extensions/cratespro/common/database/entity/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "entity" +version = "0.1.0" +edition = "2021" + +[lib] +name = "entity" +path = "src/mod.rs" + +[dependencies] +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sea-orm = { version = "1.1", features = ["sqlx-postgres"] } diff --git a/extensions/cratespro/common/database/entity/src/contributor_location.rs b/extensions/cratespro/common/database/entity/src/contributor_location.rs new file mode 100644 index 000000000..32a100cc8 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/contributor_location.rs @@ -0,0 +1,21 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "contributor_location")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub repository_id: Uuid, + pub user_id: i32, + pub is_from_china: bool, + pub common_timezone: Option, + pub analyzed_at: DateTime, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/crate_owners.rs b/extensions/cratespro/common/database/entity/src/crate_owners.rs new file mode 100644 index 000000000..4d817aed4 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/crate_owners.rs @@ -0,0 +1,54 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "crate_owners")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub crate_id: i32, + #[sea_orm(primary_key, auto_increment = false)] + pub owner_id: i32, + pub created_at: DateTimeWithTimeZone, + pub created_by: Option, + pub deleted: bool, + pub updated_at: DateTimeWithTimeZone, + #[sea_orm(primary_key, auto_increment = false)] + pub owner_kind: i32, + pub email_notifications: bool, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::crate_users::Entity", + from = "Column::CreatedBy", + to = "super::crate_users::Column::Id", + on_update = "NoAction", + on_delete = "NoAction" + )] + CrateUsers, + #[sea_orm( + belongs_to = "super::crates::Entity", + from = "Column::CrateId", + to = "super::crates::Column::Id", + on_update = "NoAction", + on_delete = "Cascade" + )] + Crates, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::CrateUsers.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Crates.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/crate_users.rs b/extensions/cratespro/common/database/entity/src/crate_users.rs new file mode 100644 index 000000000..b45bd9429 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/crate_users.rs @@ -0,0 +1,35 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "crate_users")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub gh_access_token: String, + pub gh_login: String, + pub name: Option, + pub gh_avatar: Option, + #[sea_orm(unique)] + pub gh_id: i32, + pub account_lock_reason: Option, + pub account_lock_until: Option, + pub is_admin: bool, + pub publish_notifications: bool, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm(has_many = "super::crate_owners::Entity")] + CrateOwners, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::CrateOwners.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/crates.rs b/extensions/cratespro/common/database/entity/src/crates.rs new file mode 100644 index 000000000..23230d1e6 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/crates.rs @@ -0,0 +1,35 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "crates")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub id: i32, + pub name: String, + pub updated_at: DateTimeWithTimeZone, + pub created_at: DateTimeWithTimeZone, + pub description: Option, + pub homepage: Option, + pub documentation: Option, + pub readme: Option, + pub repository: Option, + pub max_upload_size: Option, + pub max_features: Option, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm(has_many = "super::crate_owners::Entity")] + CrateOwners, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::CrateOwners.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/crates_info.rs b/extensions/cratespro/common/database/entity/src/crates_info.rs new file mode 100644 index 000000000..7919ed879 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/crates_info.rs @@ -0,0 +1,34 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "crates_info")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub description: String, + pub direct_dependency: Option, + pub indirect_dependency: Option, + pub direct_dependent: Option, + pub indirect_dependent: Option, + #[sea_orm(column_type = "Text")] + pub cves: String, + #[sea_orm(column_type = "Text")] + pub dep_cves: String, + #[sea_orm(column_type = "Text")] + pub versions: String, + #[sea_orm(column_type = "Text")] + pub license: String, + #[sea_orm(column_type = "Text")] + pub github_url: String, + #[sea_orm(column_type = "Text")] + pub doc_url: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/cves.rs b/extensions/cratespro/common/database/entity/src/cves.rs new file mode 100644 index 000000000..3aebe759c --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/cves.rs @@ -0,0 +1,21 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "cves")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub cve_id: String, + pub name: Option, + pub start_version: Option, + pub end_version: Option, + #[sea_orm(column_type = "Text", nullable)] + pub description: Option, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/dependency_cache.rs b/extensions/cratespro/common/database/entity/src/dependency_cache.rs new file mode 100644 index 000000000..293f20b08 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/dependency_cache.rs @@ -0,0 +1,20 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "dependency_cache")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + pub direct_count: Option, + pub indirect_count: Option, + #[sea_orm(column_type = "Text")] + pub dependency: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/dependent_cache.rs b/extensions/cratespro/common/database/entity/src/dependent_cache.rs new file mode 100644 index 000000000..3c339ff42 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/dependent_cache.rs @@ -0,0 +1,20 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "dependent_cache")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + pub direct_count: Option, + pub indirect_count: Option, + #[sea_orm(column_type = "Text")] + pub dependent: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/github_sync_status.rs b/extensions/cratespro/common/database/entity/src/github_sync_status.rs new file mode 100644 index 000000000..80fd245af --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/github_sync_status.rs @@ -0,0 +1,19 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "github_sync_status")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub start_date: String, + pub end_date: String, + pub sync_result: bool, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/github_user.rs b/extensions/cratespro/common/database/entity/src/github_user.rs new file mode 100644 index 000000000..92af9115d --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/github_user.rs @@ -0,0 +1,35 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "github_user")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + #[sea_orm(unique)] + pub github_id: i64, + pub login: String, + pub name: Option, + pub email: Option, + #[sea_orm(column_type = "Text", nullable)] + pub avatar_url: Option, + pub company: Option, + pub location: Option, + #[sea_orm(column_type = "Text", nullable)] + pub bio: Option, + pub public_repos: Option, + pub followers: Option, + pub following: Option, + pub created_at: DateTime, + pub updated_at: DateTime, + pub inserted_at: DateTime, + pub updated_at_local: DateTime, + pub commit_email: Option, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/graph_info.rs b/extensions/cratespro/common/database/entity/src/graph_info.rs new file mode 100644 index 000000000..359cc9ab5 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/graph_info.rs @@ -0,0 +1,18 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "graph_info")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub graph: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/license.rs b/extensions/cratespro/common/database/entity/src/license.rs new file mode 100644 index 000000000..cc0b23b44 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/license.rs @@ -0,0 +1,22 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "license")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub program_id: String, + #[sea_orm(column_type = "Text")] + pub program_name: String, + #[sea_orm(column_type = "Text")] + pub program_namespace: String, + #[sea_orm(column_type = "Text")] + pub license: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/mod.rs b/extensions/cratespro/common/database/entity/src/mod.rs new file mode 100644 index 000000000..e7d6d30b2 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/mod.rs @@ -0,0 +1,27 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +pub mod prelude; + +pub mod contributor_location; +pub mod crate_owners; +pub mod crate_users; +pub mod crates; +pub mod crates_info; +pub mod cves; +pub mod dependency_cache; +pub mod dependent_cache; +pub mod github_sync_status; +pub mod github_user; +pub mod graph_info; +pub mod license; +pub mod program_dependencies; +pub mod program_versions; +pub mod programs; +pub mod repo_sync_result; +pub mod repository_contributor; +pub mod rustsec_info; +pub mod rustsecs; +pub mod sea_orm_active_enums; +pub mod senseleak_res; +pub mod userloginfo; +pub mod version_info; diff --git a/extensions/cratespro/common/database/entity/src/prelude.rs b/extensions/cratespro/common/database/entity/src/prelude.rs new file mode 100644 index 000000000..37e71bc09 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/prelude.rs @@ -0,0 +1,24 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +pub use super::contributor_location::Entity as ContributorLocation; +pub use super::crate_owners::Entity as CrateOwners; +pub use super::crate_users::Entity as CrateUsers; +pub use super::crates::Entity as Crates; +pub use super::crates_info::Entity as CratesInfo; +pub use super::cves::Entity as Cves; +pub use super::dependency_cache::Entity as DependencyCache; +pub use super::dependent_cache::Entity as DependentCache; +pub use super::github_sync_status::Entity as GithubSyncStatus; +pub use super::github_user::Entity as GithubUser; +pub use super::graph_info::Entity as GraphInfo; +pub use super::license::Entity as License; +pub use super::program_dependencies::Entity as ProgramDependencies; +pub use super::program_versions::Entity as ProgramVersions; +pub use super::programs::Entity as Programs; +pub use super::repo_sync_result::Entity as RepoSyncResult; +pub use super::repository_contributor::Entity as RepositoryContributor; +pub use super::rustsec_info::Entity as RustsecInfo; +pub use super::rustsecs::Entity as Rustsecs; +pub use super::senseleak_res::Entity as SenseleakRes; +pub use super::userloginfo::Entity as Userloginfo; +pub use super::version_info::Entity as VersionInfo; diff --git a/extensions/cratespro/common/database/entity/src/program_dependencies.rs b/extensions/cratespro/common/database/entity/src/program_dependencies.rs new file mode 100644 index 000000000..f8a516094 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/program_dependencies.rs @@ -0,0 +1,20 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "program_dependencies")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub name_and_version: String, + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub dependency_name: String, + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub dependency_version: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/program_versions.rs b/extensions/cratespro/common/database/entity/src/program_versions.rs new file mode 100644 index 000000000..c3e38da8b --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/program_versions.rs @@ -0,0 +1,23 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "program_versions")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub id: Uuid, + pub name: String, + pub version: String, + #[sea_orm(column_type = "Text", nullable)] + pub documentation: Option, + #[sea_orm(column_type = "Text")] + pub version_type: String, + pub created_at: DateTime, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/programs.rs b/extensions/cratespro/common/database/entity/src/programs.rs new file mode 100644 index 000000000..f73b9effd --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/programs.rs @@ -0,0 +1,34 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "programs")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub id: Uuid, + pub name: String, + #[sea_orm(column_type = "Text")] + pub description: String, + pub namespace: String, + pub max_version: String, + #[sea_orm(column_type = "Text")] + pub github_url: String, + #[sea_orm(column_type = "Text")] + pub mega_url: String, + #[sea_orm(column_type = "Text")] + pub doc_url: String, + #[sea_orm(column_type = "Text")] + pub program_type: String, + pub downloads: i64, + #[sea_orm(column_type = "Text")] + pub cratesio: String, + pub github_analyzed: bool, + pub in_cratesio: bool, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/repo_sync_result.rs b/extensions/cratespro/common/database/entity/src/repo_sync_result.rs new file mode 100644 index 000000000..dcab1760f --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/repo_sync_result.rs @@ -0,0 +1,32 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use super::sea_orm_active_enums::CrateTypeEnum; +use super::sea_orm_active_enums::SyncStatusEnum; +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "repo_sync_result")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + #[sea_orm(unique)] + pub crate_name: String, + #[sea_orm(column_type = "Text", nullable)] + pub github_url: Option, + #[sea_orm(column_type = "Text")] + pub mega_url: String, + pub status: SyncStatusEnum, + pub crate_type: CrateTypeEnum, + #[sea_orm(column_type = "Text", nullable)] + pub err_message: Option, + #[sea_orm(column_type = "Text")] + pub version: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/repository_contributor.rs b/extensions/cratespro/common/database/entity/src/repository_contributor.rs new file mode 100644 index 000000000..097894b22 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/repository_contributor.rs @@ -0,0 +1,21 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "repository_contributor")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub repository_id: Uuid, + pub user_id: i32, + pub contributions: i32, + pub inserted_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/rustsec_info.rs b/extensions/cratespro/common/database/entity/src/rustsec_info.rs new file mode 100644 index 000000000..8011b15fc --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/rustsec_info.rs @@ -0,0 +1,38 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "rustsec_info")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub subtitle: String, + #[sea_orm(column_type = "Text")] + pub reported: String, + #[sea_orm(column_type = "Text")] + pub issued: String, + #[sea_orm(column_type = "Text")] + pub package: String, + #[sea_orm(column_type = "Text")] + pub r#type: String, + #[sea_orm(column_type = "Text")] + pub keywords: String, + #[sea_orm(column_type = "Text")] + pub aliases: String, + #[sea_orm(column_type = "Text")] + pub reference: String, + #[sea_orm(column_type = "Text")] + pub patched: String, + #[sea_orm(column_type = "Text")] + pub unaffected: String, + #[sea_orm(column_type = "Text")] + pub description: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/rustsecs.rs b/extensions/cratespro/common/database/entity/src/rustsecs.rs new file mode 100644 index 000000000..9d30602d3 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/rustsecs.rs @@ -0,0 +1,28 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "rustsecs")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub publish_time: String, + #[sea_orm(column_type = "Text")] + pub cratename: String, + #[sea_orm(column_type = "Text")] + pub patched: String, + #[sea_orm(column_type = "Text")] + pub aliases: String, + #[sea_orm(column_type = "Text")] + pub small_desc: String, + #[sea_orm(column_type = "Text")] + pub description: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/sea_orm_active_enums.rs b/extensions/cratespro/common/database/entity/src/sea_orm_active_enums.rs new file mode 100644 index 000000000..fe20bd7f0 --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/sea_orm_active_enums.rs @@ -0,0 +1,27 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Serialize, Deserialize)] +#[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "crate_type_enum")] +pub enum CrateTypeEnum { + #[sea_orm(string_value = "lib")] + Lib, + #[sea_orm(string_value = "application")] + Application, +} +#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Serialize, Deserialize)] +#[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "sync_status_enum")] +pub enum SyncStatusEnum { + #[sea_orm(string_value = "syncing")] + Syncing, + #[sea_orm(string_value = "succeed")] + Succeed, + #[sea_orm(string_value = "failed")] + Failed, + #[sea_orm(string_value = "analysing")] + Analysing, + #[sea_orm(string_value = "analysed")] + Analysed, +} diff --git a/extensions/cratespro/common/database/entity/src/senseleak_res.rs b/extensions/cratespro/common/database/entity/src/senseleak_res.rs new file mode 100644 index 000000000..bcd7dd54f --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/senseleak_res.rs @@ -0,0 +1,18 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "senseleak_res")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub res: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/userloginfo.rs b/extensions/cratespro/common/database/entity/src/userloginfo.rs new file mode 100644 index 000000000..9d3c415ca --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/userloginfo.rs @@ -0,0 +1,22 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "userloginfo")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub image: String, + #[sea_orm(column_type = "Text")] + pub name: String, + #[sea_orm(column_type = "Text")] + pub expires: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/entity/src/version_info.rs b/extensions/cratespro/common/database/entity/src/version_info.rs new file mode 100644 index 000000000..45031578e --- /dev/null +++ b/extensions/cratespro/common/database/entity/src/version_info.rs @@ -0,0 +1,18 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.8 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "version_info")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false, column_type = "Text")] + pub id: String, + #[sea_orm(column_type = "Text")] + pub versions: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/extensions/cratespro/common/database/src/lib.rs b/extensions/cratespro/common/database/src/lib.rs new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/extensions/cratespro/common/database/src/lib.rs @@ -0,0 +1 @@ + diff --git a/extensions/cratespro/common/model/BUCK b/extensions/cratespro/common/model/BUCK new file mode 100644 index 000000000..34eb48395 --- /dev/null +++ b/extensions/cratespro/common/model/BUCK @@ -0,0 +1,21 @@ +rust_library( + name = "model", + srcs = [ + "src/general_model.rs", + "src/github.rs", + "src/lib.rs", + "src/repo_sync_model.rs", + "src/tugraph_model.rs", + ], + crate_root = "src/lib.rs", + edition = "2021", + deps = [ + "//project/crates-pro:entity", + "//third-party:chrono", + "//third-party:sea-orm", + "//third-party:serde", + "//third-party:serde_json", + "//third-party:utoipa", + ], + visibility = ["PUBLIC"], +) diff --git a/extensions/cratespro/common/model/Cargo.toml b/extensions/cratespro/common/model/Cargo.toml new file mode 100644 index 000000000..5d983d105 --- /dev/null +++ b/extensions/cratespro/common/model/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "model" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +entity = { workspace = true } + +# third-party (第三方依赖, 不写具体版本号, 具体版本只在根目录 Cargo.toml 中出现) +chrono = "0.4" +sea-orm = "1.1" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +utoipa = { version = "5.3", features = ["actix_extras"] } diff --git a/extensions/cratespro/common/model/src/general_model.rs b/extensions/cratespro/common/model/src/general_model.rs new file mode 100644 index 000000000..d056806f2 --- /dev/null +++ b/extensions/cratespro/common/model/src/general_model.rs @@ -0,0 +1,50 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize, Hash, PartialEq, Eq, PartialOrd, Ord, Clone)] +pub struct Version { + pub name: String, + pub version: String, +} + +impl Version { + pub fn new(name: &str, version: &str) -> Self { + Self { + name: name.to_string(), + version: version.to_string(), + } + } +} + +#[derive(Debug, Deserialize, Serialize, Hash, PartialEq, Eq, PartialOrd, Ord, Clone)] +pub struct VersionWithTag { + pub name: String, + pub version: String, + pub git_url: String, + pub tag: String, +} + +impl VersionWithTag { + pub fn new(name: &str, version: &str, git_url: &str, tag: &str) -> Self { + Self { + name: name.to_string(), + version: version.to_string(), + git_url: git_url.to_string(), + tag: tag.to_string(), + } + } +} + +#[derive(Debug, Deserialize, Serialize, Hash, PartialEq, Eq, PartialOrd, Ord, Clone)] +pub struct Program { + pub name: String, + pub mega_url: String, +} + +impl Program { + pub fn new(name: &str, mega_url: &str) -> Self { + Self { + name: name.to_string(), + mega_url: mega_url.to_string(), + } + } +} diff --git a/extensions/cratespro/common/model/src/github.rs b/extensions/cratespro/common/model/src/github.rs new file mode 100644 index 000000000..b322358d5 --- /dev/null +++ b/extensions/cratespro/common/model/src/github.rs @@ -0,0 +1,192 @@ +use chrono::{DateTime, Utc}; +use entity::{contributor_location, github_user}; +use sea_orm::ActiveValue::{NotSet, Set}; +use serde::{Deserialize, Serialize}; + +// GitHub用户信息结构 +#[derive(Debug, Serialize, Deserialize, Clone, Default)] +pub struct GitHubUser { + pub id: i64, + pub login: String, + pub avatar_url: Option, + pub name: Option, + pub email: Option, + pub company: Option, + pub location: Option, + pub bio: Option, + pub public_repos: Option, + pub followers: Option, + pub following: Option, + pub created_at: DateTime, + pub updated_at: DateTime, + #[serde(rename = "type")] + pub user_type: String, +} + +impl GitHubUser { + pub fn is_bot(&self) -> bool { + self.user_type == "Bot" + } +} + +// 转换函数，用于将GitHub API返回的用户转换为数据库模型 +impl From for github_user::ActiveModel { + fn from(user: GitHubUser) -> Self { + let now = chrono::Utc::now().naive_utc(); + + Self { + id: NotSet, + github_id: Set(user.id), + login: Set(user.login), + name: Set(user.name), + email: Set(user.email), + avatar_url: Set(user.avatar_url), + company: Set(user.company), + location: Set(user.location), + bio: Set(user.bio), + public_repos: Set(user.public_repos), + followers: Set(user.followers), + following: Set(user.following), + created_at: Set(user.created_at.naive_utc()), + updated_at: Set(user.updated_at.naive_utc()), + inserted_at: Set(now), + updated_at_local: Set(now), + commit_email: NotSet, + } + } +} + +impl From for GitHubUser { + fn from(value: github_user::Model) -> Self { + Self { + id: value.github_id, + login: value.login, + avatar_url: value.avatar_url, + name: value.name, + email: value.email, + company: value.company, + location: value.location, + bio: value.bio, + public_repos: value.public_repos, + followers: value.followers, + following: value.following, + user_type: "User".to_owned(), + created_at: DateTime::::from_naive_utc_and_offset(value.created_at, Utc), + updated_at: DateTime::::from_naive_utc_and_offset(value.updated_at, Utc), + } + } +} + +pub struct AnalyzedUser { + pub user_id: i32, + pub github_id: i64, + pub login: String, + pub profile_email: Option, + pub commit_email: Option, +} + +impl From for AnalyzedUser { + fn from(value: github_user::Model) -> Self { + Self { + user_id: value.id, + github_id: value.github_id, + login: value.login, + profile_email: value.email, + commit_email: value.commit_email, + } + } +} +// 贡献者信息结构 +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct Contributor { + pub id: i64, + pub login: String, + pub avatar_url: String, + pub contributions: i32, + pub email: Option, +} + +// 贡献者分析结果 +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ContributorAnalysis { + pub has_china_timezone: bool, + pub common_timezone: String, +} + +// 转换函数，将分析结果转换为数据库模型 +impl From<&ContributorAnalysis> for contributor_location::ActiveModel { + fn from(analysis: &ContributorAnalysis) -> Self { + let now = chrono::Utc::now().naive_utc(); + + Self { + id: NotSet, + is_from_china: Set(analysis.common_timezone == "+08:00"), + common_timezone: Set(Some(analysis.common_timezone.clone())), + analyzed_at: Set(now), + ..Default::default() + } + } +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Repository { + pub name: String, + pub url: String, + pub created_at: String, +} +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct GraphQLResponse { + pub data: Option, +} +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SearchData { + pub search: SearchResult, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SearchResult { + pub edges: Vec, + pub page_info: PageInfo, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Edge { + pub node: Repository, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct PageInfo { + pub end_cursor: Option, + pub has_next_page: bool, +} + +// 解析提交数据 +#[derive(Debug, Deserialize)] +pub struct CommitAuthor { + pub login: String, + pub id: i64, + pub avatar_url: String, +} + +#[derive(Debug, Deserialize)] +pub struct CommitInfo { + pub _author: Option, + pub email: Option, +} + +#[derive(Debug, Deserialize)] +pub struct CommitDetail { + pub author: Option, +} + +#[derive(Debug, Deserialize)] +pub struct CommitData { + pub author: Option, + pub commit: CommitDetail, +} diff --git a/extensions/cratespro/common/model/src/lib.rs b/extensions/cratespro/common/model/src/lib.rs new file mode 100644 index 000000000..ba8558f79 --- /dev/null +++ b/extensions/cratespro/common/model/src/lib.rs @@ -0,0 +1,4 @@ +pub mod general_model; +pub mod repo_sync_model; +pub mod tugraph_model; +pub mod github; \ No newline at end of file diff --git a/extensions/cratespro/common/model/src/repo_sync_model.rs b/extensions/cratespro/common/model/src/repo_sync_model.rs new file mode 100644 index 000000000..3265f5eef --- /dev/null +++ b/extensions/cratespro/common/model/src/repo_sync_model.rs @@ -0,0 +1,58 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct Model { + pub id: i32, + pub crate_name: String, + pub github_url: Option, + pub mega_url: String, + #[serde(rename = "crate_type")] + pub crate_type: CrateType, + pub status: RepoSyncStatus, + pub err_message: Option, + //pub created_at: DateTime, + //pub updated_at: DateTime, + //pub message_kind: MessageKind, + //pub sourc_of_data: SourceOfData, +} + +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub enum CrateType { + Lib, + Application, +} + +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub enum RepoSyncStatus { + Syncing, + Succeed, + Failed, + Analysing, + Analysed, +} + +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub enum MessageKind { + Mega, + User, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "PascalCase")] +pub enum SourceOfData { + Cratesio, + Github, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct MessageModel { + pub db_model: Model, // 包装数据库 Model + pub message_kind: MessageKind, + pub source_of_data: SourceOfData, + pub timestamp: DateTime, // 消息发送时的时间戳 + pub extra_field: Option, // 可以添加额外字段 +} diff --git a/extensions/cratespro/common/model/src/tugraph_model.rs b/extensions/cratespro/common/model/src/tugraph_model.rs new file mode 100644 index 000000000..0b7084010 --- /dev/null +++ b/extensions/cratespro/common/model/src/tugraph_model.rs @@ -0,0 +1,233 @@ +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +/// Represents detailed information about a Rust crate. +/// +/// This structure includes metadata fields that describe a crate, such as its name, +/// current version, description, and various URLs related to its documentation, +/// repository, and license, along with a count of its dependencies. +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct Program { + /// The id + pub id: String, + /// The name of the crate. + pub name: String, + /// An optional description of the crate. + pub description: Option, + /// The namespace of the crate, such tokio-rs/tokio + pub namespace: Option, + /// The current version of the crate. + pub max_version: Option, + /// An optional URL pointing to the crate's source code repository. + pub github_url: Option, + /// mega URL + pub mega_url: Option, + /// An optional URL pointing to the crate's documentation. + pub doc_url: Option, +} + +impl Program { + /// Constructs a new `CrateInfo`. + #[allow(clippy::too_many_arguments)] + pub fn new( + id: String, + name: String, + description: Option, + namespace: Option, + max_version: Option, + github_url: Option, + mega_url: Option, + doc_url: Option, + ) -> Self { + Self { + id, + name, + description, + namespace, + max_version, + github_url, + mega_url, + doc_url, + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub enum UProgram { + Library(Library), + Application(Application), +} +impl CrateType2Idx for UProgram { + fn index(&self) -> usize { + match self { + Self::Library(_) => 0, + Self::Application(_) => 1, + } + } +} + +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct Library { + /// The id + pub id: String, + /// The name of the library. + pub name: String, + /// The number of downloads. + pub downloads: i64, + /// An optional URL pointing to the library's crates.io page. + pub cratesio: Option, +} + +impl Library { + pub fn new(id: &str, name: &str, downloads: i64, cratesio: Option<&str>) -> Self { + Library { + id: id.to_string(), + name: name.to_string(), + downloads, + cratesio: cratesio.map(|s| s.to_string()), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct Application { + /// The id + pub id: String, + /// The name of the application. + pub name: String, +} + +impl Application { + /// Creates a new `Application` instance. + pub fn new(id: String, name: &str) -> Self { + Self { + id, + name: name.to_string(), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] +pub enum UVersion { + LibraryVersion(LibraryVersion), + ApplicationVersion(ApplicationVersion), +} + +impl CrateType2Idx for UVersion { + fn index(&self) -> usize { + match self { + Self::LibraryVersion(_) => 0, + Self::ApplicationVersion(_) => 1, + } + } +} + +impl UVersion { + pub fn get_name_and_version(&self) -> String { + match self { + UVersion::LibraryVersion(l) => l.name_and_version.clone(), + UVersion::ApplicationVersion(a) => a.name_and_version.clone(), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct LibraryVersion { + /// The id + pub id: String, + /// The name and version of the crate. + pub name_and_version: String, + /// The name of the library. + pub name: String, + /// The version of the library. + pub version: String, + /// The documentation URL for the library. + pub documentation: String, +} + +impl LibraryVersion { + /// Create a new `LibraryVersion` instance. + pub fn new(id: String, name: &str, version: &str, documentation: &str) -> Self { + LibraryVersion { + id, + name_and_version: name.to_string() + "/" + version, + name: name.to_string(), + version: version.to_string(), + documentation: documentation.to_string(), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct ApplicationVersion { + /// The id + pub id: String, + /// The name and version of the crate. + pub name_and_version: String, + /// The name + pub name: String, + /// The version + pub version: String, +} + +impl ApplicationVersion { + /// Create a new `ApplicationVersion` instance. + pub fn new(id: String, name: String, version: String) -> Self { + ApplicationVersion { + id, + name_and_version: name.to_string() + "/" + &version, + name, + version, + } + } +} + +#[derive(Serialize, Deserialize, Debug, Default, Clone, ToSchema)] +pub struct Version { + /// The name and version of the crate. + pub name_and_version: String, +} + +impl Version { + /// Creates a new `Version` instance. + pub fn new(name_and_version: &str) -> Self { + Version { + name_and_version: name_and_version.to_string(), + } + } +} + +#[allow(non_snake_case)] +#[derive(Serialize, Deserialize, Debug, Default, Clone)] +pub struct HasType { + pub SRC_ID: String, + pub DST_ID: String, +} + +#[allow(non_snake_case)] +#[derive(Serialize, Deserialize, Debug, Default, Clone)] +pub struct HasVersion { + pub SRC_ID: String, + pub DST_ID: String, +} + +#[allow(non_snake_case)] +#[derive(Serialize, Deserialize, Debug, Default, Clone)] +pub struct HasDepVersion { + pub SRC_ID: String, + pub DST_ID: String, +} + +#[allow(non_snake_case)] +#[derive(Serialize, Deserialize, Debug, Default, Clone)] +pub struct DependsOn { + pub SRC_ID: String, + pub DST_ID: String, +} + +pub trait CrateType2Idx { + fn index(&self) -> usize; + fn is_library(&self) -> bool { + self.index() == 0 + } +} diff --git a/extensions/cratespro/common/repo_import/BUCK b/extensions/cratespro/common/repo_import/BUCK new file mode 100644 index 000000000..71d0a9206 --- /dev/null +++ b/extensions/cratespro/common/repo_import/BUCK @@ -0,0 +1,42 @@ +rust_library( + name = "repo_import", + srcs = [ + "src/crate_info.rs", + "src/git.rs", + "src/kafka_handler.rs", + "src/lib.rs", + "src/utils.rs", + "src/version_info.rs", + ], + crate_root = "src/lib.rs", + edition = "2021", + deps = [ + "//project/crates-pro:model", + "//third-party:bincode", + "//third-party:cargo_metadata", + "//third-party:csv", + "//third-party:dotenvy", + "//third-party:futures", + "//third-party:git2", + "//third-party:lazy_static", + "//third-party:log", + "//third-party:once_cell", + "//third-party:pretty_env_logger", + "//third-party:rayon", + "//third-party:rdkafka", + "//third-party:semver", + "//third-party:serde", + "//third-party:serde_json", + "//third-party:sqlx", + "//third-party:ssh2", + "//third-party:tempfile", + "//third-party:tokio", + "//third-party:toml", + "//third-party:tracing", + "//third-party:tracing-subscriber", + "//third-party:url", + "//third-party:uuid", + "//third-party:walkdir", + ], + visibility = ["PUBLIC"], +) diff --git a/extensions/cratespro/common/repo_import/Cargo.toml b/extensions/cratespro/common/repo_import/Cargo.toml new file mode 100644 index 000000000..dc860b268 --- /dev/null +++ b/extensions/cratespro/common/repo_import/Cargo.toml @@ -0,0 +1,44 @@ +[package] +name = "repo_import" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +model = { workspace = true } + +# third-party (第三方依赖, 不写具体版本号, 具体版本只在根目录 Cargo.toml 中出现) +bincode = "1.3" +cargo_metadata = "0.18" +csv = "1.3" +dotenvy = "0.15" +futures = "0.3" +git2 = "0.20" +lazy_static = "1.5" +log = "0.4" +once_cell = "1.21" +pretty_env_logger = "0.5" +rayon = "1.10" +rdkafka = { version = "0.37", features = ["cmake-build"] } +semver = "1.0" +serde = "1.0" +serde_json = "1.0" +sqlx = { version = "0.8", features = ["runtime-tokio", "postgres"] } +ssh2 = "0.9" +tempfile = "3.19" +tokio = { version = "1.44", features = ["full"] } +toml = "0.8" +tracing = "0.1" +tracing-subscriber = "0.3" +url = "2.5" +uuid = { version = "1.16", features = [ + "v4", # Lets you generate random UUIDs + "fast-rng", # Use a faster (but still sufficiently random) RNG + "macro-diagnostics", # Enable better diagnostics for compile-time UUIDs +] } +walkdir = "2.5" + +[dev-dependencies] +# 开发依赖可以不写在根目录 Cargo.toml 中 +mockall = "0.13" diff --git a/extensions/cratespro/common/repo_import/src/crate_info.rs b/extensions/cratespro/common/repo_import/src/crate_info.rs new file mode 100644 index 000000000..8a6574bb4 --- /dev/null +++ b/extensions/cratespro/common/repo_import/src/crate_info.rs @@ -0,0 +1,215 @@ +use crate::{ + utils::{get_namespace_by_repo_path, insert_program_by_name,extract_middle_path}, + Licenses, +}; +use model::tugraph_model::{Application, HasType, Library, Program, UProgram}; +use std::{ + fs, + path::{Path, PathBuf}, +}; +use uuid::Uuid; +use walkdir::WalkDir; + +// Given a project path, parse the metadata +pub(crate) async fn extract_info_local( + local_repo_path: PathBuf, + git_url: String, + lic: &mut Vec, +) -> Vec<(Program, HasType, UProgram)> { + let mut res = vec![]; + + // walk the directories of the project + for entry in WalkDir::new(local_repo_path.clone()) + .into_iter() + .filter_map(|x| x.ok()) + { + let entry_path = entry.path(); + + // if entry is Cargo.toml, ... + if entry_path.file_name().and_then(|n| n.to_str()) == Some("Cargo.toml") { + tracing::trace!("entry_path: {:?}", entry_path); + let crate_name_result = parse_crate_name(entry_path).await; + match crate_name_result { + Ok(name) => { + tracing::trace!("package name: {}", name); + let islib_result = is_crate_lib( + entry_path + .to_str() + .unwrap() + .strip_suffix("Cargo.toml") + .unwrap(), + ) + .await; + let islib = match islib_result { + Ok(islib) => islib, + Err(e) => { + tracing::error!("parse error: {}", e); + continue; + } + }; + + tracing::debug!("Found Crate: {}, islib: {}", name, islib); + let id = Uuid::new_v4().to_string(); + let mut program = from_cargo_toml( + local_repo_path.clone(), + entry_path.to_path_buf(), + &id, + lic, + ) + .await + .unwrap(); + + let real_git_url = extract_middle_path(&git_url).expect("Failed to parse middle_path"); + program.mega_url = Some(real_git_url.clone()); + let uprogram = if islib { + UProgram::Library(Library::new(&id.to_string(), &name, -1, None)) + } else { + UProgram::Application(Application::new(id.to_string(), &name)) + }; + + let has_type = HasType { + SRC_ID: program.id.clone(), + DST_ID: program.id.clone(), + }; + + tracing::trace!( + "program: {:?}, has_type: {:?}, uprogram: {:?}", + program, + has_type, + uprogram + ); + insert_program_by_name(name.clone(), (program.clone(), uprogram.clone())); + + res.push((program, has_type, uprogram)); + } + Err(e) => tracing::warn!("Error parsing name {}: {}", entry_path.display(), e), + } + } + } + + res +} + +async fn parse_crate_name(path: &Path) -> Result { + let content = fs::read_to_string(path).map_err(|e| e.to_string())?; + tracing::info!("finish get content:\n{}",content); + let value = content.parse::().map_err(|e| e.to_string())?; + tracing::info!("finish get value"); + // a package name, no matter lib or bin + let package_name = value + .get("package") + .and_then(|p| p.get("name")) + .and_then(|n| n.as_str()) + .ok_or("Failed to find package name, it is a workspace")? + .to_owned(); + + Ok(package_name) +} + +async fn is_crate_lib(crate_path: &str) -> Result { + let cargo_toml_path = Path::new(crate_path).join("Cargo.toml"); + let cargo_toml_content = fs::read_to_string(cargo_toml_path) + .map_err(|e| format!("Failed to read Cargo.toml: {e}"))?; + + let cargo_toml: toml::Value = cargo_toml_content + .parse::() + .map_err(|e| format!("Failed to parse Cargo.toml: {e}"))?; + + // 优先检查 Cargo.toml 中的 '[lib]' 和 '[[bin]]' + let has_lib_in_toml = cargo_toml.get("lib").is_some(); + let has_bin_in_toml = cargo_toml + .get("bin") + .is_some_and(|bins| bins.as_array().is_some_and(|b| !b.is_empty())); + + if has_lib_in_toml || has_bin_in_toml { + return Ok(has_lib_in_toml && !has_bin_in_toml); + } + + // 如果 Cargo.toml 中无明显标识，退回到检查文件 + let lib_rs_exists = Path::new(crate_path).join("src/lib.rs").exists(); + let main_rs_exists = Path::new(crate_path).join("src/main.rs").exists(); + + // 如果 'src/lib.rs' 存在，且 'src/main.rs' 不存在，更可能是库 + if lib_rs_exists && !main_rs_exists { + return Ok(true); + } + + // 如果存在 'src/main.rs'，则倾向于不是库 + if main_rs_exists { + return Ok(false); + } + + // 如果没有明显的线索，回退为默认假设不是库 + Ok(false) +} + +async fn from_cargo_toml( + local_repo_path: PathBuf, + cargo_toml_path: PathBuf, + id: &str, + lic: &mut Vec, +) -> Result> { + let content = fs::read_to_string(cargo_toml_path)?; + let parsed = content.parse::()?; + + // 处理description,将多行字符串转换为单行,用\n替换换行符 + let description = parsed["package"] + .get("decription") + .unwrap_or(&toml::Value::String(String::from("None"))) + .as_str() + .map(|desc| desc.replace('\n', "\\n").replace('\r', "")); + + let mut program = Program::new( + id.to_string(), + parsed["package"]["name"] + .as_str() + .unwrap_or_default() + .to_string(), + description, + get_namespace_by_repo_path(local_repo_path.to_str().unwrap()), + None, + parsed["package"] + .get("repository") + .unwrap_or(&toml::Value::String(String::from("None"))) + .as_str() + .map(String::from), + None, + parsed["package"] + .get("documentation") + .unwrap_or(&toml::Value::String(String::from("None"))) + .as_str() + .map(String::from), + ); + let license = parsed["package"] + .get("license") + .unwrap_or(&toml::Value::String(String::from("None"))) + .as_str() + .map(String::from); + let newlicense = Licenses { + program_id: program.id.clone(), + program_name: program.name.clone(), + program_namespace: program.namespace.clone(), + license, + }; + if program.name.is_empty() { + if let Some(ns) = program.namespace.clone() { + let new_ns = ns.as_str(); + let parts: Vec<&str> = new_ns.split('/').collect(); + if parts.len() == 2 { + program.name = parts[1].to_string(); + } + } + } + if let Some(docurl) = program.doc_url.clone() { + if docurl.is_empty() { + program.doc_url = Some("None".to_string()); + } + } + if let Some(githuburl) = program.github_url.clone() { + if githuburl.is_empty() { + program.github_url = Some("None".to_string()); + } + } + lic.push(newlicense); + Ok(program) +} diff --git a/extensions/cratespro/common/repo_import/src/git.rs b/extensions/cratespro/common/repo_import/src/git.rs new file mode 100644 index 000000000..4abd3d4f2 --- /dev/null +++ b/extensions/cratespro/common/repo_import/src/git.rs @@ -0,0 +1,181 @@ +use crate::{utils::{extract_namespace, extract_path_from_segment, insert_namespace_by_repo_path}, ImportContext}; +use git2::{ObjectType, Oid, Repository}; +use std::path::PathBuf; +use url::Url; + +impl ImportContext { + /// clone repo locally + /// 1. Get mega url from postgres + /// 2. Clone git repositories from mega, reserving the namespace as path where they are cloned + pub(crate) async fn clone_a_repo_by_url( + &mut self, + clone_dir: &str, + git_url_base: &str, + git_url_suffix: &str, + ) -> Result { + // mega_url = base + path + //tracing::info!("enter clone_a_repo_by_url"); + let git_url = { + let git_url_base = Url::parse(git_url_base) + .unwrap_or_else(|_| panic!("Failed to parse mega url base: {}", &git_url_base)); + git_url_base + .join(git_url_suffix) + .expect("Failed to join url path") + }; + tracing::info!("clone git_url:{}",git_url.clone()); + // namespace such as tokio-rs/tokio + let namespace = extract_namespace(git_url.as_ref()).expect("Failed to parse URL"); + let behind_dir = extract_path_from_segment(git_url.as_ref(),"crates").expect("Failed to parse behind_dir"); + // The path the repo will be cloned into + let path = PathBuf::from(clone_dir).join(behind_dir.clone()); + //tracing::info!("path:{:?}", path); + if !self.dont_clone { + //tracing::info!("start clone"); + clone(&path, git_url.as_ref()).await?; + //tracing::info!("finish clone"); + } + // finish cloning, store namespace ... + //tracing::info!("start insert"); + insert_namespace_by_repo_path(path.to_str().unwrap().to_string(), namespace.clone()); + //tracing::info!("finish insert"); + tracing::trace!("Finish clone all the repos\n"); + Ok(path) + } +} + +async fn clone(path: &PathBuf, url: &str) -> Result<(), git2::Error> { + if !path.is_dir() { + //tracing::info!("Start cloning repo into {:?} from URL {}", path, url); + Repository::clone(url, path)?; + //tracing::info!("Finish cloning repo into {:?}", path); + } else { + tracing::info!("Directory {:?} is not empty, skipping Clone", path); + } + Ok(()) +} + +// /// Deprecated. +// /// If it migrate from a different system, +// /// the git record will change, and this is the reset function. +/*pub(crate) async fn hard_reset_to_head(repo_path: &PathBuf) -> Result<(), git2::Error> { + let repo = Repository::open(repo_path).unwrap(); + let head = match repo.head() { + Ok(head) => head, + Err(_) => { + tracing::warn!("Repo {:?} does not have ref/heads/master", repo_path); + return Ok(()); + } + }; + let commit = repo.find_commit( + head.target() + .ok_or(git2::Error::from_str("HEAD does not point to a commit"))?, + )?; + + // commit tree + let tree = commit.tree()?; + + // Create CheckoutBuilder, set to force checkout to ensure changes to the working directory + let mut checkout_opts = CheckoutBuilder::new(); + checkout_opts.force(); + + // Correctly convert tree to Object before checking out the + let tree_obj = tree.into_object(); + repo.checkout_tree(&tree_obj as &git2::Object, Some(&mut checkout_opts))?; + tracing::info!( + "Successfully checkout Repo {:?} into ref/heads/master", + repo_path + ); + + Ok(()) +}*/ + +/// return value: (tag_name, tree_id, commit_time) +pub(crate) async fn get_all_git_tags_with_time_sorted( + repo_path: &PathBuf, +) -> Vec<(String, Oid, i64)> { + let mut tags_with_dates = Vec::new(); + + let repo = Repository::open(repo_path).unwrap(); + + // Read the repository to get all tag names + let tags: Vec = repo + .tag_names(None) + .expect("Could not retrieve tags") + .iter() + .flatten() + .map(|tag_name| tag_name.to_string()) + .collect(); + + for tag_name in tags { + let obj = repo + .revparse_single(&("refs/tags/".to_owned() + &tag_name)) + .expect("Couldn't find tag object"); + + // Convert annotated and light-weight tag into commit + let commit = if let Some(tag) = obj.as_tag() { + /*tag.target() + .expect("Couldn't get tag target") + .peel_to_commit() + .expect("Couldn't peel to commit")*/ + let tmp = tag.target().expect("Couldn't get tag target"); + match tmp.peel_to_commit() { + Ok(commit) => commit, + Err(e) => { + tracing::info!("Couldn't peel to commit: {:?}", e); + continue; + } + } + } else if let Some(commit) = obj.as_commit() { + commit.clone() + } else { + //panic!("Error!"); + tracing::info!("Error"); + continue; + }; + + let commit_time = commit.time().seconds(); + + let tree_id = match commit.tree() { + Ok(tree) => tree.id(), + _ => continue, + }; + + tags_with_dates.push((tag_name, tree_id, commit_time)); + } + tags_with_dates.sort_by(|a, b| a.2.cmp(&b.2)); + tags_with_dates +} + +pub(crate) async fn _print_all_tags(repo: &Repository, v: bool) { + let tags = repo.tag_names(None).unwrap(); + + let mut s = "".to_string(); + for tag_name in tags.iter().flatten() { + let tag_ref = repo + .find_reference(&format!("refs/tags/{tag_name}")) + .unwrap(); + + if v { + if let Ok(tag_object) = tag_ref.peel_to_tag() { + // Annotated tag + let target_commit = tag_object.target().unwrap().peel_to_commit().unwrap(); + tracing::debug!( + "Annotated Tag: {}, Commit: {}, Message: {}", + tag_name, + target_commit.id(), + tag_object.message().unwrap_or("No message") + ); + } else { + let commit_object = tag_ref.peel(ObjectType::Commit).unwrap(); + let commit = commit_object + .into_commit() + .expect("Failed to peel into commit"); + tracing::debug!("Lightweight Tag: {}, Commit: {}", tag_name, commit.id()); + } + } else { + s += &format!("{tag_name}, "); + } + } + + tracing::debug!("TAGS {:?} tags: {}", repo.path(), s); +} diff --git a/extensions/cratespro/common/repo_import/src/kafka_handler.rs b/extensions/cratespro/common/repo_import/src/kafka_handler.rs new file mode 100644 index 000000000..ea2520aef --- /dev/null +++ b/extensions/cratespro/common/repo_import/src/kafka_handler.rs @@ -0,0 +1,176 @@ +use rdkafka::config::{ClientConfig, RDKafkaLogLevel}; +use rdkafka::consumer::{BaseConsumer, CommitMode, Consumer, ConsumerContext, Rebalance}; +use rdkafka::error::{KafkaError, KafkaResult}; +use rdkafka::message::{BorrowedMessage, Headers}; +use rdkafka::producer::{BaseProducer, BaseRecord, ProducerContext}; +use rdkafka::util::Timeout; +use rdkafka::{ClientContext, Message, TopicPartitionList}; +use std::env; +use std::process::Command; +use std::time::Duration; + +#[derive(Clone)] +pub struct CustomContext; + +impl ClientContext for CustomContext {} + +impl ConsumerContext for CustomContext { + fn pre_rebalance(&self, _: &BaseConsumer, rebalance: &Rebalance) { + tracing::info!("Pre rebalance {:?}", rebalance); + } + + fn post_rebalance(&self, _: &BaseConsumer, rebalance: &Rebalance) { + tracing::info!("Post rebalance {:?}", rebalance); + } + + fn commit_callback(&self, result: KafkaResult<()>, _offsets: &TopicPartitionList) { + tracing::info!("Committing offsets: {:?}", result); + } +} + +impl ProducerContext for CustomContext { + type DeliveryOpaque = (); + + fn delivery( + &self, + _result: &rdkafka::producer::DeliveryResult, + _delivery_opaque: Self::DeliveryOpaque, + ) { + // match result { + // Ok(delivery) => tracing::info!("Delivered message to {:?}", delivery), + // Err((error, _)) => tracing::error!("Failed to deliver message: {:?}", error), + // } + } +} + +pub enum KafkaHandler { + Consumer(BaseConsumer), + Producer(BaseProducer), +} +impl KafkaHandler { + pub fn new_consumer(brokers: &str, group_id: &str, topic: &str) -> Result { + let context = CustomContext; + + let consumer: BaseConsumer = ClientConfig::new() + .set("group.id", group_id) + .set("bootstrap.servers", brokers) + .set("enable.partition.eof", "false") + .set("session.timeout.ms", "10000") + .set("heartbeat.interval.ms", "1500") + .set("max.poll.interval.ms", "3000000") + .set("enable.auto.commit", "true") + .set("auto.offset.reset", "earliest") + .set_log_level(RDKafkaLogLevel::Debug) + .create_with_context(context)?; + + consumer.subscribe(&[topic])?; + + Ok(KafkaHandler::Consumer(consumer)) + } + + pub fn new_producer(brokers: &str) -> Result { + let context = CustomContext; + + let producer: BaseProducer = ClientConfig::new() + .set("bootstrap.servers", brokers) + .create_with_context(context)?; + + Ok(KafkaHandler::Producer(producer)) + } + + pub async fn consume_once(&'_ self) -> Result, KafkaError> { + if let KafkaHandler::Consumer(consumer) = self { + tracing::debug!("Trying to consume a message"); + + match consumer.poll(Duration::from_secs(0)) { + None => { + //tracing::info!("No message received"); + Err(KafkaError::NoMessageReceived) + } + Some(m) => { + let m = m?; + tracing::debug!("{:?}", m); + if let Some(headers) = m.headers() { + for header in headers.iter() { + tracing::info!("Header {}: {:?}", header.key, header.value); + } + } + consumer.commit_message(&m, CommitMode::Async).unwrap(); + Ok(m) + } + } + } else { + unreachable!("Called consume_once on a producer"); + } + } + + pub async fn send_message(&self, topic: &str, key: &str, payload: &str) { + if let KafkaHandler::Producer(producer) = self { + let record = BaseRecord::to(topic).key(key).payload(payload); + + match producer.send(record) { + Ok(_) => { + tracing::info!("Message sent successfully"); + } + Err(e) => tracing::error!("Failed to send message: {:?}", e), + } + + producer.poll(Timeout::Never); + } else { + tracing::error!("Called send_message on a consumer"); + } + } + + /// seek to offset + pub async fn seek_to_offset(&self, offset: i64) -> Result<(), rdkafka::error::KafkaError> { + tracing::info!("Start to seek to offset: {}", offset); + if let KafkaHandler::Consumer(consumer) = self { + let topic_partitions = consumer.assignment()?; + for topic_partition in topic_partitions.elements() { + consumer.seek( + topic_partition.topic(), + topic_partition.partition(), + rdkafka::Offset::Offset(offset), + None, + )?; + } + } + Ok(()) + } +} + +/// reset the mq +pub async fn reset_kafka_offset() -> Result<(), Box> { + tracing::info!("Start to reset import kafka"); + let kafka_broker = env::var("KAFKA_BROKER").unwrap(); + let consumer_group_id = env::var("KAFKA_CONSUMER_GROUP_ID").unwrap(); + let import_topic = env::var("KAFKA_ANALYSIS_TEST_TOPIC").unwrap(); + let output = Command::new("/opt/kafka/bin/kafka-consumer-groups.sh") + .args([ + "--bootstrap-server", + &kafka_broker, + "--group", + &consumer_group_id, + "--reset-offsets", + "--to-offset", + "0", + "--execute", + "--topic", + &import_topic, + ]) + .output() + .expect("Failed to execute command"); + + if output.status.success() { + println!("Command executed successfully"); + let stdout = String::from_utf8_lossy(&output.stdout); + println!("Output: {stdout}"); + } else { + eprintln!("Command failed to execute"); + let stderr = String::from_utf8_lossy(&output.stderr); + eprintln!("Error: {stderr}"); + } + + tracing::info!("Finish to reset import kafka"); + Ok(()) +} diff --git a/extensions/cratespro/common/repo_import/src/lib.rs b/extensions/cratespro/common/repo_import/src/lib.rs new file mode 100644 index 000000000..656bda74a --- /dev/null +++ b/extensions/cratespro/common/repo_import/src/lib.rs @@ -0,0 +1,1165 @@ +mod crate_info; +mod git; +mod kafka_handler; +mod utils; +mod version_info; + +extern crate lazy_static; +extern crate pretty_env_logger; + +use crate::crate_info::extract_info_local; +use crate::kafka_handler::KafkaHandler; +use crate::utils::{ + extract_namespace, get_program_by_name, insert_namespace_by_repo_path, name_join_version, + write_into_csv,extract_path_from_segment, +}; + +//use git::hard_reset_to_head; +use git2::{ObjectType, Oid, Repository}; +use model::{repo_sync_model, tugraph_model::*}; +use rdkafka::error::KafkaError; +use rdkafka::message::BorrowedMessage; +use rdkafka::Message; +use serde::{Deserialize, Serialize}; +use tokio::process::Command; +use std::collections::{HashMap, HashSet}; +use std::env; +use std::error::Error; +use std::fs; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; +use tokio::fs::File; +use tokio::io::AsyncWriteExt; +use tokio::sync::Mutex; +use url::Url; +use version_info::VersionUpdater; + +const CLONE_CRATES_DIR: &str = "/mnt/crates/local_crates_file/"; +// const TUGRAPH_IMPORT_FILES_PG: &str = "./tugraph_import_files_mq/"; + +pub use kafka_handler::reset_kafka_offset; + +pub enum MessageKind { + Mega, + UserUpload, +} + +pub struct ImportMessage<'a> { + kind: MessageKind, + message: BorrowedMessage<'a>, +} + +pub struct ImportDriver { + pub context: ImportContext, + pub import_handler: KafkaHandler, + pub user_import_handler: KafkaHandler, + pub sender_handler: KafkaHandler, +} +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct Licenses { + pub program_id: String, + pub program_name: String, + pub program_namespace: Option, + pub license: Option, +} + +impl ImportDriver { + pub async fn new(dont_clone: bool) -> Self { + tracing::info!("Start to setup Kafka client."); + + let should_reset_kafka_offset = env::var("SHOULD_RESET_KAFKA_OFFSET").unwrap().eq("1"); + + let (import_handler, user_import_handler, sender_handler) = init_kafka_handler() + .await + .expect("Failed to initialize Kafka handlers"); + + let context = if !should_reset_kafka_offset { + // 如果不需要重置offset，则从checkpoint中恢复context + let checkpoint_dir = + env::var("CHECKPOINT_DIR").unwrap_or_else(|_| "./checkpoints".to_string()); + let checkpoint_path = format!("{checkpoint_dir}/latest.json"); + + match ImportContext::load_from_file(&checkpoint_path).await { + Ok(mut ctx) => { + // 如果有保存的 offset 且不需要重置到0，则恢复到该位置 + if let Some(offset) = ctx.kafka_offset { + tracing::info!("Restoring Kafka consumer to offset: {}", offset); + if let Err(e) = import_handler.seek_to_offset(offset).await { + tracing::error!("Failed to seek to offset {}: {}", offset, e); + } + } + ctx.write_tugraph_import_files().await; + tracing::info!("Restored context from checkpoint"); + ctx + } + Err(e) => { + tracing::warn!("Failed to load checkpoint: {}", e); + ImportContext { + dont_clone, + ..Default::default() + } + } + } + } else { + // 如果需要重置offset，则创建一个新的context + tracing::info!("Resetting Kafka offset, creating new context"); + ImportContext { + dont_clone, + ..Default::default() + } + }; + + tracing::info!("Finish to setup Kafka client."); + + Self { + context, + import_handler, + user_import_handler, + sender_handler, + } + } + + async fn consume_message(&'_ self) -> Result, KafkaError> { + // try to get data from user_import_handler + if let Ok(message) = self.user_import_handler.consume_once().await { + tracing::info!("Receive a user upload message!"); + return Ok(ImportMessage { + kind: MessageKind::UserUpload, + message, + }); + } + // if there is no data from user_import_handler，try to fetch data from import_handler + else if let Ok(message) = self.import_handler.consume_once().await { + return Ok(ImportMessage { + kind: MessageKind::Mega, + message, + }); + }; + + Err(KafkaError::NoMessageReceived) + } + #[allow(clippy::let_unit_value)] + #[allow(unused_variables)] + pub async fn import_from_mq_for_a_message(&mut self) -> Result<(), ()> { + + let kafka_analysis_topic = env::var("KAFKA_ANALYSIS_TOPIC").unwrap(); + let git_url_base = env::var("MEGA_BASE_URL").unwrap(); + + let ImportMessage { kind, message } = match self.consume_message().await { + Err(_) => { + //tracing::warn!("No message in Kafka, please check it!"); + return Err(()); + } + Ok(m) => m, + }; + let model = match serde_json::from_slice::( + message.payload().unwrap(), + ) { + Ok(m) => Some(m.clone()), + Err(e) => { + tracing::info!("Error while deserializing message payload: {:?}", e); + None + } + }; + if model.is_none(){ + return Err(()) ; + } + tracing::info!( + "Received a message, key: '{:?}', payload: '{:?}', topic: {}, partition: {}, offset: {}, timestamp: {:?}", + message.key(), + model, + message.topic(), + message.partition(), + message.offset(), + message.timestamp() + ); + + // 早一个offset，防止当前消息没解析完就结束了 + let offset = message.offset(); + self.context.kafka_offset = Some(offset); + if offset % 2000 == 0 { + tracing::info!("Reached message offset: {}", offset); + self.context.print_status().await; + } + //from mega + let url = Url::parse(&model.unwrap().db_model.mega_url).unwrap(); + let mega_url_suffix = url.path().to_string(); + + let clone_crates_dir = + env::var("NEW_CRATES_DIR").unwrap_or_else(|_| CLONE_CRATES_DIR.to_string()); + let split_crates_dir = + env::var("SPLIT_CRATES_DIR").unwrap_or_else(|_| CLONE_CRATES_DIR.to_string()); + //changes clone_or_not_clone + let git_url = { + let git_url_base = Url::parse(&git_url_base) + .unwrap_or_else(|_| panic!("Failed to parse mega url base: {}", &git_url_base)); + git_url_base + .join(&mega_url_suffix) + .expect("Failed to join url path") + }; + + let namespace = extract_namespace(git_url.as_ref()).expect("Failed to parse URL"); + let behind_dir = extract_path_from_segment(git_url.as_ref(),"crates").expect("Failed to parse behind_dir"); + let path = PathBuf::from(&clone_crates_dir).join(behind_dir.clone()); + + //changes + if !path.is_dir() { + //if user_upload, no clone + tracing::info!("dir {} not exist", path.to_str().unwrap().to_string()); + let clone_start_time = Instant::now(); + let local_repo_path = match self + .context + .clone_a_repo_by_url(&clone_crates_dir, &git_url_base, &mega_url_suffix) + .await + { + Ok(x) => x, + _ => { + tracing::error!("Failed to clone repo {}", mega_url_suffix); + return Err(()); + } + }; + let clone_need_time = clone_start_time.elapsed(); + tracing::trace!("clone need time: {:?}", clone_need_time); + + let new_versions = self + .context + .parse_a_local_repo_and_return_new_versions(local_repo_path, mega_url_suffix) + .await + .unwrap(); + } else { + tracing::info!("dir {} already exist", path.to_str().unwrap().to_string()); + let git_dir = path.join(".git"); + if git_dir.is_dir() { + tracing::info!("Performing git pull for repository: {}", path.display()); + let pull_result = std::process::Command::new("git") + .arg("pull") + .current_dir(&path) + .status() + .map_err(|e| { + tracing::error!("Failed to execute git pull: {}", e); + }); + + if let Ok(status) = pull_result { + if !status.success() { + tracing::warn!("git pull returned non-zero exit status for {}", path.display()); + } + else{ + tracing::info!("git pull success"); + } + } + } else { + tracing::warn!("{} is not a git repository, skipping pull", path.display()); + } + let insert_time = Instant::now(); + insert_namespace_by_repo_path(path.to_str().unwrap().to_string(), namespace.clone()); + let insert_need_time = insert_time.elapsed(); + tracing::trace!( + "insert_namespace_by_repo_path need time: {:?}", + insert_need_time + ); + + let new_versions = self + .context + .parse_a_local_repo_and_return_new_versions(path, mega_url_suffix) + .await + .unwrap(); + } //changes + tracing::info!("Finish to import from a message!"); + Ok(()) + } + + pub async fn save_checkpoint(&mut self) -> Result<(), Box> { + tracing::info!("Saving checkpoint..."); + let checkpoint_dir = + env::var("CHECKPOINT_DIR").unwrap_or_else(|_| "./checkpoints".to_string()); + tokio::fs::create_dir_all(&checkpoint_dir).await?; + + // 保存二进制checkpoint (如果文件存在会覆盖) + let checkpoint_path = format!("{checkpoint_dir}/latest.json"); + if tokio::fs::try_exists(&checkpoint_path).await? { + tokio::fs::remove_file(&checkpoint_path).await?; + } + tracing::info!("Saving checkpoint to {}", checkpoint_path); + self.context.save_to_file(&checkpoint_path).await?; + + // 保存人类可读的摘要 (如果文件存在会覆盖) + let summary_path = format!("{checkpoint_dir}/summary.txt"); + if tokio::fs::try_exists(&summary_path).await? { + tokio::fs::remove_file(&summary_path).await?; + } + tracing::info!("Saving summary to {}", summary_path); + tokio::fs::write(summary_path, self.context.format_status()).await?; + + tracing::info!("Checkpoint saved to {}", checkpoint_path); + + Ok(()) + } + pub async fn export_version( + &self, + repo_path: &str, + oid: &Oid, + output_path: &str, + folder_name: &str, + ) -> Result<(), Box> { + let repo = match Repository::open(repo_path) { + Ok(repo) => repo, + Err(e) => { + tracing::info!("Failed to open repository: {}", e); + return Err(e.into()); + } + }; + let output_folder = Path::new(output_path).join(folder_name); + if output_folder.exists() { + tokio::fs::remove_dir_all(&output_folder).await?; + } + tokio::fs::create_dir_all(&output_folder).await?; + let tree = repo + .find_object(*oid, Some(ObjectType::Commit))? + .peel_to_commit()? + .tree() + .unwrap(); + tree.walk(git2::TreeWalkMode::PreOrder, |root, entry| { + let entry_path = if let Some(name) = entry.name() { + PathBuf::from(root).join(name) + } else { + return 1; + }; + + let output_file_path = output_folder.join(entry_path); + if entry.kind() == Some(ObjectType::Blob) { + let blob = entry + .to_object(&repo) + .and_then(|obj| obj.peel_to_blob()) + .unwrap(); + if let Some(parent) = output_file_path.parent() { + std::fs::create_dir_all(parent).unwrap(); + } + std::fs::write(output_file_path, blob.content()).unwrap(); + } + 0 + })?; + Ok(()) + } + #[allow(clippy::manual_flatten)] + pub async fn export_tags( + &self, + repo_path: &str, + output_dir: &str, + namespace: String, + crate_name: String, + ) -> Result<(), Box> { + let repo = match Repository::open(repo_path) { + Ok(repo) => repo, + Err(e) => { + tracing::info!("Failed to open repository: {}", e); + return Err(e.into()); + } + }; + + let mut tags = Vec::new(); + if let Ok(tag_names) = repo.tag_names(None) { + for tag_name in tag_names.iter().flatten() { + if let Ok(tag_object) = repo.revparse_single(tag_name) { + let oid = match tag_object.kind() { + Some(ObjectType::Tag) => tag_object.as_tag().unwrap().target_id(), + Some(ObjectType::Commit) => tag_object.id(), + _ => { + continue; + } + }; + tags.push(( + tag_name + .replace("/", "_") + .replace("\\", "_") + .replace(":", "_"), + oid, + )); + } + } + } + if !tags.is_empty() { + let real_output_dir = output_dir.to_string() + "/" + &namespace; + for (version_name, oid) in &tags { + let real_version_name = crate_name.clone() + "-" + version_name; + match self + .export_version(repo_path, oid, real_output_dir.as_str(), &real_version_name) + .await + { + Ok(_) => {} + Err(e) => tracing::info!("Failed to export version {}: {}", version_name, e), + } + } + } + Ok(()) + } +} + +/// 根据环境变量初始化kafka handler +/// KAFKA_CONSUMER_GROUP_ID 会根据测试or部署来设置 +/// 详情见 https://github.com/crates-pro/private_docs/discussions/1#discussioncomment-12032278 +async fn init_kafka_handler() -> Result<(KafkaHandler, KafkaHandler, KafkaHandler), KafkaError> { + let kafka_broker = env::var("KAFKA_BROKER").unwrap(); + let consumer_group_id = env::var("KAFKA_CONSUMER_GROUP_ID").unwrap(); + tracing::info!("Kafka parameters: {},{}", kafka_broker, consumer_group_id); + // 创建三个kafka handler + // Data from Mega + let import_handler = KafkaHandler::new_consumer( + &kafka_broker, + &consumer_group_id, + &env::var("KAFKA_ANALYSIS_TEST_TOPIC").unwrap(), + ) + .expect("Invalid import kafka handler"); + + // Data from user-uploading + let user_import_handler = KafkaHandler::new_consumer( + &kafka_broker, + &consumer_group_id, + &env::var("KAFKA_USER_IMPORT_TOPIC").unwrap(), + ) + .expect("Invalid import kafka handler"); + + // sending for analysis + let sender_handler = + KafkaHandler::new_producer(&kafka_broker).expect("Invalid import kafka handler"); + Ok((import_handler, user_import_handler, sender_handler)) +} + +/// internal structure, +/// a context for repo parsing and importing. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct ImportContext { + pub dont_clone: bool, + + // data to write into + /// vertex + pub programs: Vec, + + pub libraries: Vec, + pub applications: Vec, + pub library_versions: Vec, + pub application_versions: Vec, + pub versions: Vec, + //pub max_versions: Arc>>, + pub licenses: Vec, + /// edge + has_lib_type: Vec, + has_app_type: Vec, + + lib_has_version: Vec, + app_has_version: Vec, + + lib_has_dep_version: Vec, + app_has_dep_version: Vec, + + pub depends_on: Vec, + + /// help is judge whether it is a new program + program_memory: HashSet, + /// help us judge whether it is a new version + version_memory: HashSet, + + pub version_updater: VersionUpdater, + + // 新增字段保存 Kafka offset + #[serde(default)] + pub kafka_offset: Option, +} + +impl ImportContext { + /// Import data from mega + /// It first clone the repositories locally from mega + pub async fn compare_versions(a: &str, b: &str) -> Result> { + let parse_version = |version: &str| { + let mut parts = version.splitn(2, ['+', '-']); // 分割主版本和附加段 + let version_part = parts.next().unwrap_or(""); + let build_metadata = parts.next(); // 可能的构建元数据或预发行信息 + + let version_numbers: Vec = version_part + .split('.') + .map(|v| v.parse::().unwrap_or(0)) // 解析每个部分 + .collect(); + + (version_numbers, build_metadata.is_some()) + }; + + let (a_parts, a_has_metadata) = parse_version(a); + let (b_parts, b_has_metadata) = parse_version(b); + + let max_length = std::cmp::max(a_parts.len(), b_parts.len()); + + for i in 0..max_length { + let a_part = a_parts.get(i).unwrap_or(&0); + let b_part = b_parts.get(i).unwrap_or(&0); + let ord = a_part.cmp(b_part); + if ord != std::cmp::Ordering::Equal { + return Ok(ord); + } + } + + // 版本部分相等，考虑构建元数据 + if a_has_metadata && !b_has_metadata { + return Ok(std::cmp::Ordering::Less); // A 有构建元数据，B 没有，A较小 + } else if !a_has_metadata && b_has_metadata { + return Ok(std::cmp::Ordering::Greater); // A 没有构建元数据，B 有，A较大 + } + + Ok(std::cmp::Ordering::Equal) + } + pub async fn update_max_version(&mut self) -> Result> { + let tmp_max_versions: Arc>> = + Arc::new(Mutex::new(HashMap::new())); + let mut tasks = vec![]; + for nv in self.versions.clone() { + let tmp_max_versions = Arc::clone(&tmp_max_versions); + let name_and_version = nv.name_and_version; + let task = tokio::spawn(async move { + if let Some((name, version)) = name_and_version.split_once('/') { + let mut tmp_max_versions2 = tmp_max_versions.lock().await; + + // 获取或插入当前版本 + let entry = tmp_max_versions2 + .entry(name.to_string()) + .or_insert(version.to_string()); + + // 进行版本比较 + match ImportContext::compare_versions(entry, version).await { + Ok(ordering) => { + // 仅在新的版本大于当前版本时更新 + if ordering == std::cmp::Ordering::Less { + *entry = version.to_string(); + } + } + Err(_) => { + //eprintln!("Error comparing versions for {}", name); + } + } + } + }); + tasks.push(task); + } + for task in tasks { + let _ = task.await; // 处理结果（如果需要） + } + let get_max_versions = tmp_max_versions.lock().await; + for (name, version) in get_max_versions.iter() { + //println!("{}/{}", name, version); + for p in &mut self.programs { + if p.name == name.clone() { + p.max_version = Some(version.clone()); + break; + } + } + } + Ok("".to_string()) + } + pub async fn max_version(&mut self, v1: &str, v2: &str) -> String { + let parts1: Vec<&str> = v1.split('.').collect(); + let parts2: Vec<&str> = v2.split('.').collect(); + let max_parts = std::cmp::max(parts1.len(), parts2.len()); + for i in 0..max_parts { + let num1 = parts1 + .get(i) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + let num2 = parts2 + .get(i) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + #[allow(clippy::comparison_chain)] + if num1 > num2 { + return v1.to_string(); + } else if num1 < num2 { + return v2.to_string(); + } + } + v1.to_string() + } + async fn parse_a_local_repo_and_return_new_versions( + &mut self, + repo_path: PathBuf, + git_url: String, + ) -> Result, String> { + let mut new_versions = vec![]; + + if repo_path.is_dir() && Path::new(&repo_path).join(".git").is_dir() { + match Repository::open(&repo_path) { + Err(e) => { + tracing::error!("Not a git repo: {:?}, Err: {}", repo_path, e); + } + Ok(_) => { + // It'a a valid git repository. Start to parse it. + tracing::info!("Processing repo: {}", repo_path.display()); + + //reset, maybe useless + /*let hard_reset_time = Instant::now(); + hard_reset_to_head(&repo_path) + .await + .map_err(|x| format!("{:?}", x))?; + let hard_reset_need_time = hard_reset_time.elapsed(); + tracing::info!("hard_reset_to_head need time: {:?}", hard_reset_need_time);*/ + let all_programs = + self.collect_and_filter_programs(&repo_path, &git_url).await; + + let all_dependencies = + self.collect_and_filter_versions(&repo_path, &git_url).await; + let proccess_time = Instant::now(); + //find max_version + + + // + + for (program, has_type, uprogram) in all_programs { + self.programs.push(program.clone()); + match uprogram { + UProgram::Library(l) => { + self.libraries.push(l.clone()); + + self.has_lib_type.push(has_type.clone()); + + } + UProgram::Application(a) => { + self.applications.push(a.clone()); + + self.has_app_type.push(has_type.clone()); + + } + }; + + // NOTE: memorize program + self.program_memory + .insert(model::general_model::Program::new( + &program.name, + &program.mega_url.clone().unwrap(), + )); + } + + for dependencies in all_dependencies.clone() { + let name = dependencies.crate_name.clone(); + let version = dependencies.version.clone(); + let git_url = dependencies.git_url.clone(); + let tag_name = dependencies.tag_name.clone(); + + new_versions.push(model::general_model::VersionWithTag::new( + &name, &version, &git_url, &tag_name, + )); + + // check whether the crate version exists. + let (program, uprogram) = match get_program_by_name(&name) { + Some((program, uprogram)) => (program, uprogram), + None => { + // continue, dont parse + continue; + } + }; + + self.version_updater.update_depends_on(&dependencies).await; + + let has_version = HasVersion { + SRC_ID: program.id.clone(), + DST_ID: name_join_version(&name, &version), //FIXME: version id undecided + }; + + let dep_version = Version { + name_and_version: name_join_version(&name, &version), + }; + + #[allow(non_snake_case)] + let SRC_ID = name_join_version(&name, &version); + #[allow(non_snake_case)] + let DST_ID = name_join_version(&name, &version); + let has_dep_version = HasDepVersion { SRC_ID, DST_ID }; + + let islib = uprogram.index() == 0; + if islib { + let version = LibraryVersion::new( + program.id.clone(), + &name.clone(), + &version.clone(), + "???", + ); + + self.library_versions.push(version.clone()); + + self.lib_has_version.push(has_version.clone()); + + self.lib_has_dep_version.push(has_dep_version.clone()); + + } else { + let version = ApplicationVersion::new( + program.id.clone(), + name.clone(), + version.clone(), + ); + + self.application_versions.push(version.clone()); + + self.app_has_version.push(has_version.clone()); + + self.app_has_dep_version.push(has_dep_version.clone()); + + } + self.versions.push(dep_version.clone()); + + + //self.depends_on + // .clone_from(&(self.version_updater.to_depends_on_edges().await)); + + // NOTE: memorize version, insert the new version into memory + self.version_memory + .insert(model::general_model::Version::new( + &dependencies.crate_name, + &dependencies.version, + )); + } + for dependency in all_dependencies.clone(){ + if let Some(tprogram) = self.programs.iter_mut() + .find(|p| p.name == dependency.clone().crate_name) + { + // 找到匹配项，修改其 max_version + let tmp_max_version = tprogram.clone().max_version.unwrap_or_else(|| "0.0.0".to_string()); + + match ImportContext::compare_versions(&tmp_max_version, &dependency.version).await { + Ok(ordering) => { + if ordering == std::cmp::Ordering::Less { + tprogram.max_version = Some(dependency.clone().version); + } + } + Err(_) => { + //eprintln!("Error comparing versions for {}", name); + } + }; + tracing::info!("Updated max_version for {} to {}", tprogram.name, dependency.version); + } else { + // 未找到匹配项时的处理（可选） + tracing::warn!("Program {} not found in self.programs",dependency.crate_name); + } + } + + let proccess_need_time = proccess_time.elapsed(); + tracing::info!("Finish processing repo: {}", repo_path.display()); + tracing::trace!("processing repo need time: {:?}", proccess_need_time); + } + } + } else { + tracing::error!("{} is not a directory", repo_path.display()); + } + Ok(new_versions) + } + + async fn collect_and_filter_programs( + &mut self, + repo_path: &Path, + git_url: &str, + ) -> Vec<(Program, HasType, UProgram)> { + tracing::info!("Start to collect_and_filter_programs {:?}", repo_path); + let collect_time = Instant::now(); + let all_programs: Vec<(Program, HasType, UProgram)> = extract_info_local( + repo_path.to_path_buf(), + git_url.to_owned(), + &mut self.licenses, + ) + .await + .into_iter() + .filter(|(p, _, _)| { + !self + .program_memory + .contains(&model::general_model::Program::new( + &p.name, + &p.mega_url.clone().unwrap(), + )) + }) + .collect(); + let collect_need_time = collect_time.elapsed(); + tracing::info!("Finish to collect_and_filter_programs {:?}", repo_path); + tracing::trace!( + "collect_and_filter_programs need time: {:?}", + collect_need_time + ); + all_programs + } + async fn collect_and_filter_versions( + &self, + repo_path: &PathBuf, + git_url: &str, + ) -> Vec { + tracing::info!("Start to collect_and_filter_versions {:?}", repo_path); + let collect_time = Instant::now(); + // get all versions and dependencies + // filter out new versions!!! + let all_dependencies: Vec = self + .parse_all_versions_of_a_repo(repo_path, git_url) + .await + .into_iter() + .filter(|x| { + !self + .version_memory + .contains(&model::general_model::Version::new( + &x.crate_name, + &x.version, + )) + }) + //.filter(|x| semver::Version::parse(&x.version).is_ok()) + .collect(); + let collect_need_time = collect_time.elapsed(); + tracing::info!("Finish to collect_and_filter_versions {:?}", repo_path); + tracing::trace!( + "collect_and_filter_versions need time: {:?}", + collect_need_time + ); + all_dependencies + } + + async fn normalize(&mut self) { + self.depends_on + .clone_from(&(self.version_updater.to_depends_on_edges().await)); + } + + /// write data base into tugraph import files + pub async fn write_tugraph_import_files(&mut self) { + + self.normalize().await; + + let write_time = Instant::now(); + let tugraph_import_files = PathBuf::from(env::var("TUGRAPH_IMPORT_FILES_PG").unwrap()); + fs::create_dir_all(tugraph_import_files.clone()) + .unwrap_or_else(|e| tracing::error!("Error: {}", e)); + + // write into csv + write_into_csv( + tugraph_import_files.join("program.csv"), + self.programs.clone(), + ) + .unwrap(); + write_into_csv( + tugraph_import_files.join("library.csv"), + self.libraries.clone(), + ) + .unwrap(); + write_into_csv( + tugraph_import_files.join("application.csv"), + self.applications.clone(), + ) + .unwrap(); + write_into_csv( + tugraph_import_files.join("library_version.csv"), + self.library_versions.clone(), + ) + .unwrap(); + write_into_csv( + tugraph_import_files.join("application_version.csv"), + self.application_versions.clone(), + ) + .unwrap(); + write_into_csv( + tugraph_import_files.join("version.csv"), + self.versions.clone(), + ) + .unwrap(); + write_into_csv( + tugraph_import_files.join("licenses.csv"), + self.licenses.clone(), + ) + .unwrap(); + + // edge + let _ = write_into_csv( + tugraph_import_files.join("has_lib_type.csv"), + self.has_lib_type.clone(), + ); + let _ = write_into_csv( + tugraph_import_files.join("has_app_type.csv"), + self.has_app_type.clone(), + ); + let _ = write_into_csv( + tugraph_import_files.join("lib_has_version.csv"), + self.lib_has_version.clone(), + ); + let _ = write_into_csv( + tugraph_import_files.join("app_has_version.csv"), + self.app_has_version.clone(), + ); + + let _ = write_into_csv( + tugraph_import_files.join("lib_has_dep_version.csv"), + self.lib_has_dep_version.clone(), + ); + let _ = write_into_csv( + tugraph_import_files.join("app_has_dep_version.csv"), + self.app_has_dep_version.clone(), + ); + let _ = write_into_csv( + tugraph_import_files.join("depends_on.csv"), + self.depends_on.clone(), + ); + + let write_need_time = write_time.elapsed(); + tracing::trace!("write need time: {:?}", write_need_time); + } + + pub async fn save_to_file(&mut self, path: &str) -> Result<(), String> { + self.normalize().await; + let serialized = + bincode::serialize(self).map_err(|e| format!("Serialization error: {e}"))?; + + let mut file = File::create(path) + .await + .map_err(|e| format!("Failed to create file: {e}"))?; + + file.write_all(&serialized) + .await + .map_err(|e| format!("Failed to write to file: {e}"))?; + + Ok(()) + } + + pub async fn load_from_file(path: &str) -> Result { + tracing::info!("Start to load from file: {}", path); + + let content = tokio::fs::read(path) + .await + .map_err(|e| format!("Failed to read file: {e}"))?; + + let context: ImportContext = + bincode::deserialize(&content).map_err(|e| format!("Deserialization error: {e}"))?; + tracing::info!( + "Context loaded successfully, there are {} programs", + context.programs.len() + ); + Ok(context) + } + + fn format_status(&self) -> String { + format!( + "Checkpoint Summary:\n\ + Time: {}\n\ + Kafka Offset: {}\n\ + \n\ + Collection Sizes:\n\ + - Programs: {}\n\ + - Libraries: {}\n\ + - Applications: {}\n\ + - Library Versions: {}\n\ + - Application Versions: {}\n\ + - Versions: {}\n\ + - Licenses: {}\n\ + \n\ + Memory Sets:\n\ + - Program Memory: {}\n\ + - Version Memory: {}\n\ + \n\ + Edge Collections:\n\ + - Has Lib Type: {}\n\ + - Has App Type: {}\n\ + - Lib Has Version: {}\n\ + - App Has Version: {}\n\ + - Lib Has Dep Version: {}\n\ + - App Has Dep Version: {}\n\ + - Depends On: {}\n", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(), + self.kafka_offset.unwrap_or(-1), + self.programs.len(), + self.libraries.len(), + self.applications.len(), + self.library_versions.len(), + self.application_versions.len(), + self.versions.len(), + self.licenses.len(), + self.program_memory.len(), + self.version_memory.len(), + self.has_lib_type.len(), + self.has_app_type.len(), + self.lib_has_version.len(), + self.app_has_version.len(), + self.lib_has_dep_version.len(), + self.app_has_dep_version.len(), + self.depends_on.len(), + ) + } + + pub async fn print_status(&mut self) { + self.normalize().await; + tracing::info!("{}", self.format_status()); + } + + pub async fn import_from_env_vars(&mut self) -> Result<(), Box> { + // 从环境变量获取配置 + let lgraph_import_path = env::var("LGRAPH_IMPORT_PATH").unwrap(); + + let config_path = env::var("LGRAPH_CONFIG_PATH").unwrap(); + + let db_dir = env::var("LGRAPH_DB_DIR").unwrap(); + + let graph_name = env::var("TUGRAPH_CRATESPRO_DB").unwrap(); + + let overwrite = true; + + // 验证lgraph_import工具是否存在 + if !Path::new(&lgraph_import_path).exists() { + return Err(format!("lgraph_import工具不存在: {lgraph_import_path}").into()); + } + + if !Path::new(&lgraph_import_path).is_file() { + return Err(format!("{lgraph_import_path} 不是有效的文件").into()); + } + + // 验证配置文件是否存在 + if !Path::new(&config_path).exists() { + return Err(format!("配置文件不存在: {config_path}").into()); + } + + if !Path::new(&config_path).is_file() { + return Err(format!("{config_path} 不是有效的文件").into()); + } + + // 确保数据库目录存在 + let db_path = Path::new(&db_dir); + if !db_path.exists() { + std::fs::create_dir_all(db_path)?; + println!("创建数据库目录: {db_dir}"); + } + + // 构建异步命令 + let mut cmd = Command::new(&lgraph_import_path); + + // 添加命令参数 + cmd.arg("-c").arg(&config_path) + .arg("--dir").arg(&db_dir) + .arg("--graph").arg(&graph_name); + + // 添加覆盖参数 + if overwrite { + cmd.arg("--overwrite").arg("true"); + } + cmd.arg("--continue_on_error").arg("true"); + // 异步执行命令并获取输出 + let output = cmd.output().await?; + + // 处理命令执行结果 + if output.status.success() { + //println!("数据导入成功！"); + + // 打印工具的标准输出 + let stdout = String::from_utf8_lossy(&output.stdout); + if !stdout.is_empty() { + println!("工具输出:\n{stdout}"); + } + Ok(()) + } else { + // 处理错误情况 + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + + tracing::info!("数据导入失败！"); + tracing::info!("退出代码: {}", output.status.code().unwrap_or(-1)); + + if !stdout.is_empty() { + tracing::info!("标准输出:\n{}", stdout); + } + + if !stderr.is_empty() { + tracing::info!("错误信息:\n{}", stderr); + } + Err("lgraph_import命令执行失败".into()) + } + } +} + +impl ImportContext { + #[allow(unused)] + fn calculate_memory_usage(&self) -> String { + use std::fmt::Write; + use std::mem; + let mut output = String::new(); + let bytes_per_gb = 1_073_741_824; // 1024^3 + + let fields = [ + ( + "Programs", + self.programs.capacity(), + mem::size_of::(), + ), + ( + "Libraries", + self.libraries.capacity(), + mem::size_of::(), + ), + ( + "Applications", + self.applications.capacity(), + mem::size_of::(), + ), + ( + "LibraryVersions", + self.library_versions.capacity(), + mem::size_of::(), + ), + ( + "ApplicationVersions", + self.application_versions.capacity(), + mem::size_of::(), + ), + ( + "Versions", + self.versions.capacity(), + mem::size_of::(), + ), + ( + "HasLibType", + self.has_lib_type.capacity(), + mem::size_of::(), + ), + ( + "HasAppType", + self.has_app_type.capacity(), + mem::size_of::(), + ), + ( + "LibHasVersion", + self.lib_has_version.capacity(), + mem::size_of::(), + ), + ( + "AppHasVersion", + self.app_has_version.capacity(), + mem::size_of::(), + ), + ( + "LibHasDepVersion", + self.lib_has_dep_version.capacity(), + mem::size_of::(), + ), + ( + "AppHasDepVersion", + self.app_has_dep_version.capacity(), + mem::size_of::(), + ), + ( + "DependsOn", + self.depends_on.capacity(), + mem::size_of::(), + ), + ( + "ProgramMemory", + self.program_memory.capacity(), + mem::size_of::(), + ), + ( + "VersionMemory", + self.version_memory.capacity(), + mem::size_of::(), + ), + ]; + + for (name, capacity, size) in &fields { + let total_size_bytes = capacity * size; + let total_size_gb = total_size_bytes as f64 / bytes_per_gb as f64; + write!(output, " [{name}: {total_size_gb:.4} GB]").unwrap(); + } + + output + } +} diff --git a/extensions/cratespro/common/repo_import/src/utils.rs b/extensions/cratespro/common/repo_import/src/utils.rs new file mode 100644 index 000000000..1e2cd666d --- /dev/null +++ b/extensions/cratespro/common/repo_import/src/utils.rs @@ -0,0 +1,197 @@ +use csv::WriterBuilder; +use lazy_static::lazy_static; +use model::tugraph_model::{Program, UProgram}; +use serde::Serialize; +use serde_json::json; +use std::collections::HashMap; +use std::error::Error; +use std::fmt::Debug; +use std::fs::OpenOptions; +use std::path::PathBuf; +use std::sync::Mutex; +use url::Url; + +lazy_static! { + pub static ref NAMESPACE_HASHMAP: Mutex> = Mutex::new(HashMap::new()); +} + +pub fn insert_namespace_by_repo_path(key: String, value: String) { + let mut map = NAMESPACE_HASHMAP.lock().unwrap(); + map.insert(key, value); +} + +pub fn get_namespace_by_repo_path(key: &str) -> Option { + let map = NAMESPACE_HASHMAP.lock().unwrap(); + map.get(key).cloned() +} + +lazy_static! { + pub static ref PROGRAM_HASHMAP: Mutex> = + Mutex::new(HashMap::new()); +} + +pub fn insert_program_by_name(key: String, value: (Program, UProgram)) { + let mut map = PROGRAM_HASHMAP.lock().unwrap(); + map.insert(key, value); +} + +pub fn get_program_by_name(key: &str) -> Option<(Program, UProgram)> { + let map = PROGRAM_HASHMAP.lock().unwrap(); + map.get(key).cloned() +} + +pub(crate) fn write_into_csv( + csv_path: PathBuf, + programs: Vec, +) -> Result<(), Box> { + let serialized = serde_json::to_value(T::default()).unwrap(); + + if let serde_json::Value::Object(map) = serialized { + let field_names: Vec<&str> = map.keys().map(|s| s.as_str()).collect(); + write_to_csv(field_names, csv_path.to_str().unwrap(), false)?; + } + + for program in &programs { + let fields = get_fields(program); + let fields = fields.iter().map(|s| s.as_str()).collect::>(); + write_to_csv(fields, csv_path.to_str().unwrap(), true)?; + } + + Ok(()) +} + +fn write_to_csv(data: Vec<&str>, file_path: &str, append: bool) -> Result<(), Box> { + let file = if append { + OpenOptions::new() + .create(true) + .append(true) + .open(file_path)? + } else { + OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(file_path)? + }; + + let mut wtr = WriterBuilder::new() + .quote_style(csv::QuoteStyle::Necessary) + .double_quote(true) + .from_writer(file); + + wtr.write_record(&data)?; + wtr.flush()?; + Ok(()) +} + +fn get_fields(item: &T) -> Vec { + let mut fields = Vec::new(); + let json = json!(item); + if let serde_json::Value::Object(map) = json { + fields = map + .values() + .map(|value| match value { + serde_json::Value::String(s) => s.clone(), + _ => value.to_string().trim_matches('"').to_owned(), + }) + .collect::>(); + } + fields +} + +/// An auxiliary function +/// +/// Extracts namespace e.g. "tokio-rs/tokio" from the git url https://www.github.com/tokio-rs/tokio +pub(crate) fn extract_namespace(url_str: &str) -> Result { + /// auxiliary function + fn remove_dot_git_suffix(input: &str) -> String { + let input = if input.ends_with('/') { + input.strip_suffix('/').unwrap() + } else { + input + }; + + let input = if input.ends_with(".git") { + input.strip_suffix(".git").unwrap().to_string() + } else { + input.to_string() + }; + input + } + + let url = Url::parse(&remove_dot_git_suffix(url_str)) + .map_err(|e| format!("Failed to parse URL {url_str}: {e}"))?; + + // /tokio-rs/tokio + let path_segments = url + .path_segments() + .ok_or("Cannot extract path segments from URL")?; + + let segments: Vec<&str> = path_segments.collect(); + //println!("{:?}", segments); + + // github URLs is of the format "/user/repo" + if segments.len() < 2 { + return Err(format!( + "URL {url_str} does not include a namespace and a repository name" + )); + } + + // join owner name and repo name + let namespace = format!( + "crates/{}", + segments[segments.len() - 2] + ); + + Ok(namespace) +} + +pub(crate) fn name_join_version(crate_name: &str, version: &str) -> String { + crate_name.to_string() + "/" + version +} + + +pub(crate) fn extract_path_from_segment(url_str: &str, start_segment: &str) -> Result { + let url = Url::parse(url_str) + .map_err(|e| format!("解析URL失败: {e}"))?; + + // 解析路径为 segments 列表（不含空字符串） + let segments: Vec<&str> = url + .path_segments() + .ok_or("无法提取路径片段")? + .collect(); + + // 查找起始段的位置 + let start_index = segments + .iter() + .position(|&s| s == start_segment) + .ok_or(format!("路径中未找到 '{start_segment}' 段"))?; + + // 截取从起始段到结尾的部分，并拼接成字符串 + let relevant_segments = &segments[start_index..]; + let path = relevant_segments.join("/"); + + Ok(path) +} + +pub(crate) fn extract_middle_path(url_str: &str) -> Result { + + // 分割路径为段（不含空字符串） + let segments: Vec<&str> = url_str + .split('/') + .filter(|s| !s.is_empty()) // 过滤掉空字符串 + .collect(); + + // 检查是否有足够的段（至少需要2个段才能移除最后一个） + if segments.len() < 2 { + return Err("URL路径段不足，无法提取中间部分".to_string()); + } + + // 移除最后一个段（版本号） + let segments_without_version = &segments[..segments.len() - 1]; + + // 拼接成带前缀/的路径 + let result = format!("/{}", segments_without_version.join("/")); + + Ok(result) +} \ No newline at end of file diff --git a/extensions/cratespro/common/repo_import/src/version_info.rs b/extensions/cratespro/common/repo_import/src/version_info.rs new file mode 100644 index 000000000..d77ec3726 --- /dev/null +++ b/extensions/cratespro/common/repo_import/src/version_info.rs @@ -0,0 +1,461 @@ +use crate::git::get_all_git_tags_with_time_sorted; +use crate::utils::name_join_version; +use crate::ImportContext; +use git2::{Oid, Repository}; +use git2::{TreeWalkMode, TreeWalkResult}; +use model::tugraph_model::DependsOn; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::mem; +use std::path::PathBuf; +use toml::Value; + +/// A representation for the info +/// extracted from a `cargo.toml` file +#[allow(unused)] +#[derive(Debug, Default, Clone)] +pub struct Dependencies { + pub(crate) crate_name: String, + pub(crate) version: String, + pub(crate) dependencies: Vec<(String, String)>, + + pub(crate) git_url: String, + pub(crate) tag_name: String, +} + +impl ImportContext { + /// a git repo contains different crates + #[allow(clippy::type_complexity)] + pub(crate) async fn parse_all_versions_of_a_repo( + &self, + repo_path: &PathBuf, + git_url: &str, + ) -> Vec { + let mut crate_version_map: HashMap<(String, String), Dependencies> = HashMap::default(); + + let versions = get_all_git_tags_with_time_sorted(repo_path).await; //tag id time + + // parse each version of a repository with an order of time, walk all the packages of it + for (tag_name, tree, _) in versions.iter() { + let all_packages_dependencies = self + .parse_a_repo_of_a_version(repo_path, git_url, tag_name, *tree) + .await; + + // NOTE: At certain time, a version in Cargo.toml will exists in several tags, + // while a tag corresponds to a unique Cargo.toml version. + // So, I use a map to select the lastest tag which contains the version. + for dependencies in all_packages_dependencies { + let name = dependencies.crate_name.clone(); + let version = dependencies.version.clone(); + crate_version_map.insert((name.clone(), version.clone()), dependencies); + } + } + + crate_version_map.into_values().collect() + } + + /// for a given commit(version), walk all the package + async fn parse_a_repo_of_a_version( + &self, + repo_path: &PathBuf, + git_url: &str, + tag_name: &str, + tree: Oid, + ) -> Vec { + let mut res = Vec::new(); + + // Lock the repository and tree for reading + let repo = Repository::open(repo_path).unwrap(); + let tree = repo.find_tree(tree).expect("Failed to find tree"); + + // Walk the tree to find Cargo.toml + tree.walk(TreeWalkMode::PostOrder, |_, entry| { + if entry.name() == Some("Cargo.toml") { + // for each Cargo.toml in repo of given commit + let obj = entry + .to_object(&repo) + .expect("Failed to convert TreeEntry to Object"); + let blob = obj.as_blob().expect("Failed to interpret object as blob"); + let content = std::str::from_utf8(blob.content()) + .expect("Cargo.toml content is not valid UTF-8"); + + if let Some(dependencies) = + self.parse_a_package_of_a_version(content, git_url, tag_name) + { + res.push(dependencies); + } + } + + TreeWalkResult::Ok + }) + .unwrap(); + + res + } + + fn parse_a_package_of_a_version( + &self, + cargo_toml_content: &str, + git_url: &str, + tag_name: &str, + ) -> Option { + match cargo_toml_content.parse::() { + Ok(toml) => { + if let Some(package) = toml.get("package") { + if let Some(crate_name) = package.get("name") { + let crate_name = crate_name.as_str()?.to_string(); + let version = package.get("version")?.as_str()?.to_string(); + + // e.g. 0.1.53a2 is invalid version number. + if semver::Version::parse(&version).is_err() { + return None; + } + + // dedup + if self + .version_updater + .version_parser + .exists(&crate_name, &version) + { + return None; + } + + let mut dependencies = vec![]; + + if let Some(dep_table) = toml.get("dependencies") { + if let Some(deps_table) = dep_table.as_table() { + for (name, val) in deps_table { + if let Some(version) = val.as_str() { + dependencies.push((name.clone(), version.to_owned())); + } else if let Some(ver_tab) = val.as_table() { + if let Some(val) = ver_tab.get("version") { + if let Some(version) = val.as_str() { + dependencies + .push((name.clone(), version.to_owned())); + } + } + } + } + } + } + + let dependencies = Dependencies { + crate_name, + version, + dependencies, + git_url: git_url.to_string(), + tag_name: tag_name.to_string(), + }; + + return Some(dependencies); + } + } + } + Err(_) => tracing::error!("Failed to parse Cargo.toml for {:?}", cargo_toml_content), + } + None + } +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +pub struct VersionUpdater { + /// a reverse record: who depends on the key? + pub reverse_depends_on_map: HashMap>, + + /// a actual map: a crate **actually** depends on which? + /// it is used to build `depends_on` edges. + pub actually_depends_on_map: + HashMap>, + + pub version_parser: VersionParser, +} + +impl VersionUpdater { + pub async fn to_depends_on_edges(&self) -> Vec { + let mut edges = vec![]; + for (src, dsts) in &self.actually_depends_on_map { + for dst in dsts { + #[allow(non_snake_case)] + let SRC_ID = name_join_version(&src.name, &src.version); + + #[allow(non_snake_case)] + let DST_ID = name_join_version(&dst.name, &dst.version); + let depends_on = DependsOn { SRC_ID, DST_ID }; + edges.push(depends_on); + } + } + edges + } + + /// Given a dependency list, + pub async fn update_depends_on(&mut self, info: &Dependencies) { + self.version_parser + .insert_version(&info.crate_name, &info.version) + .await; + let cur_release = model::general_model::Version::new(&info.crate_name, &info.version); + self.ensure_dependencies(&cur_release, info).await; + self.ensure_dependents(&cur_release).await; + } + + async fn ensure_dependencies( + &mut self, + cur_release: &model::general_model::Version, + info: &Dependencies, + ) { + for (name, version) in &info.dependencies { + //let dep = model::general_model::Version::new(&name, &version); + self.insert_reverse_dep(name, version, &cur_release.name, &cur_release.version) + .await; + } + + // a new version should not exist before. + assert!(!self.actually_depends_on_map.contains_key(cur_release)); + let cur_dependencies = self.search_dependencies(info).await; + self.actually_depends_on_map + .insert(cur_release.clone(), cur_dependencies); + } + + async fn search_dependencies(&self, info: &Dependencies) -> Vec { + let mut res: Vec = vec![]; + for (dependency_name, dependency_version) in &info.dependencies { + let version_option = self + .version_parser + .find_latest_matching_version(dependency_name, dependency_version) + .await; + + if let Some(dependency_actual_version) = &version_option { + let dependency = + model::general_model::Version::new(dependency_name, dependency_actual_version); + res.push(dependency); + } + } + res + } + + async fn ensure_dependents(&mut self, cur_release: &model::general_model::Version) { + let sem_ver = semver::Version::parse(&cur_release.version) + .unwrap_or_else(|_| panic!("failed to parse version {:?}", &cur_release)); + let wrapped_reverse_map = self.reverse_depends_on_map.get(&cur_release.name); + if let Some(reverse_map) = wrapped_reverse_map { + for (required_version, reverse_dep) in reverse_map { + let requirement = match semver::VersionReq::parse(required_version) { + Ok(req) => req, + Err(_) => { + tracing::error!("failed to transform to VersionReq"); + continue; + } + }; + + if requirement.matches(&sem_ver) { + if let Some(v) = self.actually_depends_on_map.get_mut(reverse_dep) { + let mut found = false; + let mut exist = false; + for x in &mut *v { + if x.name == cur_release.name { + found = true; + let prev_sem_ver = semver::Version::parse(&x.version).unwrap(); + if sem_ver == prev_sem_ver { + exist = true; + //replace + //x.version.clone_from(&cur_release.version); + } + //found break; + //break; + } + } + #[allow(clippy::if_same_then_else)] + if !found { + v.push(model::general_model::Version::new( + &cur_release.name, + &cur_release.version, + )); + } else if !exist { + v.push(model::general_model::Version::new( + &cur_release.name, + &cur_release.version, + )); + } + } else { + // No vec + self.actually_depends_on_map.insert( + reverse_dep.clone(), + vec![model::general_model::Version::new( + &cur_release.name, + &cur_release.version, + )], + ); + } + } + } + } + } + + /// insert (dependency, dependent) + /// notice that: dependent is unique, but dependency should be newest. + pub async fn insert_reverse_dep( + &mut self, + dependency_name: &str, + dependency_version: &str, + dependent_name: &str, + dependent_version: &str, + ) { + //let dependency = model::general_model::Version::new(dependency_name, dependency_version); + let dependent = model::general_model::Version::new(dependent_name, dependent_version); + self.reverse_depends_on_map + .entry(dependency_name.to_string()) + .or_default() + .push((dependency_version.to_string(), dependent)); + } +} + +#[derive(Default, Debug, Clone, Serialize, Deserialize)] +pub struct VersionParser { + version_map: HashMap>, +} + +impl VersionParser { + pub async fn insert_version(&mut self, crate_name: &str, version: &str) { + self.version_map + .entry(crate_name.to_string()) + .or_default() + .push(version.to_string()); + } + + pub(crate) fn exists(&self, name: &str, version: &str) -> bool { + if let Some(map) = self.version_map.get(name) { + return map.contains(&version.to_string()); + } + false + } + + pub(crate) async fn _remove(&mut self, name: &str) { + self.version_map.remove(name); + } + + pub async fn find_latest_matching_version( + &self, + target_lib: &str, + target_version: &str, + ) -> Option { + if let Some(lib_map) = self.version_map.get(target_lib) { + // if the lib exists + let req_str = if target_version.contains('.') { + format!("^{target_version}") + } else { + format!("{target_version}.*") + }; + + let requirement = match semver::VersionReq::parse(&req_str) { + Ok(req) => req, + Err(_) => return None, // 如果无法解析为有效的版本请求，则返回 None + }; + + let mut matching_versions: Vec = lib_map + .iter() + .filter_map(|ver| semver::Version::parse(ver).ok()) // 将所有版本字符串解析为 Version 对象 + .filter(|ver| requirement.matches(ver)) + .collect(); + + // Sort the matched versions and return the last (largest) one + matching_versions.sort(); + return matching_versions.last().map(|v| v.to_string()); + } + None + } +} + +impl VersionUpdater { + #[allow(unused)] + pub fn calculate_memory_usage(&self) -> String { + let stack_size = mem::size_of_val(self); + + let mut heap_size = 0; + + // Calculate heap size for reverse_depends_on_map + for (key, value) in &self.reverse_depends_on_map { + heap_size += key.capacity() * mem::size_of::(); // String capacity + heap_size += + value.capacity() * mem::size_of::<(String, model::general_model::Version)>(); + for (s, _) in value { + heap_size += s.capacity() * mem::size_of::(); // String capacity + } + } + + // Calculate heap size for actually_depends_on_map + for value in self.actually_depends_on_map.values() { + heap_size += mem::size_of::(); // Key size + heap_size += value.capacity() * mem::size_of::(); + } + + format!(" [Version Updater: {}] ", stack_size + heap_size) + + &self.version_parser.calculate_memory_usage() + } +} +impl VersionParser { + #[allow(unused)] + fn calculate_memory_usage(&self) -> String { + let stack_size = mem::size_of_val(self); + + let mut heap_size = 0; + + // Calculate heap size for version_map + for (key, value) in &self.version_map { + heap_size += key.capacity() * mem::size_of::(); // String capacity + heap_size += value.capacity() * mem::size_of::(); + for s in value { + heap_size += s.capacity() * mem::size_of::(); // String capacity + } + } + + format!(" [Version Parser: {}] ", stack_size + heap_size) + } +} + +#[cfg(test)] +mod tests { + use super::VersionParser; + + #[tokio::test] + async fn test_insert_and_find_version() { + let mut parser = VersionParser::default(); + parser.insert_version("crate_a", "1.0.1").await; + parser.insert_version("crate_a", "1.1.1").await; + parser.insert_version("crate_a", "1.2.1").await; + parser.insert_version("crate_a", "1.2.2").await; + + // Test finding the latest exact version + assert_eq!( + parser.find_latest_matching_version("crate_a", "1.2").await, + Some("1.2.2".to_string()) + ); + assert_eq!( + parser.find_latest_matching_version("crate_a", "1").await, + Some("1.2.2".to_string()) + ); + + // Test finding versions when there's no match + assert_eq!( + parser.find_latest_matching_version("crate_a", "2.0").await, + None + ); + + // Test finding versions with a precise match + parser.insert_version("crate_b", "2.0.0").await; + parser.insert_version("crate_b", "2.0.1").await; + assert_eq!( + parser + .find_latest_matching_version("crate_b", "2.0.1") + .await, + Some("2.0.1".to_string()) + ); + + assert_eq!( + parser.find_latest_matching_version("crate_b", "2").await, + Some("2.0.1".to_string()) + ); + assert_eq!( + parser.find_latest_matching_version("crate_c", "2").await, + None + ); + } +} diff --git a/extensions/cratespro/common/search/BUCK b/extensions/cratespro/common/search/BUCK new file mode 100644 index 000000000..082b7d8db --- /dev/null +++ b/extensions/cratespro/common/search/BUCK @@ -0,0 +1,26 @@ +filegroup( + name = "search-0.1.0.crate", + srcs = [ + "src/ai.rs", + "src/crates_search.rs", + "src/embedding.rs", + "src/lib.rs", + "src/search_prepare.rs" + ], +) + +rust_library( + name = "search", + srcs = [":search-0.1.0.crate"], + crate_root = "search-0.1.0.crate/src/lib.rs", + edition = "2021", + deps = [ + "//third-party:pgvector", + "//third-party:reqwest", + "//third-party:serde", + "//third-party:semver", + "//third-party:serde_json", + "//third-party:tokio-postgres", + ], + visibility = ["PUBLIC"], +) \ No newline at end of file diff --git a/extensions/cratespro/common/search/Cargo.toml b/extensions/cratespro/common/search/Cargo.toml new file mode 100644 index 000000000..28ee727d2 --- /dev/null +++ b/extensions/cratespro/common/search/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "search" +version = "0.1.0" +edition = "2021" + +[dependencies] +# third-party (第三方依赖, 不写具体版本号, 具体版本只在根目录 Cargo.toml 中出现) +pgvector = { version = "0.4", features = ["postgres"] } +reqwest = { version = "0.12", features = ["json"] } +semver = "1.0" +serde = "1.0" +serde_json = "1.0" +tokio-postgres = "0.7" diff --git a/extensions/cratespro/common/search/src/ai.rs b/extensions/cratespro/common/search/src/ai.rs new file mode 100644 index 000000000..c758a930b --- /dev/null +++ b/extensions/cratespro/common/search/src/ai.rs @@ -0,0 +1,144 @@ +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use std::env; +use tokio_postgres::Client as PgClient; + +/// 模型请求体结构体 +#[derive(Serialize)] +struct RequestBody { + model: String, + messages: Vec, +} + +/// 消息结构体 +#[derive(Serialize, Deserialize, Clone)] +struct Message { + role: String, + content: String, +} + +/// 响应选择结构体 +#[derive(Deserialize)] +struct ResponseChoice { + message: Message, +} + +/// 响应体结构体 +#[derive(Deserialize)] +struct ResponseBody { + choices: Vec, +} + +/// 聊天上下文结构体 +struct ChatContext { + messages: Vec, +} + +impl ChatContext { + /// 创建新的聊天上下文 + fn new() -> Self { + ChatContext { + messages: Vec::new(), + } + } + + /// 添加消息到聊天上下文 + fn add_message(&mut self, role: &str, content: &str) { + self.messages.push(Message { + role: role.to_string(), + content: content.to_string(), + }); + } + + /// 获取聊天上下文中的所有消息 + fn get_messages(&self) -> &Vec { + &self.messages + } +} + +/// AI 聊天结构体 +pub struct AIChat<'a> { + context: ChatContext, + client: &'a PgClient, +} + +impl<'a> AIChat<'a> { + /// 创建新的 AI 聊天实例 + pub fn new(client: &'a PgClient) -> Self { + let mut ret = AIChat { + context: ChatContext::new(), + client, + }; + ret.context.add_message( + "system", + "You are an experienced rust programmer, know all the major crates.You should help user and answer the question.", + ); + ret + } + + /// 处理用户消息并返回 AI 响应 + pub async fn chat(&mut self, user_message: &str) -> Result> { + self.context.add_message("user", user_message); + let answer = query_openai(&self.context).await?; + self.context.add_message("assistant", &answer); + Ok(answer) + } + + /// 处理带有嵌入信息的用户消息并返回 AI 响应 + pub async fn chat_with_embedding( + &mut self, + user_message: &str, + ) -> Result> { + let addition_infomation = get_crate_info_with_embedding(self.client, user_message).await?; + let user_message = format!( + "Here are some revelvant crates to refer:{addition_infomation}. Question:{user_message}" + ); + self.context.add_message("user", user_message.as_str()); + let answer = query_openai(&self.context).await?; + self.context.add_message("assistant", &answer); + Ok(answer) + } +} + +/// 查询 OpenAI API 并返回响应 +async fn query_openai(context: &ChatContext) -> Result { + let api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set"); + let client = Client::new(); + let open_ai_chat_url = env::var("OPEN_AI_CHAT_URL").expect("OPEN_AI_CHAT_URL not set"); + + let request_body = RequestBody { + model: "gpt-3.5-turbo".to_string(), + messages: context.get_messages().clone(), + }; + + let response = client + .post(&open_ai_chat_url) + .header("Authorization", format!("Bearer {api_key}")) + .json(&request_body) + .send() + .await?; + + let response_body: ResponseBody = response.json().await?; + + Ok(response_body.choices[0].message.content.clone()) +} + +/// 获取带有嵌入信息的 crate 信息 +async fn get_crate_info_with_embedding( + client: &PgClient, + question: &str, +) -> Result> { + let question_embedding = crate::embedding::get_one_text_embedding(question).await?; + let top_n = 5; + let results: Vec<(i32, String, String)> = + crate::embedding::search_crates_by_embedding(client, &question_embedding, top_n).await?; + + let mut results_text = String::new(); + for (id, name, description) in results { + results_text.push_str(&format!( + "ID: {id}, Name: {name}, Description: {description}\n" + )); + } + + Ok(results_text) +} diff --git a/extensions/cratespro/common/search/src/crates_search.rs b/extensions/cratespro/common/search/src/crates_search.rs new file mode 100644 index 000000000..032a95524 --- /dev/null +++ b/extensions/cratespro/common/search/src/crates_search.rs @@ -0,0 +1,156 @@ +use semver::Version; +use std::env; +use tokio_postgres::Client as PgClient; + +pub struct SearchModule<'a> { + pg_client: &'a PgClient, + table_name: String, +} + +pub enum SearchSortCriteria { + Comprehensive, + Relavance, + Downloads, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct RecommendCrate { + pub id: String, + pub name: String, + pub description: String, + pub downloads: i64, + pub namespace: String, + pub max_version: String, + pub rank: f32, +} + +impl<'a> SearchModule<'a> { + pub async fn new(pg_client: &'a PgClient) -> Self { + let table_name = env::var("TABLE_NAME").unwrap_or_else(|_| "crates".to_string()); + SearchModule { + pg_client, + table_name, + } + } + + pub async fn search_crate( + &self, + keyword: &str, + sort_by: SearchSortCriteria, + ) -> Result, Box> { + let mut crates = search_crate_without_ai(self.pg_client, &self.table_name, keyword).await?; + sort_crates(&mut crates, sort_by); + rearrange_crates(&mut crates, keyword); + Ok(crates) + } +} + +fn version_cmp(a: &RecommendCrate, b: &RecommendCrate) -> std::cmp::Ordering { + let version_a = Version::parse(&a.max_version); + let version_b = Version::parse(&b.max_version); + + match (version_a, version_b) { + (Ok(ver_a), Ok(ver_b)) => ver_b.cmp(&ver_a), + (Ok(_), Err(_)) => std::cmp::Ordering::Less, + (Err(_), Ok(_)) => std::cmp::Ordering::Greater, + (Err(_), Err(_)) => std::cmp::Ordering::Equal, + } +} + +fn sort_crates_by_relevance(crate_vec: &mut [RecommendCrate]) { + crate_vec.sort_by(|a, b| { + b.rank + .partial_cmp(&a.rank) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.name.cmp(&b.name)) + .then_with(|| version_cmp(a, b)) + }); +} + +fn normalize_downloads(downloads: i64, max_downloads: i64) -> f32 { + downloads as f32 / max_downloads as f32 +} + +fn sort_crates_by_downloads(crate_vec: &mut [RecommendCrate]) { + let max_downloads = crate_vec.iter().map(|c| c.downloads).max().unwrap_or(1); + + crate_vec.sort_by(|a, b| { + let normalized_downloads_a = normalize_downloads(a.downloads, max_downloads); + let normalized_downloads_b = normalize_downloads(b.downloads, max_downloads); + + let score_a = 0.3 * a.rank + 0.7 * normalized_downloads_a; + let score_b = 0.3 * b.rank + 0.7 * normalized_downloads_b; + + score_b + .partial_cmp(&score_a) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.name.cmp(&b.name)) + .then_with(|| version_cmp(a, b)) + }); +} + +fn sort_crates(crate_vec: &mut [RecommendCrate], sort_by: SearchSortCriteria) { + match sort_by { + SearchSortCriteria::Comprehensive | SearchSortCriteria::Relavance => { + sort_crates_by_relevance(crate_vec); + } + SearchSortCriteria::Downloads => { + sort_crates_by_downloads(crate_vec); + } + } +} + +fn rearrange_crates(crates: &mut Vec, keyword: &str) { + let mut matching_crates: Vec = Vec::new(); + crates.retain(|c| { + if c.name == keyword { + matching_crates.push(c.clone()); + false + } else { + true + } + }); + sort_crates(&mut matching_crates, SearchSortCriteria::Relavance); + crates.splice(0..0, matching_crates); +} +#[allow(clippy::uninlined_format_args)] +async fn search_crate_without_ai( + client: &PgClient, + table_name: &str, + keyword: &str, +) -> Result, Box> { + let tsquery_keyword = keyword.replace(" ", " & "); + let query = format!("{tsquery_keyword}:*"); + + let statement = format!( + "SELECT {0}.id::text, {0}.name, {0}.description, ts_rank({0}.tsv, to_tsquery($1)) AS rank,{0}.downloads,{0}.namespace,{0}.max_version + FROM {0} + WHERE {0}.tsv @@ to_tsquery($1) + ORDER BY rank DESC", + table_name + ); + let rows = client.query(statement.as_str(), &[&query]).await?; + let mut recommend_crates = Vec::::new(); + + for row in rows.iter() { + let id: Option = row.get("id"); + let name: Option = row.get("name"); + let description: Option = row.get("description"); + let downloads: Option = row.get("downloads"); + let namespace: Option = row.get("namespace"); + let max_version: Option = row.get("max_version"); + let rank: Option = row.get("rank"); + + recommend_crates.push(RecommendCrate { + id: id.unwrap_or_default(), + name: name.unwrap_or_default(), + description: description.unwrap_or_default(), + downloads: downloads.unwrap_or(0), + namespace: namespace.unwrap_or_default(), + max_version: max_version.unwrap_or_default(), + rank: rank.unwrap_or(0.0), + }); + } + + Ok(recommend_crates) +} diff --git a/extensions/cratespro/common/search/src/embedding.rs b/extensions/cratespro/common/search/src/embedding.rs new file mode 100644 index 000000000..285bac53f --- /dev/null +++ b/extensions/cratespro/common/search/src/embedding.rs @@ -0,0 +1,134 @@ +use pgvector::Vector; +use reqwest::Client; +use serde::Deserialize; +use serde_json::json; +use std::env; +use tokio_postgres::Client as PgClient; + +#[derive(Deserialize)] +struct EmbeddingData { + embedding: Vec, +} + +#[derive(Deserialize)] +struct EmbeddingResponse { + data: Vec, +} + +pub async fn get_one_text_embedding(text: &str) -> Result, Box> { + get_texts_embedding(&[text]) + .await + .map(|embeddings| embeddings[0].clone()) +} + +//TODO 1: 优化get_texts_embedding函数，使其使用batch API +async fn get_texts_embedding(texts: &[&str]) -> Result>, Box> { + let api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set"); + let client = Client::new(); + let open_ai_embedding_url = + env::var("OPEN_AI_EMBEDDING_URL").expect("OPEN_AI_EMBEDDING_URL not set"); + + // let url = "https://api.xty.app/v1/embeddings"; + let request_body = json!({ + "input": texts, + "model": "text-embedding-3-small" + }); + let response = client + .post(open_ai_embedding_url) + .header("Content-Type", "application/json") + .header("AUTHORIZATION", format!("Bearer {api_key}")) + .json(&request_body) + .send() + .await? + .json::() + .await?; + + let embeddings: Vec> = response.data.into_iter().map(|d| d.embedding).collect(); + + Ok(embeddings) +} + +pub async fn update_crate_embeddings( + client: &PgClient, + crate_id: &str, +) -> Result<(), Box> { + let table_name = env::var("TABLE_NAME").expect("TABLE_NAME not set"); + let query = format!("SELECT name, description FROM {table_name} WHERE id = $1"); + let row = client + .query_one(&query, &[&crate_id]) + .await + .expect("Failed to get crate"); + let name: String = row.get("name"); + let description: String = row.get("description"); + let text = format!("crate name:{name}, crate description:{description}"); + let embedding = get_one_text_embedding(&text).await.unwrap(); + let embedding = Vector::from(embedding); + let update_query = format!("UPDATE {table_name} SET embedding = $1 WHERE id = $2"); + client + .execute(&update_query, &[&embedding, &crate_id]) + .await + .expect("Failed to update crate embedding"); + Ok(()) +} + +pub async fn update_all_crate_embeddings( + client: &PgClient, +) -> Result<(), Box> { + let table_name = env::var("TABLE_NAME").expect("TABLE_NAME not set"); + let query = format!("SELECT id, name, description FROM {table_name}"); + let rows = client.query(&query, &[]).await?; + let texts: Vec = rows + .iter() + .map(|row| { + let name: Option<&str> = row.get("name"); + let description: Option<&str> = row.get("description"); + format!( + "crate name:{}, crate description:{}", + name.unwrap_or(""), + description.unwrap_or("") + ) + }) + .collect(); + let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect(); + let mut all_embeddings = Vec::new(); + for chunk in text_refs.chunks(32) { + let embeddings = get_texts_embedding(chunk).await?; + println!("Got embeddings for {} texts", embeddings.len()); + all_embeddings.extend(embeddings); + } + + // println!("{:?}", all_embeddings); + for (i, row) in rows.iter().enumerate() { + let id: &str = row.get("id"); + let embedding = Vector::from(all_embeddings[i].clone()); + let update_query = format!("UPDATE {table_name} SET embedding = $1 WHERE id = $2"); + client.execute(&update_query, &[&embedding, &id]).await?; + println!("Updated embedding for crate {id}"); + } + Ok(()) +} + +// 新增函数 +pub(crate) async fn search_crates_by_embedding( + client: &PgClient, + embedding: &[f32], + n: i64, +) -> Result, Box> { + let table_name = env::var("TABLE_NAME").expect("TABLE_NAME not set"); + let query = format!( + "SELECT id, name, description FROM {table_name} ORDER BY embedding <=> $1 LIMIT $2" + ); + let rows = client.query(&query, &[&embedding, &n]).await?; + + let results = rows + .iter() + .map(|row| { + let id: i32 = row.get("id"); + let name: String = row.get("name"); + let description: String = row.get("description"); + (id, name, description) + }) + .collect(); + + Ok(results) +} diff --git a/extensions/cratespro/common/search/src/lib.rs b/extensions/cratespro/common/search/src/lib.rs new file mode 100644 index 000000000..e6f5c28fb --- /dev/null +++ b/extensions/cratespro/common/search/src/lib.rs @@ -0,0 +1,4 @@ +pub mod ai; +pub mod crates_search; +pub mod embedding; +pub mod search_prepare; diff --git a/extensions/cratespro/common/search/src/search_prepare.rs b/extensions/cratespro/common/search/src/search_prepare.rs new file mode 100644 index 000000000..5ce8a05cf --- /dev/null +++ b/extensions/cratespro/common/search/src/search_prepare.rs @@ -0,0 +1,201 @@ +use std::env; +use tokio_postgres::Client as PgClient; + +use crate::embedding; + +pub struct SearchPrepare<'a> { + pg_client: &'a PgClient, + table_name: String, +} +impl<'a> SearchPrepare<'a> { + pub async fn new(pgclient: &'a PgClient) -> Self { + let table_name = env::var("TABLE_NAME").unwrap_or_else(|_| "crates".to_string()); + SearchPrepare { + pg_client: pgclient, + table_name, + } + } + //检查crates表是否存在 + + pub async fn prepare_tsv(&self) -> Result<(), Box> { + let table_exists = self.crates_table_exists().await?; + if !table_exists { + return Err("crates table not exists".into()); + } + self.add_tsv_column().await?; + self.set_tsv_column().await?; + self.create_tsv_index().await?; + Ok(()) + } + + pub async fn prepare_embedding(&self) -> Result<(), Box> { + let table_exists = self.crates_table_exists().await?; + if !table_exists { + return Err("crates table not exists".into()); + } + self.add_pgvector_extension().await?; + self.add_embedding_column().await?; + self.set_embedding_column().await?; + self.create_embedding_index().await?; + Ok(()) + } + + pub async fn crates_table_exists(&self) -> Result> { + let query = format!( + "SELECT EXISTS ( + SELECT FROM information_schema.tables + WHERE table_name = '{}' + )", + self.table_name + ); + let rows = self.pg_client.query(&query, &[]).await?; + Ok(rows[0].get(0)) + } + + // 功能二：添加tsv列 + pub async fn add_tsv_column(&self) -> Result<(), Box> { + let query = format!( + "ALTER TABLE {} ADD COLUMN IF NOT EXISTS tsv tsvector", + self.table_name + ); + self.pg_client.execute(&query, &[]).await?; + Ok(()) + } + + // 功能三：为tsv列创建索引 + pub async fn create_tsv_index(&self) -> Result<(), Box> { + let query = format!( + "CREATE INDEX IF NOT EXISTS idx_{}_tsv ON {} USING gin(tsv)", + self.table_name, self.table_name + ); + self.pg_client.execute(&query, &[]).await?; + Ok(()) + } + + // 功能四：设置tsv为crates中name属性和description属性的全文搜索tsvector + pub async fn set_tsv_column(&self) -> Result<(), Box> { + let query = format!( + "UPDATE {} SET tsv = + setweight(to_tsvector('english', coalesce(name, '')), 'A') || + setweight(to_tsvector('english', coalesce(description, '')), 'B')", + self.table_name + ); + self.pg_client.execute(&query, &[]).await?; + Ok(()) + } + + // 功能五：添加embedding列 + pub async fn add_embedding_column(&self) -> Result<(), Box> { + let query = format!( + "ALTER TABLE {} ADD COLUMN IF NOT EXISTS embedding vector(1536)", + self.table_name + ); + self.pg_client.execute(&query, &[]).await?; + Ok(()) + } + + pub async fn set_embedding_column(&self) -> Result<(), Box> { + embedding::update_all_crate_embeddings(self.pg_client).await + } + // 功能六：为embedding列创建索引 + pub async fn create_embedding_index(&self) -> Result<(), Box> { + let query = format!( + "CREATE INDEX IF NOT EXISTS idx_{}_embedding ON {} USING ivfflat (embedding)", + self.table_name, self.table_name + ); + self.pg_client.execute(&query, &[]).await?; + Ok(()) + } + // 功能七:检查是否有crates表,crates表是否有tsv列,是否有embedding列,是否有索引,是否有数据,如果都有,返回true,否则返回false + pub async fn check_ok(&self) -> bool { + let table_exists = self.crates_table_exists().await.unwrap_or(false); + if !table_exists { + return false; + } + + let tsv_column_exists = self + .pg_client + .query( + &format!( + "SELECT EXISTS ( + SELECT FROM information_schema.columns + WHERE table_name = '{}' AND column_name = 'tsv' + )", + self.table_name + ), + &[], + ) + .await + .map(|rows| rows[0].get(0)) + .unwrap_or(false); + if !tsv_column_exists { + return false; + } + + let embedding_column_exists = self + .pg_client + .query( + &format!( + "SELECT EXISTS ( + SELECT FROM information_schema.columns + WHERE table_name = '{}' AND column_name = 'embedding' + )", + self.table_name + ), + &[], + ) + .await + .map(|rows| rows[0].get(0)) + .unwrap_or(false); + if !embedding_column_exists { + return false; + } + + let tsv_index_exists = self + .pg_client + .query( + &format!( + "SELECT EXISTS ( + SELECT FROM pg_indexes + WHERE tablename = '{}' AND indexname = 'idx_{}_tsv' + )", + self.table_name, self.table_name + ), + &[], + ) + .await + .map(|rows| rows[0].get(0)) + .unwrap_or(false); + if !tsv_index_exists { + return false; + } + + let embedding_index_exists = self + .pg_client + .query( + &format!( + "SELECT EXISTS ( + SELECT FROM pg_indexes + WHERE tablename = '{}' AND indexname = 'idx_{}_embedding' + )", + self.table_name, self.table_name + ), + &[], + ) + .await + .map(|rows| rows[0].get(0)) + .unwrap_or(false); + if !embedding_index_exists { + return false; + } + + true + } + + // 功能八：为crates表添加pgvector扩展 + pub async fn add_pgvector_extension(&self) -> Result<(), Box> { + let query = "CREATE EXTENSION IF NOT EXISTS vector"; + self.pg_client.execute(query, &[]).await?; + Ok(()) + } +} diff --git a/extensions/cratespro/common/tudriver/BUCK b/extensions/cratespro/common/tudriver/BUCK new file mode 100644 index 000000000..ed1e721f0 --- /dev/null +++ b/extensions/cratespro/common/tudriver/BUCK @@ -0,0 +1,17 @@ +rust_library( + name = "tudriver", + srcs = [ + "src/lib.rs", + "src/tugraph_client.rs", + ], + crate_root = "src/lib.rs", + edition = "2021", + deps = [ + "//third-party:base64", + "//third-party:neo4rs", + "//third-party:serde_json", + "//third-party:tokio", + "//third-party:tracing", + ], + visibility = ["PUBLIC"], +) diff --git a/extensions/cratespro/common/tudriver/Cargo.toml b/extensions/cratespro/common/tudriver/Cargo.toml new file mode 100644 index 000000000..0213c50c1 --- /dev/null +++ b/extensions/cratespro/common/tudriver/Cargo.toml @@ -0,0 +1,12 @@ +[package] +edition = "2021" +name = "tudriver" +version = "0.1.0" + +[dependencies] +# third-party (第三方依赖, 不写具体版本号, 具体版本只在根目录 Cargo.toml 中出现) +base64 = "0.22" +neo4rs = "0.8" +serde_json = "1.0" +tokio = { version = "1.44", features = ["full"] } +tracing = "0.1" diff --git a/extensions/cratespro/common/tudriver/src/lib.rs b/extensions/cratespro/common/tudriver/src/lib.rs new file mode 100644 index 000000000..1877771ce --- /dev/null +++ b/extensions/cratespro/common/tudriver/src/lib.rs @@ -0,0 +1,81 @@ +pub mod tugraph_client; + +#[cfg(test)] +mod tests { + + use neo4rs::*; + use std::env; + + /// This is the test to test whether the Tugraph is setup. + #[tokio::test] + async fn test_tugraph_setup() { + let tugraph_bolt_url = env::var("TUGRAPH_BOLT_URL").unwrap(); + let tugraph_user_name = env::var("TUGRAPH_USER_NAME").unwrap(); + let tugraph_user_password = env::var("TUGRAPH_USER_PASSWORD").unwrap(); + + // build bolt config + let default_graph_config = ConfigBuilder::default() + .uri(&tugraph_bolt_url) + .user(&tugraph_user_name) + .password(&tugraph_user_password) + .db("default") + .build() + .unwrap(); + + // connect the database + let default_graph = Graph::connect(default_graph_config).await.unwrap(); + + let _ = default_graph + .run(query( + "CALL dbms.graph.createGraph('graph_for_test', 'description', 2045)", + )) + .await; + + let config = ConfigBuilder::default() + .uri(&tugraph_bolt_url) + .user(&tugraph_user_name) + .password(&tugraph_user_password) + .db("graph_for_test") + .build() + .unwrap(); + + let graph = Graph::connect(config).await.unwrap(); + + graph.run(query("CALL db.dropDB()")).await.unwrap(); + graph.run(query("CALL db.createVertexLabel('person', 'id' , 'id' ,INT32, false, 'name' ,STRING, false)")).await.unwrap(); + graph + .run(query( + "CALL db.createEdgeLabel('is_friend','[[\"person\",\"person\"]]')", + )) + .await + .unwrap(); + graph + .run(query( + "create (n1:person {name:'jack',id:1}), (n2:person {name:'lucy',id:2})", + )) + .await + .unwrap(); + graph + .run(query( + "match (n1:person {id:1}), (n2:person {id:2}) create (n1)-[r:is_friend]->(n2)", + )) + .await + .unwrap(); + let mut result = graph + .execute(query("match (n)-[r]->(m) return n,r,m")) + .await + .unwrap(); + + if let Ok(Some(row)) = result.next().await { + let n: Node = row.get("n").unwrap(); + assert_eq!(n.id(), 0); + let r: Relation = row.get("r").unwrap(); + assert_eq!(r.start_node_id(), 0); + assert_eq!(r.end_node_id(), 1); + let m: Node = row.get("m").unwrap(); + assert_eq!(m.id(), 1); + } else { + panic!("Error no result"); + } + } +} diff --git a/extensions/cratespro/common/tudriver/src/tugraph_client.rs b/extensions/cratespro/common/tudriver/src/tugraph_client.rs new file mode 100644 index 000000000..bc361732e --- /dev/null +++ b/extensions/cratespro/common/tudriver/src/tugraph_client.rs @@ -0,0 +1,380 @@ +use neo4rs::*; +use serde_json::{Map, Value}; +use std::collections::HashMap; +use std::error::Error; +use std::fs::File; +use std::io::{self, Read}; + +/// doc: https://tugraph-db.readthedocs.io/zh-cn/latest/5.developer-manual/6.interface/1.query/1.cypher.html +/// https://github.com/TuGraph-family/tugraph-db/blob/master/src/cypher/procedure/procedure.h +#[derive(Clone)] +pub struct TuGraphClient { + pub graph: Graph, +} + +impl TuGraphClient { + /// Initialize TuGraph Client + /// Arguments: + /// * `url`: url for bolt + /// * `user`: user name + /// * `password`: password for user + /// * `db`: graph, default is 'default' + pub async fn new( + uri: &str, + user: &str, + password: &str, + db: &str, + ) -> Result> { + let graph_name = if db.is_empty() { "default" } else { db }; + let config = ConfigBuilder::default() + .uri(uri) + .user(user) + .password(password) + .max_connections(1000) + .fetch_size(10000) + .db(graph_name) + .build()?; + tracing::info!( + "Begin to connect to Tugraph, uri: {uri}, user: {user}, password: {password}, db: {db}" + ); + + let graph = Graph::connect(config).await.unwrap(); + tracing::info!("Success to connect to Tugraph"); + Ok(TuGraphClient { graph }) + } + + /// Reset the database, be carefully + #[allow(dead_code)] + pub(crate) async unsafe fn drop_database(&self) -> Result<(), Box> { + self.graph.run(query("CALL db.dropDB()")).await?; + Ok(()) + } + + pub async fn test_ping(&self) { + let ping_query = "RETURN 1"; + let result = self.exec_query(ping_query).await; + match result { + Ok(_) => tracing::info!("Connection to Neo4j is successful."), + Err(e) => tracing::error!("Failed to connect to Neo4j: {}", e), + } + } + + pub async fn exec_query(&self, q: &str) -> Result, Box> { + let mut labels = vec![]; + //tracing::info!("start query"); + let mut result = self.graph.execute(query(q)).await?; + //tracing::info!("end query"); + while let Some(row) = result.next().await? { + let value: Value = row.to().unwrap(); // 打印出 row 的内容以调试 + //println!("{:#?}", value); + labels.push(serde_json::to_string(&value).unwrap()); + } + Ok(labels) + } + + pub async fn list_edge_labels(&self) -> Result> { + let mut labels = String::default(); + let mut result = self.graph.execute(query("CALL db.edgeLabels()")).await?; + while let Some(row) = result.next().await? { + labels = row.to().unwrap(); + } + Ok(labels) + } + + /// Creates a vertex label in the database. + /// + /// Arguments: + /// * `label_name`: name of vertex label + /// * `primary`: primary field of vertex label + /// * `field_specs`: A slice of tuples where each tuple represents field_spec in the form of (property_name, property_type, is_option). + /// + /// Returns: + /// * `Result<(), Box>` - Ok(()) if successful, or an error wrapped in Box otherwise. + /// Example usage: Create a `person` vertex label with an ID of type INT32 and additional properties for `name` and `age`. + /// ```ignore + /// client.create_vertex_label("person", "id", &[("name".to_string(), "STRING".to_string(), false), ("age".to_string(), "INT32".to_string(), false)]).await.unwrap(); + /// ``` + pub async fn create_vertex_label( + &self, + label_name: &str, + primary_field: &str, + field_specs: &[(String, String, bool)], + ) -> Result<(), Box> { + let mut fields_string = field_specs + .iter() + .map(|(name, type_, option)| format!("'{name}', {type_}, {option}")) + .collect::>() + .join(", "); + + fields_string = + if !fields_string.is_empty() { ", " } else { "" }.to_string() + &fields_string; + + let query_string = format!( + "CALL db.createVertexLabel('{label_name}', '{primary_field}'{fields_string})" + ); + println!("Query: {query_string}"); + self.graph.run(query(&query_string)).await?; + Ok(()) + } + + pub async fn create_subgraph(&self, graph_name: &str) -> Result<(), Box> { + let query_string = format!("CALL dbms.graph.createGraph('{graph_name}')"); + println!("Query: {query_string}"); + self.graph.run(query(&query_string)).await?; + Ok(()) + } + + /// Creates an edge label in the database. + /// + /// Arguments: + /// * `label_name`: Name of the edge label. + /// * `edge_constraints`: Vec of tuple pairs, each representing a valid start and end vertex label for the edge. + /// * `field_specs`: A slice of tuples where each tuple represents a field_spec in the form of (field_name, field_type, is_optional). + /// + /// Returns: + /// * `Result<(), Box>` - Ok(()) if successful, or an error wrapped in Box otherwise. + /// + /// Example usage: Create a `KNOWS` edge label with constraints that it can only exist between `Person` and `Person` or `Person` and `Organization`, and with an optional `name` property of type `int32`. + /// ```ignore + /// client.create_edge_label( + /// "KNOWS".to_string(), + /// &[("Person".to_string(), "Person".to_string()), ("Person".to_string(), "Organization".to_string())], + /// &[("name".to_string(), "INT32".to_string(), true)] + /// ).await.unwrap(); + /// ``` + pub async fn create_edge_label( + &self, + label_name: String, + edge_constraints: &[(String, String)], + field_specs: &[(String, String, bool)], + ) -> Result<(), Box> { + let constraint_strings = edge_constraints + .iter() + .map(|(start_label, end_label)| format!("[\"{start_label}\", \"{end_label}\"]")) + .collect::>() + .join(", "); + + let mut fields_string = field_specs + .iter() + .map(|(name, type_, optional)| { + let option_str = if *optional { "true" } else { "false" }; + format!("'{name}', '{type_}', {option_str}") + }) + .collect::>() + .join(", "); + fields_string = + if !fields_string.is_empty() { ", " } else { "" }.to_string() + &fields_string; + + let query_string = format!( + "CALL db.createEdgeLabel('{label_name}', '[{constraint_strings}]'{fields_string})" + ); + + println!("Query: {query_string}"); + self.graph.run(query(&query_string)).await?; + Ok(()) + } + + /// Loads a plugin into the database. + /// + /// Arguments: + /// * `plugin_name`: The name of the plugin as a STRING. + /// * `plugin_so_path`: The path of the plugin as a STRING. + /// + /// Returns: + /// * `Result<(), Box>` - Ok(()) if successful, or an error wrapped in Box otherwise. + /// + /// Example usage: Load a custom `HelloWorld` plugin. + /// ```ignore + /// client.load_plugin("trace_dependencies", "../trace_dependencies.so").await.unwrap(); + /// ``` + pub async fn load_plugin( + &self, + plugin_name: &str, + plugin_so_path: &str, + ) -> Result<(), Box> { + let plugin_type: &str = "CPP"; + let plugin_description: &str = "plugin"; + let read_only: bool = true; + let version: &str = "v1"; + let code_type: &str = "SO"; + + let mut file = File::open(plugin_so_path)?; + let mut buffer = Vec::new(); + + file.read_to_end(&mut buffer)?; + + use base64::engine::general_purpose::STANDARD; + use base64::Engine; + let engine = STANDARD; + + let plugin_content: &str = &engine.encode(buffer); + + let query_string = format!( + "CALL db.plugin.loadPlugin('{plugin_type}', '{plugin_name}', '{plugin_content}', '{code_type}', '{plugin_description}', {read_only}, '{version}')" + ); + + self.graph.run(query(&query_string)).await.unwrap(); + + println!("load plugin {plugin_name}"); + Ok(()) + } + + /// list the info of loaded plugins + /// + /// # param + /// * `plugin_type` - + /// * `plugin_version` - + pub async fn list_plugin( + &self, + plugin_type: &str, + plugin_version: &str, + ) -> Result, Box> { + let query_string = format!( + "CALL db.plugin.listPlugin('{plugin_type}', '{plugin_version}')" + ); + + let mut result = self.graph.execute(query(&query_string)).await?; + // println!("{:?}", result.next().await); + let mut plugins = Vec::new(); + + while let Some(row) = result.next().await? { + let desc: Map<_, _> = row.get("plugin_description").unwrap(); + let name = desc + .get("name") + .unwrap() + .to_string() + .trim_matches('"') + .to_string(); + plugins.push(name); + } + Ok(plugins) + } + + /// Deletes a plugin based on its type and name. + /// + /// # Parameters + /// * `plugin_type` - The type of the plugin to delete (e.g., "CPP", "PY"). + /// * `plugin_name` - The name of the plugin to delete. + /// + /// # Returns + /// Returns `Ok(())` if the plugin was successfully deleted, or an error if the operation fails. + /// + /// # Errors + /// Returns an error if the delete operation cannot be executed or if the database responds with an error. + pub async fn delete_plugin( + &self, + plugin_type: &str, + plugin_name: &str, + ) -> Result<(), Box> { + let query_string = format!( + "CALL db.plugin.deletePlugin('{plugin_type}', '{plugin_name}')" + ); + + // Executes the query. + // In a real-world scenario, you should handle potential errors properly, + // e.g., if the plugin does not exist or if the arguments are invalid. + self.graph.run(query(&query_string)).await?; + println!("delete plugin {plugin_name}"); + Ok(()) + } + + /// List all the subgraphs in the database + pub async fn list_graphs(&self) -> Result, Box> { + let mut result = self + .graph + .execute(query( + "CALL dbms.graph.listGraphs() YIELD graph_name RETURN graph_name", + )) + .await?; + + let mut names = Vec::new(); + while let Ok(Some(row)) = result.next().await { + let name: String = row.get("graph_name")?; + names.push(name); + } + + Ok(names) + } + + /// Retrieves information about a plugin based on its type and name. + /// + /// # Parameters + /// * `plugin_type` - The type of the plugin to retrieve information for. + /// * `plugin_name` - The name of the plugin to retrieve information for. + /// * `show_code` - A boolean flag to indicate whether to include code in the plugin information. + /// + /// # Returns + /// Returns a `Result` with a map containing plugin descriptions if the operation is successful, or an error if it fails. + /// + /// # Errors + /// Returns an error if the retrieval operation fails or if the database returns an error. + pub async fn get_plugin_info( + &self, + plugin_type: &str, + plugin_name: &str, + show_code: bool, + ) -> Result, Box> { + let show_code_param = if show_code { "true" } else { "false" }; + let query_string = format!( + "CALL db.plugin.getPluginInfo('{plugin_type}', '{plugin_name}', {show_code_param})" + ); + + // Executes the query. + // Proper error handling should be implemented in a real-world scenario. + let mut result = self.graph.execute(query(&query_string)).await?; + + if let Some(row) = result.next().await? { + let desc: Map<_, _> = row.get("plugin_description").unwrap(); + let mut map: HashMap = HashMap::default(); + for (k, v) in &desc { + map.insert(k.clone(), v.to_string()); + } + + return Ok(map); + } + + Err(io::Error::new(io::ErrorKind::NotFound, "No data found").into()) + } + + /// Calls a plugin with specified parameters. + /// + /// # Parameters + /// * `plugin_type` - The type of the plugin to call. + /// * `plugin_name` - The name of the plugin to call. + /// * `param` - Additional parameter for the plugin call. + /// * `timeout` - Timeout value for the plugin call. + /// * `in_process` - A flag indicating whether to process the plugin internally. + /// + /// # Returns + /// Returns a `Result` with a tuple containing success status and result if the call is successful, or an error if it fails. + /// + /// # Errors + /// Returns an error if the plugin call fails or encounters an issue during execution. + pub async fn call_plugin( + &self, + plugin_type: &str, + plugin_name: &str, + param: &str, + timeout: f64, + in_process: bool, + ) -> Result<(bool, String), Box> { + let in_process_param = if in_process { "true" } else { "false" }; + let query_string = format!( + "CALL db.plugin.callPlugin('{plugin_type}', '{plugin_name}', '{param}', {timeout}, {in_process_param})" + ); + + // Executes the query. + // Proper error handling should be implemented in a real-world scenario. + let mut result = self.graph.execute(query(&query_string)).await?; + + if let Some(row) = result.next().await? { + let success: bool = true; + // row.get("success").unwrap(); + let result_str: String = row.get("result").unwrap(); + + return Ok((success, result_str)); + } + + Err(io::Error::new(io::ErrorKind::NotFound, "No data found").into()) + } +} diff --git a/extensions/cratespro/import_tugraph/.env b/extensions/cratespro/import_tugraph/.env new file mode 100644 index 000000000..5f70f29fd --- /dev/null +++ b/extensions/cratespro/import_tugraph/.env @@ -0,0 +1,44 @@ + +CLONE_CRATES_DIR="target/local_crates_file/" +TUGRAPH_IMPORT_FILES_PG= "target/tugraph_import_files_mq/" +NEW_CRATES_DIR="target/new_crates_file" +SPLIT_CRATES_DIR="target/split_crates_file" + +CHECKPOINT_DIR="target/checkpoints" + +KAFKA_BROKER="172.17.0.1:30092" +KAFKA_IMPORT_TOPIC="REPO_SYNC_STATUS.dev.0102" +KAFKA_USER_IMPORT_TOPIC="USER_IMPORT" +KAFKA_ANALYSIS_TOPIC="ANALYSIS.full.20250902-1747" +KAFKA_ANALYSIS_TEST_TOPIC="ANALYSIS.demo.20250902" +KAFKA_CONSUMER_GROUP_ID="test_lhw39" + +POSTGRES_HOST_IP="172.17.0.1" +POSTGRES_HOST_PORT="31432" +POSTGRES_USER_NAME="mega" +POSTGRES_USER_PASSWORD="mega" +POSTGRES_CRATESPRO_DB="cratespro" +POSTGRES_CRATESIO_DB="cratesio" + +TUGRAPH_BOLT_URL="bolt://172.17.0.1:30687" +TUGRAPH_USER_NAME="admin" +TUGRAPH_USER_PASSWORD="rust@2024" +TUGRAPH_CRATESPRO_DB="cratespro" +LGRAPH_IMPORT_PATH="/usr/local/bin/lgraph_import" +LGRAPH_CONFIG_PATH="/home/rust/lhw/import_tugraph/import.config" +LGRAPH_DB_DIR="/home/rust/tugraph/lgraph_db" + +REDIS_HOST = "172.17.0.1" +REDIS_PASSWORD = "rust@2024" + +MEGA_BASE_URL="http://git.gitmega.nju:8000" + +TABLE_NAME="programs" +OPENAI_API_KEY="" + +RUST_LOG=info + +CRATES_PRO_IMPORT=1 +CRATES_PRO_ANALYSIS=1 +CRATES_PRO_PACKAGE=1 +SHOULD_RESET_KAFKA_OFFSET=0 diff --git a/extensions/cratespro/import_tugraph/BUCK b/extensions/cratespro/import_tugraph/BUCK new file mode 100644 index 000000000..c88306efb --- /dev/null +++ b/extensions/cratespro/import_tugraph/BUCK @@ -0,0 +1,30 @@ +load("@prelude//rust:cargo_package.bzl", "cargo") + +filegroup( + name = "import-tugraph-0.1.0.crate", + srcs = [ + "src/cli.rs", + "src/core_controller.rs", + "src/main.rs", + ], +) + +pkg_deps = [ + "//project/crates-pro:repo_import", + "//third-party:dotenvy", + "//third-party:futures-util", + "//third-party:serde_json", + "//third-party:structopt", + "//third-party:tokio", + "//third-party:tracing", + "//third-party:tracing-subscriber", +] + +cargo.rust_binary( + name = "import_tugraph", + srcs = [":import-tugraph-0.1.0.crate"], + crate_root = "import-tugraph-0.1.0.crate/src/main.rs", + edition = "2024", + deps = pkg_deps, + visibility = ["PUBLIC"], +) diff --git a/extensions/cratespro/import_tugraph/Cargo.toml b/extensions/cratespro/import_tugraph/Cargo.toml new file mode 100644 index 000000000..535abb087 --- /dev/null +++ b/extensions/cratespro/import_tugraph/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "import_tugraph" +version = "0.1.0" +edition = "2021" +default-run = "import_tugraph" + +[[bin]] +name = "import_tugraph" +path = "src/main.rs" + + +[dependencies] +model = { workspace = true } +repo_import = { workspace = true } + +# third-party (第三方依赖, 不写具体版本号, 具体版本只在根目录 Cargo.toml 中出现) +dotenvy = "0.15" +futures = "0.3" +futures-util = "0.3" +neo4rs = "0.8" +rdkafka = { version = "0.37", features = ["cmake-build"] } +serde_json = "1.0" +structopt = "0.3" +tokio = { version = "1.44", features = ["full"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +[dev-dependencies] +serial_test = "0.5" \ No newline at end of file diff --git a/extensions/cratespro/import_tugraph/Dockerfile b/extensions/cratespro/import_tugraph/Dockerfile new file mode 100644 index 000000000..9b085ac9a --- /dev/null +++ b/extensions/cratespro/import_tugraph/Dockerfile @@ -0,0 +1,33 @@ +FROM c9-ssl32-gh19 + +# Install dependencies +# libgfortran: TuGraph 4.5.2 (Dockerfile) +RUN dnf update -y \ + && dnf install -y libgfortran + +# Install TuGraph 4.5.2 +RUN curl -LO https://github.com/TuGraph-family/tugraph-db/releases/download/v4.5.2/tugraph-4.5.2-1.el8.x86_64.rpm \ + && dnf install -y tugraph-4.5.2-1.el8.x86_64.rpm \ + && rm -f tugraph-4.5.2-1.el8.x86_64.rpm +ENV LD_LIBRARY_PATH="/usr/local/lib64/lgraph:$LD_LIBRARY_PATH" + +ARG USERNAME="rust" +USER $USERNAME + +# Create and set permissions for workdir directory +USER root +RUN mkdir -p /workdir && chown $USERNAME:$USERNAME /workdir +USER $USERNAME + +WORKDIR /workdir + +# Copy artifacts +COPY ./import_tugraph ./import_tugraph +COPY ./import.config ./import.config +COPY ./.env ./.env + +RUN mkdir log +RUN mkdir target + +ENV RUST_BACKTRACE=1 +CMD ["./import_tugraph"] diff --git a/extensions/cratespro/import_tugraph/deploy-import-tugraph.sh b/extensions/cratespro/import_tugraph/deploy-import-tugraph.sh new file mode 100755 index 000000000..b282cdda5 --- /dev/null +++ b/extensions/cratespro/import_tugraph/deploy-import-tugraph.sh @@ -0,0 +1,57 @@ +#!/bin/bash +set -euxo pipefail + +# source +CODE_DIR=$(pwd) +CRATESPRO_DIR=/home/rust/crates-pro +INFRA_DIR=/home/rust/workspace/crates-pro-infra +MEGA_DIR=/home/rust/lhw/mega +# deployment +NAMESPACE=crates-pro +DEPLOYMENT=import-tugraph +KAFKA_HOST=172.17.0.1:30092 +KAFKA_CONSUMER_GROUP=test_lhw38 +# build +BUILD_DIR=$(mktemp -d) +TIMESTAMP=$(date +%Y%m%d-%H%M) +IMAGE_NAME=localhost:30500/crates-pro-import-tugraph:local-$TIMESTAMP + +### Preparation: Sync source directories +rsync --delete --archive $CRATESPRO_DIR/ $INFRA_DIR/project/crates-pro --exclude="/.git" --exclude="/buck-out" --exclude="/build" --exclude="/target" +rsync --delete --archive $CODE_DIR/ $INFRA_DIR/project/crates-pro/import_tugraph +rsync --delete --archive $MEGA_DIR/extensions/cratespro/common/repo_import/ $INFRA_DIR/project/crates-pro/repo_import/ +rsync --delete --archive $MEGA_DIR/extensions/cratespro/common/database/ $INFRA_DIR/project/crates-pro/database/ +rsync --delete --archive $MEGA_DIR/extensions/cratespro/common/model/ $INFRA_DIR/project/crates-pro/model/ +### Step 1: Compile, then copy artifacts to $BUILD_DIR +cd $INFRA_DIR +buck2 build //project/crates-pro/import_tugraph:import_tugraph --out $BUILD_DIR/import_tugraph +cp $CODE_DIR/import.config $BUILD_DIR/import.config +cp $CODE_DIR/.env $BUILD_DIR/.env +cd $CODE_DIR + +### Step 2: Build Docker images +docker build -t $IMAGE_NAME -f $CODE_DIR/Dockerfile $BUILD_DIR + +### Step 3: Push Docker images +docker push $IMAGE_NAME + +### Step 4: Stop current containers +# Scale deployment to 0 replicas +kubectl scale deployment $DEPLOYMENT -n $NAMESPACE --replicas=0 + +# Wait until all pods are terminated +while kubectl get pods -n $NAMESPACE | grep $DEPLOYMENT > /dev/null; do + sleep 5 +done + +### Step 5: Set new images +kubectl set image deployment/$DEPLOYMENT -n $NAMESPACE container-0=$IMAGE_NAME + +# Wait until all kafka consumers are removed +while docker run --rm -t bitnami/kafka -- kafka-consumer-groups.sh --bootstrap-server $KAFKA_HOST --group $KAFKA_CONSUMER_GROUP --describe | grep rdkafka > /dev/null; do + sleep 5 +done + +### Step 6: Run new containers +# Scale deployment back to 1 replica +kubectl scale deployment $DEPLOYMENT -n $NAMESPACE --replicas=1 diff --git a/extensions/cratespro/import_tugraph/deployment.yaml b/extensions/cratespro/import_tugraph/deployment.yaml new file mode 100644 index 000000000..3e6092f31 --- /dev/null +++ b/extensions/cratespro/import_tugraph/deployment.yaml @@ -0,0 +1,54 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: import-tugraph + namespace: crates-pro +spec: + replicas: 1 + selector: + matchLabels: + app: import-tugraph + template: + metadata: + labels: + app: import-tugraph + spec: + securityContext: + runAsUser: 1000 + runAsGroup: 1000 + fsGroup: 1000 + containers: + - name: container-0 + image: localhost:30500/cratespro-import-tugraph:local-20250903-0740 + imagePullPolicy: IfNotPresent + env: + - name: MEGA_BASE_URL + value: "http://git.gitmega.nju:8000" + # - name: KAFKA_CONSUMER_GROUP_ID + # value: "instance-test1-import-tugraph-group" + - name: POSTGRES_HOST_IP + value: "postgresql-test1.postgresql.svc.cluster.local" + - name: POSTGRES_HOST_PORT + value: "5432" + - name: TUGRAPH_BOLT_URL + value: "bolt://tugraph-test1.tugraph.svc.cluster.local:7687" + - name: LGRAPH_IMPORT_PATH + value: "/usr/local/bin/lgraph_import" + - name: LGRAPH_CONFIG_PATH + value: "/workdir/import.config" + - name: LGRAPH_DB_DIR + value: "/var/lgraph_db" + volumeMounts: + - name: output-volume + mountPath: /workdir/target + - name: lgraph-db-volume + mountPath: /var/lgraph_db + volumes: + - name: output-volume + hostPath: + path: /mnt/devops/instances/test1/cratespro-backend/output + type: Directory + - name: lgraph-db-volume + hostPath: + path: /home/rust/tugraph/lgraph_db + type: Directory diff --git a/extensions/cratespro/import_tugraph/import.config b/extensions/cratespro/import_tugraph/import.config new file mode 100644 index 000000000..268fe7d21 --- /dev/null +++ b/extensions/cratespro/import_tugraph/import.config @@ -0,0 +1,201 @@ +{ + "schema": [ + { + "label": "program", + "type": "VERTEX", + "properties": [ + { "name": "id", "type": "STRING" }, + { "name": "name", "type": "STRING" }, + { "name": "description", "type": "STRING" }, + { "name": "namespace", "type": "STRING" }, + { "name": "max_version", "type": "STRING" }, + { "name": "github_url", "type": "STRING" }, + { "name": "mega_url", "type": "STRING" }, + { "name": "doc_url", "type": "STRING" } + ], + "primary": "id" + }, + { + "label": "library", + "type": "VERTEX", + "properties": [ + { "name": "id", "type": "STRING" }, + { "name": "name", "type": "STRING" }, + { "name": "downloads", "type": "INT64" }, + { "name": "cratesio", "type": "STRING", "optional": true } + ], + "primary": "id" + }, + { + "label": "application", + "type": "VERTEX", + "properties": [ + { "name": "id", "type": "STRING" }, + { "name": "name", "type": "STRING" } + ], + "primary": "id" + }, + { + "label": "library_version", + "type": "VERTEX", + "properties": [ + { "name": "id", "type": "STRING" }, + { "name": "name_and_version", "type": "STRING" }, + { "name": "name", "type": "STRING" }, + { "name": "version", "type": "STRING" }, + { "name": "documentation", "type": "STRING" } + ], + "primary": "name_and_version" + }, + { + "label": "application_version", + "type": "VERTEX", + "properties": [ + { "name": "id", "type": "STRING" }, + { "name": "name_and_version", "type": "STRING" }, + { "name": "name", "type": "STRING" }, + { "name": "version", "type": "STRING" } + ], + "primary": "name_and_version" + }, + { + "label": "version", + "type": "VERTEX", + "properties": [ + { "name": "name_and_version", "type": "STRING" } + ], + "primary": "name_and_version" + }, + { + "label": "has_type", + "type": "EDGE", + "properties": [], + "constraints": [["program", "library"], ["program", "application"]] + }, + { + "label": "has_version", + "type": "EDGE", + "properties": [], + "constraints": [["library", "library_version"], ["application", "application_version"]] + }, + { + "label": "has_dep_version", + "type": "EDGE", + "properties": [], + "constraints": [["library_version", "version"], ["application_version", "version"]] + }, + { + "label": "depends_on", + "type": "EDGE", + "properties": [], + "constraints": [["version", "version"]] + } + ], + "files": [ + { + "path": "/workdir/target/tugraph_import_files_mq/program.csv", + "header": 1, + "format": "CSV", + "label": "program", + "columns": ["description","doc_url","github_url","id","max_version","mega_url","name","namespace"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/library.csv", + "header": 1, + "format": "CSV", + "label": "library", + "columns": ["cratesio","downloads","id","name"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/application.csv", + "header": 1, + "format": "CSV", + "label": "application", + "columns": ["id", "name"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/library_version.csv", + "header": 1, + "format": "CSV", + "label": "library_version", + "columns": ["documentation","id","name","name_and_version","version"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/application_version.csv", + "header": 1, + "format": "CSV", + "label": "application_version", + "columns": ["id","name","name_and_version","version"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/version.csv", + "header": 1, + "format": "CSV", + "label": "version", + "columns": ["name_and_version"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/has_lib_type.csv", + "header": 1, + "format": "CSV", + "label": "has_type", + "SRC_ID": "program", + "DST_ID": "library", + "columns": ["DST_ID", "SRC_ID"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/has_app_type.csv", + "header": 1, + "format": "CSV", + "label": "has_type", + "SRC_ID": "program", + "DST_ID": "application", + "columns": ["DST_ID", "SRC_ID"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/lib_has_version.csv", + "header": 1, + "format": "CSV", + "label": "has_version", + "SRC_ID": "library", + "DST_ID": "library_version", + "columns": ["DST_ID", "SRC_ID"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/app_has_version.csv", + "header": 1, + "format": "CSV", + "label": "has_version", + "SRC_ID": "application", + "DST_ID": "application_version", + "columns": ["DST_ID", "SRC_ID"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/lib_has_dep_version.csv", + "header": 1, + "format": "CSV", + "label": "has_dep_version", + "SRC_ID": "library_version", + "DST_ID": "version", + "columns": ["DST_ID", "SRC_ID"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/app_has_dep_version.csv", + "header": 1, + "format": "CSV", + "label": "has_dep_version", + "SRC_ID": "application_version", + "DST_ID": "version", + "columns": ["DST_ID", "SRC_ID"] + }, + { + "path": "/workdir/target/tugraph_import_files_mq/depends_on.csv", + "header": 1, + "format": "CSV", + "label": "depends_on", + "SRC_ID": "version", + "DST_ID": "version", + "columns": ["DST_ID", "SRC_ID"] + } + ] +} \ No newline at end of file diff --git a/extensions/cratespro/import_tugraph/src/cli.rs b/extensions/cratespro/import_tugraph/src/cli.rs new file mode 100644 index 000000000..e249b7671 --- /dev/null +++ b/extensions/cratespro/import_tugraph/src/cli.rs @@ -0,0 +1,19 @@ +use structopt::StructOpt; + +#[derive(StructOpt, Debug, Default, Clone)] +pub struct CratesProCli { + #[structopt(subcommand)] + pub(crate) _command: Option, + + #[structopt(short, long)] + pub(crate) _mega_base: Option, + + #[structopt(short, long)] + pub(crate) dont_clone: bool, +} + +#[derive(StructOpt, Debug, Default, Clone)] +pub enum Command { + #[default] + Mega, +} diff --git a/extensions/cratespro/import_tugraph/src/core_controller.rs b/extensions/cratespro/import_tugraph/src/core_controller.rs new file mode 100644 index 000000000..1b6405689 --- /dev/null +++ b/extensions/cratespro/import_tugraph/src/core_controller.rs @@ -0,0 +1,115 @@ +//! the core controller module +//! It receives messages from Kafka MQ one by one, +//! parse them, and store it into tugraph, and notify +//! other processes. + +//use analysis::analyse_once; +#[allow(unused_imports)] +use repo_import::ImportDriver; + +use crate::cli::CratesProCli; +use futures_util::future::FutureExt; +//use std::process::Command; +//use std::process::Command; +use std::sync::atomic::{AtomicBool, Ordering}; +#[allow(unused_imports)] +use std::{env, fs, sync::Arc, time::Duration}; +use tokio::signal::unix::{signal, SignalKind}; +use tokio::sync::Mutex; + +pub struct CoreController { + pub cli: CratesProCli, + + pub import: bool, +} +struct SharedState { + is_packaging: bool, +} + +impl CoreController { + pub async fn new(cli: CratesProCli) -> Self { + let import = env::var("CRATES_PRO_IMPORT").unwrap().eq("1"); + Self { + cli, + import, + } + } + + pub async fn run(&self) { + let import = self.import; + + let shared_state: Arc> = + Arc::new(Mutex::new(SharedState { + is_packaging: false, + })); + + let dont_clone = self.cli.dont_clone; + let state_clone1: Arc> = Arc::clone(&shared_state); + let import_task = tokio::spawn(async move { + if import { + let should_reset_kafka_offset = + env::var("SHOULD_RESET_KAFKA_OFFSET").unwrap().eq("1"); + if should_reset_kafka_offset { + repo_import::reset_kafka_offset() + .await + .unwrap_or_else(|x| panic!("{}", x)); + } + + let mut import_driver = ImportDriver::new(dont_clone).await; + let mut count = 0; + let is_importing = Arc::new(AtomicBool::new(false)); + let is_importing_clone = Arc::clone(&is_importing); + let mut term_signal = signal(SignalKind::terminate()).unwrap(); + let mut received_term = false; + + loop { + let mut state = state_clone1.lock().await; + while state.is_packaging { + drop(state); + tokio::time::sleep(Duration::from_secs(1)).await; + state = state_clone1.lock().await; + } + + is_importing_clone.store(true, Ordering::SeqCst); + let result = import_driver.import_from_mq_for_a_message().await; + is_importing_clone.store(false, Ordering::SeqCst); + + if !received_term && term_signal.recv().now_or_never().is_some() { + tracing::info!( + "Import task received SIGTERM, will exit after current message" + ); + received_term = true; + } + + if received_term { + tracing::info!("Import task saving final checkpoint..."); + if let Err(e) = import_driver.save_checkpoint().await { + tracing::error!("Failed to save final checkpoint: {}", e); + } else { + tracing::info!("Final checkpoint saved successfully"); + } + std::process::exit(0); + } + + count += 1; + if count == 10000 { + import_driver.context.write_tugraph_import_files().await; + let _ =import_driver.context.import_from_env_vars().await; + count = 0; + } + drop(state); + + if result.is_err() { + tokio::time::sleep(Duration::from_secs(0)).await; + } + } + } + }); + + // 只等待实际运行的任务 + if import { + import_task.await.unwrap(); + tracing::info!("Import task completed"); + } + } +} diff --git a/extensions/cratespro/import_tugraph/src/main.rs b/extensions/cratespro/import_tugraph/src/main.rs new file mode 100644 index 000000000..291afecd6 --- /dev/null +++ b/extensions/cratespro/import_tugraph/src/main.rs @@ -0,0 +1,41 @@ +mod cli; +mod core_controller; + +use cli::CratesProCli; +use core_controller::CoreController; +use std::fs::{self, File}; +use std::time::{SystemTime, UNIX_EPOCH}; +use structopt::StructOpt; +use tracing_subscriber::EnvFilter; + +#[tokio::main] +async fn main() { + dotenvy::dotenv().ok(); + + // 获取当前时间戳 + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + // 创建带时间戳的日志文件 + let log_path = format!("log/log_{timestamp}.ans"); + let parent_dir = std::path::Path::new(&log_path) + .parent() + .ok_or_else(|| std::io::Error::other("Invalid log path")).unwrap(); + + fs::create_dir_all(parent_dir).unwrap(); + let file = File::create(&log_path).expect("Unable to create log file"); + + // 设置日志记录器 + tracing_subscriber::fmt() + .with_writer(file) + .with_env_filter(EnvFilter::from_default_env()) + .init(); + + tracing::info!("Starting with log file: {}", log_path); + + let cli = CratesProCli::from_args(); + let core_controller = CoreController::new(cli).await; + core_controller.run().await; +}