1use std::{
2 collections::HashMap,
3 pin::Pin,
4 sync::{Arc, LazyLock, Mutex},
5};
6
7use arrow::ipc::writer::IpcWriteOptions;
8use arrow_array::RecordBatch;
9use arrow_flight::{
10 FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse,
11 IpcMessage, SchemaAsIpc, Ticket,
12 decode::FlightRecordBatchStream,
13 encode::FlightDataEncoderBuilder,
14 error::FlightError,
15 sql::{
16 ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
17 ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetDbSchemas,
18 CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandPreparedStatementQuery,
19 CommandStatementQuery, CommandStatementUpdate, DoPutPreparedStatementResult,
20 ProstMessageExt, SqlInfo, TicketStatementQuery,
21 metadata::{SqlInfoData, SqlInfoDataBuilder},
22 server::{FlightSqlService, PeekableFlightDataStream},
23 },
24};
25use futures::{Stream, TryStreamExt};
26use graphar_falkordb::FalkorDbLoader;
27use prost::Message;
28use tonic::{Request, Response, Status, Streaming};
29
30use crate::{
31 auth::{AuthConfig, Identity},
32 error::FlightSqlError,
33 falkor::FalkorExecutor,
34 registry::{SchemaRegistry, parse_select_from},
35};
36
37static SQL_INFO: LazyLock<SqlInfoData> = LazyLock::new(|| {
41 let mut b = SqlInfoDataBuilder::new();
42 b.append(
43 SqlInfo::FlightSqlServerName,
44 "knut graphar-flight (Cypher over FalkorDB)",
45 );
46 b.append(SqlInfo::FlightSqlServerVersion, env!("CARGO_PKG_VERSION"));
47 b.append(SqlInfo::FlightSqlServerArrowVersion, "1.3");
49 b.append(SqlInfo::FlightSqlServerReadOnly, true);
50 b.append(SqlInfo::SqlDdlCatalog, false);
52 b.append(SqlInfo::SqlIdentifierQuoteChar, "\"");
53 b.build().expect("static SqlInfoData builds")
54});
55
56type BoxStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
57
58const SCHEMA_CACHE_CAP: usize = 16;
64
65#[derive(Clone, Default)]
74struct SchemaCache {
75 inner: Arc<Mutex<HashMap<String, RecordBatch>>>,
76}
77
78impl SchemaCache {
79 fn insert(&self, cypher: String, batch: RecordBatch) {
83 let mut map = self.inner.lock().expect("schema cache mutex");
84 if map.len() >= SCHEMA_CACHE_CAP && !map.contains_key(&cypher) {
85 if let Some(victim) = map.keys().next().cloned() {
88 map.remove(&victim);
89 }
90 }
91 map.insert(cypher, batch);
92 }
93
94 fn take(&self, cypher: &str) -> Option<RecordBatch> {
96 self.inner
97 .lock()
98 .expect("schema cache mutex")
99 .remove(cypher)
100 }
101}
102
103#[derive(Clone)]
114pub struct CypherFlightService {
115 executor: FalkorExecutor,
116 registry: Arc<SchemaRegistry>,
117 redis_url: Arc<String>,
118 graph: Arc<String>,
119 auth: Arc<AuthConfig>,
120 schema_cache: SchemaCache,
122}
123
124impl CypherFlightService {
125 pub async fn connect(
126 redis_url: impl Into<String>,
127 graph: impl Into<String>,
128 registry: Arc<SchemaRegistry>,
129 ) -> crate::error::Result<Self> {
130 let redis_url = Arc::new(redis_url.into());
131 let graph = Arc::new(graph.into());
132 let executor = FalkorExecutor::connect(redis_url.as_str(), graph.as_str()).await?;
133 Ok(Self {
134 executor,
135 registry,
136 redis_url,
137 graph,
138 auth: Arc::new(AuthConfig::default()),
139 schema_cache: SchemaCache::default(),
140 })
141 }
142
143 pub fn with_auth(mut self, auth: Arc<AuthConfig>) -> Self {
146 self.auth = auth;
147 self
148 }
149
150 async fn new_loader(&self) -> crate::error::Result<FalkorDbLoader> {
151 FalkorDbLoader::connect(&self.redis_url, self.graph.as_str())
152 .await
153 .map_err(FlightSqlError::from)
154 }
155
156 fn target_cypher(&self, incoming: &str) -> String {
163 if let Some(query) = self.registry.resolve_table(incoming) {
164 return query;
165 }
166 if let Some(ident) = parse_select_from(incoming)
167 && let Some(query) = self.registry.resolve_table(&ident)
168 {
169 return query;
170 }
171 incoming.to_string()
172 }
173
174 async fn schema_for(&self, cypher: &str) -> Result<arrow_schema::SchemaRef, Status> {
182 if let Some(s) = self.registry.get(cypher) {
183 Ok(s)
184 } else {
185 let batch = self
186 .executor
187 .clone()
188 .query_auto(cypher)
189 .await
190 .map_err(|e| Status::internal(e.to_string()))?;
191 let schema = batch.schema();
192 self.schema_cache.insert(cypher.to_string(), batch);
193 Ok(schema)
194 }
195 }
196
197 fn authz_name(&self, cypher: &str) -> String {
201 self.registry
202 .name_of(cypher)
203 .unwrap_or_else(|| cypher.to_string())
204 }
205
206 fn authorize<T>(&self, request: &Request<T>, cypher: &str) -> Result<(), Status> {
212 let identity = request
213 .extensions()
214 .get::<Identity>()
215 .cloned()
216 .unwrap_or(Identity::Anonymous);
217 self.auth.authorize(&identity, &self.authz_name(cypher))
218 }
219
220 async fn execute_cypher_stream(&self, cypher: &str) -> Result<Response<BoxStream>, Status> {
230 let mut exec = self.executor.clone();
231 let batch = if let Some(schema) = self.registry.get(cypher) {
232 exec.query(cypher, &schema)
233 .await
234 .map_err(|e| Status::internal(e.to_string()))?
235 } else if let Some(cached) = self.schema_cache.take(cypher) {
236 cached
237 } else {
238 exec.query_auto(cypher)
239 .await
240 .map_err(|e| Status::internal(e.to_string()))?
241 };
242 let schema = batch.schema();
243 let stream: BoxStream = Box::pin(
244 FlightDataEncoderBuilder::new()
245 .with_schema(schema)
246 .build(futures::stream::iter([Ok(batch)]))
247 .map_err(|e: FlightError| Status::internal(e.to_string())),
248 );
249 Ok(Response::new(stream))
250 }
251}
252
253#[tonic::async_trait]
254impl FlightSqlService for CypherFlightService {
255 type FlightService = CypherFlightService;
256
257 async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {}
260
261 async fn do_handshake(
270 &self,
271 request: Request<Streaming<HandshakeRequest>>,
272 ) -> Result<
273 Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
274 Status,
275 > {
276 let header = request
277 .metadata()
278 .get("authorization")
279 .and_then(|v| v.to_str().ok());
280 self.auth.check_header(header)?;
282
283 let token = self.auth.issued_token().unwrap_or("").to_string();
284 let response_msg = HandshakeResponse {
285 protocol_version: 0,
286 payload: token.clone().into(),
287 };
288 let output = futures::stream::once(async move { Ok(response_msg) });
289 let mut response: Response<Pin<Box<dyn Stream<Item = _> + Send>>> =
290 Response::new(Box::pin(output));
291 if !token.is_empty() {
292 let bearer = format!("Bearer {token}")
293 .parse()
294 .map_err(|_| Status::internal("token not header-safe"))?;
295 response.metadata_mut().insert("authorization", bearer);
296 }
297 Ok(response)
298 }
299
300 async fn get_flight_info_statement(
308 &self,
309 query: CommandStatementQuery,
310 request: Request<arrow_flight::FlightDescriptor>,
311 ) -> Result<Response<FlightInfo>, Status> {
312 let cypher = self.target_cypher(&query.query);
315 self.authorize(&request, &cypher)?;
316 let schema = self.schema_for(&cypher).await?;
317
318 let ticket_bytes = TicketStatementQuery {
319 statement_handle: cypher.into_bytes().into(),
320 }
321 .encode_to_vec();
322
323 let info = FlightInfo::new()
324 .try_with_schema(schema.as_ref())
325 .map_err(|e| Status::internal(e.to_string()))?
326 .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(ticket_bytes)))
327 .with_total_records(-1)
328 .with_total_bytes(-1)
329 .with_ordered(false);
330
331 Ok(Response::new(info))
332 }
333
334 async fn do_get_statement(
337 &self,
338 ticket: TicketStatementQuery,
339 request: Request<Ticket>,
340 ) -> Result<Response<BoxStream>, Status> {
341 let query = String::from_utf8(ticket.statement_handle.to_vec())
342 .map_err(|e| Status::invalid_argument(format!("invalid ticket: {e}")))?;
343 self.authorize(&request, &query)?;
344 self.execute_cypher_stream(&query).await
345 }
346
347 async fn do_put_statement_update(
350 &self,
351 cmd: CommandStatementUpdate,
352 request: Request<PeekableFlightDataStream>,
353 ) -> Result<i64, Status> {
354 let directive = cmd.query.trim().to_string();
355 let action =
356 parse_put_directive(&directive).map_err(|e| Status::invalid_argument(e.to_string()))?;
357
358 let raw = request.into_inner().into_inner(); let batches = FlightRecordBatchStream::new_from_flight_data(
361 raw.map_err(|s| FlightError::Tonic(Box::new(s))),
362 )
363 .try_collect::<Vec<_>>()
364 .await
365 .map_err(|e| Status::internal(e.to_string()))?;
366
367 if batches.is_empty() {
368 return Ok(0);
369 }
370
371 let mut loader = self
372 .new_loader()
373 .await
374 .map_err(|e| Status::internal(e.to_string()))?;
375
376 let n = match action {
377 PutAction::Vertices { label } => loader
378 .insert_vertices(&label, &batches)
379 .await
380 .map_err(|e| Status::internal(e.to_string()))?,
381 PutAction::Edges {
382 src,
383 edge_type,
384 dst,
385 } => loader
386 .insert_edges(&src, &edge_type, &dst, &batches)
387 .await
388 .map_err(|e| Status::internal(e.to_string()))?,
389 };
390
391 Ok(n as i64)
392 }
393
394 async fn get_flight_info_catalogs(
400 &self,
401 query: CommandGetCatalogs,
402 request: Request<FlightDescriptor>,
403 ) -> Result<Response<FlightInfo>, Status> {
404 let descriptor = request.into_inner();
405 let ticket = Ticket::new(query.as_any().encode_to_vec());
406 metadata_info(descriptor, ticket, query.into_builder().schema())
407 }
408
409 async fn get_flight_info_schemas(
410 &self,
411 query: CommandGetDbSchemas,
412 request: Request<FlightDescriptor>,
413 ) -> Result<Response<FlightInfo>, Status> {
414 let descriptor = request.into_inner();
415 let ticket = Ticket::new(query.as_any().encode_to_vec());
416 metadata_info(descriptor, ticket, query.into_builder().schema())
417 }
418
419 async fn get_flight_info_tables(
420 &self,
421 query: CommandGetTables,
422 request: Request<FlightDescriptor>,
423 ) -> Result<Response<FlightInfo>, Status> {
424 let descriptor = request.into_inner();
425 let ticket = Ticket::new(query.as_any().encode_to_vec());
426 metadata_info(descriptor, ticket, query.into_builder().schema())
427 }
428
429 async fn get_flight_info_table_types(
430 &self,
431 query: CommandGetTableTypes,
432 request: Request<FlightDescriptor>,
433 ) -> Result<Response<FlightInfo>, Status> {
434 let descriptor = request.into_inner();
435 let ticket = Ticket::new(query.as_any().encode_to_vec());
436 metadata_info(descriptor, ticket, query.into_builder().schema())
437 }
438
439 async fn get_flight_info_sql_info(
440 &self,
441 query: CommandGetSqlInfo,
442 request: Request<FlightDescriptor>,
443 ) -> Result<Response<FlightInfo>, Status> {
444 let descriptor = request.into_inner();
445 let ticket = Ticket::new(query.as_any().encode_to_vec());
446 metadata_info(descriptor, ticket, query.into_builder(&SQL_INFO).schema())
447 }
448
449 async fn do_get_catalogs(
450 &self,
451 query: CommandGetCatalogs,
452 _request: Request<Ticket>,
453 ) -> Result<Response<BoxStream>, Status> {
454 let mut builder = query.into_builder();
455 builder.append(self.graph.as_str());
456 let schema = builder.schema();
457 Ok(Response::new(metadata_stream(schema, builder.build())))
458 }
459
460 async fn do_get_schemas(
461 &self,
462 query: CommandGetDbSchemas,
463 _request: Request<Ticket>,
464 ) -> Result<Response<BoxStream>, Status> {
465 let mut builder = query.into_builder();
466 builder.append(self.graph.as_str(), DB_SCHEMA);
467 let schema = builder.schema();
468 Ok(Response::new(metadata_stream(schema, builder.build())))
469 }
470
471 async fn do_get_tables(
472 &self,
473 query: CommandGetTables,
474 _request: Request<Ticket>,
475 ) -> Result<Response<BoxStream>, Status> {
476 let mut builder = query.into_builder();
477 for (name, table_schema) in self.registry.tables() {
478 builder
479 .append(
480 self.graph.as_str(),
481 DB_SCHEMA,
482 &name,
483 "TABLE",
484 table_schema.as_ref(),
485 )
486 .map_err(|e| Status::internal(e.to_string()))?;
487 }
488 let schema = builder.schema();
489 Ok(Response::new(metadata_stream(schema, builder.build())))
490 }
491
492 async fn do_get_table_types(
493 &self,
494 query: CommandGetTableTypes,
495 _request: Request<Ticket>,
496 ) -> Result<Response<BoxStream>, Status> {
497 let mut builder = query.into_builder();
498 builder.append("TABLE");
499 let schema = builder.schema();
500 Ok(Response::new(metadata_stream(schema, builder.build())))
501 }
502
503 async fn do_get_sql_info(
504 &self,
505 query: CommandGetSqlInfo,
506 _request: Request<Ticket>,
507 ) -> Result<Response<BoxStream>, Status> {
508 let builder = query.into_builder(&SQL_INFO);
509 let schema = builder.schema();
510 Ok(Response::new(metadata_stream(schema, builder.build())))
511 }
512
513 async fn do_action_create_prepared_statement(
518 &self,
519 query: ActionCreatePreparedStatementRequest,
520 request: Request<arrow_flight::Action>,
521 ) -> Result<ActionCreatePreparedStatementResult, Status> {
522 let cypher = self.target_cypher(&query.query);
523 self.authorize(&request, &cypher)?;
524 let schema = self.schema_for(&cypher).await?;
525 let IpcMessage(schema_bytes) =
526 SchemaAsIpc::new(schema.as_ref(), &IpcWriteOptions::default())
527 .try_into()
528 .map_err(|e: arrow::error::ArrowError| Status::internal(e.to_string()))?;
529 Ok(ActionCreatePreparedStatementResult {
530 prepared_statement_handle: cypher.into_bytes().into(),
531 dataset_schema: schema_bytes,
532 parameter_schema: Default::default(), })
534 }
535
536 async fn do_action_close_prepared_statement(
537 &self,
538 _query: ActionClosePreparedStatementRequest,
539 _request: Request<arrow_flight::Action>,
540 ) -> Result<(), Status> {
541 Ok(()) }
543
544 async fn get_flight_info_prepared_statement(
545 &self,
546 cmd: CommandPreparedStatementQuery,
547 request: Request<FlightDescriptor>,
548 ) -> Result<Response<FlightInfo>, Status> {
549 let cypher = String::from_utf8(cmd.prepared_statement_handle.to_vec())
550 .map_err(|e| Status::invalid_argument(format!("invalid handle: {e}")))?;
551 self.authorize(&request, &cypher)?;
552 let schema = self.schema_for(&cypher).await?;
553
554 let ticket_bytes = TicketStatementQuery {
555 statement_handle: cypher.into_bytes().into(),
556 }
557 .encode_to_vec();
558
559 let info = FlightInfo::new()
560 .try_with_schema(schema.as_ref())
561 .map_err(|e| Status::internal(e.to_string()))?
562 .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(ticket_bytes)))
563 .with_total_records(-1)
564 .with_total_bytes(-1)
565 .with_ordered(false);
566 Ok(Response::new(info))
567 }
568
569 async fn do_get_prepared_statement(
570 &self,
571 cmd: CommandPreparedStatementQuery,
572 request: Request<Ticket>,
573 ) -> Result<Response<BoxStream>, Status> {
574 let cypher = String::from_utf8(cmd.prepared_statement_handle.to_vec())
575 .map_err(|e| Status::invalid_argument(format!("invalid handle: {e}")))?;
576 self.authorize(&request, &cypher)?;
577 self.execute_cypher_stream(&cypher).await
578 }
579
580 async fn do_put_prepared_statement_query(
581 &self,
582 query: CommandPreparedStatementQuery,
583 _request: Request<PeekableFlightDataStream>,
584 ) -> Result<DoPutPreparedStatementResult, Status> {
585 Ok(DoPutPreparedStatementResult {
587 prepared_statement_handle: Some(query.prepared_statement_handle),
588 })
589 }
590}
591
592const DB_SCHEMA: &str = "public";
593
594#[cfg(test)]
597pub(crate) fn sql_info_batch() -> arrow_array::RecordBatch {
598 CommandGetSqlInfo::default()
599 .into_builder(&SQL_INFO)
600 .build()
601 .expect("sql info batch builds")
602}
603
604fn metadata_info(
607 descriptor: FlightDescriptor,
608 ticket: Ticket,
609 schema: arrow_schema::SchemaRef,
610) -> Result<Response<FlightInfo>, Status> {
611 let info = FlightInfo::new()
612 .try_with_schema(schema.as_ref())
613 .map_err(|e| Status::internal(e.to_string()))?
614 .with_endpoint(FlightEndpoint::new().with_ticket(ticket))
615 .with_descriptor(descriptor);
616 Ok(Response::new(info))
617}
618
619fn metadata_stream(
622 schema: arrow_schema::SchemaRef,
623 batch: std::result::Result<arrow_array::RecordBatch, FlightError>,
624) -> BoxStream {
625 Box::pin(
626 FlightDataEncoderBuilder::new()
627 .with_schema(schema)
628 .build(futures::stream::once(async move { batch }))
629 .map_err(|e: FlightError| Status::internal(e.to_string())),
630 )
631}
632
633enum PutAction {
636 Vertices {
637 label: String,
638 },
639 Edges {
640 src: String,
641 edge_type: String,
642 dst: String,
643 },
644}
645
646fn parse_put_directive(s: &str) -> Result<PutAction, FlightSqlError> {
647 if let Some(rest) = s.strip_prefix("VERTICES:") {
648 return Ok(PutAction::Vertices {
649 label: rest.to_string(),
650 });
651 }
652 if let Some(rest) = s.strip_prefix("EDGES:") {
653 let parts: Vec<&str> = rest.splitn(3, ':').collect();
654 if parts.len() == 3 {
655 return Ok(PutAction::Edges {
656 src: parts[0].to_string(),
657 edge_type: parts[1].to_string(),
658 dst: parts[2].to_string(),
659 });
660 }
661 }
662 Err(FlightSqlError::UnknownPutCommand(s.to_string()))
663}