Refactor gRPC reflection!

This commit is contained in:
Gregory Schier
2024-06-20 12:49:58 -07:00
parent 8f06a834c8
commit c18d30b89f
15 changed files with 169 additions and 123 deletions

7
src-tauri/Cargo.lock generated
View File

@@ -2290,6 +2290,7 @@ dependencies = [
"hyper 0.14.29",
"hyper-rustls 0.24.2",
"log",
"md5",
"prost",
"prost-reflect",
"prost-types",
@@ -3250,6 +3251,12 @@ dependencies = [
"digest",
]
[[package]]
name = "md5"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771"
[[package]]
name = "memchr"
version = "2.7.2"

View File

@@ -20,3 +20,4 @@ hyper-rustls = { version = "0.24.0", features = ["http2"] }
uuid = { version = "1.7.0", features = ["v4"] }
tauri = { version = "2.0.0-beta" }
tauri-plugin-shell = "2.0.0-beta"
md5 = "0.7.0"

View File

@@ -16,7 +16,9 @@ use tonic::transport::Uri;
use tonic::{IntoRequest, IntoStreamingRequest, Request, Response, Status, Streaming};
use crate::codec::DynamicCodec;
use crate::proto::{fill_pool, fill_pool_from_files, get_transport, method_desc_to_path};
use crate::proto::{
fill_pool_from_files, fill_pool_from_reflection, get_transport, method_desc_to_path,
};
use crate::{json_schema, MethodDefinition, ServiceDefinition};
#[derive(Clone)]
@@ -182,39 +184,36 @@ impl GrpcHandle {
}
impl GrpcHandle {
pub async fn services_from_files(
&mut self,
id: &str,
paths: Vec<PathBuf>,
) -> Result<Vec<ServiceDefinition>, String> {
let pool_key = format!(
"{}-{}",
id,
paths
.iter()
.map(|p| p.to_string_lossy().to_string())
.collect::<Vec<String>>()
.join(":")
);
let pool = fill_pool_from_files(&self.app_handle, paths).await?;
self.pools.insert(pool_key, pool.clone());
Ok(self.services_from_pool(&pool))
}
pub async fn services_from_reflection(
pub async fn reflect(
&mut self,
id: &str,
uri: &str,
) -> Result<Vec<ServiceDefinition>, String> {
// Short-circuit if no URL is set
if uri.is_empty() {
return Ok(Vec::new());
}
proto_files: &Vec<PathBuf>,
) -> Result<(), String> {
let pool = if proto_files.is_empty() {
let full_uri = uri_from_str(uri)?;
fill_pool_from_reflection(&full_uri).await
} else {
fill_pool_from_files(&self.app_handle, proto_files).await
}?;
let uri = uri_from_str(uri)?;
let pool = fill_pool(&uri).await?;
let pool_key = format!("{}-{}", id, uri);
self.pools.insert(pool_key, pool.clone());
self.pools
.insert(make_pool_key(id, uri, proto_files), pool.clone());
Ok(())
}
pub async fn services(
&mut self,
id: &str,
uri: &str,
proto_files: &Vec<PathBuf>,
) -> Result<Vec<ServiceDefinition>, String> {
// Ensure reflection is up-to-date
self.reflect(id, uri, proto_files).await?;
let pool = self
.get_pool(id, uri, proto_files)
.ok_or("Failed to get pool".to_string())?;
Ok(self.services_from_pool(&pool))
}
@@ -247,25 +246,26 @@ impl GrpcHandle {
&mut self,
id: &str,
uri: &str,
proto_files: Vec<PathBuf>,
proto_files: &Vec<PathBuf>,
) -> Result<GrpcConnection, String> {
let uri = uri_from_str(uri)?;
let pool = match self.pools.get(id) {
Some(p) => p.clone(),
None => match proto_files.len() {
0 => fill_pool(&uri).await?,
_ => {
let pool = fill_pool_from_files(&self.app_handle, proto_files).await?;
self.pools.insert(id.to_string(), pool.clone());
pool
}
},
};
self.reflect(id, uri, proto_files).await?;
let pool = self
.get_pool(id, uri, proto_files)
.ok_or("Failed to get pool")?;
let uri = uri_from_str(uri)?;
let conn = get_transport();
let connection = GrpcConnection { pool, conn, uri };
let connection = GrpcConnection {
pool: pool.clone(),
conn,
uri,
};
Ok(connection)
}
fn get_pool(&self, id: &str, uri: &str, proto_files: &Vec<PathBuf>) -> Option<&DescriptorPool> {
self.pools.get(make_pool_key(id, uri, proto_files).as_str())
}
}
fn decorate_req<T>(metadata: HashMap<String, String>, req: &mut Request<T>) -> Result<(), String> {
@@ -287,3 +287,18 @@ fn uri_from_str(uri_str: &str) -> Result<Uri, String> {
}
}
}
fn make_pool_key(id: &str, uri: &str, proto_files: &Vec<PathBuf>) -> String {
let pool_key = format!(
"{}::{}::{}",
id,
uri,
proto_files
.iter()
.map(|p| p.to_string_lossy().to_string())
.collect::<Vec<String>>()
.join(":")
);
format!("{:x}", md5::compute(pool_key))
}

View File

@@ -28,7 +28,7 @@ use tonic_reflection::pb::ServerReflectionRequest;
pub async fn fill_pool_from_files(
app_handle: &AppHandle,
paths: Vec<PathBuf>,
paths: &Vec<PathBuf>,
) -> Result<DescriptorPool, String> {
let mut pool = DescriptorPool::new();
let random_file_name = format!("{}.desc", uuid::Uuid::new_v4());
@@ -121,7 +121,7 @@ pub async fn fill_pool_from_files(
Ok(pool)
}
pub async fn fill_pool(uri: &Uri) -> Result<DescriptorPool, String> {
pub async fn fill_pool_from_reflection(uri: &Uri) -> Result<DescriptorPool, String> {
let mut pool = DescriptorPool::new();
let mut client = ServerReflectionClient::with_origin(get_transport(), uri.clone());

View File

@@ -6,7 +6,7 @@ extern crate objc;
use std::collections::HashMap;
use std::env::current_dir;
use std::fs;
use std::fs::{create_dir_all, File, read_to_string};
use std::fs::{create_dir_all, read_to_string, File};
use std::path::PathBuf;
use std::process::exit;
use std::str::FromStr;
@@ -17,45 +17,45 @@ use fern::colors::ColoredLevelConfig;
use log::{debug, error, info, warn};
use rand::random;
use serde_json::{json, Value};
use sqlx::{Pool, Sqlite, SqlitePool};
use sqlx::migrate::Migrator;
use sqlx::sqlite::SqliteConnectOptions;
use sqlx::types::Json;
use tauri::{AppHandle, LogicalSize, RunEvent, State, WebviewUrl, WebviewWindow};
use tauri::{Manager, WindowEvent};
use sqlx::{Pool, Sqlite, SqlitePool};
use tauri::path::BaseDirectory;
#[cfg(target_os = "macos")]
use tauri::TitleBarStyle;
use tauri::{AppHandle, LogicalSize, RunEvent, State, WebviewUrl, WebviewWindow};
use tauri::{Manager, WindowEvent};
use tauri_plugin_log::{fern, Target, TargetKind};
use tauri_plugin_shell::ShellExt;
use tokio::sync::Mutex;
use ::grpc::{Code, deserialize_message, serialize_message, ServiceDefinition};
use ::grpc::manager::{DynamicMessage, GrpcHandle};
use ::grpc::{deserialize_message, serialize_message, Code, ServiceDefinition};
use crate::analytics::{AnalyticsAction, AnalyticsResource};
use crate::grpc::metadata_to_map;
use crate::http_request::send_http_request;
use crate::models::{
cancel_pending_grpc_connections, cancel_pending_responses, CookieJar,
create_http_response, delete_all_grpc_connections, delete_all_http_responses, delete_cookie_jar,
delete_environment, delete_folder, delete_grpc_connection, delete_grpc_request,
delete_http_request, delete_http_response, delete_workspace, duplicate_grpc_request,
duplicate_http_request, Environment, EnvironmentVariable, Folder, generate_model_id,
get_cookie_jar, get_environment, get_folder, get_grpc_connection,
cancel_pending_grpc_connections, cancel_pending_responses, create_http_response,
delete_all_grpc_connections, delete_all_http_responses, delete_cookie_jar, delete_environment,
delete_folder, delete_grpc_connection, delete_grpc_request, delete_http_request,
delete_http_response, delete_workspace, duplicate_grpc_request, duplicate_http_request,
generate_model_id, get_cookie_jar, get_environment, get_folder, get_grpc_connection,
get_grpc_request, get_http_request, get_http_response, get_key_value_raw,
get_or_create_settings, get_workspace, get_workspace_export_resources, GrpcConnection, GrpcEvent,
GrpcEventType, GrpcRequest, HttpRequest, HttpResponse, KeyValue,
list_cookie_jars, list_environments, list_folders, list_grpc_connections, list_grpc_events,
list_grpc_requests, list_http_requests, list_responses, list_workspaces, ModelType,
set_key_value_raw, Settings, update_response_if_id, update_settings, upsert_cookie_jar, upsert_environment,
upsert_folder, upsert_grpc_connection, upsert_grpc_event, upsert_grpc_request, upsert_http_request, upsert_workspace, Workspace,
get_or_create_settings, get_workspace, get_workspace_export_resources, list_cookie_jars,
list_environments, list_folders, list_grpc_connections, list_grpc_events, list_grpc_requests,
list_http_requests, list_responses, list_workspaces, set_key_value_raw, update_response_if_id,
update_settings, upsert_cookie_jar, upsert_environment, upsert_folder, upsert_grpc_connection,
upsert_grpc_event, upsert_grpc_request, upsert_http_request, upsert_workspace, CookieJar,
Environment, EnvironmentVariable, Folder, GrpcConnection, GrpcEvent, GrpcEventType,
GrpcRequest, HttpRequest, HttpResponse, KeyValue, ModelType, Settings, Workspace,
WorkspaceExportResources,
};
use crate::notifications::YaakNotifier;
use crate::plugin::{
find_plugins, get_plugin, ImportResult, PluginCapability,
run_plugin_export_curl, run_plugin_filter, run_plugin_import,
find_plugins, get_plugin, run_plugin_export_curl, run_plugin_filter, run_plugin_import,
ImportResult, PluginCapability,
};
use crate::render::{render_request, variables_from_environment};
use crate::updates::{UpdateMode, YaakUpdater};
@@ -135,26 +135,18 @@ async fn cmd_grpc_reflect(
.await
.map_err(|e| e.to_string())?;
let uri = safe_uri(req.url.as_str());
if proto_files.len() > 0 {
grpc_handle
.lock()
.await
.services_from_files(
&req.id,
proto_files
.iter()
.map(|p| PathBuf::from_str(p).unwrap())
.collect(),
)
.await
} else {
grpc_handle
.lock()
.await
.services_from_reflection(&req.id, uri.as_str())
.await
}
grpc_handle
.lock()
.await
.services(
&req.id,
&req.url,
&proto_files
.iter()
.map(|p| PathBuf::from_str(p).unwrap())
.collect(),
)
.await
}
#[tauri::command]
@@ -231,6 +223,7 @@ async fn cmd_grpc_go(
workspace_id: req.workspace_id,
request_id: req.id,
status: -1,
elapsed: 0,
url: req.url.clone(),
..Default::default()
},
@@ -238,6 +231,7 @@ async fn cmd_grpc_go(
.await
.map_err(|e| e.to_string())?
};
let conn_id = conn.id.clone();
let base_msg = GrpcEvent {
@@ -270,12 +264,29 @@ async fn cmd_grpc_go(
.connect(
&req.clone().id,
uri.as_str(),
proto_files
&proto_files
.iter()
.map(|p| PathBuf::from_str(p).unwrap())
.collect(),
)
.await?;
.await;
let connection = match connection {
Ok(c) => c,
Err(err) => {
upsert_grpc_connection(
&w,
&GrpcConnection {
elapsed: start.elapsed().as_millis() as i64,
error: Some(err.clone()),
..conn.clone()
},
)
.await
.map_err(|e| e.to_string())?;
return Ok(conn_id);
}
};
let method_desc = connection
.method(&service, &method)

View File

@@ -1,6 +1,6 @@
{
"productName": "yaak",
"version": "__YAAK_VERSION__",
"version": "0.0.0",
"identifier": "app.yaak.desktop",
"build": {
"beforeBuildCommand": "npm run build",