Skip to main content

nexql_tools/
completions.rs

1// SPDX-License-Identifier: GPL-3.0-only
2// Copyright (C) 2026 NexQL-OSS Team
3
4//! Minimal `completions/complete` for `ref` tool arguments from the schema index.
5
6use 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
45/// Suggests `schema.name` refs from indexed shards for tool `ref` arguments.
46pub struct CompletionsProvider {
47    store: IndexStore,
48}
49
50impl CompletionsProvider {
51    pub fn new(store: IndexStore) -> Self {
52        Self { store }
53    }
54
55    /// Complete when `argument_name` looks like a schema object ref.
56    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}