sql_lsp/
schema.rs

1use dashmap::DashMap;
2use serde::{Deserialize, Serialize};
3use std::str::FromStr;
4use std::sync::Arc;
5use uuid::Uuid;
6
7/// Schema ID,用于隔离不同的 schema 数据
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub struct SchemaId(pub Uuid);
10
11impl SchemaId {
12    pub fn new() -> Self {
13        Self(Uuid::new_v4())
14    }
15}
16
17impl FromStr for SchemaId {
18    type Err = uuid::Error;
19
20    fn from_str(s: &str) -> Result<Self, Self::Err> {
21        Ok(Self(Uuid::parse_str(s)?))
22    }
23}
24
25impl Default for SchemaId {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31/// 数据库 Schema 信息
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct Schema {
34    /// Schema ID
35    pub id: SchemaId,
36    /// 数据库名称
37    pub database: String,
38    /// 表列表
39    pub tables: Vec<Table>,
40    /// 函数列表
41    pub functions: Vec<Function>,
42    /// Schema 定义文件的 URI (可选)
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub source_uri: Option<String>,
45}
46
47/// 表信息
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct Table {
50    /// 表名
51    pub name: String,
52    /// 列列表
53    pub columns: Vec<Column>,
54    /// 表注释
55    pub comment: Option<String>,
56    /// 表定义位置的 URI 和行号 (可选)
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub source_location: Option<(String, u32)>,
59}
60
61/// 列信息
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct Column {
64    /// 列名
65    pub name: String,
66    /// 数据类型
67    pub data_type: String,
68    /// 是否可空
69    pub nullable: bool,
70    /// 列注释
71    pub comment: Option<String>,
72    /// 列定义位置的 URI 和行号 (可选)
73    #[serde(skip_serializing_if = "Option::is_none")]
74    pub source_location: Option<(String, u32)>,
75}
76
77/// 函数信息
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct Function {
80    /// 函数名
81    pub name: String,
82    /// 参数列表
83    pub parameters: Vec<FunctionParameter>,
84    /// 返回类型
85    pub return_type: String,
86    /// 函数描述
87    pub description: Option<String>,
88}
89
90/// 函数参数信息
91#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct FunctionParameter {
93    /// 参数名
94    pub name: String,
95    /// 参数类型
96    pub data_type: String,
97    /// 是否可选
98    pub optional: bool,
99}
100
101/// Schema 管理器,用于管理和隔离不同的 schema
102#[derive(Debug, Clone)]
103pub struct SchemaManager {
104    /// Schema 存储,使用 DashMap 实现线程安全的并发访问
105    schemas: Arc<DashMap<SchemaId, Schema>>,
106}
107
108impl SchemaManager {
109    pub fn new() -> Self {
110        Self {
111            schemas: Arc::new(DashMap::new()),
112        }
113    }
114
115    /// 注册一个新的 schema
116    pub fn register(&self, schema: Schema) -> SchemaId {
117        let id = schema.id;
118        self.schemas.insert(id, schema);
119        id
120    }
121
122    /// 获取指定的 schema
123    pub fn get(&self, id: SchemaId) -> Option<Schema> {
124        self.schemas.get(&id).map(|s| s.clone())
125    }
126
127    /// 更新 schema
128    pub fn update(&self, id: SchemaId, schema: Schema) -> bool {
129        if self.schemas.contains_key(&id) {
130            self.schemas.insert(id, schema);
131            true
132        } else {
133            false
134        }
135    }
136
137    /// 删除 schema
138    pub fn remove(&self, id: SchemaId) -> bool {
139        self.schemas.remove(&id).is_some()
140    }
141
142    /// 列出所有 schema ID
143    pub fn list_ids(&self) -> Vec<SchemaId> {
144        self.schemas.iter().map(|entry| *entry.key()).collect()
145    }
146
147    /// 清空所有 schema
148    pub fn clear(&self) {
149        self.schemas.clear();
150    }
151}
152
153impl Default for SchemaManager {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_schema_id() {
165        let id1 = SchemaId::new();
166        let id2 = SchemaId::new();
167        assert_ne!(id1, id2);
168
169        let id_str = id1.0.to_string();
170        let id3 = SchemaId::from_str(&id_str).unwrap();
171        assert_eq!(id1, id3);
172    }
173
174    #[test]
175    fn test_schema_manager() {
176        let manager = SchemaManager::new();
177
178        let schema = Schema {
179            id: SchemaId::new(),
180            database: "test_db".to_string(),
181            tables: vec![],
182            functions: vec![],
183            source_uri: None,
184        };
185
186        let id = manager.register(schema.clone());
187        assert_eq!(id, schema.id);
188
189        let retrieved = manager.get(id).unwrap();
190        assert_eq!(retrieved.database, "test_db");
191
192        manager.remove(id);
193        assert!(manager.get(id).is_none());
194    }
195
196    #[tokio::test]
197    async fn test_schema_manager_concurrent() {
198        let manager = SchemaManager::new();
199        let manager_clone = manager.clone();
200
201        let schema1 = Schema {
202            id: SchemaId::new(),
203            database: "db1".to_string(),
204            tables: vec![],
205            functions: vec![],
206            source_uri: None,
207        };
208
209        let schema2 = Schema {
210            id: SchemaId::new(),
211            database: "db2".to_string(),
212            tables: vec![],
213            functions: vec![],
214            source_uri: None,
215        };
216
217        let id1 = manager.register(schema1);
218        let id2 = manager_clone.register(schema2);
219
220        assert_eq!(manager.get(id1).unwrap().database, "db1");
221        assert_eq!(manager_clone.get(id2).unwrap().database, "db2");
222
223        assert_ne!(id1, id2);
224    }
225}