flow/batching_mode/
frontend_client.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Frontend client to run flow as batching task which is time-window-aware normal query triggered every tick set by user
16
17use std::collections::HashMap;
18use std::sync::{Arc, Mutex, RwLock, Weak};
19
20use api::v1::greptime_request::Request;
21use api::v1::query_request::Query;
22use api::v1::{CreateTableExpr, QueryRequest, RowInsertRequests};
23use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics};
24use common_error::ext::BoxedError;
25use common_grpc::channel_manager::{ChannelConfig, ChannelManager, load_client_tls_config};
26use common_meta::peer::{Peer, PeerDiscovery};
27use common_query::{Output, OutputData};
28use common_telemetry::warn;
29use common_time::Timezone;
30use futures::stream::{FuturesUnordered, StreamExt};
31use meta_client::client::MetaClient;
32use query::datafusion::QUERY_PARALLELISM_HINT;
33use query::metrics::terminal_recordbatch_metrics_from_plan;
34use query::options::{FlowQueryExtensions, QueryOptions};
35use rand::rng;
36use rand::seq::SliceRandom;
37use servers::query_handler::grpc::GrpcQueryHandler;
38use session::context::{Channel, QueryContextBuilder, QueryContextRef};
39use session::hints::READ_PREFERENCE_HINT;
40use snafu::{OptionExt, ResultExt};
41use tokio::sync::SetOnce;
42
43use crate::Error;
44use crate::batching_mode::BatchingModeOptions;
45use crate::error::{
46    CreateSinkTableSnafu, ExternalSnafu, InvalidClientConfigSnafu, InvalidRequestSnafu,
47    NoAvailableFrontendSnafu, UnexpectedSnafu,
48};
49
50/// Adapter trait for [`GrpcQueryHandler`] that boxes the underlying error into [`BoxedError`].
51///
52/// This is mainly used by flownode to invoke a frontend instance in standalone mode.
53#[async_trait::async_trait]
54pub trait GrpcQueryHandlerWithBoxedError: Send + Sync + 'static {
55    async fn do_query(
56        &self,
57        query: Request,
58        ctx: QueryContextRef,
59    ) -> std::result::Result<Output, BoxedError>;
60}
61
62/// auto impl
63#[async_trait::async_trait]
64impl<T: GrpcQueryHandler + Send + Sync + 'static> GrpcQueryHandlerWithBoxedError for T {
65    async fn do_query(
66        &self,
67        query: Request,
68        ctx: QueryContextRef,
69    ) -> std::result::Result<Output, BoxedError> {
70        self.do_query(query, ctx).await.map_err(BoxedError::new)
71    }
72}
73
74#[derive(Debug, Clone)]
75pub struct HandlerMutable {
76    handler: Arc<Mutex<Option<Weak<dyn GrpcQueryHandlerWithBoxedError>>>>,
77    is_initialized: Arc<SetOnce<()>>,
78}
79
80impl HandlerMutable {
81    pub async fn set_handler(&self, handler: Weak<dyn GrpcQueryHandlerWithBoxedError>) {
82        *self.handler.lock().unwrap() = Some(handler);
83        // Ignore the error, as we allow the handler to be set multiple times.
84        let _ = self.is_initialized.set(());
85    }
86}
87
88/// A simple frontend client able to execute sql using grpc protocol
89///
90/// This is for computation-heavy query which need to offload computation to frontend, lifting the load from flownode
91#[derive(Debug, Clone)]
92pub enum FrontendClient {
93    Distributed {
94        meta_client: Arc<MetaClient>,
95        query_channel_manager: ChannelManager,
96        control_channel_manager: ChannelManager,
97        query: QueryOptions,
98        batch_opts: BatchingModeOptions,
99    },
100    Standalone {
101        /// for the sake of simplicity still use grpc even in standalone mode
102        /// notice the client here should all be lazy, so that can wait after frontend is booted then make conn
103        database_client: HandlerMutable,
104        query: QueryOptions,
105    },
106}
107
108impl FrontendClient {
109    /// Create a new empty frontend client, with a `HandlerMutable` to set the grpc handler later
110    pub fn from_empty_grpc_handler(query: QueryOptions) -> (Self, HandlerMutable) {
111        let is_initialized = Arc::new(SetOnce::new());
112        let handler = HandlerMutable {
113            handler: Arc::new(Mutex::new(None)),
114            is_initialized,
115        };
116        (
117            Self::Standalone {
118                database_client: handler.clone(),
119                query,
120            },
121            handler,
122        )
123    }
124
125    /// Waits until the frontend client is initialized.
126    pub async fn wait_initialized(&self) {
127        if let FrontendClient::Standalone {
128            database_client, ..
129        } = self
130        {
131            database_client.is_initialized.wait().await;
132        }
133    }
134
135    pub fn from_meta_client(
136        meta_client: Arc<MetaClient>,
137        query: QueryOptions,
138        batch_opts: BatchingModeOptions,
139    ) -> Result<Self, Error> {
140        common_telemetry::info!("Frontend client build without auth");
141        let cfg = ChannelConfig::new()
142            .connect_timeout(batch_opts.grpc_conn_timeout)
143            .timeout(Some(batch_opts.query_timeout));
144        let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone())
145            .context(InvalidClientConfigSnafu)?;
146        // Keep separate TLS config handles and pools for query and control traffic.
147        let query_channel_manager = ChannelManager::with_config(cfg.clone(), tls_config.clone());
148        let control_channel_manager = ChannelManager::with_config(cfg, tls_config);
149
150        Ok(Self::Distributed {
151            meta_client,
152            query_channel_manager,
153            control_channel_manager,
154            query,
155            batch_opts,
156        })
157    }
158
159    pub fn from_grpc_handler(
160        grpc_handler: Weak<dyn GrpcQueryHandlerWithBoxedError>,
161        query: QueryOptions,
162    ) -> Self {
163        let is_initialized = Arc::new(SetOnce::new_with(Some(())));
164        let handler = HandlerMutable {
165            handler: Arc::new(Mutex::new(Some(grpc_handler))),
166            is_initialized: is_initialized.clone(),
167        };
168
169        Self::Standalone {
170            database_client: handler,
171            query,
172        }
173    }
174}
175
176#[derive(Debug, Clone)]
177pub struct DatabaseWithPeer {
178    pub database: Database,
179    pub peer: Peer,
180}
181
182impl DatabaseWithPeer {
183    fn new(database: Database, peer: Peer) -> Self {
184        Self { database, peer }
185    }
186
187    /// Try sending a "SELECT 1" to the database
188    async fn try_select_one(&self) -> Result<(), Error> {
189        // notice here use `sql` for `SELECT 1` return 1 row
190        let output = self
191            .database
192            .sql("SELECT 1")
193            .await
194            .with_context(|_| InvalidRequestSnafu {
195                context: format!("Failed to handle `SELECT 1` request at {:?}", self.peer),
196            })?;
197
198        if let OutputData::Stream(stream) = output.data {
199            common_recordbatch::util::collect(stream)
200                .await
201                .map_err(BoxedError::new)
202                .context(ExternalSnafu)?;
203        }
204
205        Ok(())
206    }
207}
208
209impl FrontendClient {
210    /// scan for available frontend from metadata
211    pub(crate) async fn scan_for_frontend(&self) -> Result<Vec<Peer>, Error> {
212        let Self::Distributed { meta_client, .. } = self else {
213            return Ok(vec![]);
214        };
215
216        meta_client
217            .active_frontends()
218            .await
219            .map(|nodes| nodes.into_iter().map(|node| node.peer).collect())
220            .map_err(BoxedError::new)
221            .context(ExternalSnafu)
222    }
223
224    /// Probes all discovered frontends without auth.
225    ///
226    /// Returns non-auth failures to allow callers to retry transient connectivity
227    /// errors. Authentication failures are returned immediately because they mean
228    /// a frontend advertised an auth-protected endpoint to flownodes.
229    pub(crate) async fn check_all_frontends_without_auth(
230        &self,
231        frontends: &[Peer],
232    ) -> Result<Vec<String>, Error> {
233        let Self::Distributed {
234            query_channel_manager,
235            control_channel_manager,
236            batch_opts,
237            ..
238        } = self
239        else {
240            return Ok(vec![]);
241        };
242
243        let probe_timeout = batch_opts.grpc_conn_timeout;
244        let mut probes = frontends
245            .iter()
246            .map(|peer| {
247                let addr = peer.addr.clone();
248                let query_channel_manager = query_channel_manager.clone();
249                let control_channel_manager = control_channel_manager.clone();
250
251                async move {
252                    let client = Client::with_query_and_control_managers(
253                        query_channel_manager,
254                        control_channel_manager,
255                        vec![addr.clone()],
256                    );
257                    let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client);
258
259                    match tokio::time::timeout(probe_timeout, database.sql("SELECT 1")).await {
260                        Ok(Ok(_)) => Ok(None),
261                        Ok(Err(err)) if err.tonic_code() == Some(tonic::Code::Unauthenticated) => {
262                            Err(err).context(InvalidRequestSnafu {
263                                context: format!(
264                                    "Frontend {addr} rejected unauthenticated flownode probe; ensure frontend internal_grpc is advertised to metasrv"
265                                ),
266                            })
267                        }
268                        Ok(Err(err)) => Ok(Some(format!("{addr}: {err}"))),
269                        Err(_) => Ok(Some(format!(
270                            "{addr}: health check timed out after {probe_timeout:?}"
271                        ))),
272                    }
273                }
274            })
275            .collect::<FuturesUnordered<_>>();
276
277        let mut failures = Vec::new();
278        while let Some(probe_result) = probes.next().await {
279            if let Some(failure) = probe_result? {
280                failures.push(failure);
281            }
282        }
283
284        Ok(failures)
285    }
286
287    /// Get a frontend discovered by metasrv and verified with a query probe.
288    async fn get_random_active_frontend(
289        &self,
290        catalog: &str,
291        schema: &str,
292    ) -> Result<DatabaseWithPeer, Error> {
293        let Self::Distributed {
294            meta_client: _,
295            query_channel_manager,
296            control_channel_manager,
297            query: _,
298            batch_opts,
299            ..
300        } = self
301        else {
302            return UnexpectedSnafu {
303                reason: "Expect distributed mode",
304            }
305            .fail();
306        };
307
308        let mut interval = tokio::time::interval(batch_opts.grpc_conn_timeout);
309        interval.tick().await;
310        for retry in 0..batch_opts.experimental_grpc_max_retries {
311            let mut frontends = self.scan_for_frontend().await?;
312            // shuffle the frontends to avoid always pick the same one
313            frontends.shuffle(&mut rng());
314
315            for peer in frontends {
316                let addr = peer.addr.clone();
317                let client = Client::with_query_and_control_managers(
318                    query_channel_manager.clone(),
319                    control_channel_manager.clone(),
320                    vec![addr.clone()],
321                );
322                let database = Database::new(catalog, schema, client);
323                let db = DatabaseWithPeer::new(database, peer);
324                match db.try_select_one().await {
325                    Ok(_) => return Ok(db),
326                    Err(e) => {
327                        warn!(
328                            "Failed to connect to frontend {} on retry={}: \n{e:?}",
329                            addr, retry
330                        );
331                    }
332                }
333            }
334            // no available frontend
335            // sleep and retry
336            interval.tick().await;
337        }
338
339        NoAvailableFrontendSnafu {
340            timeout: batch_opts.grpc_conn_timeout,
341            context: "No available frontend found that is able to process query",
342        }
343        .fail()
344    }
345
346    pub async fn create(
347        &self,
348        create: CreateTableExpr,
349        catalog: &str,
350        schema: &str,
351    ) -> Result<u32, Error> {
352        self.handle(
353            Request::Ddl(api::v1::DdlRequest {
354                expr: Some(api::v1::ddl_request::Expr::CreateTable(create.clone())),
355            }),
356            catalog,
357            schema,
358            &mut None,
359        )
360        .await
361        .map_err(BoxedError::new)
362        .with_context(|_| CreateSinkTableSnafu {
363            create: create.clone(),
364        })
365    }
366
367    /// Executes SQL with an optional timezone override.
368    pub async fn sql(
369        &self,
370        catalog: &str,
371        schema: &str,
372        sql: &str,
373        timezone: Option<Timezone>,
374    ) -> Result<Output, Error> {
375        match self {
376            FrontendClient::Distributed { .. } => {
377                let mut db = self.get_random_active_frontend(catalog, schema).await?;
378                if let Some(timezone) = timezone {
379                    db.database.set_timezone(timezone.to_string());
380                }
381                db.database
382                    .sql(sql)
383                    .await
384                    .map_err(BoxedError::new)
385                    .context(ExternalSnafu)
386            }
387            FrontendClient::Standalone {
388                database_client, ..
389            } => {
390                // Each request gets fresh query state, including read snapshots.
391                let mut ctx = QueryContextBuilder::default()
392                    .current_catalog(catalog.to_string())
393                    .current_schema(schema.to_string());
394                if let Some(timezone) = timezone {
395                    ctx = ctx.timezone(timezone);
396                }
397                let ctx = Arc::new(ctx.build());
398                let database_client = {
399                    database_client
400                        .handler
401                        .lock()
402                        .map_err(|e| {
403                            UnexpectedSnafu {
404                                reason: format!("Failed to lock database client: {e}"),
405                            }
406                            .build()
407                        })?
408                        .as_ref()
409                        .context(UnexpectedSnafu {
410                            reason: "Standalone's frontend instance is not set",
411                        })?
412                        .upgrade()
413                        .context(UnexpectedSnafu {
414                            reason: "Failed to upgrade database client",
415                        })?
416                };
417                let req = Request::Query(QueryRequest {
418                    query: Some(Query::Sql(sql.to_string())),
419                });
420                database_client
421                    .do_query(req, ctx)
422                    .await
423                    .map_err(BoxedError::new)
424                    .context(ExternalSnafu)
425            }
426        }
427    }
428
429    /// Execute row inserts on the frontend.
430    pub async fn row_inserts(
431        &self,
432        catalog: &str,
433        schema: &str,
434        requests: RowInsertRequests,
435        hints: &[(&str, &str)],
436    ) -> Result<u32, Error> {
437        match self {
438            FrontendClient::Distributed { .. } => {
439                let db = self.get_random_active_frontend(catalog, schema).await?;
440                db.database
441                    .row_inserts_with_hints(requests, hints)
442                    .await
443                    .with_context(|_| InvalidRequestSnafu {
444                        context: format!("Failed to handle row inserts at {:?}", db.peer),
445                    })
446            }
447            FrontendClient::Standalone {
448                database_client, ..
449            } => {
450                let extensions = HashMap::from_iter(
451                    hints
452                        .iter()
453                        .map(|(key, value)| ((*key).to_string(), (*value).to_string())),
454                );
455                let ctx = QueryContextBuilder::default()
456                    .current_catalog(catalog.to_string())
457                    .current_schema(schema.to_string())
458                    .extensions(extensions)
459                    .build();
460                let ctx = Arc::new(ctx);
461                {
462                    let database_client = {
463                        database_client
464                            .handler
465                            .lock()
466                            .unwrap()
467                            .as_ref()
468                            .context(UnexpectedSnafu {
469                                reason: "Standalone's frontend instance is not set",
470                            })?
471                            .upgrade()
472                            .context(UnexpectedSnafu {
473                                reason: "Failed to upgrade database client",
474                            })?
475                    };
476                    let resp: common_query::Output = database_client
477                        .do_query(Request::RowInserts(requests), ctx)
478                        .await
479                        .map_err(BoxedError::new)
480                        .context(ExternalSnafu)?;
481                    match resp.data {
482                        OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| {
483                            UnexpectedSnafu {
484                                reason: format!("Failed to convert rows to u32: {}", rows),
485                            }
486                            .build()
487                        })?),
488                        _ => UnexpectedSnafu {
489                            reason: "Unexpected output data",
490                        }
491                        .fail(),
492                    }
493                }
494            }
495        }
496    }
497
498    /// Execute a flow query and return terminal metrics. `snapshot_seqs` are
499    /// optional read upper bounds used only by snapshot-fenced repair chunks.
500    pub(crate) async fn query_with_terminal_metrics(
501        &self,
502        catalog: &str,
503        schema: &str,
504        request: QueryRequest,
505        extensions: &[(&str, &str)],
506        snapshot_seqs: &HashMap<u64, u64>,
507        peer_desc: &mut Option<PeerDesc>,
508    ) -> Result<OutputWithMetrics, Error> {
509        let flow_extensions = build_flow_extensions(extensions)?;
510        match self {
511            FrontendClient::Distributed {
512                query, batch_opts, ..
513            } => {
514                let query_parallelism = query.parallelism.to_string();
515                let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
516                if query.parallelism > 0 {
517                    hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
518                }
519                let db = self.get_random_active_frontend(catalog, schema).await?;
520                *peer_desc = Some(PeerDesc::Dist {
521                    peer: db.peer.clone(),
522                });
523                db.database
524                    .flight_request()
525                    .with_hints(&hints)
526                    .with_flow_extensions(extensions)
527                    .with_snapshot_seqs(snapshot_seqs)
528                    .with_timeout(batch_opts.experimental_flight_do_get_timeout)
529                    .query_with_terminal_metrics(request)
530                    .await
531                    .map_err(BoxedError::new)
532                    .context(ExternalSnafu)
533            }
534            FrontendClient::Standalone {
535                database_client,
536                query,
537                ..
538            } => {
539                *peer_desc = Some(PeerDesc::Standalone);
540                let mut extensions_map = HashMap::new();
541                if query.parallelism > 0 {
542                    extensions_map.insert(
543                        QUERY_PARALLELISM_HINT.to_string(),
544                        query.parallelism.to_string(),
545                    );
546                }
547                for (key, value) in extensions {
548                    extensions_map.insert((*key).to_string(), (*value).to_string());
549                }
550                let ctx = QueryContextBuilder::default()
551                    .current_catalog(catalog.to_string())
552                    .current_schema(schema.to_string())
553                    .extensions(extensions_map)
554                    .channel(Channel::Internal)
555                    .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs.clone())))
556                    .build();
557                let ctx = Arc::new(ctx);
558                let database_client = {
559                    database_client
560                        .handler
561                        .lock()
562                        .map_err(|e| {
563                            UnexpectedSnafu {
564                                reason: format!("Failed to lock database client: {e}"),
565                            }
566                            .build()
567                        })?
568                        .as_ref()
569                        .context(UnexpectedSnafu {
570                            reason: "Standalone's frontend instance is not set",
571                        })?
572                        .upgrade()
573                        .context(UnexpectedSnafu {
574                            reason: "Failed to upgrade database client",
575                        })?
576                };
577                database_client
578                    .do_query(Request::Query(request), ctx.clone())
579                    .await
580                    .map(|output| {
581                        wrap_standalone_output_with_terminal_metrics(output, &flow_extensions)
582                    })
583                    .map_err(BoxedError::new)
584                    .context(ExternalSnafu)
585            }
586        }
587    }
588
589    /// Handle an insert request with one attempt.
590    ///
591    /// Unlike [`Self::handle`], this does not use the batching retry policy. It
592    /// is intended for stateless streaming sinks, where retrying an insert can
593    /// duplicate rows.
594    pub(crate) async fn handle_insert_once(
595        &self,
596        req: api::v1::greptime_request::Request,
597        catalog: &str,
598        schema: &str,
599        peer_desc: &mut Option<PeerDesc>,
600    ) -> Result<u32, Error> {
601        match self {
602            FrontendClient::Distributed { .. } => {
603                let db = self.get_random_active_frontend(catalog, schema).await?;
604
605                *peer_desc = Some(PeerDesc::Dist {
606                    peer: db.peer.clone(),
607                });
608
609                db.database
610                    .handle(req.clone())
611                    .await
612                    .with_context(|_| InvalidRequestSnafu {
613                        context: format!("Failed to handle request at {:?}: {:?}", db.peer, req),
614                    })
615            }
616            FrontendClient::Standalone {
617                database_client,
618                query,
619                ..
620            } => {
621                let ctx = QueryContextBuilder::default()
622                    .current_catalog(catalog.to_string())
623                    .current_schema(schema.to_string())
624                    .channel(Channel::Internal)
625                    .extensions(if query.parallelism > 0 {
626                        HashMap::from([(
627                            QUERY_PARALLELISM_HINT.to_string(),
628                            query.parallelism.to_string(),
629                        )])
630                    } else {
631                        HashMap::new()
632                    })
633                    .build();
634                let ctx = Arc::new(ctx);
635                let database_client = {
636                    database_client
637                        .handler
638                        .lock()
639                        .unwrap()
640                        .as_ref()
641                        .context(UnexpectedSnafu {
642                            reason: "Standalone's frontend instance is not set",
643                        })?
644                        .upgrade()
645                        .context(UnexpectedSnafu {
646                            reason: "Failed to upgrade database client",
647                        })?
648                };
649                let resp: common_query::Output = database_client
650                    .do_query(req, ctx)
651                    .await
652                    .map_err(BoxedError::new)
653                    .context(ExternalSnafu)?;
654                match resp.data {
655                    common_query::OutputData::AffectedRows(rows) => rows.try_into().map_err(|_| {
656                        UnexpectedSnafu {
657                            reason: format!("Failed to convert rows to u32: {}", rows),
658                        }
659                        .build()
660                    }),
661                    _ => UnexpectedSnafu {
662                        reason: "Unexpected output data",
663                    }
664                    .fail(),
665                }
666            }
667        }
668    }
669
670    /// Handle a request to frontend
671    pub(crate) async fn handle(
672        &self,
673        req: api::v1::greptime_request::Request,
674        catalog: &str,
675        schema: &str,
676        peer_desc: &mut Option<PeerDesc>,
677    ) -> Result<u32, Error> {
678        match self {
679            FrontendClient::Distributed {
680                query, batch_opts, ..
681            } => {
682                let db = self.get_random_active_frontend(catalog, schema).await?;
683
684                *peer_desc = Some(PeerDesc::Dist {
685                    peer: db.peer.clone(),
686                });
687
688                let query_parallelism = query.parallelism.to_string();
689                let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
690                if query.parallelism > 0 {
691                    hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
692                }
693                db.database
694                    .handle_with_retry(
695                        req.clone(),
696                        batch_opts.experimental_grpc_max_retries,
697                        &hints,
698                    )
699                    .await
700                    .with_context(|_| InvalidRequestSnafu {
701                        context: format!("Failed to handle request at {:?}: {:?}", db.peer, req),
702                    })
703            }
704            FrontendClient::Standalone {
705                database_client,
706                query,
707                ..
708            } => {
709                let ctx = QueryContextBuilder::default()
710                    .current_catalog(catalog.to_string())
711                    .current_schema(schema.to_string())
712                    .extensions(if query.parallelism > 0 {
713                        HashMap::from([(
714                            QUERY_PARALLELISM_HINT.to_string(),
715                            query.parallelism.to_string(),
716                        )])
717                    } else {
718                        HashMap::new()
719                    })
720                    .build();
721                let ctx = Arc::new(ctx);
722                {
723                    let database_client = {
724                        database_client
725                            .handler
726                            .lock()
727                            .unwrap()
728                            .as_ref()
729                            .context(UnexpectedSnafu {
730                                reason: "Standalone's frontend instance is not set",
731                            })?
732                            .upgrade()
733                            .context(UnexpectedSnafu {
734                                reason: "Failed to upgrade database client",
735                            })?
736                    };
737                    let resp: common_query::Output = database_client
738                        .do_query(req, ctx)
739                        .await
740                        .map_err(BoxedError::new)
741                        .context(ExternalSnafu)?;
742                    match resp.data {
743                        OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| {
744                            UnexpectedSnafu {
745                                reason: format!("Failed to convert rows to u32: {}", rows),
746                            }
747                            .build()
748                        })?),
749                        _ => UnexpectedSnafu {
750                            reason: "Unexpected output data",
751                        }
752                        .fail(),
753                    }
754                }
755            }
756        }
757    }
758}
759
760fn build_flow_extensions(extensions: &[(&str, &str)]) -> Result<FlowQueryExtensions, Error> {
761    let flow_extensions = HashMap::from_iter(
762        extensions
763            .iter()
764            .map(|(key, value)| ((*key).to_string(), (*value).to_string())),
765    );
766    FlowQueryExtensions::parse_flow_extensions(&flow_extensions)
767        .map_err(BoxedError::new)
768        .context(ExternalSnafu)
769        .map(|extensions| extensions.unwrap_or_default())
770}
771
772fn wrap_standalone_output_with_terminal_metrics(
773    output: Output,
774    flow_extensions: &FlowQueryExtensions,
775) -> OutputWithMetrics {
776    let should_collect_region_watermark = flow_extensions.should_collect_region_watermark();
777    let terminal_metrics =
778        if should_collect_region_watermark && !matches!(&output.data, OutputData::Stream(_)) {
779            output
780                .meta
781                .plan
782                .clone()
783                .and_then(terminal_recordbatch_metrics_from_plan)
784        } else {
785            None
786        };
787    let result = OutputWithMetrics::from_output(output);
788    if let Some(metrics) = terminal_metrics {
789        result.metrics.update(Some(metrics));
790    }
791    result
792}
793
794/// Describe a peer of frontend
795#[derive(Debug, Default, Clone)]
796pub enum PeerDesc {
797    /// The query failed before a frontend peer was selected.
798    #[default]
799    Unknown,
800    /// Distributed mode's frontend peer address
801    Dist {
802        /// frontend peer address
803        peer: Peer,
804    },
805    /// Standalone mode
806    Standalone,
807}
808
809impl std::fmt::Display for PeerDesc {
810    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
811        match self {
812            PeerDesc::Unknown => write!(f, "unknown"),
813            PeerDesc::Dist { peer } => write!(f, "{}", peer.addr),
814            PeerDesc::Standalone => write!(f, "standalone"),
815        }
816    }
817}
818
819#[cfg(test)]
820mod tests {
821    use std::pin::Pin;
822    use std::sync::atomic::{AtomicUsize, Ordering};
823    use std::task::{Context, Poll};
824    use std::time::Duration;
825
826    use api::v1::query_request::Query;
827    use arrow_flight::flight_service_server::FlightServiceServer;
828    use arrow_flight::{FlightData, Ticket};
829    use common_grpc::flight::FlightEncoder;
830    use common_query::{Output, OutputData};
831    use common_recordbatch::adapter::RecordBatchMetrics;
832    use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream};
833    use datatypes::arrow::datatypes::Schema as ArrowSchema;
834    use datatypes::prelude::{ConcreteDataType, VectorRef};
835    use datatypes::schema::{ColumnSchema, Schema};
836    use datatypes::vectors::Int32Vector;
837    use futures::StreamExt;
838    use servers::grpc::flight::{FlightCraft, FlightCraftWrapper, TonicStream};
839    use tokio::net::TcpListener;
840    use tokio::task::JoinHandle;
841    use tokio::time::timeout;
842    use tokio_stream::wrappers::TcpListenerStream;
843    use tonic::{Request as TonicRequest, Response as TonicResponse, Status};
844
845    use super::*;
846
847    #[derive(Debug)]
848    struct NoopHandler;
849
850    struct MockMetricsStream {
851        schema: datatypes::schema::SchemaRef,
852        batch: Option<RecordBatch>,
853        metrics: RecordBatchMetrics,
854        terminal_metrics_only: bool,
855    }
856
857    impl futures::Stream for MockMetricsStream {
858        type Item = common_recordbatch::error::Result<RecordBatch>;
859
860        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
861            Poll::Ready(self.batch.take().map(Ok))
862        }
863
864        fn size_hint(&self) -> (usize, Option<usize>) {
865            (
866                usize::from(self.batch.is_some()),
867                Some(usize::from(self.batch.is_some())),
868            )
869        }
870    }
871
872    impl RecordBatchStream for MockMetricsStream {
873        fn name(&self) -> &str {
874            "MockMetricsStream"
875        }
876
877        fn schema(&self) -> datatypes::schema::SchemaRef {
878            self.schema.clone()
879        }
880
881        fn output_ordering(&self) -> Option<&[OrderOption]> {
882            None
883        }
884
885        fn metrics(&self) -> Option<RecordBatchMetrics> {
886            if self.terminal_metrics_only && self.batch.is_some() {
887                return None;
888            }
889            Some(self.metrics.clone())
890        }
891    }
892
893    #[derive(Debug)]
894    struct MetricsHandler;
895
896    #[derive(Debug)]
897    struct InsertOnceHandler {
898        calls: Arc<AtomicUsize>,
899    }
900
901    #[derive(Debug)]
902    struct ExtensionAwareHandler;
903
904    #[derive(Debug)]
905    struct SnapshotBindingHandler;
906
907    #[derive(Debug)]
908    struct RejectUnauthenticatedFlight;
909
910    #[derive(Debug)]
911    struct SlowFlight;
912
913    struct DelayedEofFlight {
914        schema_sent: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
915        release: Arc<tokio::sync::Notify>,
916    }
917
918    #[derive(Debug)]
919    struct LateStreamErrorFlight;
920
921    struct WaitForConcurrentFlight {
922        barrier: Arc<tokio::sync::Barrier>,
923    }
924
925    #[async_trait::async_trait]
926    impl GrpcQueryHandlerWithBoxedError for NoopHandler {
927        async fn do_query(
928            &self,
929            _query: Request,
930            _ctx: QueryContextRef,
931        ) -> std::result::Result<Output, BoxedError> {
932            Ok(Output::new_with_affected_rows(0))
933        }
934    }
935
936    #[async_trait::async_trait]
937    impl GrpcQueryHandlerWithBoxedError for InsertOnceHandler {
938        async fn do_query(
939            &self,
940            _query: Request,
941            ctx: QueryContextRef,
942        ) -> std::result::Result<Output, BoxedError> {
943            assert_eq!(ctx.channel(), Channel::Internal);
944            self.calls.fetch_add(1, Ordering::SeqCst);
945            Ok(Output::new_with_affected_rows(1))
946        }
947    }
948
949    #[async_trait::async_trait]
950    impl GrpcQueryHandlerWithBoxedError for MetricsHandler {
951        async fn do_query(
952            &self,
953            _query: Request,
954            _ctx: QueryContextRef,
955        ) -> std::result::Result<Output, BoxedError> {
956            let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
957                "v",
958                ConcreteDataType::int32_datatype(),
959                false,
960            )]));
961            let batch = RecordBatch::new(
962                schema.clone(),
963                vec![Arc::new(Int32Vector::from_slice([1, 2])) as VectorRef],
964            )
965            .unwrap();
966            Ok(Output::new_with_stream(Box::pin(MockMetricsStream {
967                schema,
968                batch: Some(batch),
969                metrics: RecordBatchMetrics {
970                    region_watermarks: vec![common_recordbatch::adapter::RegionWatermarkEntry {
971                        region_id: 42,
972                        watermark: Some(99),
973                    }],
974                    ..Default::default()
975                },
976                terminal_metrics_only: true,
977            })))
978        }
979    }
980
981    #[async_trait::async_trait]
982    impl GrpcQueryHandlerWithBoxedError for ExtensionAwareHandler {
983        async fn do_query(
984            &self,
985            _query: Request,
986            ctx: QueryContextRef,
987        ) -> std::result::Result<Output, BoxedError> {
988            assert_eq!(ctx.extension("flow.return_region_seq"), Some("true"));
989            assert_eq!(ctx.channel(), Channel::Internal);
990            Ok(Output::new_with_affected_rows(1))
991        }
992    }
993
994    #[async_trait::async_trait]
995    impl GrpcQueryHandlerWithBoxedError for SnapshotBindingHandler {
996        async fn do_query(
997            &self,
998            _query: Request,
999            ctx: QueryContextRef,
1000        ) -> std::result::Result<Output, BoxedError> {
1001            assert_eq!(ctx.extension("flow.return_region_seq"), Some("true"));
1002            assert_eq!(ctx.get_snapshot(1), Some(10));
1003            assert_eq!(ctx.get_snapshot(2), Some(20));
1004            ctx.set_snapshot(42, 99);
1005            Ok(Output::new_with_affected_rows(1))
1006        }
1007    }
1008
1009    #[async_trait::async_trait]
1010    impl FlightCraft for RejectUnauthenticatedFlight {
1011        async fn do_get(
1012            &self,
1013            _request: TonicRequest<Ticket>,
1014        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1015            Err(Status::unauthenticated("auth failed"))
1016        }
1017    }
1018
1019    #[async_trait::async_trait]
1020    impl FlightCraft for SlowFlight {
1021        async fn do_get(
1022            &self,
1023            _request: TonicRequest<Ticket>,
1024        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1025            tokio::time::sleep(Duration::from_secs(60)).await;
1026            Err(Status::unavailable("slow response"))
1027        }
1028    }
1029
1030    #[async_trait::async_trait]
1031    impl FlightCraft for DelayedEofFlight {
1032        async fn do_get(
1033            &self,
1034            _request: TonicRequest<Ticket>,
1035        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1036            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
1037            let schema_sent = self.schema_sent.lock().unwrap().take();
1038            let schema_stream = futures::stream::once(async move {
1039                if let Some(schema_sent) = schema_sent {
1040                    let _ = schema_sent.send(());
1041                }
1042                Ok(schema)
1043            });
1044            let release = self.release.clone();
1045            let delayed_eof = futures::stream::unfold(release, |release| async move {
1046                release.notified().await;
1047                None::<(std::result::Result<FlightData, Status>, _)>
1048            });
1049
1050            Ok(TonicResponse::new(Box::pin(
1051                schema_stream.chain(delayed_eof),
1052            )))
1053        }
1054    }
1055
1056    #[async_trait::async_trait]
1057    impl FlightCraft for LateStreamErrorFlight {
1058        async fn do_get(
1059            &self,
1060            _request: TonicRequest<Ticket>,
1061        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1062            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
1063            let stream =
1064                futures::stream::iter([Ok(schema), Err(Status::unavailable("late stream error"))]);
1065            Ok(TonicResponse::new(Box::pin(stream)))
1066        }
1067    }
1068
1069    #[async_trait::async_trait]
1070    impl FlightCraft for WaitForConcurrentFlight {
1071        async fn do_get(
1072            &self,
1073            _request: TonicRequest<Ticket>,
1074        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1075            self.barrier.wait().await;
1076            Err(Status::unavailable("probe started concurrently"))
1077        }
1078    }
1079
1080    async fn start_flight_server<T: FlightCraft>(handler: T) -> (String, JoinHandle<()>) {
1081        let listener = TcpListener::bind("127.0.0.1:0")
1082            .await
1083            .expect("bind test flight server");
1084        let addr = listener.local_addr().expect("local addr").to_string();
1085        let server = tokio::spawn(async move {
1086            tonic::transport::Server::builder()
1087                .add_service(FlightServiceServer::new(FlightCraftWrapper(handler)))
1088                .serve_with_incoming(TcpListenerStream::new(listener))
1089                .await
1090                .expect("serve test flight server");
1091        });
1092
1093        (addr, server)
1094    }
1095
1096    #[tokio::test]
1097    async fn test_sql_timezone_and_fresh_context() {
1098        struct ContextHandler {
1099            expected_timezone: Timezone,
1100            contexts: Mutex<Vec<QueryContextRef>>,
1101        }
1102
1103        #[async_trait::async_trait]
1104        impl GrpcQueryHandlerWithBoxedError for ContextHandler {
1105            async fn do_query(
1106                &self,
1107                request: Request,
1108                ctx: QueryContextRef,
1109            ) -> std::result::Result<Output, BoxedError> {
1110                assert_eq!(ctx.current_catalog(), "test_catalog");
1111                assert_eq!(ctx.current_schema(), "test_schema");
1112                assert_eq!(ctx.timezone(), self.expected_timezone);
1113                assert!(ctx.snapshots().is_empty());
1114                ctx.set_snapshot(1, 42);
1115                self.contexts.lock().unwrap().push(ctx);
1116                assert!(matches!(request, Request::Query(QueryRequest {
1117                    query: Some(Query::Sql(sql)),
1118                }) if sql == "SELECT to_unixtime('2026-01-01 00:00:00')"));
1119                Ok(Output::new_with_affected_rows(0))
1120            }
1121        }
1122
1123        for timezone in [None, Some("UTC"), Some("+08:00")] {
1124            let timezone = timezone.map(|tz| Timezone::from_tz_string(tz).unwrap());
1125            let handler = Arc::new(ContextHandler {
1126                expected_timezone: timezone
1127                    .clone()
1128                    .unwrap_or_else(|| QueryContextBuilder::default().build().timezone()),
1129                contexts: Mutex::new(Vec::new()),
1130            });
1131            let grpc_handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = handler.clone();
1132            let client = FrontendClient::from_grpc_handler(
1133                Arc::downgrade(&grpc_handler),
1134                QueryOptions::default(),
1135            );
1136            for _ in 0..2 {
1137                client
1138                    .sql(
1139                        "test_catalog",
1140                        "test_schema",
1141                        "SELECT to_unixtime('2026-01-01 00:00:00')",
1142                        timezone.clone(),
1143                    )
1144                    .await
1145                    .unwrap();
1146            }
1147            let contexts = handler.contexts.lock().unwrap();
1148            assert_eq!(contexts.len(), 2);
1149            assert!(!Arc::ptr_eq(&contexts[0], &contexts[1]));
1150        }
1151    }
1152
1153    #[tokio::test]
1154    async fn wait_initialized() {
1155        let (client, handler_mut) =
1156            FrontendClient::from_empty_grpc_handler(QueryOptions::default());
1157
1158        assert!(
1159            timeout(Duration::from_millis(50), client.wait_initialized())
1160                .await
1161                .is_err()
1162        );
1163
1164        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1165        handler_mut.set_handler(Arc::downgrade(&handler)).await;
1166
1167        timeout(Duration::from_secs(1), client.wait_initialized())
1168            .await
1169            .expect("wait_initialized should complete after handler is set");
1170
1171        timeout(Duration::from_millis(10), client.wait_initialized())
1172            .await
1173            .expect("wait_initialized should be a no-op once initialized");
1174
1175        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1176        let client =
1177            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1178        assert!(
1179            timeout(Duration::from_millis(10), client.wait_initialized())
1180                .await
1181                .is_ok()
1182        );
1183
1184        let meta_client = Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend));
1185        let client = FrontendClient::from_meta_client(
1186            meta_client,
1187            QueryOptions::default(),
1188            BatchingModeOptions::default(),
1189        )
1190        .unwrap();
1191        assert!(
1192            timeout(Duration::from_millis(10), client.wait_initialized())
1193                .await
1194                .is_ok()
1195        );
1196    }
1197
1198    #[tokio::test]
1199    async fn test_handle_insert_once_calls_standalone_handler_once() {
1200        let calls = Arc::new(AtomicUsize::new(0));
1201        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(InsertOnceHandler {
1202            calls: calls.clone(),
1203        });
1204        let client =
1205            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1206        let mut peer_desc = None;
1207
1208        let affected_rows = client
1209            .handle_insert_once(
1210                Request::RowInserts(api::v1::RowInsertRequests { inserts: vec![] }),
1211                "greptime",
1212                "public",
1213                &mut peer_desc,
1214            )
1215            .await
1216            .unwrap();
1217
1218        assert_eq!(affected_rows, 1);
1219        assert_eq!(calls.load(Ordering::SeqCst), 1);
1220        assert!(peer_desc.is_none());
1221    }
1222
1223    #[tokio::test]
1224    async fn test_query_with_terminal_metrics_tracks_watermark_in_standalone_mode() {
1225        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(MetricsHandler);
1226        let client =
1227            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1228        let mut peer_desc = None;
1229
1230        let result = client
1231            .query_with_terminal_metrics(
1232                "greptime",
1233                "public",
1234                QueryRequest {
1235                    query: Some(Query::Sql("select 1".to_string())),
1236                },
1237                &[],
1238                &HashMap::new(),
1239                &mut peer_desc,
1240            )
1241            .await
1242            .unwrap();
1243        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1244
1245        let terminal_metrics = result.metrics.clone();
1246        assert!(!result.metrics.is_ready());
1247        assert!(terminal_metrics.get().is_none());
1248
1249        let OutputData::Stream(mut stream) = result.output.data else {
1250            panic!("expected stream output");
1251        };
1252        while stream.next().await.is_some() {}
1253
1254        assert!(terminal_metrics.is_ready());
1255        assert_eq!(
1256            terminal_metrics.region_watermark_map(),
1257            Some(HashMap::from([(42_u64, 99_u64)]))
1258        );
1259    }
1260
1261    #[tokio::test]
1262    async fn test_query_with_terminal_metrics_forwards_flow_extensions_in_standalone_mode() {
1263        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(ExtensionAwareHandler);
1264        let client =
1265            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1266        let mut peer_desc = None;
1267
1268        let result = client
1269            .query_with_terminal_metrics(
1270                "greptime",
1271                "public",
1272                QueryRequest {
1273                    query: Some(Query::Sql("insert into t select 1".to_string())),
1274                },
1275                &[("flow.return_region_seq", "true")],
1276                &HashMap::new(),
1277                &mut peer_desc,
1278            )
1279            .await
1280            .unwrap();
1281        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1282
1283        assert!(result.metrics.is_ready());
1284        assert!(result.region_watermark_map().is_none());
1285    }
1286
1287    #[tokio::test]
1288    async fn test_query_with_terminal_metrics_uses_standalone_snapshot_bounds() {
1289        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(SnapshotBindingHandler);
1290        let client =
1291            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1292        let mut peer_desc = None;
1293
1294        let result = client
1295            .query_with_terminal_metrics(
1296                "greptime",
1297                "public",
1298                QueryRequest {
1299                    query: Some(Query::Sql("insert into t select * from src".to_string())),
1300                },
1301                &[("flow.return_region_seq", "true")],
1302                &HashMap::from([(1, 10), (2, 20)]),
1303                &mut peer_desc,
1304            )
1305            .await
1306            .unwrap();
1307        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1308
1309        assert!(result.metrics.is_ready());
1310        assert_eq!(result.region_watermark_map(), None);
1311    }
1312
1313    #[tokio::test]
1314    async fn test_query_with_terminal_metrics_rejects_invalid_flow_extensions() {
1315        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1316        let client =
1317            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1318        let mut peer_desc = None;
1319
1320        let err = client
1321            .query_with_terminal_metrics(
1322                "greptime",
1323                "public",
1324                QueryRequest {
1325                    query: Some(Query::Sql("select 1".to_string())),
1326                },
1327                &[("flow.return_region_seq", "not-a-bool")],
1328                &HashMap::new(),
1329                &mut peer_desc,
1330            )
1331            .await
1332            .unwrap_err();
1333
1334        assert!(format!("{err:?}").contains("Invalid value for flow.return_region_seq"));
1335    }
1336
1337    #[tokio::test]
1338    async fn test_try_select_one_waits_for_stream_eof() {
1339        let (schema_sent, schema_sent_rx) = tokio::sync::oneshot::channel();
1340        let release = Arc::new(tokio::sync::Notify::new());
1341        let (addr, server) = start_flight_server(DelayedEofFlight {
1342            schema_sent: Mutex::new(Some(schema_sent)),
1343            release: release.clone(),
1344        })
1345        .await;
1346        let database = Database::new(
1347            DEFAULT_CATALOG_NAME,
1348            DEFAULT_SCHEMA_NAME,
1349            Client::with_query_and_control_managers(
1350                ChannelManager::new(),
1351                ChannelManager::new(),
1352                [addr.as_str()],
1353            ),
1354        );
1355        let db = DatabaseWithPeer::new(
1356            database,
1357            Peer {
1358                id: 1,
1359                addr: addr.clone(),
1360            },
1361        );
1362        let mut probe = tokio::spawn(async move { db.try_select_one().await });
1363
1364        timeout(Duration::from_secs(1), schema_sent_rx)
1365            .await
1366            .expect("server should send the schema")
1367            .expect("schema signal should be sent");
1368        assert!(
1369            timeout(Duration::from_millis(100), &mut probe)
1370                .await
1371                .is_err(),
1372            "SELECT 1 must wait for the delayed stream tail and EOF"
1373        );
1374
1375        release.notify_one();
1376        timeout(Duration::from_secs(1), &mut probe)
1377            .await
1378            .expect("SELECT 1 should complete after EOF")
1379            .expect("probe task should not panic")
1380            .expect("SELECT 1 should succeed after EOF");
1381        server.abort();
1382    }
1383
1384    #[tokio::test]
1385    async fn test_try_select_one_propagates_late_stream_error() {
1386        let (addr, server) = start_flight_server(LateStreamErrorFlight).await;
1387        let database = Database::new(
1388            DEFAULT_CATALOG_NAME,
1389            DEFAULT_SCHEMA_NAME,
1390            Client::with_query_and_control_managers(
1391                ChannelManager::new(),
1392                ChannelManager::new(),
1393                [addr.as_str()],
1394            ),
1395        );
1396        let db = DatabaseWithPeer::new(
1397            database,
1398            Peer {
1399                id: 1,
1400                addr: addr.clone(),
1401            },
1402        );
1403
1404        let err = db.try_select_one().await.unwrap_err();
1405        server.abort();
1406
1407        assert!(format!("{err:?}").contains("late stream error"));
1408    }
1409
1410    #[tokio::test]
1411    async fn test_check_all_frontends_without_auth_fails_fast_on_unauthenticated_frontend() {
1412        let (addr, server) = start_flight_server(RejectUnauthenticatedFlight).await;
1413        let client = FrontendClient::from_meta_client(
1414            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1415            QueryOptions::default(),
1416            BatchingModeOptions::default(),
1417        )
1418        .unwrap();
1419
1420        let err = client
1421            .check_all_frontends_without_auth(&[Peer {
1422                id: 1,
1423                addr: addr.clone(),
1424            }])
1425            .await
1426            .unwrap_err();
1427        server.abort();
1428
1429        let Error::InvalidRequest {
1430            context, source, ..
1431        } = err
1432        else {
1433            panic!("expected InvalidRequest, got {err:?}");
1434        };
1435        assert!(context.contains(&addr));
1436        assert!(context.contains("rejected unauthenticated flownode probe"));
1437        assert_eq!(source.tonic_code(), Some(tonic::Code::Unauthenticated));
1438    }
1439
1440    #[tokio::test]
1441    async fn test_check_all_frontends_without_auth_uses_grpc_connection_timeout() {
1442        let (addr, server) = start_flight_server(SlowFlight).await;
1443        let client = FrontendClient::from_meta_client(
1444            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1445            QueryOptions::default(),
1446            BatchingModeOptions {
1447                grpc_conn_timeout: Duration::from_millis(50),
1448                ..Default::default()
1449            },
1450        )
1451        .unwrap();
1452
1453        let failures = client
1454            .check_all_frontends_without_auth(&[Peer {
1455                id: 1,
1456                addr: addr.clone(),
1457            }])
1458            .await
1459            .unwrap();
1460        server.abort();
1461
1462        assert_eq!(failures.len(), 1);
1463        assert!(failures[0].contains(&addr));
1464        assert!(failures[0].contains("health check timed out"));
1465    }
1466
1467    #[tokio::test]
1468    async fn test_check_all_frontends_without_auth_checks_frontends_concurrently() {
1469        let barrier = Arc::new(tokio::sync::Barrier::new(2));
1470        let (addr1, server1) = start_flight_server(WaitForConcurrentFlight {
1471            barrier: barrier.clone(),
1472        })
1473        .await;
1474        let (addr2, server2) = start_flight_server(WaitForConcurrentFlight { barrier }).await;
1475        let client = FrontendClient::from_meta_client(
1476            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1477            QueryOptions::default(),
1478            BatchingModeOptions {
1479                grpc_conn_timeout: Duration::from_millis(500),
1480                ..Default::default()
1481            },
1482        )
1483        .unwrap();
1484
1485        let failures = timeout(
1486            Duration::from_secs(2),
1487            client.check_all_frontends_without_auth(&[
1488                Peer {
1489                    id: 1,
1490                    addr: addr1.clone(),
1491                },
1492                Peer {
1493                    id: 2,
1494                    addr: addr2.clone(),
1495                },
1496            ]),
1497        )
1498        .await
1499        .expect("concurrent probes should complete before per-peer timeouts")
1500        .unwrap();
1501        server1.abort();
1502        server2.abort();
1503
1504        assert_eq!(failures.len(), 2);
1505        assert!(failures.iter().any(|failure| failure.contains(&addr1)));
1506        assert!(failures.iter().any(|failure| failure.contains(&addr2)));
1507        assert!(
1508            failures
1509                .iter()
1510                .all(|failure| !failure.contains("health check timed out")),
1511            "sequential probes would time out before both requests reach the barrier: {failures:?}"
1512        );
1513    }
1514}
flow/batching_mode/frontend_client.rs

flow/batching_mode/
frontend_client.rs
