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::{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                    .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs.clone())))
525                    .build();
526                let ctx = Arc::new(ctx);
527                let database_client = {
528                    database_client
529                        .handler
530                        .lock()
531                        .map_err(|e| {
532                            UnexpectedSnafu {
533                                reason: format!("Failed to lock database client: {e}"),
534                            }
535                            .build()
536                        })?
537                        .as_ref()
538                        .context(UnexpectedSnafu {
539                            reason: "Standalone's frontend instance is not set",
540                        })?
541                        .upgrade()
542                        .context(UnexpectedSnafu {
543                            reason: "Failed to upgrade database client",
544                        })?
545                };
546                database_client
547                    .do_query(Request::Query(request), ctx.clone())
548                    .await
549                    .map(|output| {
550                        wrap_standalone_output_with_terminal_metrics(output, &flow_extensions)
551                    })
552                    .map_err(BoxedError::new)
553                    .context(ExternalSnafu)
554            }
555        }
556    }
557
558    /// Handle a request to frontend
559    pub(crate) async fn handle(
560        &self,
561        req: api::v1::greptime_request::Request,
562        catalog: &str,
563        schema: &str,
564        peer_desc: &mut Option<PeerDesc>,
565    ) -> Result<u32, Error> {
566        match self {
567            FrontendClient::Distributed {
568                query, batch_opts, ..
569            } => {
570                let db = self.get_random_active_frontend(catalog, schema).await?;
571
572                *peer_desc = Some(PeerDesc::Dist {
573                    peer: db.peer.clone(),
574                });
575
576                db.database
577                    .handle_with_retry(
578                        req.clone(),
579                        batch_opts.experimental_grpc_max_retries,
580                        &[
581                            (QUERY_PARALLELISM_HINT, &query.parallelism.to_string()),
582                            (READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref()),
583                        ],
584                    )
585                    .await
586                    .with_context(|_| InvalidRequestSnafu {
587                        context: format!("Failed to handle request at {:?}: {:?}", db.peer, req),
588                    })
589            }
590            FrontendClient::Standalone {
591                database_client,
592                query,
593            } => {
594                let ctx = QueryContextBuilder::default()
595                    .current_catalog(catalog.to_string())
596                    .current_schema(schema.to_string())
597                    .extensions(HashMap::from([(
598                        QUERY_PARALLELISM_HINT.to_string(),
599                        query.parallelism.to_string(),
600                    )]))
601                    .build();
602                let ctx = Arc::new(ctx);
603                {
604                    let database_client = {
605                        database_client
606                            .handler
607                            .lock()
608                            .unwrap()
609                            .as_ref()
610                            .context(UnexpectedSnafu {
611                                reason: "Standalone's frontend instance is not set",
612                            })?
613                            .upgrade()
614                            .context(UnexpectedSnafu {
615                                reason: "Failed to upgrade database client",
616                            })?
617                    };
618                    let resp: common_query::Output = database_client
619                        .do_query(req, ctx)
620                        .await
621                        .map_err(BoxedError::new)
622                        .context(ExternalSnafu)?;
623                    match resp.data {
624                        OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| {
625                            UnexpectedSnafu {
626                                reason: format!("Failed to convert rows to u32: {}", rows),
627                            }
628                            .build()
629                        })?),
630                        _ => UnexpectedSnafu {
631                            reason: "Unexpected output data",
632                        }
633                        .fail(),
634                    }
635                }
636            }
637        }
638    }
639}
640
641fn build_flow_extensions(extensions: &[(&str, &str)]) -> Result<FlowQueryExtensions, Error> {
642    let flow_extensions = HashMap::from_iter(
643        extensions
644            .iter()
645            .map(|(key, value)| ((*key).to_string(), (*value).to_string())),
646    );
647    FlowQueryExtensions::parse_flow_extensions(&flow_extensions)
648        .map_err(BoxedError::new)
649        .context(ExternalSnafu)
650        .map(|extensions| extensions.unwrap_or_default())
651}
652
653fn wrap_standalone_output_with_terminal_metrics(
654    output: Output,
655    flow_extensions: &FlowQueryExtensions,
656) -> OutputWithMetrics {
657    let should_collect_region_watermark = flow_extensions.should_collect_region_watermark();
658    let terminal_metrics =
659        if should_collect_region_watermark && !matches!(&output.data, OutputData::Stream(_)) {
660            output
661                .meta
662                .plan
663                .clone()
664                .and_then(terminal_recordbatch_metrics_from_plan)
665        } else {
666            None
667        };
668    let result = OutputWithMetrics::from_output(output);
669    if let Some(metrics) = terminal_metrics {
670        result.metrics.update(Some(metrics));
671    }
672    result
673}
674
675/// Describe a peer of frontend
676#[derive(Debug, Default, Clone)]
677pub(crate) enum PeerDesc {
678    /// The query failed before a frontend peer was selected.
679    #[default]
680    Unknown,
681    /// Distributed mode's frontend peer address
682    Dist {
683        /// frontend peer address
684        peer: Peer,
685    },
686    /// Standalone mode
687    Standalone,
688}
689
690impl std::fmt::Display for PeerDesc {
691    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
692        match self {
693            PeerDesc::Unknown => write!(f, "unknown"),
694            PeerDesc::Dist { peer } => write!(f, "{}", peer.addr),
695            PeerDesc::Standalone => write!(f, "standalone"),
696        }
697    }
698}
699
700#[cfg(test)]
701mod tests {
702    use std::pin::Pin;
703    use std::task::{Context, Poll};
704    use std::time::Duration;
705
706    use api::v1::query_request::Query;
707    use arrow_flight::flight_service_server::FlightServiceServer;
708    use arrow_flight::{FlightData, Ticket};
709    use common_grpc::flight::FlightEncoder;
710    use common_query::{Output, OutputData};
711    use common_recordbatch::adapter::RecordBatchMetrics;
712    use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream};
713    use datatypes::arrow::datatypes::Schema as ArrowSchema;
714    use datatypes::prelude::{ConcreteDataType, VectorRef};
715    use datatypes::schema::{ColumnSchema, Schema};
716    use datatypes::vectors::Int32Vector;
717    use futures::StreamExt;
718    use servers::grpc::flight::{FlightCraft, FlightCraftWrapper, TonicStream};
719    use tokio::net::TcpListener;
720    use tokio::task::JoinHandle;
721    use tokio::time::timeout;
722    use tokio_stream::wrappers::TcpListenerStream;
723    use tonic::{Request as TonicRequest, Response as TonicResponse, Status};
724
725    use super::*;
726
727    #[derive(Debug)]
728    struct NoopHandler;
729
730    struct MockMetricsStream {
731        schema: datatypes::schema::SchemaRef,
732        batch: Option<RecordBatch>,
733        metrics: RecordBatchMetrics,
734        terminal_metrics_only: bool,
735    }
736
737    impl futures::Stream for MockMetricsStream {
738        type Item = common_recordbatch::error::Result<RecordBatch>;
739
740        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
741            Poll::Ready(self.batch.take().map(Ok))
742        }
743
744        fn size_hint(&self) -> (usize, Option<usize>) {
745            (
746                usize::from(self.batch.is_some()),
747                Some(usize::from(self.batch.is_some())),
748            )
749        }
750    }
751
752    impl RecordBatchStream for MockMetricsStream {
753        fn name(&self) -> &str {
754            "MockMetricsStream"
755        }
756
757        fn schema(&self) -> datatypes::schema::SchemaRef {
758            self.schema.clone()
759        }
760
761        fn output_ordering(&self) -> Option<&[OrderOption]> {
762            None
763        }
764
765        fn metrics(&self) -> Option<RecordBatchMetrics> {
766            if self.terminal_metrics_only && self.batch.is_some() {
767                return None;
768            }
769            Some(self.metrics.clone())
770        }
771    }
772
773    #[derive(Debug)]
774    struct MetricsHandler;
775
776    #[derive(Debug)]
777    struct ExtensionAwareHandler;
778
779    #[derive(Debug)]
780    struct SnapshotBindingHandler;
781
782    #[derive(Debug)]
783    struct RejectUnauthenticatedFlight;
784
785    #[derive(Debug)]
786    struct SlowFlight;
787
788    struct DelayedEofFlight {
789        schema_sent: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
790        release: Arc<tokio::sync::Notify>,
791    }
792
793    #[derive(Debug)]
794    struct LateStreamErrorFlight;
795
796    struct WaitForConcurrentFlight {
797        barrier: Arc<tokio::sync::Barrier>,
798    }
799
800    #[async_trait::async_trait]
801    impl GrpcQueryHandlerWithBoxedError for NoopHandler {
802        async fn do_query(
803            &self,
804            _query: Request,
805            _ctx: QueryContextRef,
806        ) -> std::result::Result<Output, BoxedError> {
807            Ok(Output::new_with_affected_rows(0))
808        }
809    }
810
811    #[async_trait::async_trait]
812    impl GrpcQueryHandlerWithBoxedError for MetricsHandler {
813        async fn do_query(
814            &self,
815            _query: Request,
816            _ctx: QueryContextRef,
817        ) -> std::result::Result<Output, BoxedError> {
818            let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
819                "v",
820                ConcreteDataType::int32_datatype(),
821                false,
822            )]));
823            let batch = RecordBatch::new(
824                schema.clone(),
825                vec![Arc::new(Int32Vector::from_slice([1, 2])) as VectorRef],
826            )
827            .unwrap();
828            Ok(Output::new_with_stream(Box::pin(MockMetricsStream {
829                schema,
830                batch: Some(batch),
831                metrics: RecordBatchMetrics {
832                    region_watermarks: vec![common_recordbatch::adapter::RegionWatermarkEntry {
833                        region_id: 42,
834                        watermark: Some(99),
835                    }],
836                    ..Default::default()
837                },
838                terminal_metrics_only: true,
839            })))
840        }
841    }
842
843    #[async_trait::async_trait]
844    impl GrpcQueryHandlerWithBoxedError for ExtensionAwareHandler {
845        async fn do_query(
846            &self,
847            _query: Request,
848            ctx: QueryContextRef,
849        ) -> std::result::Result<Output, BoxedError> {
850            assert_eq!(ctx.extension("flow.return_region_seq"), Some("true"));
851            Ok(Output::new_with_affected_rows(1))
852        }
853    }
854
855    #[async_trait::async_trait]
856    impl GrpcQueryHandlerWithBoxedError for SnapshotBindingHandler {
857        async fn do_query(
858            &self,
859            _query: Request,
860            ctx: QueryContextRef,
861        ) -> std::result::Result<Output, BoxedError> {
862            assert_eq!(ctx.extension("flow.return_region_seq"), Some("true"));
863            assert_eq!(ctx.get_snapshot(1), Some(10));
864            assert_eq!(ctx.get_snapshot(2), Some(20));
865            ctx.set_snapshot(42, 99);
866            Ok(Output::new_with_affected_rows(1))
867        }
868    }
869
870    #[async_trait::async_trait]
871    impl FlightCraft for RejectUnauthenticatedFlight {
872        async fn do_get(
873            &self,
874            _request: TonicRequest<Ticket>,
875        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
876            Err(Status::unauthenticated("auth failed"))
877        }
878    }
879
880    #[async_trait::async_trait]
881    impl FlightCraft for SlowFlight {
882        async fn do_get(
883            &self,
884            _request: TonicRequest<Ticket>,
885        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
886            tokio::time::sleep(Duration::from_secs(60)).await;
887            Err(Status::unavailable("slow response"))
888        }
889    }
890
891    #[async_trait::async_trait]
892    impl FlightCraft for DelayedEofFlight {
893        async fn do_get(
894            &self,
895            _request: TonicRequest<Ticket>,
896        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
897            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
898            let schema_sent = self.schema_sent.lock().unwrap().take();
899            let schema_stream = futures::stream::once(async move {
900                if let Some(schema_sent) = schema_sent {
901                    let _ = schema_sent.send(());
902                }
903                Ok(schema)
904            });
905            let release = self.release.clone();
906            let delayed_eof = futures::stream::unfold(release, |release| async move {
907                release.notified().await;
908                None::<(std::result::Result<FlightData, Status>, _)>
909            });
910
911            Ok(TonicResponse::new(Box::pin(
912                schema_stream.chain(delayed_eof),
913            )))
914        }
915    }
916
917    #[async_trait::async_trait]
918    impl FlightCraft for LateStreamErrorFlight {
919        async fn do_get(
920            &self,
921            _request: TonicRequest<Ticket>,
922        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
923            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
924            let stream =
925                futures::stream::iter([Ok(schema), Err(Status::unavailable("late stream error"))]);
926            Ok(TonicResponse::new(Box::pin(stream)))
927        }
928    }
929
930    #[async_trait::async_trait]
931    impl FlightCraft for WaitForConcurrentFlight {
932        async fn do_get(
933            &self,
934            _request: TonicRequest<Ticket>,
935        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
936            self.barrier.wait().await;
937            Err(Status::unavailable("probe started concurrently"))
938        }
939    }
940
941    async fn start_flight_server<T: FlightCraft>(handler: T) -> (String, JoinHandle<()>) {
942        let listener = TcpListener::bind("127.0.0.1:0")
943            .await
944            .expect("bind test flight server");
945        let addr = listener.local_addr().expect("local addr").to_string();
946        let server = tokio::spawn(async move {
947            tonic::transport::Server::builder()
948                .add_service(FlightServiceServer::new(FlightCraftWrapper(handler)))
949                .serve_with_incoming(TcpListenerStream::new(listener))
950                .await
951                .expect("serve test flight server");
952        });
953
954        (addr, server)
955    }
956
957    #[tokio::test]
958    async fn wait_initialized() {
959        let (client, handler_mut) =
960            FrontendClient::from_empty_grpc_handler(QueryOptions::default());
961
962        assert!(
963            timeout(Duration::from_millis(50), client.wait_initialized())
964                .await
965                .is_err()
966        );
967
968        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
969        handler_mut.set_handler(Arc::downgrade(&handler)).await;
970
971        timeout(Duration::from_secs(1), client.wait_initialized())
972            .await
973            .expect("wait_initialized should complete after handler is set");
974
975        timeout(Duration::from_millis(10), client.wait_initialized())
976            .await
977            .expect("wait_initialized should be a no-op once initialized");
978
979        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
980        let client =
981            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
982        assert!(
983            timeout(Duration::from_millis(10), client.wait_initialized())
984                .await
985                .is_ok()
986        );
987
988        let meta_client = Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend));
989        let client = FrontendClient::from_meta_client(
990            meta_client,
991            QueryOptions::default(),
992            BatchingModeOptions::default(),
993        )
994        .unwrap();
995        assert!(
996            timeout(Duration::from_millis(10), client.wait_initialized())
997                .await
998                .is_ok()
999        );
1000    }
1001
1002    #[tokio::test]
1003    async fn test_query_with_terminal_metrics_tracks_watermark_in_standalone_mode() {
1004        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(MetricsHandler);
1005        let client =
1006            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1007        let mut peer_desc = None;
1008
1009        let result = client
1010            .query_with_terminal_metrics(
1011                "greptime",
1012                "public",
1013                QueryRequest {
1014                    query: Some(Query::Sql("select 1".to_string())),
1015                },
1016                &[],
1017                &HashMap::new(),
1018                &mut peer_desc,
1019            )
1020            .await
1021            .unwrap();
1022        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1023
1024        let terminal_metrics = result.metrics.clone();
1025        assert!(!result.metrics.is_ready());
1026        assert!(terminal_metrics.get().is_none());
1027
1028        let OutputData::Stream(mut stream) = result.output.data else {
1029            panic!("expected stream output");
1030        };
1031        while stream.next().await.is_some() {}
1032
1033        assert!(terminal_metrics.is_ready());
1034        assert_eq!(
1035            terminal_metrics.region_watermark_map(),
1036            Some(HashMap::from([(42_u64, 99_u64)]))
1037        );
1038    }
1039
1040    #[tokio::test]
1041    async fn test_query_with_terminal_metrics_forwards_flow_extensions_in_standalone_mode() {
1042        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(ExtensionAwareHandler);
1043        let client =
1044            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1045        let mut peer_desc = None;
1046
1047        let result = client
1048            .query_with_terminal_metrics(
1049                "greptime",
1050                "public",
1051                QueryRequest {
1052                    query: Some(Query::Sql("insert into t select 1".to_string())),
1053                },
1054                &[("flow.return_region_seq", "true")],
1055                &HashMap::new(),
1056                &mut peer_desc,
1057            )
1058            .await
1059            .unwrap();
1060        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1061
1062        assert!(result.metrics.is_ready());
1063        assert!(result.region_watermark_map().is_none());
1064    }
1065
1066    #[tokio::test]
1067    async fn test_query_with_terminal_metrics_uses_standalone_snapshot_bounds() {
1068        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(SnapshotBindingHandler);
1069        let client =
1070            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1071        let mut peer_desc = None;
1072
1073        let result = client
1074            .query_with_terminal_metrics(
1075                "greptime",
1076                "public",
1077                QueryRequest {
1078                    query: Some(Query::Sql("insert into t select * from src".to_string())),
1079                },
1080                &[("flow.return_region_seq", "true")],
1081                &HashMap::from([(1, 10), (2, 20)]),
1082                &mut peer_desc,
1083            )
1084            .await
1085            .unwrap();
1086        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1087
1088        assert!(result.metrics.is_ready());
1089        assert_eq!(result.region_watermark_map(), None);
1090    }
1091
1092    #[tokio::test]
1093    async fn test_query_with_terminal_metrics_rejects_invalid_flow_extensions() {
1094        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1095        let client =
1096            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1097        let mut peer_desc = None;
1098
1099        let err = client
1100            .query_with_terminal_metrics(
1101                "greptime",
1102                "public",
1103                QueryRequest {
1104                    query: Some(Query::Sql("select 1".to_string())),
1105                },
1106                &[("flow.return_region_seq", "not-a-bool")],
1107                &HashMap::new(),
1108                &mut peer_desc,
1109            )
1110            .await
1111            .unwrap_err();
1112
1113        assert!(format!("{err:?}").contains("Invalid value for flow.return_region_seq"));
1114    }
1115
1116    #[tokio::test]
1117    async fn test_try_select_one_waits_for_stream_eof() {
1118        let (schema_sent, schema_sent_rx) = tokio::sync::oneshot::channel();
1119        let release = Arc::new(tokio::sync::Notify::new());
1120        let (addr, server) = start_flight_server(DelayedEofFlight {
1121            schema_sent: Mutex::new(Some(schema_sent)),
1122            release: release.clone(),
1123        })
1124        .await;
1125        let database = Database::new(
1126            DEFAULT_CATALOG_NAME,
1127            DEFAULT_SCHEMA_NAME,
1128            Client::with_urls([addr.as_str()]),
1129        );
1130        let db = DatabaseWithPeer::new(
1131            database,
1132            Peer {
1133                id: 1,
1134                addr: addr.clone(),
1135            },
1136        );
1137        let mut probe = tokio::spawn(async move { db.try_select_one().await });
1138
1139        timeout(Duration::from_secs(1), schema_sent_rx)
1140            .await
1141            .expect("server should send the schema")
1142            .expect("schema signal should be sent");
1143        assert!(
1144            timeout(Duration::from_millis(100), &mut probe)
1145                .await
1146                .is_err(),
1147            "SELECT 1 must wait for the delayed stream tail and EOF"
1148        );
1149
1150        release.notify_one();
1151        timeout(Duration::from_secs(1), &mut probe)
1152            .await
1153            .expect("SELECT 1 should complete after EOF")
1154            .expect("probe task should not panic")
1155            .expect("SELECT 1 should succeed after EOF");
1156        server.abort();
1157    }
1158
1159    #[tokio::test]
1160    async fn test_try_select_one_propagates_late_stream_error() {
1161        let (addr, server) = start_flight_server(LateStreamErrorFlight).await;
1162        let database = Database::new(
1163            DEFAULT_CATALOG_NAME,
1164            DEFAULT_SCHEMA_NAME,
1165            Client::with_urls([addr.as_str()]),
1166        );
1167        let db = DatabaseWithPeer::new(
1168            database,
1169            Peer {
1170                id: 1,
1171                addr: addr.clone(),
1172            },
1173        );
1174
1175        let err = db.try_select_one().await.unwrap_err();
1176        server.abort();
1177
1178        assert!(format!("{err:?}").contains("late stream error"));
1179    }
1180
1181    #[tokio::test]
1182    async fn test_check_all_frontends_without_auth_fails_fast_on_unauthenticated_frontend() {
1183        let (addr, server) = start_flight_server(RejectUnauthenticatedFlight).await;
1184        let client = FrontendClient::from_meta_client(
1185            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1186            QueryOptions::default(),
1187            BatchingModeOptions::default(),
1188        )
1189        .unwrap();
1190
1191        let err = client
1192            .check_all_frontends_without_auth(&[Peer {
1193                id: 1,
1194                addr: addr.clone(),
1195            }])
1196            .await
1197            .unwrap_err();
1198        server.abort();
1199
1200        let Error::InvalidRequest {
1201            context, source, ..
1202        } = err
1203        else {
1204            panic!("expected InvalidRequest, got {err:?}");
1205        };
1206        assert!(context.contains(&addr));
1207        assert!(context.contains("rejected unauthenticated flownode probe"));
1208        assert_eq!(source.tonic_code(), Some(tonic::Code::Unauthenticated));
1209    }
1210
1211    #[tokio::test]
1212    async fn test_check_all_frontends_without_auth_uses_grpc_connection_timeout() {
1213        let (addr, server) = start_flight_server(SlowFlight).await;
1214        let client = FrontendClient::from_meta_client(
1215            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1216            QueryOptions::default(),
1217            BatchingModeOptions {
1218                grpc_conn_timeout: Duration::from_millis(50),
1219                ..Default::default()
1220            },
1221        )
1222        .unwrap();
1223
1224        let failures = client
1225            .check_all_frontends_without_auth(&[Peer {
1226                id: 1,
1227                addr: addr.clone(),
1228            }])
1229            .await
1230            .unwrap();
1231        server.abort();
1232
1233        assert_eq!(failures.len(), 1);
1234        assert!(failures[0].contains(&addr));
1235        assert!(failures[0].contains("health check timed out"));
1236    }
1237
1238    #[tokio::test]
1239    async fn test_check_all_frontends_without_auth_checks_frontends_concurrently() {
1240        let barrier = Arc::new(tokio::sync::Barrier::new(2));
1241        let (addr1, server1) = start_flight_server(WaitForConcurrentFlight {
1242            barrier: barrier.clone(),
1243        })
1244        .await;
1245        let (addr2, server2) = start_flight_server(WaitForConcurrentFlight { barrier }).await;
1246        let client = FrontendClient::from_meta_client(
1247            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1248            QueryOptions::default(),
1249            BatchingModeOptions {
1250                grpc_conn_timeout: Duration::from_millis(500),
1251                ..Default::default()
1252            },
1253        )
1254        .unwrap();
1255
1256        let failures = timeout(
1257            Duration::from_secs(2),
1258            client.check_all_frontends_without_auth(&[
1259                Peer {
1260                    id: 1,
1261                    addr: addr1.clone(),
1262                },
1263                Peer {
1264                    id: 2,
1265                    addr: addr2.clone(),
1266                },
1267            ]),
1268        )
1269        .await
1270        .expect("concurrent probes should complete before per-peer timeouts")
1271        .unwrap();
1272        server1.abort();
1273        server2.abort();
1274
1275        assert_eq!(failures.len(), 2);
1276        assert!(failures.iter().any(|failure| failure.contains(&addr1)));
1277        assert!(failures.iter().any(|failure| failure.contains(&addr2)));
1278        assert!(
1279            failures
1280                .iter()
1281                .all(|failure| !failure.contains("health check timed out")),
1282            "sequential probes would time out before both requests reach the barrier: {failures:?}"
1283        );
1284    }
1285}
flow/batching_mode/frontend_client.rs

flow/batching_mode/
frontend_client.rs
