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

flow/batching_mode/
frontend_client.rs
