Skip to main content

wowlab_sentinel/mcp/tools/query/
types.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5const DEFAULT_QUERY_LIMIT: u64 = 100;
6
7#[derive(Clone, Debug, Default, Deserialize, JsonSchema)]
8#[schemars(
9    description = "Single table query. Run `get_schema` first to see valid tables and columns.",
10    inline
11)]
12pub(crate) struct TableQuery {
13    #[schemars(description = "Table name from `get_schema` (e.g. `game.spells`, `game.items`)")]
14    pub table: String,
15    #[serde(default)]
16    #[schemars(description = "Filter conditions, AND-ed together. Omit for unfiltered results.")]
17    pub filters: Vec<Filter>,
18    #[serde(default)]
19    #[schemars(description = "Column to sort by. Must exist in the table.")]
20    pub order_by: Option<String>,
21    #[serde(default)]
22    #[schemars(description = "Sort descending. Default: false (ascending).")]
23    pub order_desc: bool,
24    #[serde(default = "default_limit")]
25    #[schemars(description = "Max rows to return. Default: 100, max: 1000.")]
26    pub limit: u64,
27    #[serde(default)]
28    #[schemars(description = "Rows to skip (for pagination).")]
29    pub offset: Option<u64>,
30}
31
32const fn default_limit() -> u64 {
33    DEFAULT_QUERY_LIMIT
34}
35
36#[derive(Clone, Debug, Default, Deserialize, JsonSchema)]
37#[schemars(description = "Batch of queries. Each uses the same schema as `query`.")]
38pub(crate) struct TableBatchQuery {
39    #[schemars(description = "Array of TableQuery objects (max 50). Each runs independently.")]
40    pub queries: Vec<TableQuery>,
41}
42
43#[derive(Clone, Debug, JsonSchema, Serialize)]
44#[schemars(description = "Result for a single query in a batch request")]
45pub(in crate::mcp::tools) struct BatchQueryResult {
46    #[schemars(description = "Zero-based index of this query in the batch input")]
47    pub index: usize,
48    #[schemars(description = "Table that was queried")]
49    pub table: String,
50    #[schemars(description = "Number of rows returned for this query")]
51    pub row_count: usize,
52    #[schemars(description = "Rows returned for this query. Null when this query failed.")]
53    pub results: Option<Vec<Value>>,
54    #[schemars(description = "Query error message when this query failed, otherwise null")]
55    pub error: Option<String>,
56    #[serde(skip)]
57    #[schemars(skip)]
58    pub error_type: Option<crate::telemetry::McpErrorType>,
59}
60
61#[derive(Clone, Debug, Deserialize, JsonSchema)]
62#[schemars(
63    description = "A single filter condition to apply to the query",
64    inline
65)]
66pub(crate) struct Filter {
67    #[schemars(description = "Column name to filter on (must exist in the table)")]
68    pub column: String,
69    #[schemars(description = "Comparison operation to apply")]
70    pub op: FilterOp,
71    #[serde(default)]
72    #[schemars(
73        description = "Value to compare against (type must match column: int, text, bool, or json for jsonb ops). Omit for is_null/is_not_null."
74    )]
75    pub value: Option<Value>,
76}
77
78#[derive(Clone, Debug, Deserialize, JsonSchema)]
79#[serde(rename_all = "snake_case")]
80#[schemars(description = "Filter comparison operation", inline)]
81#[non_exhaustive]
82pub(crate) enum FilterOp {
83    #[schemars(description = "Equals (works with int, text, bool)")]
84    Eq,
85    #[schemars(description = "Not equals (works with int, text, bool)")]
86    Ne,
87    #[schemars(description = "Column value is SQL NULL (omit value)")]
88    IsNull,
89    #[schemars(description = "Column value is not SQL NULL (omit value)")]
90    IsNotNull,
91    #[schemars(description = "Greater than (works with int)")]
92    Gt,
93    #[schemars(description = "Greater than or equal (works with int)")]
94    Gte,
95    #[schemars(description = "Less than (works with int)")]
96    Lt,
97    #[schemars(description = "Less than or equal (works with int)")]
98    Lte,
99    #[schemars(description = "Text contains substring (works with text columns only)")]
100    Contains,
101    #[schemars(description = "Text starts with prefix (works with text columns only)")]
102    StartsWith,
103    #[schemars(description = "Match any value in a list (works with int, float, text, bool)")]
104    In,
105    #[schemars(
106        description = "JSONB contains (json columns only). Value is matched with @> operator. Example: {\"effect\": 68}"
107    )]
108    JsonbContains,
109    #[schemars(description = "JSONB has key (json columns only). Value is a string key name.")]
110    JsonbHasKey,
111    #[schemars(
112        description = "JSONB array contains element (json columns only). Value is an object to match within the array. Wraps value in [] automatically."
113    )]
114    JsonbArrayContains,
115    #[schemars(
116        description = "Integer array contains element (int_array columns only). Value is a single integer. Matches rows where the column array contains the value."
117    )]
118    ArrayContains,
119}
120
121#[derive(Clone, Debug, Default, Deserialize, JsonSchema)]
122#[schemars(
123    description = "Count rows matching filters. Same filters as `query` but returns only the count."
124)]
125pub(crate) struct CountQuery {
126    #[schemars(description = "Table name from `get_schema` (e.g. `game.spells`, `game.items`)")]
127    pub table: String,
128    #[serde(default)]
129    #[schemars(description = "Filter conditions, AND-ed together. Omit to count all rows.")]
130    pub filters: Vec<Filter>,
131}
132
133#[derive(Debug, JsonSchema, Serialize)]
134pub(in crate::mcp::tools) struct CountResult {
135    pub table: String,
136    pub count: i64,
137}
138
139#[cfg(test)]
140mod tests {
141    use googletest::prelude::*;
142    use rmcp::{
143        handler::server::{tool::schema_for_type, wrapper::Parameters},
144        model::JsonObject,
145    };
146    use schemars::JsonSchema;
147    use serde_json::{Value, json};
148
149    use super::{CountQuery, TableBatchQuery};
150
151    fn input_schema<T>() -> Value
152    where
153        T: JsonSchema + 'static,
154    {
155        let schema: JsonObject = schema_for_type::<Parameters<T>>().as_ref().clone();
156
157        Value::Object(schema)
158    }
159
160    #[gtest]
161    fn count_query_schema_exposes_filter_fields_without_references() -> Result<()> {
162        let schema = input_schema::<CountQuery>();
163
164        verify_that!(
165            schema.pointer("/properties/filters/items/properties/column/type"),
166            some(eq(&json!("string")))
167        )?;
168        let operators = schema
169            .pointer("/properties/filters/items/properties/op/oneOf")
170            .and_then(Value::as_array)
171            .into_iter()
172            .flatten()
173            .filter_map(|operator| operator.get("const").and_then(Value::as_str))
174            .collect::<Vec<_>>();
175
176        verify_that!(
177            operators,
178            contains_each![
179                eq(&"eq"),
180                eq(&"is_null"),
181                eq(&"is_not_null"),
182                eq(&"contains"),
183                eq(&"jsonb_array_contains")
184            ]
185        )?;
186
187        verify_that!(schema.pointer("/properties/filters/items/$ref"), none())
188    }
189
190    #[gtest]
191    fn batch_query_schema_exposes_nested_query_and_filter_fields() -> Result<()> {
192        let schema = input_schema::<TableBatchQuery>();
193
194        verify_that!(
195            schema.pointer("/properties/queries/items/properties/table/type"),
196            some(eq(&json!("string")))
197        )?;
198        verify_that!(
199            schema.pointer(
200                "/properties/queries/items/properties/filters/items/properties/value/description"
201            ),
202            some(anything())
203        )?;
204
205        verify_that!(schema.pointer("/properties/queries/items/$ref"), none())
206    }
207}