mod ariadne_diagnostics;
mod db_stats;
mod query_records;
mod query_source_cache;
mod query_source_keys;
use {
crate::{
Partitions,
TRACER,
protocol::{
lsp::LanguageServer,
task::RpcTask,
},
scheduler::task::TaskContext,
},
otel::span,
serde_json::Value,
};
#[cfg(feature = "testing-commands")]
use crate::protocol::jsonrpc;
async fn handle_custom_command<P: Partitions, T: LanguageServer<P>>(
command: &str,
arguments: &[Value],
ctx: &mut TaskContext<P, T>,
) -> jsonrpc::Result<Option<Value>> {
match command {
| "laburnum/dbStats" => db_stats::handle_db_stats(arguments, ctx).await,
| "laburnum/querySourceCache" => {
query_source_cache::handle_query_source_cache(arguments, ctx).await
},
| "laburnum/queryRecords" => {
query_records::handle_query_records(arguments, ctx).await
},
| "laburnum/querySourceKeys" => {
query_source_keys::handle_query_source_keys(arguments, ctx).await
},
| "laburnum/ariadneDiagnostics" => {
ariadne_diagnostics::handle_ariadne_diagnostics(arguments, ctx).await
},
| _ => Err(jsonrpc::Error::method_not_found()),
}
}
impl<P: Partitions, T: LanguageServer<P>> RpcTask<P, T> {
#[cfg(feature = "testing-commands")]
pub(super) fn try_handle_laburnum_command(
&mut self,
method: &str,
params: &serde_json::Value,
request_id: jsonrpc::Id,
) -> bool {
use crate::scheduler::lanes::SYNC_LANE;
if method != "workspace/executeCommand" {
return false;
}
let command = match params.get("command").and_then(|v| v.as_str()) {
| Some(cmd) if cmd.starts_with("laburnum/") => cmd.to_string(),
| _ => return false,
};
let arguments = params
.get("arguments")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
let cx = span!(
^ @TRACER,
"laburnum.commands.try_handle",
"command" = command.clone(),
"method" = method.to_string(),
"request_id" = request_id.to_string()
);
let tx = self.response_tx.clone();
let request_id_clone = request_id.clone();
self.ctx.spawn_task(
move |mut ctx| -> _ {
use opentelemetry::trace::FutureExt;
async move {
let result =
handle_custom_command(&command, &arguments, &mut ctx).await;
let response = match result {
| Ok(value) => {
jsonrpc::Response::from_ok(
request_id_clone.clone(),
value.unwrap_or(serde_json::Value::Null),
)
},
| Err(error) => {
jsonrpc::Response::from_error(request_id_clone.clone(), error)
},
};
if let Err(_e) = tx.send(response).await {}
None
}
.with_context(cx)
},
SYNC_LANE,
);
true
}
}