diff --git a/.gitignore b/.gitignore index 9033a94e6..05714bd77 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,6 @@ Cargo.lock # Mac .DS_Store .VSCodeCounter# See https://help.github.com/articles/ignoring-files/ for more about ignoring files. + +# zmt agent data file +ztm_agent.db \ No newline at end of file diff --git a/gateway/Cargo.toml b/gateway/Cargo.toml index 63f688e38..36a5a07d2 100644 --- a/gateway/Cargo.toml +++ b/gateway/Cargo.toml @@ -15,6 +15,7 @@ jupiter = { workspace = true } callisto = { workspace = true } ceres = { workspace = true } gemini = { workspace = true } +vault = { workspace = true } venus = { workspace = true } mercury = { workspace = true } anyhow = { workspace = true } diff --git a/gateway/src/https_server.rs b/gateway/src/https_server.rs index 457868bce..292096b0f 100644 --- a/gateway/src/https_server.rs +++ b/gateway/src/https_server.rs @@ -14,7 +14,7 @@ use axum::Router; use axum_server::tls_rustls::RustlsConfig; use clap::Args; use common::enums::ZtmType; -use gemini::ztm::run_ztm_client; +use gemini::ztm::{run_ztm_client, LocalZTM}; use regex::Regex; use tower::ServiceBuilder; use tower_http::cors::{Any, CorsLayer}; @@ -196,7 +196,10 @@ async fn post_method_router( return lfs::lfs_delete_lock(state, &lfs_config, uri.path(), req).await; } else if Regex::new(r"/objects/batch$").unwrap().is_match(uri.path()) { return lfs::lfs_process_batch(state, &lfs_config, req).await; - } else if Regex::new(r"objects/chunkids$").unwrap().is_match(uri.path()) { + } else if Regex::new(r"objects/chunkids$") + .unwrap() + .is_match(uri.path()) + { return lfs::lfs_fetch_chunk_ids(state, &lfs_config, req).await; } // Routing git services. @@ -257,7 +260,6 @@ pub fn check_run_with_ztm(config: Config, common: CommonOptions) { match ztm_type { ZtmType::Agent => { //Mega server join a ztm mesh - let config = config.ztm; let bootstrap_node = match common.bootstrap_node { Some(n) => n, None => { @@ -265,8 +267,10 @@ pub fn check_run_with_ztm(config: Config, common: CommonOptions) { return; } }; - let peer_id = "123".to_string(); - tokio::spawn(async move { run_ztm_client(bootstrap_node, config, peer_id).await }); + let (peer_id, _) = vault::init(); + let ztm: LocalZTM = LocalZTM { agent_port: 7778 }; + ztm.clone().start_ztm_agent(); + tokio::spawn(async move { run_ztm_client(bootstrap_node, config, peer_id, ztm).await }); } ZtmType::Relay => { //Start a sub thread to run relay server diff --git a/gateway/src/relay_server.rs b/gateway/src/relay_server.rs index a12c2d4cd..cf2542eef 100644 --- a/gateway/src/relay_server.rs +++ b/gateway/src/relay_server.rs @@ -2,13 +2,14 @@ use std::net::SocketAddr; use std::str::FromStr; use axum::body::Body; -use axum::extract::{Query, State}; +use axum::extract::{FromRequest, Query, State}; use axum::http::{Request, Response, StatusCode, Uri}; use axum::routing::get; -use axum::Router; +use axum::{Json, Router}; +use callisto::{ztm_node, ztm_repo_info}; use common::config::{Config, ZTMConfig}; -use gemini::ztm::{RemoteZTM, ZTM}; -use gemini::RelayGetParams; +use gemini::ztm::{RemoteZTM, ZTMAgent, ZTMCA}; +use gemini::{Node, RelayGetParams, RelayResultRes, RepoInfo}; use jupiter::context::Context; use regex::Regex; use tower::ServiceBuilder; @@ -63,10 +64,7 @@ pub async fn app(config: Config, host: String, port: u16) -> Router { // add TraceLayer for log record // add CorsLayer to add cors header Router::new() - .nest( - "/api/v1", - api_router::routers().with_state(api_state), - ) + .nest("/api/v1", api_router::routers().with_state(api_state)) .route( "/*path", get(get_method_router).post(post_method_router), @@ -88,6 +86,12 @@ async fn get_method_router( return hello_relay(params).await; } else if Regex::new(r"/certificate$").unwrap().is_match(uri.path()) { return certificate(ztm_config, params).await; + } else if Regex::new(r"/ping$").unwrap().is_match(uri.path()) { + return ping(state, params).await; + } else if Regex::new(r"/node_list$").unwrap().is_match(uri.path()) { + return node_list(state, params).await; + } else if Regex::new(r"/repo_list$").unwrap().is_match(uri.path()) { + return repo_list(state, params).await; } Err(( StatusCode::NOT_FOUND, @@ -97,10 +101,13 @@ async fn get_method_router( async fn post_method_router( state: State, - _uri: Uri, - _req: Request, + uri: Uri, + req: Request, ) -> Result, (StatusCode, String)> { let _ztm_config = state.context.config.ztm.clone(); + if Regex::new(r"/repo_provide$").unwrap().is_match(uri.path()) { + return repo_provide(state, req).await; + } Err(( StatusCode::NOT_FOUND, String::from("Operation not supported\n"), @@ -137,6 +144,101 @@ pub async fn certificate( .unwrap()) } +pub async fn ping( + state: State, + params: RelayGetParams, +) -> Result, (StatusCode, String)> { + let storage = state.context.services.ztm_storage.clone(); + let node: ztm_node::Model = match params.try_into() { + Ok(n) => n, + Err(_) => { + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "invalid paras".to_string(), + )); + } + }; + match storage.insert_or_update_node(node).await { + Ok(_) => { + let res = serde_json::to_string(&RelayResultRes { success: true }).unwrap(); + Ok(Response::builder() + .header("Content-Type", "application/json") + .body(Body::from(res)) + .unwrap()) + } + Err(_) => Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "invalid paras".to_string(), + )), + } +} + +pub async fn node_list( + state: State, + _params: RelayGetParams, +) -> Result, (StatusCode, String)> { + let storage = state.context.services.ztm_storage.clone(); + let nodelist: Vec = storage + .get_all_node() + .await + .unwrap() + .into_iter() + .map(|x| x.into()) + .collect(); + let json_string = serde_json::to_string(&nodelist).unwrap(); + Ok(Response::builder() + .header("Content-Type", "application/json") + .body(Body::from(json_string)) + .unwrap()) +} + +pub async fn repo_provide( + state: State, + req: Request, +) -> Result, (StatusCode, String)> { + let storage = state.context.services.ztm_storage.clone(); + let request = Json::from_request(req, &state) + .await + .unwrap_or_else(|_| Json(RepoInfo::default())); + let repo_info: RepoInfo = request.0; + if repo_info.identifier.is_empty() { + return Err((StatusCode::BAD_REQUEST, "paras invalid".to_string())); + } + let repo_info_model: ztm_repo_info::Model = repo_info.into(); + match storage.insert_or_update_repo_info(repo_info_model).await { + Ok(_) => { + let res = serde_json::to_string(&RelayResultRes { success: true }).unwrap(); + Ok(Response::builder() + .header("Content-Type", "application/json") + .body(Body::from(res)) + .unwrap()) + } + Err(_) => Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "invalid paras".to_string(), + )), + } +} + +pub async fn repo_list( + state: State, + _params: RelayGetParams, +) -> Result, (StatusCode, String)> { + let storage = state.context.services.ztm_storage.clone(); + let repo_info_list: Vec = storage + .get_all_repo_info() + .await + .unwrap() + .into_iter() + .map(|x| x.into()) + .collect(); + let json_string = serde_json::to_string(&repo_info_list).unwrap(); + Ok(Response::builder() + .header("Content-Type", "application/json") + .body(Body::from(json_string)) + .unwrap()) +} + pub async fn relay_connect_ztm(config: ZTMConfig, relay_port: u16) -> Result { // 1. generate a permit for relay let name = "relay".to_string(); diff --git a/gemini/Cargo.toml b/gemini/Cargo.toml index 08f916d36..cdd712775 100644 --- a/gemini/Cargo.toml +++ b/gemini/Cargo.toml @@ -16,4 +16,9 @@ axum = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } reqwest = { workspace = true } -tracing = { workspace = true } \ No newline at end of file +tracing = { workspace = true } +tokio = { workspace = true, features = ["net"] } +chrono = { workspace = true } +jupiter = { workspace = true } +callisto = { workspace = true } +rust-ztm = { git = "https://github.com/flomesh-io/rust-ztm.git" } \ No newline at end of file diff --git a/gemini/src/lib.rs b/gemini/src/lib.rs index 198635f61..37171820d 100644 --- a/gemini/src/lib.rs +++ b/gemini/src/lib.rs @@ -1,9 +1,148 @@ -use serde::Deserialize; +use std::fmt; + +use callisto::{ztm_node, ztm_repo_info}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; pub mod http; pub mod ztm; #[derive(Deserialize, Debug)] pub struct RelayGetParams { + pub peer_id: Option, + pub hub: Option, pub name: Option, + pub agent_name: Option, + pub service_name: Option, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct RelayResultRes { + pub success: bool, +} + +#[derive(Debug)] +pub enum MegaType { + Agent, + Relay, +} + +impl fmt::Display for MegaType { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + MegaType::Agent => write!(f, "Agent"), + MegaType::Relay => write!(f, "Relay"), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct Node { + pub peer_id: String, + pub hub: String, + pub agent_name: String, + pub service_name: String, + pub mega_type: String, + pub online: bool, + pub last_online_time: i64, +} + +#[derive(Debug)] +pub enum ConversionError { + InvalidParas, +} + +impl TryFrom for Node { + type Error = ConversionError; + + fn try_from(paras: RelayGetParams) -> Result { + if paras.peer_id.is_none() + || paras.hub.is_none() + || paras.agent_name.is_none() + || paras.service_name.is_none() + { + return Err(ConversionError::InvalidParas); + } + let now = Utc::now().timestamp_millis(); + Ok(Node { + peer_id: paras.peer_id.unwrap(), + hub: paras.hub.unwrap(), + agent_name: paras.agent_name.unwrap(), + service_name: paras.service_name.unwrap(), + mega_type: MegaType::Agent.to_string(), + online: true, + last_online_time: now, + }) + } +} + +impl TryFrom for ztm_node::Model { + type Error = ConversionError; + + fn try_from(paras: RelayGetParams) -> Result { + if paras.peer_id.is_none() + || paras.hub.is_none() + || paras.agent_name.is_none() + || paras.service_name.is_none() + { + return Err(ConversionError::InvalidParas); + } + let now = Utc::now().timestamp_millis(); + Ok(ztm_node::Model { + peer_id: paras.peer_id.unwrap(), + hub: paras.hub.unwrap(), + agent_name: paras.agent_name.unwrap(), + service_name: paras.service_name.unwrap(), + r#type: MegaType::Agent.to_string(), + online: true, + last_online_time: now, + }) + } +} + +impl From for Node { + fn from(n: ztm_node::Model) -> Self { + Node { + peer_id: n.peer_id, + hub: n.hub, + agent_name: n.agent_name, + service_name: n.service_name, + mega_type: n.r#type, + online: n.online, + last_online_time: n.last_online_time, + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct RepoInfo { + pub name: String, + pub identifier: String, + pub origin: String, + pub update_time: i64, + pub commit: String, +} + +impl From for ztm_repo_info::Model { + fn from(r: RepoInfo) -> Self { + ztm_repo_info::Model { + identifier: r.identifier, + name: r.name, + origin: r.origin, + update_time: r.update_time, + commit: r.commit, + } + } +} + +impl From for RepoInfo { + fn from(r: ztm_repo_info::Model) -> Self { + RepoInfo { + identifier: r.identifier, + name: r.name, + origin: r.origin, + update_time: r.update_time, + commit: r.commit, + } + } } diff --git a/gemini/src/ztm/mod.rs b/gemini/src/ztm/mod.rs index e6a984850..1e649df59 100644 --- a/gemini/src/ztm/mod.rs +++ b/gemini/src/ztm/mod.rs @@ -1,16 +1,18 @@ +use std::{collections::HashMap, thread::sleep, time::Duration}; + use axum::async_trait; -use common::config::ZTMConfig; +use common::config::{Config, ZTMConfig}; use reqwest::{header::CONTENT_TYPE, Client}; use serde::{Deserialize, Serialize}; -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct ZTMUserPermit { pub ca: String, pub agent: CertAgent, pub bootstraps: Vec, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct CertAgent { pub name: String, pub certificate: String, @@ -18,7 +20,7 @@ pub struct CertAgent { pub private_key: String, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct ZTMMesh { pub name: String, pub ca: String, @@ -28,7 +30,7 @@ pub struct ZTMMesh { pub errors: Vec, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct Agent { pub id: String, pub name: String, @@ -36,7 +38,7 @@ pub struct Agent { pub certificate: String, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct ZTMEndPoint { pub id: String, pub name: String, @@ -44,18 +46,18 @@ pub struct ZTMEndPoint { pub online: bool, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct ZTMServiceReq { pub host: String, pub port: u16, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct ZTMPortReq { pub target: ZTMPortService, } -#[derive(Deserialize, Serialize, Debug)] +#[derive(Deserialize, Serialize, Debug, Clone)] pub struct ZTMPortService { pub service: String, } @@ -63,10 +65,13 @@ pub struct ZTMPortService { const MESH_NAME: &str = "relay_mesh"; #[async_trait] -pub trait ZTM { +pub trait ZTMCA { async fn create_ztm_certificate(&self, name: String) -> Result; - // fn create_ztm_certificate(&self, name: String) -> Result; async fn delete_ztm_certificate(&self, name: String) -> Result; +} + +#[async_trait] +pub trait ZTMAgent { async fn connect_ztm_hub(&self, permit: ZTMUserPermit) -> Result; async fn create_ztm_service( &self, @@ -80,8 +85,6 @@ pub trait ZTM { service_name: String, port: u16, ) -> Result; - - async fn search_ztm_endpoint(&self, name: String) -> Result; } pub struct RemoteZTM { @@ -89,7 +92,7 @@ pub struct RemoteZTM { } #[async_trait] -impl ZTM for RemoteZTM { +impl ZTMCA for RemoteZTM { async fn create_ztm_certificate(&self, name: String) -> Result { let ca_address = self.config.ca.clone(); let hub_address = self.config.hub.clone(); @@ -160,7 +163,10 @@ impl ZTM for RemoteZTM { }; Ok(s) } +} +#[async_trait] +impl ZTMAgent for RemoteZTM { async fn connect_ztm_hub(&self, permit: ZTMUserPermit) -> Result { // POST {agent}/api/meshes/${meshName} let permit_string = serde_json::to_string(&permit).unwrap(); @@ -255,12 +261,37 @@ impl ZTM for RemoteZTM { }; Ok(response_text) } +} - async fn search_ztm_endpoint(&self, name: String) -> Result { - // GET {agent}/api/meshes/{mesh}/endpoints - let agent_address = self.config.agent.clone(); - let url = format!("{agent_address}/api/meshes/{MESH_NAME}/endpoints"); - let request_result = reqwest::get(url).await; +#[derive(Debug, Clone)] +pub struct LocalZTM { + pub agent_port: u16, +} + +impl LocalZTM { + pub fn start_ztm_agent(self) { + tokio::spawn(async move { + rust_ztm::start_agent("ztm_agent.db", self.agent_port); + }); + } +} + +#[async_trait] +impl ZTMAgent for LocalZTM { + async fn connect_ztm_hub(&self, permit: ZTMUserPermit) -> Result { + // POST {agent}/api/meshes/${meshName} + let permit_string = serde_json::to_string(&permit).unwrap(); + tracing::info!("permit_string:{permit_string}"); + let agent_port = self.agent_port; + let agent_address = format!("http://127.0.0.1:{agent_port}"); + let url = format!("{agent_address}/api/meshes/{MESH_NAME}"); + let client = Client::new(); + let request_result = client + .post(url) + .header(CONTENT_TYPE, "application/json") + .body(permit_string) + .send() + .await; let response_text = match handle_ztm_response(request_result).await { Ok(s) => s, Err(s) => { @@ -268,25 +299,93 @@ impl ZTM for RemoteZTM { } }; - let endpoint_list: Vec = match serde_json::from_slice(response_text.as_bytes()) - { + let mesh: ZTMMesh = match serde_json::from_slice(response_text.as_bytes()) { Ok(p) => p, Err(e) => { return Err(e.to_string()); } }; - for endpoint in endpoint_list { - if endpoint.name == name { - return Ok(endpoint); + Ok(mesh) + } + + async fn create_ztm_service( + &self, + ep_id: String, + service_name: String, + port: u16, + ) -> Result { + // create a ZTM service + // POST {agent}/api/meshes/${mesh.name}/endpoints/${ep.id}/services/${svcName} + let agent_port = self.agent_port; + let agent_address = format!("http://127.0.0.1:{agent_port}"); + let url = format!( + "{agent_address}/api/meshes/{MESH_NAME}/endpoints/{ep_id}/services/tcp/{service_name}" + ); + let client = Client::new(); + let req = ZTMServiceReq { + host: "127.0.0.1".to_string(), + port, + }; + let req_string = serde_json::to_string(&req).unwrap(); + let request_result = client + .post(url) + .header(CONTENT_TYPE, "application/json") + .body(req_string) + .send() + .await; + let response_text = match handle_ztm_response(request_result).await { + Ok(s) => s, + Err(s) => { + return Err(s); + } + }; + Ok(response_text) + } + + async fn create_ztm_port( + &self, + ep_id: String, + service_name: String, + port: u16, + ) -> Result { + //POST {agent}/api/meshes/${mesh.name}/endpoints/${ep.id}/ports/127.0.0.1/tcp/{port} + // request body {"service":service_name} + let agent_port = self.agent_port; + let agent_address = format!("http://127.0.0.1:{agent_port}"); + let url = format!( + "{agent_address}/api/meshes/{MESH_NAME}/endpoints/{ep_id}/ports/127.0.0.1/tcp/{port}" + ); + let client = Client::new(); + let req = ZTMPortReq { + target: ZTMPortService { + service: service_name, + }, + }; + let req_string = serde_json::to_string(&req).unwrap(); + let request_result = client + .post(url) + .header(CONTENT_TYPE, "application/json") + .body(req_string) + .send() + .await; + let response_text = match handle_ztm_response(request_result).await { + Ok(s) => s, + Err(s) => { + return Err(s); } - } - return Err("endpoint not found".to_string()); + }; + Ok(response_text) } } -pub async fn run_ztm_client(bootstrap_node: String, config: ZTMConfig, peer_id: String) { - let name = peer_id; - let ztm: RemoteZTM = RemoteZTM { config }; +pub async fn run_ztm_client( + bootstrap_node: String, + _config: Config, + peer_id: String, + agent: LocalZTM, +) { + let name = peer_id.clone(); + // let _context = Context::new(config.clone()).await; // 1. to get permit json from bootstrap_node // GET {bootstrap_node}/api/v1/certificate?name={name} @@ -308,7 +407,7 @@ pub async fn run_ztm_client(bootstrap_node: String, config: ZTMConfig, peer_id: }; // 2. join ztm mesh - let mesh = match ztm.connect_ztm_hub(permit).await { + let mesh = match agent.connect_ztm_hub(permit.clone()).await { Ok(m) => m, Err(s) => { tracing::error!(s); @@ -317,20 +416,9 @@ pub async fn run_ztm_client(bootstrap_node: String, config: ZTMConfig, peer_id: }; tracing::info!("connect to ztm hub successfully"); - // // 3. find ztm relay service - // let relay_endpoint = match ztm.search_ztm_endpoint("relay".to_string()).await { - // Ok(endpoint) => endpoint, - // Err(s) => { - // tracing::error!("find relay ztm endpoint failed, {s}"); - // return; - // } - // }; - // let endpoint_id = relay_endpoint.id; - // tracing::info!("find relay ztm endpoint successfully, id = {endpoint_id}"); - - // 4. create a ztm port + // 3. create a ztm port for relay let ztm_port = 8002; - match ztm + match agent .create_ztm_port(mesh.agent.id, "relay".to_string(), ztm_port) .await { @@ -341,6 +429,35 @@ pub async fn run_ztm_client(bootstrap_node: String, config: ZTMConfig, peer_id: } } tracing::info!("create a ztm port successfully, port:{ztm_port}"); + let peer_id_clone = peer_id.clone(); + loop { + ping( + peer_id_clone.clone(), + permit.bootstraps.first().unwrap().to_string(), + ztm_port, + ) + .await; + sleep(Duration::from_secs(15)); + } +} + +pub async fn ping(peer_id: String, hub: String, ztm_port: u16) { + let url = format!("http://127.0.0.1:{ztm_port}/ping"); + let mut params = HashMap::new(); + params.insert("peer_id", peer_id.clone()); + params.insert("hub", hub); + params.insert("agent_name", peer_id.clone()); + params.insert("service_name", peer_id.clone()); + let client = reqwest::Client::new(); + let response = client.get(url.clone()).query(¶ms).send().await; + let response_text = match handle_ztm_response(response).await { + Ok(s) => s, + Err(s) => { + tracing::error!("GET {url} failed,{s}"); + return; + } + }; + tracing::info!("Get {url}, response: {response_text}"); } pub async fn handle_ztm_response( diff --git a/jupiter/callisto/src/lib.rs b/jupiter/callisto/src/lib.rs index 1271d6da4..445fd8195 100644 --- a/jupiter/callisto/src/lib.rs +++ b/jupiter/callisto/src/lib.rs @@ -13,6 +13,7 @@ pub mod git_tree; pub mod import_refs; pub mod lfs_locks; pub mod lfs_objects; +pub mod lfs_split_relation; pub mod mega_blob; pub mod mega_commit; pub mod mega_issue; @@ -23,4 +24,5 @@ pub mod mega_refs; pub mod mega_tag; pub mod mega_tree; pub mod raw_blob; -pub mod lfs_split_relation; \ No newline at end of file +pub mod ztm_node; +pub mod ztm_repo_info; diff --git a/jupiter/callisto/src/ztm_node.rs b/jupiter/callisto/src/ztm_node.rs new file mode 100644 index 000000000..6e29ca519 --- /dev/null +++ b/jupiter/callisto/src/ztm_node.rs @@ -0,0 +1,21 @@ +//! `SeaORM` Entity. Generated by sea-orm-codegen 0.11.3 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] +#[sea_orm(table_name = "ztm_node")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub peer_id: String, + pub hub: String, + pub agent_name: String, + pub service_name: String, + pub r#type: String, + pub online: bool, + pub last_online_time: i64, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/jupiter/callisto/src/ztm_repo_info.rs b/jupiter/callisto/src/ztm_repo_info.rs new file mode 100644 index 000000000..b4112ddda --- /dev/null +++ b/jupiter/callisto/src/ztm_repo_info.rs @@ -0,0 +1,19 @@ +//! `SeaORM` Entity. Generated by sea-orm-codegen 0.11.3 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] +#[sea_orm(table_name = "ztm_repo_info")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub identifier: String, + pub name: String, + pub origin: String, + pub update_time: i64, + pub commit: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/jupiter/src/context.rs b/jupiter/src/context.rs index c70401951..8d98df037 100644 --- a/jupiter/src/context.rs +++ b/jupiter/src/context.rs @@ -4,7 +4,7 @@ use common::config::Config; use crate::storage::{ git_db_storage::GitDbStorage, init::database_connection, lfs_storage::LfsStorage, - mega_storage::MegaStorage, + mega_storage::MegaStorage, ztm_storage::ZTMStorage, }; #[derive(Clone)] @@ -33,6 +33,7 @@ pub struct Service { pub mega_storage: Arc, pub git_db_storage: Arc, pub lfs_storage: Arc, + pub ztm_storage: Arc, } impl Service { @@ -46,6 +47,7 @@ impl Service { GitDbStorage::new(connection.clone(), config.storage.clone()).await, ), lfs_storage: Arc::new(LfsStorage::new(connection.clone()).await), + ztm_storage: Arc::new(ZTMStorage::new(connection.clone()).await), } } @@ -58,6 +60,7 @@ impl Service { mega_storage: Arc::new(MegaStorage::mock()), git_db_storage: Arc::new(GitDbStorage::mock()), lfs_storage: Arc::new(LfsStorage::mock()), + ztm_storage: Arc::new(ZTMStorage::mock()), }) } } diff --git a/jupiter/src/storage/mod.rs b/jupiter/src/storage/mod.rs index 88e89b94c..204e4a9fa 100644 --- a/jupiter/src/storage/mod.rs +++ b/jupiter/src/storage/mod.rs @@ -3,6 +3,7 @@ pub mod git_fs_storage; pub mod init; pub mod lfs_storage; pub mod mega_storage; +pub mod ztm_storage; use async_trait::async_trait; diff --git a/jupiter/src/storage/ztm_storage.rs b/jupiter/src/storage/ztm_storage.rs new file mode 100644 index 000000000..c804301d8 --- /dev/null +++ b/jupiter/src/storage/ztm_storage.rs @@ -0,0 +1,166 @@ +use std::sync::Arc; + +use callisto::{ztm_node, ztm_repo_info}; +use sea_orm::{DatabaseConnection, EntityTrait, InsertResult, IntoActiveModel, Set}; + +use common::errors::MegaError; + +#[derive(Clone)] +pub struct ZTMStorage { + pub connection: Arc, +} + +impl ZTMStorage { + pub fn get_connection(&self) -> &DatabaseConnection { + &self.connection + } + + pub async fn new(connection: Arc) -> Self { + ZTMStorage { connection } + } + + pub fn mock() -> Self { + ZTMStorage { + connection: Arc::new(DatabaseConnection::default()), + } + } + + pub async fn get_node_by_id( + &self, + peer_id: &str, + ) -> Result, MegaError> { + let result = ztm_node::Entity::find_by_id(peer_id) + .one(self.get_connection()) + .await + .unwrap(); + Ok(result) + } + + pub async fn get_all_node(&self) -> Result, MegaError> { + Ok(ztm_node::Entity::find() + .all(self.get_connection()) + .await + .unwrap()) + } + + pub async fn insert_node( + &self, + node: ztm_node::Model, + ) -> Result, MegaError> { + Ok(ztm_node::Entity::insert(node.into_active_model()) + .exec(self.get_connection()) + .await + .unwrap()) + } + + pub async fn update_node(&self, node: ztm_node::Model) -> Result { + Ok(ztm_node::Entity::update(node.into_active_model()) + .exec(self.get_connection()) + .await + .unwrap()) + } + + pub async fn insert_or_update_node( + &self, + node: ztm_node::Model, + ) -> Result { + match self.get_node_by_id(&node.peer_id).await.unwrap() { + Some(_) => { + let mut active_model: ztm_node::ActiveModel = node.clone().into_active_model(); + active_model.hub = Set(node.hub); + active_model.agent_name = Set(node.agent_name); + active_model.service_name = Set(node.service_name); + active_model.r#type = Set(node.r#type); + active_model.online = Set(node.online); + active_model.last_online_time = Set(node.last_online_time); + Ok(ztm_node::Entity::update(active_model) + .exec(self.get_connection()) + .await + .unwrap()) + } + None => { + ztm_node::Entity::insert(node.clone().into_active_model()) + .exec(self.get_connection()) + .await + .unwrap(); + Ok(node) + } + } + } + + pub async fn delete_node_by_id(&self, id: String) { + ztm_node::Entity::delete_by_id(id) + .exec(self.get_connection()) + .await + .unwrap(); + } + + pub async fn get_repo_info_by_id( + &self, + identifier: &str, + ) -> Result, MegaError> { + let result = ztm_repo_info::Entity::find_by_id(identifier) + .one(self.get_connection()) + .await + .unwrap(); + Ok(result) + } + + pub async fn get_all_repo_info(&self) -> Result, MegaError> { + Ok(ztm_repo_info::Entity::find() + .all(self.get_connection()) + .await + .unwrap()) + } + + pub async fn insert_repo_info( + &self, + repo_info: ztm_repo_info::Model, + ) -> Result, MegaError> { + Ok(ztm_repo_info::Entity::insert(repo_info.into_active_model()) + .exec(self.get_connection()) + .await + .unwrap()) + } + + pub async fn update_repo_info( + &self, + repo_info: ztm_repo_info::Model, + ) -> Result { + Ok(ztm_repo_info::Entity::update(repo_info.into_active_model()) + .exec(self.get_connection()) + .await + .unwrap()) + } + + pub async fn insert_or_update_repo_info( + &self, + repo_info: ztm_repo_info::Model, + ) -> Result { + match self + .get_repo_info_by_id(&repo_info.identifier) + .await + .unwrap() + { + Some(_) => { + let mut active_model: ztm_repo_info::ActiveModel = + repo_info.clone().into_active_model(); + active_model.name = Set(repo_info.name); + active_model.origin = Set(repo_info.origin); + active_model.update_time = Set(repo_info.update_time); + active_model.commit = Set(repo_info.commit); + Ok(ztm_repo_info::Entity::update(active_model) + .exec(self.get_connection()) + .await + .unwrap()) + } + None => { + ztm_repo_info::Entity::insert(repo_info.clone().into_active_model()) + .exec(self.get_connection()) + .await + .unwrap(); + Ok(repo_info) + } + } + } +} diff --git a/sql/postgres/pg_20240205__init.sql b/sql/postgres/pg_20240205__init.sql index cde3b8c02..aad29dd22 100644 --- a/sql/postgres/pg_20240205__init.sql +++ b/sql/postgres/pg_20240205__init.sql @@ -224,4 +224,23 @@ CREATE TABLE IF NOT EXISTS "lfs_split_relations" ( "offset" BIGINT NOT NULL, "size" BIGINT NOT NULL, PRIMARY KEY ("ori_oid", "sub_oid", "offset") -) \ No newline at end of file +) + + +CREATE TABLE IF NOT EXISTS "ztm_node" ( + "peer_id" VARCHAR(64) PRIMARY KEY, + "hub" VARCHAR(64), + "agent_name" VARCHAR(64), + "service_name" VARCHAR(64), + "type" VARCHAR(64), + "online" BOOLEAN NOT NULL, + "last_online_time" BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS "ztm_repo_info" ( + "identifier" VARCHAR(128) PRIMARY KEY, + "name" VARCHAR(64), + "origin" VARCHAR(64), + "update_time" BIGINT NOT NULL, + "commit" VARCHAR(64) +); \ No newline at end of file