nexql_tools/
completions.rs1use nexql_index::IndexStore;
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9
10const MAX_SUGGESTIONS: usize = 50;
11
12#[derive(Debug, Error)]
13pub enum CompletionError {
14 #[error("{0}")]
15 InvalidParams(String),
16 #[error("{0}")]
17 Internal(String),
18}
19
20impl CompletionError {
21 pub fn code(&self) -> i32 {
22 match self {
23 Self::InvalidParams(_) => -32602,
24 Self::Internal(_) => -32603,
25 }
26 }
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct CompletionValue {
31 pub value: String,
32 #[serde(skip_serializing_if = "Option::is_none")]
33 pub description: Option<String>,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct CompletionResult {
38 pub values: Vec<CompletionValue>,
39 #[serde(rename = "total", skip_serializing_if = "Option::is_none")]
40 pub total: Option<usize>,
41 #[serde(rename = "hasMore", skip_serializing_if = "Option::is_none")]
42 pub has_more: Option<bool>,
43}
44
45pub struct CompletionsProvider {
47 store: IndexStore,
48}
49
50impl CompletionsProvider {
51 pub fn new(store: IndexStore) -> Self {
52 Self { store }
53 }
54
55 pub fn complete_ref(
57 &self,
58 connection_id: &str,
59 database: &str,
60 argument_name: &str,
61 value_prefix: &str,
62 ) -> Result<CompletionResult, CompletionError> {
63 if !is_ref_argument(argument_name) {
64 return Ok(CompletionResult {
65 values: Vec::new(),
66 total: Some(0),
67 has_more: Some(false),
68 });
69 }
70
71 let base = self.store.base_dir(connection_id, database);
72 let Some(manifest) = self
73 .store
74 .read_manifest(&base)
75 .map_err(|e| CompletionError::Internal(e.to_string()))?
76 else {
77 return Ok(CompletionResult {
78 values: Vec::new(),
79 total: Some(0),
80 has_more: Some(false),
81 });
82 };
83
84 let overrides = self
85 .store
86 .read_overrides(&base)
87 .map_err(|e| CompletionError::Internal(e.to_string()))?;
88
89 let prefix_lower = value_prefix.to_ascii_lowercase();
90 let mut refs = Vec::new();
91 for shard in &manifest.shards {
92 let Some(entries) = self
93 .store
94 .read_shard_entries(&base, &shard.file)
95 .map_err(|e| CompletionError::Internal(e.to_string()))?
96 else {
97 continue;
98 };
99 for (ref_, entry) in entries {
100 if entry.excluded == Some(true) {
101 continue;
102 }
103 if let Some(objects) = overrides.as_ref().and_then(|o| o.objects.as_ref())
104 && objects.get(&ref_).and_then(|o| o.excluded) == Some(true)
105 {
106 continue;
107 }
108 if !prefix_lower.is_empty() && !ref_.to_ascii_lowercase().starts_with(&prefix_lower)
109 {
110 continue;
111 }
112 refs.push((ref_, entry.kind.as_str().to_owned()));
113 }
114 }
115 refs.sort_by(|a, b| a.0.cmp(&b.0));
116 let total = refs.len();
117 let has_more = total > MAX_SUGGESTIONS;
118 refs.truncate(MAX_SUGGESTIONS);
119
120 Ok(CompletionResult {
121 values: refs
122 .into_iter()
123 .map(|(value, kind)| CompletionValue {
124 value,
125 description: Some(kind),
126 })
127 .collect(),
128 total: Some(total),
129 has_more: Some(has_more),
130 })
131 }
132}
133
134fn is_ref_argument(name: &str) -> bool {
135 matches!(name, "ref" | "table" | "from" | "to" | "a" | "b" | "object")
136 || name.ends_with("_ref")
137 || name.ends_with("Ref")
138}
139
140#[cfg(test)]
141mod tests {
142 use super::*;
143
144 #[test]
145 fn ref_argument_detection() {
146 assert!(is_ref_argument("ref"));
147 assert!(is_ref_argument("table"));
148 assert!(!is_ref_argument("sql"));
149 assert!(!is_ref_argument("limit"));
150 }
151}