catalog/system_schema/information_schema/
region_peers.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
15use core::pin::pin;
16use std::sync::{Arc, Weak};
17
18use arrow_schema::SchemaRef as ArrowSchemaRef;
19use common_catalog::consts::INFORMATION_SCHEMA_REGION_PEERS_TABLE_ID;
20use common_error::ext::BoxedError;
21use common_meta::rpc::router::RegionRoute;
22use common_recordbatch::adapter::RecordBatchStreamAdapter;
23use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
24use datafusion::common::HashMap;
25use datafusion::execution::TaskContext;
26use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
27use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
28use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
29use datatypes::prelude::{ConcreteDataType, ScalarVectorBuilder, VectorRef};
30use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
31use datatypes::value::Value;
32use datatypes::vectors::{Int64VectorBuilder, StringVectorBuilder, UInt64VectorBuilder};
33use futures::{StreamExt, TryStreamExt};
34use snafu::{OptionExt, ResultExt};
35use store_api::storage::{RegionId, ScanRequest, TableId};
36use table::metadata::TableType;
37
38use crate::CatalogManager;
39use crate::error::{
40    CreateRecordBatchSnafu, FindRegionRoutesSnafu, InternalSnafu, Result,
41    UpgradeWeakCatalogManagerRefSnafu,
42};
43use crate::kvbackend::KvBackendCatalogManager;
44use crate::system_schema::SystemTable;
45use crate::system_schema::information_schema::{
46    InformationSchemaTableScope, InformationTable, Predicates, REGION_PEERS,
47};
48
49pub const TABLE_CATALOG: &str = "table_catalog";
50pub const TABLE_SCHEMA: &str = "table_schema";
51pub const TABLE_NAME: &str = "table_name";
52pub const REGION_ID: &str = "region_id";
53pub const PEER_ID: &str = "peer_id";
54const PEER_ADDR: &str = "peer_addr";
55pub const IS_LEADER: &str = "is_leader";
56const STATUS: &str = "status";
57const DOWN_SECONDS: &str = "down_seconds";
58const INIT_CAPACITY: usize = 42;
59
60/// The `REGION_PEERS` table provides information about the region distribution and routes. Including fields:
61///
62/// - `table_catalog`: the table catalog name
63/// - `table_schema`: the table schema name
64/// - `table_name`: the table name
65/// - `region_id`: the region id
66/// - `peer_id`: the region storage datanode peer id
67/// - `peer_addr`: the region storage datanode gRPC peer address
68/// - `is_leader`: whether the peer is the leader
69/// - `status`: the region status, `ALIVE` or `DOWNGRADED`.
70/// - `down_seconds`: the duration of being offline, in seconds.
71#[derive(Debug)]
72pub(super) struct InformationSchemaRegionPeers {
73    schema: SchemaRef,
74    catalog_name: String,
75    catalog_manager: Weak<dyn CatalogManager>,
76}
77
78impl InformationSchemaRegionPeers {
79    pub(super) fn new(catalog_name: String, catalog_manager: Weak<dyn CatalogManager>) -> Self {
80        Self {
81            schema: Self::schema(),
82            catalog_name,
83            catalog_manager,
84        }
85    }
86
87    pub(crate) fn schema() -> SchemaRef {
88        Arc::new(Schema::new(vec![
89            ColumnSchema::new(TABLE_CATALOG, ConcreteDataType::string_datatype(), false),
90            ColumnSchema::new(TABLE_SCHEMA, ConcreteDataType::string_datatype(), false),
91            ColumnSchema::new(TABLE_NAME, ConcreteDataType::string_datatype(), false),
92            ColumnSchema::new(REGION_ID, ConcreteDataType::uint64_datatype(), false),
93            ColumnSchema::new(PEER_ID, ConcreteDataType::uint64_datatype(), true),
94            ColumnSchema::new(PEER_ADDR, ConcreteDataType::string_datatype(), true),
95            ColumnSchema::new(IS_LEADER, ConcreteDataType::string_datatype(), true),
96            ColumnSchema::new(STATUS, ConcreteDataType::string_datatype(), true),
97            ColumnSchema::new(DOWN_SECONDS, ConcreteDataType::int64_datatype(), true),
98        ]))
99    }
100
101    fn builder(&self) -> InformationSchemaRegionPeersBuilder {
102        InformationSchemaRegionPeersBuilder::new(
103            self.schema.clone(),
104            self.catalog_name.clone(),
105            self.catalog_manager.clone(),
106        )
107    }
108}
109
110impl InformationTable for InformationSchemaRegionPeers {
111    fn scope(&self) -> InformationSchemaTableScope {
112        InformationSchemaTableScope::Catalog
113    }
114}
115
116impl SystemTable for InformationSchemaRegionPeers {
117    fn table_id(&self) -> TableId {
118        INFORMATION_SCHEMA_REGION_PEERS_TABLE_ID
119    }
120
121    fn table_name(&self) -> &'static str {
122        REGION_PEERS
123    }
124
125    fn schema(&self) -> SchemaRef {
126        self.schema.clone()
127    }
128
129    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
130        let schema = self.schema.arrow_schema().clone();
131        let mut builder = self.builder();
132        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
133            schema,
134            futures::stream::once(async move {
135                builder
136                    .make_region_peers(Some(request))
137                    .await
138                    .map(|x| x.into_df_record_batch())
139                    .map_err(Into::into)
140            }),
141        ));
142        Ok(Box::pin(
143            RecordBatchStreamAdapter::try_new(stream)
144                .map_err(BoxedError::new)
145                .context(InternalSnafu)?,
146        ))
147    }
148}
149
150struct InformationSchemaRegionPeersBuilder {
151    schema: SchemaRef,
152    catalog_name: String,
153    catalog_manager: Weak<dyn CatalogManager>,
154
155    table_catalogs: StringVectorBuilder,
156    table_schemas: StringVectorBuilder,
157    table_names: StringVectorBuilder,
158    region_ids: UInt64VectorBuilder,
159    peer_ids: UInt64VectorBuilder,
160    peer_addrs: StringVectorBuilder,
161    is_leaders: StringVectorBuilder,
162    statuses: StringVectorBuilder,
163    down_seconds: Int64VectorBuilder,
164}
165
166impl InformationSchemaRegionPeersBuilder {
167    fn new(
168        schema: SchemaRef,
169        catalog_name: String,
170        catalog_manager: Weak<dyn CatalogManager>,
171    ) -> Self {
172        Self {
173            schema,
174            catalog_name,
175            catalog_manager,
176            table_catalogs: StringVectorBuilder::with_capacity(INIT_CAPACITY),
177            table_schemas: StringVectorBuilder::with_capacity(INIT_CAPACITY),
178            table_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
179            region_ids: UInt64VectorBuilder::with_capacity(INIT_CAPACITY),
180            peer_ids: UInt64VectorBuilder::with_capacity(INIT_CAPACITY),
181            peer_addrs: StringVectorBuilder::with_capacity(INIT_CAPACITY),
182            is_leaders: StringVectorBuilder::with_capacity(INIT_CAPACITY),
183            statuses: StringVectorBuilder::with_capacity(INIT_CAPACITY),
184            down_seconds: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
185        }
186    }
187
188    /// Construct the `information_schema.region_peers` virtual table
189    async fn make_region_peers(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
190        let catalog_name = self.catalog_name.clone();
191        let catalog_manager = self
192            .catalog_manager
193            .upgrade()
194            .context(UpgradeWeakCatalogManagerRefSnafu)?;
195
196        let partition_manager = catalog_manager
197            .as_any()
198            .downcast_ref::<KvBackendCatalogManager>()
199            .map(|catalog_manager| catalog_manager.partition_manager());
200
201        let predicates = Predicates::from_scan_request(&request);
202
203        for schema_name in catalog_manager.schema_names(&catalog_name, None).await? {
204            let table_stream = catalog_manager
205                .tables(&catalog_name, &schema_name, None)
206                .try_filter_map(|t| async move {
207                    let table_info = t.table_info();
208                    if table_info.table_type == TableType::Temporary {
209                        Ok(None)
210                    } else {
211                        Ok(Some((table_info.ident.table_id, table_info.name.clone())))
212                    }
213                });
214
215            const BATCH_SIZE: usize = 128;
216
217            // Split tables into chunks
218            let mut table_chunks = pin!(table_stream.ready_chunks(BATCH_SIZE));
219
220            while let Some(tables) = table_chunks.next().await {
221                let tables = tables.into_iter().collect::<Result<HashMap<_, _>>>()?;
222                let table_ids = tables.keys().cloned().collect::<Vec<_>>();
223
224                let table_routes = if let Some(partition_manager) = &partition_manager {
225                    partition_manager
226                        .batch_find_region_routes(&table_ids)
227                        .await
228                        .context(FindRegionRoutesSnafu)?
229                } else {
230                    table_ids.into_iter().map(|id| (id, vec![])).collect()
231                };
232
233                for (table_id, routes) in table_routes {
234                    // Safety: table_id is guaranteed to be in the map
235                    let table_name = tables.get(&table_id).unwrap();
236                    self.add_region_peers(
237                        &catalog_name,
238                        &schema_name,
239                        table_name,
240                        &predicates,
241                        table_id,
242                        &routes,
243                    );
244                }
245            }
246        }
247
248        self.finish()
249    }
250
251    fn add_region_peers(
252        &mut self,
253        table_catalog: &str,
254        table_schema: &str,
255        table_name: &str,
256        predicates: &Predicates,
257        table_id: TableId,
258        routes: &[RegionRoute],
259    ) {
260        for route in routes {
261            let region_id = RegionId::new(table_id, route.region.id.region_number()).as_u64();
262            let peer_id = route.leader_peer.clone().map(|p| p.id);
263            let peer_addr = route.leader_peer.clone().map(|p| p.addr);
264            let state = if let Some(state) = route.leader_state {
265                Some(state.as_ref().to_string())
266            } else {
267                // Alive by default
268                Some("ALIVE".to_string())
269            };
270
271            let row = [
272                (TABLE_CATALOG, &Value::from(table_catalog)),
273                (TABLE_SCHEMA, &Value::from(table_schema)),
274                (TABLE_NAME, &Value::from(table_name)),
275                (REGION_ID, &Value::from(region_id)),
276            ];
277
278            if !predicates.eval(&row) {
279                continue;
280            }
281
282            self.table_catalogs.push(Some(table_catalog));
283            self.table_schemas.push(Some(table_schema));
284            self.table_names.push(Some(table_name));
285            self.region_ids.push(Some(region_id));
286            self.peer_ids.push(peer_id);
287            self.peer_addrs.push(peer_addr.as_deref());
288            self.is_leaders.push(Some("Yes"));
289            self.statuses.push(state.as_deref());
290            self.down_seconds
291                .push(route.leader_down_millis().map(|m| m / 1000));
292
293            for follower in &route.follower_peers {
294                self.table_catalogs.push(Some(table_catalog));
295                self.table_schemas.push(Some(table_schema));
296                self.table_names.push(Some(table_name));
297                self.region_ids.push(Some(region_id));
298                self.peer_ids.push(Some(follower.id));
299                self.peer_addrs.push(Some(follower.addr.as_str()));
300                self.is_leaders.push(Some("No"));
301                self.statuses.push(None);
302                self.down_seconds.push(None);
303            }
304        }
305    }
306
307    fn finish(&mut self) -> Result<RecordBatch> {
308        let columns: Vec<VectorRef> = vec![
309            Arc::new(self.table_catalogs.finish()),
310            Arc::new(self.table_schemas.finish()),
311            Arc::new(self.table_names.finish()),
312            Arc::new(self.region_ids.finish()),
313            Arc::new(self.peer_ids.finish()),
314            Arc::new(self.peer_addrs.finish()),
315            Arc::new(self.is_leaders.finish()),
316            Arc::new(self.statuses.finish()),
317            Arc::new(self.down_seconds.finish()),
318        ];
319        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
320    }
321}
322
323impl DfPartitionStream for InformationSchemaRegionPeers {
324    fn schema(&self) -> &ArrowSchemaRef {
325        self.schema.arrow_schema()
326    }
327
328    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
329        let schema = self.schema.arrow_schema().clone();
330        let mut builder = self.builder();
331        Box::pin(DfRecordBatchStreamAdapter::new(
332            schema,
333            futures::stream::once(async move {
334                builder
335                    .make_region_peers(None)
336                    .await
337                    .map(|x| x.into_df_record_batch())
338                    .map_err(Into::into)
339            }),
340        ))
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use api::v1::meta::Peer;
347    use arrow::array::AsArray;
348    use common_meta::rpc::router::{Region, RegionRoute};
349    use datafusion::common::ScalarValue;
350    use datafusion::logical_expr::{BinaryExpr, Expr, Operator, col};
351    use store_api::storage::{RegionId, ScanRequest};
352
353    use super::*;
354
355    fn new_region_route(table_id: u32, region_number: u32, peer_id: u64) -> RegionRoute {
356        RegionRoute {
357            region: Region {
358                id: RegionId::new(table_id, region_number),
359                ..Default::default()
360            },
361            leader_peer: Some(Peer {
362                id: peer_id,
363                addr: format!("127.0.0.1:{}", 3000 + peer_id),
364            }),
365            follower_peers: vec![],
366            leader_state: None,
367            leader_down_since: None,
368            write_route_policy: None,
369        }
370    }
371
372    #[test]
373    fn test_add_region_peers_predicate_filters_correctly() {
374        let schema = InformationSchemaRegionPeers::schema();
375        let mut builder = InformationSchemaRegionPeersBuilder::new(
376            schema,
377            "greptime".to_string(),
378            Weak::<KvBackendCatalogManager>::new(),
379        );
380
381        let table_id = 1;
382        // 3 regions: region_number 0, 1, 2
383        let routes = vec![
384            new_region_route(table_id, 0, 1),
385            new_region_route(table_id, 1, 2),
386            new_region_route(table_id, 2, 3),
387        ];
388
389        // Build a predicate that matches only the last region (region_number=2).
390        // With the old `return` bug, encountering the first non-matching region
391        // (region_number=0) would exit add_region_peers entirely, so region_number=2
392        // would never be found.
393        let target_region_id = RegionId::new(table_id, 2).as_u64();
394        let filter = Expr::BinaryExpr(BinaryExpr::new(
395            Box::new(col(REGION_ID)),
396            Operator::Eq,
397            Box::new(Expr::Literal(
398                ScalarValue::UInt64(Some(target_region_id)),
399                None,
400            )),
401        ));
402        let request = ScanRequest {
403            filters: vec![filter],
404            ..Default::default()
405        };
406        let predicates = Predicates::from_scan_request(&Some(request));
407
408        builder.add_region_peers(
409            "greptime",
410            "public",
411            "test_table",
412            &predicates,
413            table_id,
414            &routes,
415        );
416
417        let batch = builder.finish().unwrap();
418        // Should have exactly 1 row for the matching region
419        assert_eq!(batch.num_rows(), 1);
420        // Verify it's the correct region
421        let region_id_col = batch
422            .column(3)
423            .as_primitive::<arrow::datatypes::UInt64Type>();
424        assert_eq!(region_id_col.value(0), target_region_id);
425    }
426}
catalog/system_schema/information_schema/region_peers.rs

catalog/system_schema/information_schema/
region_peers.rs
