kql_language_tools/
schema.rs

1//! Schema types for semantic validation
2//!
3//! These types represent the database schema that can be used for
4//! schema-aware validation. The schema includes tables, columns,
5//! and user-defined functions.
6
7use serde::{Deserialize, Serialize};
8
9/// Database schema for semantic validation
10///
11/// Contains definitions of tables, columns, and functions that
12/// the KQL validator should be aware of when performing semantic
13/// analysis.
14#[derive(Debug, Clone, Default, Serialize, Deserialize)]
15pub struct Schema {
16    /// Database name (optional)
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub database: Option<String>,
19
20    /// Tables in the schema
21    #[serde(default)]
22    pub tables: Vec<Table>,
23
24    /// User-defined functions
25    #[serde(default)]
26    pub functions: Vec<Function>,
27}
28
29impl Schema {
30    /// Create a new empty schema
31    #[must_use]
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    /// Create a schema with a database name
37    #[must_use]
38    pub fn with_database(database: impl Into<String>) -> Self {
39        Self {
40            database: Some(database.into()),
41            ..Self::default()
42        }
43    }
44
45    /// Add a table to the schema
46    pub fn add_table(&mut self, table: Table) -> &mut Self {
47        self.tables.push(table);
48        self
49    }
50
51    /// Add a function to the schema
52    pub fn add_function(&mut self, function: Function) -> &mut Self {
53        self.functions.push(function);
54        self
55    }
56
57    /// Builder method to add a table
58    #[must_use]
59    pub fn table(mut self, table: Table) -> Self {
60        self.tables.push(table);
61        self
62    }
63
64    /// Builder method to add a function
65    #[must_use]
66    pub fn function(mut self, function: Function) -> Self {
67        self.functions.push(function);
68        self
69    }
70
71    /// Check if the schema is empty
72    #[must_use]
73    pub fn is_empty(&self) -> bool {
74        self.tables.is_empty() && self.functions.is_empty()
75    }
76
77    /// Get a table by name
78    #[must_use]
79    pub fn get_table(&self, name: &str) -> Option<&Table> {
80        self.tables.iter().find(|t| t.name.eq_ignore_ascii_case(name))
81    }
82
83    /// Get a function by name
84    #[must_use]
85    pub fn get_function(&self, name: &str) -> Option<&Function> {
86        self.functions
87            .iter()
88            .find(|f| f.name.eq_ignore_ascii_case(name))
89    }
90}
91
92/// Table definition
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct Table {
95    /// Table name
96    pub name: String,
97
98    /// Table columns
99    #[serde(default)]
100    pub columns: Vec<Column>,
101
102    /// Optional table description
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub description: Option<String>,
105}
106
107impl Table {
108    /// Create a new table with the given name
109    #[must_use]
110    pub fn new(name: impl Into<String>) -> Self {
111        Self {
112            name: name.into(),
113            columns: Vec::new(),
114            description: None,
115        }
116    }
117
118    /// Add a column to the table
119    pub fn add_column(&mut self, column: Column) -> &mut Self {
120        self.columns.push(column);
121        self
122    }
123
124    /// Builder method to add a column
125    #[must_use]
126    pub fn column(mut self, column: Column) -> Self {
127        self.columns.push(column);
128        self
129    }
130
131    /// Builder method to add a column with name and type
132    #[must_use]
133    pub fn with_column(mut self, name: impl Into<String>, data_type: impl Into<String>) -> Self {
134        self.columns.push(Column::new(name, data_type));
135        self
136    }
137
138    /// Set the description
139    #[must_use]
140    pub fn description(mut self, desc: impl Into<String>) -> Self {
141        self.description = Some(desc.into());
142        self
143    }
144
145    /// Get a column by name
146    #[must_use]
147    pub fn get_column(&self, name: &str) -> Option<&Column> {
148        self.columns
149            .iter()
150            .find(|c| c.name.eq_ignore_ascii_case(name))
151    }
152}
153
154/// Column definition
155#[derive(Debug, Clone, Serialize, Deserialize)]
156pub struct Column {
157    /// Column name
158    pub name: String,
159
160    /// KQL data type (string, long, datetime, dynamic, etc.)
161    pub data_type: String,
162
163    /// Optional column description
164    #[serde(skip_serializing_if = "Option::is_none")]
165    pub description: Option<String>,
166}
167
168impl Column {
169    /// Create a new column
170    #[must_use]
171    pub fn new(name: impl Into<String>, data_type: impl Into<String>) -> Self {
172        Self {
173            name: name.into(),
174            data_type: data_type.into(),
175            description: None,
176        }
177    }
178
179    /// Set the description
180    #[must_use]
181    pub fn description(mut self, desc: impl Into<String>) -> Self {
182        self.description = Some(desc.into());
183        self
184    }
185
186    /// Create a string column
187    #[must_use]
188    pub fn string(name: impl Into<String>) -> Self {
189        Self::new(name, "string")
190    }
191
192    /// Create a long column
193    #[must_use]
194    pub fn long(name: impl Into<String>) -> Self {
195        Self::new(name, "long")
196    }
197
198    /// Create a real column
199    #[must_use]
200    pub fn real(name: impl Into<String>) -> Self {
201        Self::new(name, "real")
202    }
203
204    /// Create a bool column
205    #[must_use]
206    pub fn bool(name: impl Into<String>) -> Self {
207        Self::new(name, "bool")
208    }
209
210    /// Create a datetime column
211    #[must_use]
212    pub fn datetime(name: impl Into<String>) -> Self {
213        Self::new(name, "datetime")
214    }
215
216    /// Create a timespan column
217    #[must_use]
218    pub fn timespan(name: impl Into<String>) -> Self {
219        Self::new(name, "timespan")
220    }
221
222    /// Create a guid column
223    #[must_use]
224    pub fn guid(name: impl Into<String>) -> Self {
225        Self::new(name, "guid")
226    }
227
228    /// Create a dynamic column
229    #[must_use]
230    pub fn dynamic(name: impl Into<String>) -> Self {
231        Self::new(name, "dynamic")
232    }
233}
234
235/// User-defined function definition
236#[derive(Debug, Clone, Serialize, Deserialize)]
237pub struct Function {
238    /// Function name
239    pub name: String,
240
241    /// Parameter definitions
242    #[serde(default)]
243    pub parameters: Vec<Parameter>,
244
245    /// Return type
246    pub return_type: String,
247
248    /// Optional function body (KQL expression)
249    #[serde(skip_serializing_if = "Option::is_none")]
250    pub body: Option<String>,
251
252    /// Optional description
253    #[serde(skip_serializing_if = "Option::is_none")]
254    pub description: Option<String>,
255}
256
257impl Function {
258    /// Create a new function
259    #[must_use]
260    pub fn new(name: impl Into<String>, return_type: impl Into<String>) -> Self {
261        Self {
262            name: name.into(),
263            parameters: Vec::new(),
264            return_type: return_type.into(),
265            body: None,
266            description: None,
267        }
268    }
269
270    /// Add a parameter
271    pub fn add_parameter(&mut self, param: Parameter) -> &mut Self {
272        self.parameters.push(param);
273        self
274    }
275
276    /// Builder method to add a parameter
277    #[must_use]
278    pub fn param(mut self, name: impl Into<String>, data_type: impl Into<String>) -> Self {
279        self.parameters.push(Parameter::new(name, data_type));
280        self
281    }
282
283    /// Set the function body
284    #[must_use]
285    pub fn body(mut self, body: impl Into<String>) -> Self {
286        self.body = Some(body.into());
287        self
288    }
289
290    /// Set the description
291    #[must_use]
292    pub fn description(mut self, desc: impl Into<String>) -> Self {
293        self.description = Some(desc.into());
294        self
295    }
296}
297
298/// Function parameter definition
299#[derive(Debug, Clone, Serialize, Deserialize)]
300pub struct Parameter {
301    /// Parameter name
302    pub name: String,
303
304    /// Parameter data type
305    pub data_type: String,
306
307    /// Optional default value
308    #[serde(skip_serializing_if = "Option::is_none")]
309    pub default_value: Option<String>,
310}
311
312impl Parameter {
313    /// Create a new parameter
314    #[must_use]
315    pub fn new(name: impl Into<String>, data_type: impl Into<String>) -> Self {
316        Self {
317            name: name.into(),
318            data_type: data_type.into(),
319            default_value: None,
320        }
321    }
322
323    /// Set a default value
324    #[must_use]
325    pub fn default(mut self, value: impl Into<String>) -> Self {
326        self.default_value = Some(value.into());
327        self
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334
335    #[test]
336    fn test_schema_builder() {
337        let schema = Schema::with_database("SecurityDB")
338            .table(
339                Table::new("SecurityEvent")
340                    .with_column("TimeGenerated", "datetime")
341                    .with_column("Account", "string")
342                    .with_column("EventID", "long")
343                    .with_column("Computer", "string"),
344            )
345            .table(
346                Table::new("SigninLogs")
347                    .with_column("TimeGenerated", "datetime")
348                    .with_column("UserPrincipalName", "string")
349                    .with_column("IPAddress", "string")
350                    .with_column("ResultType", "string"),
351            );
352
353        assert_eq!(schema.database, Some("SecurityDB".to_string()));
354        assert_eq!(schema.tables.len(), 2);
355        assert_eq!(schema.tables[0].columns.len(), 4);
356    }
357
358    #[test]
359    fn test_schema_serialization() {
360        let schema = Schema::new().table(
361            Table::new("Test")
362                .with_column("Id", "long")
363                .with_column("Name", "string"),
364        );
365
366        let json = serde_json::to_string(&schema).unwrap();
367        let parsed: Schema = serde_json::from_str(&json).unwrap();
368
369        assert_eq!(parsed.tables.len(), 1);
370        assert_eq!(parsed.tables[0].name, "Test");
371        assert_eq!(parsed.tables[0].columns.len(), 2);
372    }
373}
kql_language_tools/schema.rs

kql_language_tools/
schema.rs
