1use std::convert::Infallible;
2use std::sync::Arc;
3use std::time::{Duration, Instant};
4
5use alopex_cluster::{
6 CatalogTableRef, CatalogTableSnapshot, PlacementCatalog, PlanId, QueryRouter,
7 QueryRoutingRequest, QueryTableReference, QueryTableReferenceAccess, QueryTableReferenceSource,
8 RoutingDecisionKind, RoutingDiagnostics, TableLifecycleEffect, TableRef,
9};
10use alopex_core::kv::async_adapter::AsyncKVTransactionAdapter;
11use alopex_core::kv::{KVStore, KVTransaction};
12use alopex_core::storage::format::bincode_config;
13use alopex_core::types::TxnMode;
14use alopex_sql::catalog::persistent::{PersistedTableMeta, TableFqn, TABLES_PREFIX};
15use alopex_sql::catalog::TableMetadata;
16use alopex_sql::planner::{
17 LogicalPlan, PlannedStatement, TableReference, TableReferenceAccess, TableReferenceSource,
18};
19use alopex_sql::storage::async_storage::AsyncTxnBridge;
20use alopex_sql::storage::AsyncSqlTransaction;
21use alopex_sql::AlopexDialect;
22use axum::extract::Extension;
23use axum::response::{IntoResponse, Response};
24use axum::Json;
25use bincode::Options;
26use futures::StreamExt;
27use serde::{Deserialize, Serialize};
28use tokio::sync::mpsc;
29use tokio_stream::wrappers::ReceiverStream;
30
31use crate::error::{Result, ServerError};
32use crate::http::{error_response, json_response, RequestContext};
33use crate::ops::memory::MemoryControlPolicy;
34use crate::server::ServerState;
35use crate::session::{CatalogRollbackEffect, SessionId, TxnHandle};
36
37#[derive(Debug, Deserialize)]
38pub struct SqlRequest {
39 pub sql: String,
40 pub session_id: Option<String>,
41 #[serde(default)]
42 pub streaming: bool,
43}
44
45#[derive(Debug, Serialize)]
46pub struct ColumnInfoResponse {
47 pub name: String,
48 pub data_type: String,
49}
50
51#[derive(Debug, Serialize)]
52pub struct SqlResponse {
53 pub columns: Vec<ColumnInfoResponse>,
54 pub rows: Vec<Vec<alopex_sql::storage::SqlValue>>,
55 pub affected_rows: Option<u64>,
56 #[serde(skip_serializing_if = "Vec::is_empty")]
57 pub routing_diagnostics: Vec<RoutingDiagnostics>,
58}
59
60#[derive(Debug, Serialize)]
61struct StreamItem {
62 row: Option<Vec<alopex_sql::storage::SqlValue>>,
63 error: Option<StreamError>,
64 done: bool,
65}
66
67#[derive(Debug, Serialize)]
68struct StreamError {
69 code: String,
70 message: String,
71 correlation_id: String,
72}
73
74type AsyncTxn = AsyncTxnBridge<'static, AsyncKVTransactionAdapter>;
75
76enum StreamSource {
77 Txn(AsyncTxn),
78 Handle(TxnHandle),
79}
80
81struct RoutingPlan {
82 planned: Vec<PlannedStatement>,
83 diagnostics: Vec<RoutingDiagnostics>,
84}
85
86#[derive(Clone)]
87struct TableLifecycleState {
88 table_ref: TableRef,
89 table_id: u32,
90 table: TableMetadata,
91}
92
93enum TableLifecycleCandidate {
94 Created {
95 table_name: String,
96 before: Option<TableLifecycleState>,
97 },
98 Dropped {
99 table_name: String,
100 before: Option<TableLifecycleState>,
101 },
102 CreateIndex {
103 index_name: String,
104 index_existed_before: bool,
105 },
106 DropIndex {
107 index_name: String,
108 before: Option<TableLifecycleState>,
109 },
110}
111
112pub async fn handle(
113 Extension(state): Extension<Arc<ServerState>>,
114 Extension(ctx): Extension<RequestContext>,
115 Json(request): Json<SqlRequest>,
116) -> Response {
117 if request.sql.trim().is_empty() {
118 return error_response(
119 ServerError::BadRequest("sql must not be empty".into()),
120 &ctx,
121 );
122 }
123
124 if request.streaming {
125 return stream_response(state, request, &ctx);
126 }
127
128 let result = execute_non_streaming(state.clone(), &request, &ctx).await;
129 match result {
130 Ok(response) => json_response(response, state.config.max_response_size, &ctx),
131 Err(err) => error_response(err, &ctx),
132 }
133}
134
135pub(crate) async fn execute_non_streaming(
141 state: Arc<ServerState>,
142 request: &SqlRequest,
143 ctx: &RequestContext,
144) -> Result<SqlResponse> {
145 let start = Instant::now();
146 let sql = request.sql.as_str();
147 let is_ddl = is_ddl(sql);
148 if is_write_sql(sql) {
149 state.lifecycle_state.check_write_allowed()?;
150 }
151
152 let exec_result: Result<(
153 alopex_sql::executor::ExecutionResult,
154 Vec<RoutingDiagnostics>,
155 )> = async {
156 if let Some(session_id) = &request.session_id {
157 let session_id = session_id
158 .parse::<SessionId>()
159 .map_err(|_| ServerError::BadRequest("invalid session_id".into()))?;
160 execute_session_statement_with_routing(
161 &state,
162 &session_id,
163 sql,
164 &ctx.correlation_id,
165 state.config.query_timeout,
166 )
167 .await
168 } else {
169 execute_non_session_statement_with_routing(
170 &state,
171 sql,
172 &ctx.correlation_id,
173 state.config.query_timeout,
174 )
175 .await
176 }
177 }
178 .await;
179 let exec_result = match exec_result {
180 Ok(result) => result,
181 Err(err) => {
182 state.metrics.record_query(start.elapsed(), false);
183 return Err(err);
184 }
185 };
186
187 if state.config.audit_log_enabled && is_ddl {
188 state
189 .audit
190 .log_ddl(sql, ctx.actor.as_deref(), &ctx.correlation_id);
191 }
192
193 if is_ddl && request.session_id.is_none() {
194 sync_catalog_to_store(&state)?;
195 }
196
197 state.metrics.record_query(start.elapsed(), true);
198
199 Ok(map_execution_result(exec_result.0, exec_result.1))
200}
201
202pub(crate) async fn execute_session_statement_with_routing(
203 state: &ServerState,
204 session_id: &SessionId,
205 sql: &str,
206 correlation_id: &str,
207 timeout: Duration,
208) -> Result<(
209 alopex_sql::executor::ExecutionResult,
210 Vec<RoutingDiagnostics>,
211)> {
212 let handle = state.session_manager.get_transaction(session_id).await?;
213 let routing_plan = route_session_sql(state, &handle, sql, correlation_id).await?;
214 if let Some(diagnostic) = future_distributed_diagnostic(&routing_plan.diagnostics) {
215 return Err(future_distributed_error(diagnostic));
216 }
217 let lifecycle_candidates = table_lifecycle_candidates(state, &routing_plan.planned)?;
218 let result = tokio::time::timeout(timeout, handle.execute(sql))
219 .await
220 .map_err(|_| ServerError::Timeout("query timeout".into()))?
221 .map_err(|err| ServerError::Sql(err.into()))?;
222 let (lifecycle_effects, catalog_rollback_effects) =
223 statement_effects_after_execution(state, lifecycle_candidates)?;
224 handle
225 .buffer_table_lifecycle_effects(lifecycle_effects)
226 .await;
227 handle
228 .buffer_catalog_rollback_effects(catalog_rollback_effects)
229 .await;
230 Ok((result, routing_plan.diagnostics))
231}
232
233pub(crate) async fn execute_non_session_statement_with_routing(
234 state: &ServerState,
235 sql: &str,
236 correlation_id: &str,
237 timeout: Duration,
238) -> Result<(
239 alopex_sql::executor::ExecutionResult,
240 Vec<RoutingDiagnostics>,
241)> {
242 let mut txn = state.begin_sql_txn().await?;
243 let routing_plan = match route_non_session_sql(state, &txn, sql, correlation_id).await {
244 Ok(plan) => plan,
245 Err(err) => {
246 let _ = txn.async_rollback().await;
247 return Err(err);
248 }
249 };
250 if let Some(diagnostic) = future_distributed_diagnostic(&routing_plan.diagnostics) {
251 let _ = txn.async_rollback().await;
252 return Err(future_distributed_error(diagnostic));
253 }
254 let lifecycle_candidates = table_lifecycle_candidates(state, &routing_plan.planned)?;
255 let fut = match tokio::time::timeout(timeout, txn.async_execute(sql)).await {
256 Ok(result) => result,
257 Err(_) => {
258 let _ = txn.async_rollback().await;
259 return Err(ServerError::Timeout("query timeout".into()));
260 }
261 };
262 match fut {
263 Ok(result) => {
264 let (lifecycle_effects, _) =
265 statement_effects_after_execution(state, lifecycle_candidates)?;
266 txn.async_commit()
267 .await
268 .map_err(|err| ServerError::Sql(err.into()))?;
269 state.apply_table_lifecycle_effects(lifecycle_effects)?;
270 Ok((result, routing_plan.diagnostics))
271 }
272 Err(err) => {
273 let _ = txn.async_rollback().await;
274 Err(ServerError::Sql(err.into()))
275 }
276 }
277}
278
279pub(crate) async fn route_session_statement_for_execution(
280 state: &ServerState,
281 handle: &TxnHandle,
282 sql: &str,
283 correlation_id: &str,
284) -> Result<Vec<RoutingDiagnostics>> {
285 let routing_plan = route_session_sql(state, handle, sql, correlation_id).await?;
286 if let Some(diagnostic) = future_distributed_diagnostic(&routing_plan.diagnostics) {
287 return Err(future_distributed_error(diagnostic));
288 }
289 Ok(routing_plan.diagnostics)
290}
291
292async fn route_non_session_sql(
293 state: &ServerState,
294 txn: &AsyncTxn,
295 sql: &str,
296 correlation_id: &str,
297) -> Result<RoutingPlan> {
298 let planned = txn
299 .async_plan_for_routing(sql)
300 .await
301 .map_err(|err| ServerError::Sql(err.into()))?;
302 route_planned_sql(state, planned, correlation_id)
303}
304
305async fn route_session_sql(
306 state: &ServerState,
307 handle: &TxnHandle,
308 sql: &str,
309 correlation_id: &str,
310) -> Result<RoutingPlan> {
311 let planned = handle
312 .plan_for_routing(sql)
313 .await
314 .map_err(|err| ServerError::Sql(err.into()))?;
315 route_planned_sql(state, planned, correlation_id)
316}
317
318fn route_planned_sql(
319 state: &ServerState,
320 planned: Vec<PlannedStatement>,
321 correlation_id: &str,
322) -> Result<RoutingPlan> {
323 let cluster_snapshot = state.cluster_status_snapshot()?;
324 let placement_catalog = PlacementCatalog::from_view(cluster_snapshot.placement);
325 let membership = cluster_snapshot.membership;
326 let catalog_snapshot = catalog_table_snapshot(state, cluster_snapshot.identity.update_epoch)?;
327 let router = QueryRouter::new(&placement_catalog, &membership);
328
329 let mut diagnostics = Vec::with_capacity(planned.len());
330 for (index, statement) in planned.iter().enumerate() {
331 let plan_id = PlanId::new(format!("{correlation_id}:{index}"));
332 let request = query_routing_request(plan_id, statement, &catalog_snapshot);
333 diagnostics.push(router.route(request));
334 }
335 Ok(RoutingPlan {
336 planned,
337 diagnostics,
338 })
339}
340
341fn future_distributed_diagnostic(
342 diagnostics: &[RoutingDiagnostics],
343) -> Option<&RoutingDiagnostics> {
344 diagnostics.iter().find(|diagnostic| {
345 diagnostic.decision == RoutingDecisionKind::FutureDistributedExecutionRequired
346 })
347}
348
349fn future_distributed_error(diagnostic: &RoutingDiagnostics) -> ServerError {
350 ServerError::FutureDistributedExecutionRequired(format!(
351 "routing decision {:?} for plan {}: {:?}",
352 diagnostic.decision,
353 diagnostic.plan_id.as_str(),
354 diagnostic.reason
355 ))
356}
357
358fn query_routing_request(
359 plan_id: PlanId,
360 statement: &PlannedStatement,
361 catalog_snapshot: &CatalogTableSnapshot,
362) -> QueryRoutingRequest {
363 let table_references = statement
364 .table_references()
365 .iter()
366 .map(|reference| query_table_reference(reference, catalog_snapshot))
367 .collect();
368
369 QueryRoutingRequest::new(plan_id, catalog_snapshot.clone(), table_references)
370}
371
372fn query_table_reference(
373 reference: &TableReference,
374 catalog_snapshot: &CatalogTableSnapshot,
375) -> QueryTableReference {
376 QueryTableReference::new(
377 table_ref_for_reference(&reference.table_name, catalog_snapshot),
378 query_table_reference_access(reference.access),
379 query_table_reference_source(reference.source),
380 )
381}
382
383fn query_table_reference_access(access: TableReferenceAccess) -> QueryTableReferenceAccess {
384 match access {
385 TableReferenceAccess::Read => QueryTableReferenceAccess::Read,
386 TableReferenceAccess::Write => QueryTableReferenceAccess::Write,
387 TableReferenceAccess::Create => QueryTableReferenceAccess::Create,
388 TableReferenceAccess::Drop => QueryTableReferenceAccess::Drop,
389 TableReferenceAccess::Metadata => QueryTableReferenceAccess::Metadata,
390 }
391}
392
393fn query_table_reference_source(source: TableReferenceSource) -> QueryTableReferenceSource {
394 match source {
395 TableReferenceSource::TopLevelPlanTableName => {
396 QueryTableReferenceSource::TopLevelPlanTableName
397 }
398 TableReferenceSource::LogicalPlanScan => QueryTableReferenceSource::LogicalPlanScan,
399 TableReferenceSource::LogicalPlanMutationTarget => {
400 QueryTableReferenceSource::LogicalPlanMutationTarget
401 }
402 TableReferenceSource::LogicalPlanDdlTarget => {
403 QueryTableReferenceSource::LogicalPlanDdlTarget
404 }
405 TableReferenceSource::LogicalPlanIndexTarget => {
406 QueryTableReferenceSource::LogicalPlanIndexTarget
407 }
408 TableReferenceSource::TypedExprSubquery => QueryTableReferenceSource::TypedExprSubquery,
409 }
410}
411
412fn catalog_table_snapshot(state: &ServerState, update_epoch: u64) -> Result<CatalogTableSnapshot> {
413 let guard = state
414 .catalog
415 .read()
416 .map_err(|_| ServerError::Internal("catalog lock poisoned".into()))?;
417 let tables = guard
418 .list_tables()
419 .iter()
420 .map(|table| CatalogTableRef::new(table_fqn_string(table), table.table_id))
421 .collect();
422 Ok(CatalogTableSnapshot::from_tables(update_epoch, tables))
423}
424
425fn table_ref_for_reference(table_name: &str, snapshot: &CatalogTableSnapshot) -> TableRef {
426 snapshot
427 .tables
428 .iter()
429 .find(|table| {
430 table.table_ref.as_str() == table_name
431 || table.table_ref.as_str().rsplit('.').next() == Some(table_name)
432 })
433 .map(|table| table.table_ref.clone())
434 .unwrap_or_else(|| TableRef::new(default_table_ref(table_name)))
435}
436
437fn table_fqn_string(table: &TableMetadata) -> String {
438 let fqn = TableFqn::from(table);
439 format!("{}.{}.{}", fqn.catalog, fqn.namespace, fqn.table)
440}
441
442fn table_lifecycle_candidates(
443 state: &ServerState,
444 planned: &[PlannedStatement],
445) -> Result<Vec<TableLifecycleCandidate>> {
446 let mut candidates = Vec::new();
447 for statement in planned {
448 match &statement.plan {
449 LogicalPlan::CreateTable { table, .. } => {
450 candidates.push(TableLifecycleCandidate::Created {
451 table_name: table.name.clone(),
452 before: table_lifecycle_state(state, &table.name)?,
453 });
454 }
455 LogicalPlan::DropTable { name, .. } => {
456 candidates.push(TableLifecycleCandidate::Dropped {
457 table_name: name.clone(),
458 before: table_lifecycle_state(state, name)?,
459 });
460 }
461 LogicalPlan::CreateIndex { index, .. } => {
462 candidates.push(TableLifecycleCandidate::CreateIndex {
463 index_name: index.name.clone(),
464 index_existed_before: index_exists(state, &index.name)?,
465 });
466 }
467 LogicalPlan::DropIndex { name, .. } => {
468 candidates.push(TableLifecycleCandidate::DropIndex {
469 index_name: name.clone(),
470 before: index_table_lifecycle_state(state, name)?,
471 });
472 }
473 _ => {}
474 }
475 }
476 Ok(candidates)
477}
478
479fn statement_effects_after_execution(
480 state: &ServerState,
481 candidates: Vec<TableLifecycleCandidate>,
482) -> Result<(Vec<TableLifecycleEffect>, Vec<CatalogRollbackEffect>)> {
483 let mut lifecycle_effects = Vec::new();
484 let mut catalog_rollback_effects = Vec::new();
485 for candidate in candidates {
486 match candidate {
487 TableLifecycleCandidate::Created { table_name, before } => {
488 let after = table_lifecycle_state(state, &table_name)?;
489 if let Some(after) = after {
490 let changed = match before.as_ref() {
491 Some(before) => before.table_id != after.table_id,
492 None => true,
493 };
494 if changed {
495 lifecycle_effects.push(TableLifecycleEffect::Created {
496 table_ref: after.table_ref,
497 table_id: after.table_id,
498 });
499 catalog_rollback_effects.push(CatalogRollbackEffect::DropTable {
500 table_name: after.table.name,
501 });
502 }
503 }
504 }
505 TableLifecycleCandidate::Dropped { table_name, before } => {
506 let after = table_lifecycle_state(state, &table_name)?;
507 if let Some(before) = before {
508 let changed = match after.as_ref() {
509 Some(after) => after.table_id != before.table_id,
510 None => true,
511 };
512 if changed {
513 lifecycle_effects.push(TableLifecycleEffect::Dropped {
514 table_ref: before.table_ref.clone(),
515 table_id: before.table_id,
516 });
517 catalog_rollback_effects.push(CatalogRollbackEffect::CreateTable {
518 table: Box::new(before.table),
519 });
520 }
521 }
522 }
523 TableLifecycleCandidate::CreateIndex {
524 index_name,
525 index_existed_before,
526 } => {
527 if !index_existed_before {
528 if let Some(after) = index_table_lifecycle_state(state, &index_name)? {
529 lifecycle_effects.push(TableLifecycleEffect::SchemaChanged {
530 table_ref: after.table_ref,
531 table_id: after.table_id,
532 });
533 }
534 }
535 }
536 TableLifecycleCandidate::DropIndex { index_name, before } => {
537 if let Some(before) = before {
538 if !index_exists(state, &index_name)? {
539 lifecycle_effects.push(TableLifecycleEffect::SchemaChanged {
540 table_ref: before.table_ref,
541 table_id: before.table_id,
542 });
543 }
544 }
545 }
546 }
547 }
548 Ok((lifecycle_effects, catalog_rollback_effects))
549}
550
551fn table_lifecycle_state(
552 state: &ServerState,
553 table_name: &str,
554) -> Result<Option<TableLifecycleState>> {
555 let guard = state
556 .catalog
557 .read()
558 .map_err(|_| ServerError::Internal("catalog lock poisoned".into()))?;
559 Ok(guard
560 .get_table(table_name)
561 .map(|table| TableLifecycleState {
562 table_ref: TableRef::new(table_fqn_string(table)),
563 table_id: table.table_id,
564 table: table.clone(),
565 }))
566}
567
568fn index_table_lifecycle_state(
569 state: &ServerState,
570 index_name: &str,
571) -> Result<Option<TableLifecycleState>> {
572 let guard = state
573 .catalog
574 .read()
575 .map_err(|_| ServerError::Internal("catalog lock poisoned".into()))?;
576 let Some(index) = guard.get_index(index_name) else {
577 return Ok(None);
578 };
579 Ok(guard
580 .get_table(&index.table)
581 .map(|table| TableLifecycleState {
582 table_ref: TableRef::new(table_fqn_string(table)),
583 table_id: table.table_id,
584 table: table.clone(),
585 }))
586}
587
588fn index_exists(state: &ServerState, index_name: &str) -> Result<bool> {
589 let guard = state
590 .catalog
591 .read()
592 .map_err(|_| ServerError::Internal("catalog lock poisoned".into()))?;
593 Ok(guard.get_index(index_name).is_some())
594}
595
596fn default_table_ref(table_name: &str) -> String {
597 if table_name.matches('.').count() >= 2 {
598 table_name.to_string()
599 } else {
600 format!("default.default.{table_name}")
601 }
602}
603
604pub(crate) fn sync_catalog_to_store(state: &ServerState) -> Result<()> {
605 let guard = state
606 .catalog
607 .read()
608 .map_err(|_| ServerError::Internal("catalog lock poisoned".into()))?;
609 let tables = guard.list_tables();
610 let mut txn = state.store.begin(TxnMode::ReadWrite)?;
611 delete_prefix(&mut txn, TABLES_PREFIX)?;
612 for table in tables {
613 let persisted = PersistedTableMeta::from(&table);
614 let value = bincode_config()
615 .serialize(&persisted)
616 .map_err(|err| ServerError::Internal(err.to_string()))?;
617 txn.put(
618 table_key(&table.catalog_name, &table.namespace_name, &table.name),
619 value,
620 )?;
621 }
622 txn.commit_self()?;
623 Ok(())
624}
625
626fn delete_prefix<'a, T: KVTransaction<'a>>(txn: &mut T, prefix: &[u8]) -> Result<()> {
627 let mut keys = Vec::new();
628 for (key, _) in txn.scan_prefix(prefix)? {
629 keys.push(key);
630 }
631 for key in keys {
632 txn.delete(key)?;
633 }
634 Ok(())
635}
636
637fn table_key(catalog_name: &str, namespace_name: &str, table_name: &str) -> Vec<u8> {
638 let mut key = TABLES_PREFIX.to_vec();
639 key.extend_from_slice(catalog_name.as_bytes());
640 key.push(b'/');
641 key.extend_from_slice(namespace_name.as_bytes());
642 key.push(b'/');
643 key.extend_from_slice(table_name.as_bytes());
644 key
645}
646
647fn stream_response(state: Arc<ServerState>, request: SqlRequest, ctx: &RequestContext) -> Response {
648 if is_write_sql(&request.sql) {
649 if let Err(err) = state.lifecycle_state.check_write_allowed() {
650 return error_response(err, ctx);
651 }
652 }
653 let (sender, receiver) = mpsc::channel(32);
654 let sql = request.sql.clone();
655 let correlation_id = ctx.correlation_id.clone();
656 let max_response_size = state.config.max_response_size;
657 let timeout = state.config.query_timeout;
658 let memory_policy = MemoryControlPolicy::from_env();
659 let metrics = state.metrics.clone();
660 let mut audit = None;
661 if state.config.audit_log_enabled && is_ddl(&sql) {
662 audit = Some(state.audit.clone());
663 }
664
665 let session_id = request.session_id.clone();
666 let state_clone = state.clone();
667 let memory_policy = memory_policy.clone();
668 tokio::spawn(async move {
669 let start = Instant::now();
670 let mut bytes_sent = 0usize;
671 let mut success = true;
672 let mut source = match session_id {
673 Some(id) => {
674 let parsed = match id.parse::<SessionId>() {
675 Ok(id) => id,
676 Err(_) => {
677 let _ = sender
678 .send(stream_item_error(
679 ServerError::BadRequest("invalid session_id".into()),
680 &correlation_id,
681 ))
682 .await;
683 return;
684 }
685 };
686 match state_clone.session_manager.get_transaction(&parsed).await {
687 Ok(handle) => {
688 match route_session_statement_for_execution(
689 &state_clone,
690 &handle,
691 &sql,
692 &correlation_id,
693 )
694 .await
695 {
696 Ok(_) => {}
697 Err(err) => {
698 let _ = sender.send(stream_item_error(err, &correlation_id)).await;
699 return;
700 }
701 }
702 StreamSource::Handle(handle)
703 }
704 Err(err) => {
705 let _ = sender.send(stream_item_error(err, &correlation_id)).await;
706 return;
707 }
708 }
709 }
710 None => match state_clone.begin_sql_txn().await {
711 Ok(txn) => StreamSource::Txn(txn),
712 Err(err) => {
713 let _ = sender.send(stream_item_error(err, &correlation_id)).await;
714 return;
715 }
716 },
717 };
718
719 let mut stream = match &mut source {
720 StreamSource::Handle(handle) => handle.query(&sql),
721 StreamSource::Txn(txn) => txn.async_query(&sql),
722 };
723 let deadline = start + timeout;
724 loop {
725 let remaining = deadline.saturating_duration_since(Instant::now());
726 if remaining.is_zero() {
727 let _ = sender
728 .send(stream_item_error(
729 ServerError::Timeout("query timeout".into()),
730 &correlation_id,
731 ))
732 .await;
733 success = false;
734 break;
735 }
736
737 tokio::select! {
738 _ = sender.closed() => {
739 success = false;
740 break;
741 }
742 item = tokio::time::timeout(remaining, stream.next()) => {
743 let next = match item {
744 Ok(value) => value,
745 Err(_) => {
746 let _ = sender
747 .send(stream_item_error(
748 ServerError::Timeout("query timeout".into()),
749 &correlation_id,
750 ))
751 .await;
752 success = false;
753 break;
754 }
755 };
756
757 match next {
758 Some(Ok(row)) => {
759 let item = StreamItem {
760 row: Some(row.values),
761 error: None,
762 done: false,
763 };
764 match serde_json::to_vec(&item) {
765 Ok(bytes) => {
766 bytes_sent += bytes.len();
767 if let Err(err) =
768 memory_policy.enforce_output_bytes(bytes_sent as u64)
769 {
770 let _ = sender
771 .send(stream_item_error(err, &correlation_id))
772 .await;
773 success = false;
774 break;
775 }
776 if bytes_sent > max_response_size {
777 let _ = sender
778 .send(stream_item_error(
779 ServerError::PayloadTooLarge(
780 "response size exceeds limit".into(),
781 ),
782 &correlation_id,
783 ))
784 .await;
785 success = false;
786 break;
787 }
788 }
789 Err(err) => {
790 let _ = sender
791 .send(stream_item_error(
792 ServerError::Internal(err.to_string()),
793 &correlation_id,
794 ))
795 .await;
796 success = false;
797 break;
798 }
799 }
800 match sender.try_send(item) {
801 Ok(()) => {}
802 Err(mpsc::error::TrySendError::Full(item)) => {
803 metrics.record_backpressure();
804 if sender.send(item).await.is_err() {
805 success = false;
806 break;
807 }
808 }
809 Err(mpsc::error::TrySendError::Closed(_)) => {
810 success = false;
811 break;
812 }
813 }
814 }
815 Some(Err(err)) => {
816 let _ = sender
817 .send(stream_item_error(
818 ServerError::Sql(err.into()),
819 &correlation_id,
820 ))
821 .await;
822 success = false;
823 break;
824 }
825 None => break,
826 }
827 }
828 }
829 }
830
831 drop(stream);
832 if let StreamSource::Txn(txn) = source {
833 let _ = txn.async_rollback().await;
834 }
835 if let Some(logger) = audit {
836 logger.log_ddl(&sql, None, &correlation_id);
837 }
838 metrics.record_query(start.elapsed(), success);
839 let _ = sender
840 .send(StreamItem {
841 row: None,
842 error: None,
843 done: true,
844 })
845 .await;
846 });
847
848 let stream = ReceiverStream::new(receiver).map(|item| {
849 let json = serde_json::to_string(&item).unwrap_or_else(|_| "{}".to_string());
850 Ok::<axum::body::Bytes, Infallible>(axum::body::Bytes::from(json + "\n"))
851 });
852
853 let body = axum::body::Body::from_stream(stream);
854 axum::response::Response::builder()
855 .status(axum::http::StatusCode::OK)
856 .header(axum::http::header::CONTENT_TYPE, "application/jsonl")
857 .body(body)
858 .unwrap_or_else(|_| axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response())
859}
860
861fn stream_item_error(err: ServerError, correlation_id: &str) -> StreamItem {
862 StreamItem {
863 row: None,
864 error: Some(StreamError {
865 code: err.error_code(),
866 message: err.to_string(),
867 correlation_id: correlation_id.to_string(),
868 }),
869 done: false,
870 }
871}
872
873fn map_execution_result(
874 exec_result: alopex_sql::executor::ExecutionResult,
875 routing_diagnostics: Vec<RoutingDiagnostics>,
876) -> SqlResponse {
877 match exec_result {
878 alopex_sql::executor::ExecutionResult::Query(query) => SqlResponse {
879 columns: query
880 .columns
881 .into_iter()
882 .map(|col| ColumnInfoResponse {
883 name: col.name,
884 data_type: type_to_string(&col.data_type),
885 })
886 .collect(),
887 rows: query.rows,
888 affected_rows: None,
889 routing_diagnostics,
890 },
891 alopex_sql::executor::ExecutionResult::RowsAffected(rows) => SqlResponse {
892 columns: Vec::new(),
893 rows: Vec::new(),
894 affected_rows: Some(rows),
895 routing_diagnostics,
896 },
897 alopex_sql::executor::ExecutionResult::Success => SqlResponse {
898 columns: Vec::new(),
899 rows: Vec::new(),
900 affected_rows: None,
901 routing_diagnostics,
902 },
903 }
904}
905
906fn type_to_string(data_type: &alopex_sql::planner::ResolvedType) -> String {
907 match data_type {
908 alopex_sql::planner::ResolvedType::Integer => "INTEGER".to_string(),
909 alopex_sql::planner::ResolvedType::BigInt => "BIGINT".to_string(),
910 alopex_sql::planner::ResolvedType::Float => "FLOAT".to_string(),
911 alopex_sql::planner::ResolvedType::Double => "DOUBLE".to_string(),
912 alopex_sql::planner::ResolvedType::Text => "TEXT".to_string(),
913 alopex_sql::planner::ResolvedType::Blob => "BLOB".to_string(),
914 alopex_sql::planner::ResolvedType::Boolean => "BOOLEAN".to_string(),
915 alopex_sql::planner::ResolvedType::Timestamp => "TIMESTAMP".to_string(),
916 alopex_sql::planner::ResolvedType::Vector { dimension, metric } => {
917 format!("VECTOR({dimension}, {metric:?})")
918 }
919 alopex_sql::planner::ResolvedType::Null => "NULL".to_string(),
920 }
921}
922
923fn is_ddl(sql: &str) -> bool {
924 let Ok(statements) = alopex_sql::parser::Parser::parse_sql(&AlopexDialect, sql) else {
925 return false;
926 };
927 statements.iter().any(|stmt| match &stmt.kind {
928 alopex_sql::ast::StatementKind::CreateTable(_)
929 | alopex_sql::ast::StatementKind::DropTable(_)
930 | alopex_sql::ast::StatementKind::CreateIndex(_)
931 | alopex_sql::ast::StatementKind::DropIndex(_) => true,
932 alopex_sql::ast::StatementKind::Select(_)
933 | alopex_sql::ast::StatementKind::Insert(_)
934 | alopex_sql::ast::StatementKind::Update(_)
935 | alopex_sql::ast::StatementKind::Delete(_) => false,
936 })
937}
938
939fn is_write_sql(sql: &str) -> bool {
940 let Ok(statements) = alopex_sql::parser::Parser::parse_sql(&AlopexDialect, sql) else {
941 return false;
942 };
943 statements
944 .iter()
945 .any(|stmt| !matches!(stmt.kind, alopex_sql::ast::StatementKind::Select(_)))
946}