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

flow/batching_mode/
frontend_client.rs
