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

flow/batching_mode/
frontend_client.rs
