nexql_tools/
completions.rs

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

nexql_tools/
completions.rs
