linkrs_server/graph_service/
batch_exec.rs

1use std::collections::HashMap;
2
3use linkrs_api::api_core::QueryResult;
4
5use crate::query::parser::ast::Stmt;
6use crate::storage::{
7    StorageClient, StorageOperationContextOps, StorageSchemaContextOps, StorageSyncContextOps,
8};
9
10use super::error::GraphServiceError;
11use super::GraphService;
12
13impl<
14        S: StorageClient
15            + StorageSchemaContextOps
16            + StorageSyncContextOps
17            + StorageOperationContextOps
18            + Clone
19            + crate::storage::AutoCommitBatchOps
20            + 'static,
21    > GraphService<S>
22{
23    /// Execute a batch of auto-commit DML statements inside one shared
24    /// auto-commit batch window.
25    pub async fn execute_batch(
26        &self,
27        session_id: i64,
28        statements: &[String],
29        parameters: Option<HashMap<String, linkrs_core::Value>>,
30        session_variables: Option<HashMap<String, linkrs_core::Value>>,
31    ) -> Vec<Result<QueryResult, GraphServiceError>> {
32        let Some(session) = self.session_manager.find_session(session_id) else {
33            return statements
34                .iter()
35                .map(|_| {
36                    Err(GraphServiceError::new(format!(
37                        "Invalid session ID: {session_id}"
38                    )))
39                })
40                .collect();
41        };
42        session.charge();
43
44        let mut results: Vec<Option<Result<QueryResult, GraphServiceError>>> =
45            vec![None; statements.len()];
46        let mut batch_indices: Vec<usize> = Vec::new();
47        let mut batch_statements: Vec<String> = Vec::new();
48        for (index, stmt) in statements.iter().enumerate() {
49            match Self::parse_command(stmt) {
50                Err(parse_error) => results[index] = Some(Err(parse_error)),
51                Ok(Some(parsed)) => match parsed.ast.stmt() {
52                    Stmt::AssignVariable(_) => {
53                        results[index] = Some(
54                            self.execute_with_params(
55                                session_id,
56                                stmt,
57                                parameters.clone(),
58                                session_variables.clone(),
59                            )
60                            .await,
61                        );
62                    }
63                    _ => {
64                        results[index] = Some(Err(GraphServiceError::new(
65                            "Transaction commands are not supported in batch execution",
66                        )));
67                    }
68                },
69                Ok(None) => {
70                    batch_indices.push(index);
71                    batch_statements.push(stmt.clone());
72                }
73            }
74        }
75
76        if session.current_transaction().is_some() || !session.is_auto_commit() {
77            for (index, stmt) in batch_statements.iter().enumerate() {
78                results[batch_indices[index]] = Some(
79                    self.execute_with_params(
80                        session_id,
81                        stmt,
82                        parameters.clone(),
83                        session_variables.clone(),
84                    )
85                    .await,
86                );
87            }
88            return finalize_batch_outcomes(results);
89        }
90
91        let space_id = session.space().map(|s| s.id as i64).unwrap_or(0);
92        let username = session.user();
93        let (permitted, denied) = self.partition_permitted(&username, space_id, &batch_statements);
94
95        let query_request = linkrs_api::api_core::QueryRequest {
96            isolation_level: None,
97            space_id: session.space().map(|s| s.id),
98            space_name: session.space().map(|s| s.name),
99            auto_commit: true,
100            transaction_id: None,
101            parameters,
102            session_variables,
103            query_id: None,
104            parsed_statement: None,
105            consistency: Default::default(),
106        };
107        let outcomes = self
108            .query_api
109            .write()
110            .execute_batch(&permitted, query_request);
111
112        merge_batch_outcomes(&mut results, &batch_indices, &denied, outcomes);
113        finalize_batch_outcomes(results)
114    }
115
116    /// Split batch statements into permitted and denied groups.
117    fn partition_permitted(
118        &self,
119        username: &str,
120        space_id: i64,
121        batch_statements: &[String],
122    ) -> (Vec<String>, Vec<(usize, String)>) {
123        use super::context::QueryExecutionContext;
124        let mut denied: Vec<(usize, String)> = Vec::new();
125        let mut permitted: Vec<String> = Vec::with_capacity(batch_statements.len());
126        for (index, stmt) in batch_statements.iter().enumerate() {
127            let context = QueryExecutionContext {
128                stmt,
129                parsed_ast: None,
130                space_id,
131                parameters: None,
132                session_variables: None,
133            };
134            if let Err(e) = self.check_query_permission(username, &context) {
135                denied.push((index, e.message().to_string()));
136                continue;
137            }
138            permitted.push(stmt.clone());
139        }
140        (permitted, denied)
141    }
142}
143
144impl<
145        S: StorageClient
146            + StorageSchemaContextOps
147            + StorageSyncContextOps
148            + StorageOperationContextOps
149            + Clone
150            + crate::storage::AutoCommitBatchOps
151            + crate::storage::AutoCommitGroupOps
152            + 'static,
153    > GraphService<S>
154{
155    /// Execute a batch of auto-commit DML statements inside shared
156    /// group-commit windows.
157    pub async fn execute_batch_grouped(
158        &self,
159        session_id: i64,
160        statements: &[String],
161        group_size: usize,
162    ) -> Vec<Result<QueryResult, GraphServiceError>> {
163        let Some(session) = self.session_manager.find_session(session_id) else {
164            return statements
165                .iter()
166                .map(|_| {
167                    Err(GraphServiceError::new(format!(
168                        "Invalid session ID: {session_id}"
169                    )))
170                })
171                .collect();
172        };
173        session.charge();
174
175        let mut results: Vec<Option<Result<QueryResult, GraphServiceError>>> =
176            vec![None; statements.len()];
177        let mut batch_indices: Vec<usize> = Vec::new();
178        let mut batch_statements: Vec<String> = Vec::new();
179        for (index, stmt) in statements.iter().enumerate() {
180            match Self::parse_command(stmt) {
181                Err(parse_error) => results[index] = Some(Err(parse_error)),
182                Ok(Some(parsed)) => match parsed.ast.stmt() {
183                    Stmt::AssignVariable(_) => {
184                        results[index] = Some(self.execute(session_id, stmt).await);
185                    }
186                    _ => {
187                        results[index] = Some(Err(GraphServiceError::new(
188                            "Transaction commands are not supported in batch execution",
189                        )));
190                    }
191                },
192                Ok(None) => {
193                    batch_indices.push(index);
194                    batch_statements.push(stmt.clone());
195                }
196            }
197        }
198
199        if session.current_transaction().is_some() || !session.is_auto_commit() {
200            for (index, stmt) in batch_statements.iter().enumerate() {
201                results[batch_indices[index]] = Some(self.execute(session_id, stmt).await);
202            }
203            return finalize_batch_outcomes(results);
204        }
205
206        let space_id = session.space().map(|s| s.id as i64).unwrap_or(0);
207        let username = session.user();
208        let (permitted, denied) = self.partition_permitted(&username, space_id, &batch_statements);
209
210        let query_request = linkrs_api::api_core::QueryRequest {
211            isolation_level: None,
212            space_id: session.space().map(|s| s.id),
213            space_name: session.space().map(|s| s.name),
214            auto_commit: true,
215            transaction_id: None,
216            parameters: None,
217            session_variables: None,
218            query_id: None,
219            parsed_statement: None,
220            consistency: Default::default(),
221        };
222        let outcomes =
223            self.query_api
224                .write()
225                .execute_batch_grouped(&permitted, query_request, group_size);
226
227        merge_batch_outcomes(&mut results, &batch_indices, &denied, outcomes);
228        finalize_batch_outcomes(results)
229    }
230}
231
232/// Convert the per-slot batch results into the final ordered outcome vector.
233fn finalize_batch_outcomes(
234    results: Vec<Option<Result<QueryResult, GraphServiceError>>>,
235) -> Vec<Result<QueryResult, GraphServiceError>> {
236    results
237        .into_iter()
238        .map(|slot| slot.unwrap_or_else(|| Err(GraphServiceError::new("Batch outcome missing"))))
239        .collect()
240}
241
242/// Merge the batch-window outcomes back into the per-slot results.
243fn merge_batch_outcomes(
244    results: &mut [Option<Result<QueryResult, GraphServiceError>>],
245    batch_indices: &[usize],
246    denied: &[(usize, String)],
247    outcomes: Vec<Result<linkrs_api::api_core::QueryResult, linkrs_api::api_core::CoreError>>,
248) {
249    let mut permitted_outcomes = outcomes.into_iter();
250    for (batch_pos, original_index) in batch_indices.iter().enumerate() {
251        if let Some((_, error)) = denied.iter().find(|(i, _)| *i == batch_pos) {
252            results[*original_index] = Some(Err(GraphServiceError::new(error.clone())));
253            continue;
254        }
255        match permitted_outcomes.next() {
256            Some(Ok(result)) => {
257                results[*original_index] = Some(Ok(result));
258            }
259            Some(Err(error)) => {
260                results[*original_index] = Some(Err(GraphServiceError::from_core_error(error)));
261            }
262            None => {
263                results[*original_index] =
264                    Some(Err(GraphServiceError::new("Batch outcome missing")));
265            }
266        }
267    }
268}
linkrs_server/graph_service/batch_exec.rs

linkrs_server/graph_service/
batch_exec.rs
