datanode/region_server/
catalog.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 std::collections::HashMap;
16use std::sync::Arc;
17
18use datafusion::catalog::{
19    CatalogProvider, CatalogProviderList, MemTable, SchemaProvider, TableProvider,
20};
21use datafusion::datasource::provider_as_source;
22use datafusion::error as df_error;
23use datafusion::error::Result as DfResult;
24use datafusion_common::DataFusionError;
25use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
26use datafusion_expr::{LogicalPlan, TableSource};
27use futures::TryStreamExt;
28use session::context::QueryContextRef;
29use snafu::{OptionExt, ResultExt};
30use store_api::region_info::RegionInfoEntry;
31use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry};
32use store_api::storage::RegionId;
33
34use crate::error::{DataFusionSnafu, ListStorageSstsSnafu, Result, UnexpectedSnafu};
35use crate::region_server::RegionServer;
36
37/// Reserved internal table kinds used.
38/// These are recognized by reserved table names and mapped to providers.
39#[allow(clippy::enum_variant_names)]
40#[derive(Clone, Debug, PartialEq, Eq, Hash, Copy)]
41enum InternalTableKind {
42    InspectSstManifest,
43    InspectSstStorage,
44    InspectSstIndexMeta,
45    InspectRegionInfo,
46}
47
48impl InternalTableKind {
49    /// Determine if the name is a reserved internal table (case-insensitive).
50    pub fn from_table_name(name: &str) -> Option<Self> {
51        if name.eq_ignore_ascii_case(ManifestSstEntry::reserved_table_name_for_inspection()) {
52            return Some(Self::InspectSstManifest);
53        }
54        if name.eq_ignore_ascii_case(StorageSstEntry::reserved_table_name_for_inspection()) {
55            return Some(Self::InspectSstStorage);
56        }
57        if name.eq_ignore_ascii_case(PuffinIndexMetaEntry::reserved_table_name_for_inspection()) {
58            return Some(Self::InspectSstIndexMeta);
59        }
60        if name.eq_ignore_ascii_case(RegionInfoEntry::reserved_table_name_for_inspection()) {
61            return Some(Self::InspectRegionInfo);
62        }
63        None
64    }
65
66    /// Return the `TableProvider` for the internal table.
67    pub async fn table_provider(&self, server: &RegionServer) -> Result<Arc<dyn TableProvider>> {
68        match self {
69            Self::InspectSstManifest => server.inspect_sst_manifest_provider().await,
70            Self::InspectSstStorage => server.inspect_sst_storage_provider().await,
71            Self::InspectSstIndexMeta => server.inspect_sst_index_meta_provider().await,
72            Self::InspectRegionInfo => server.inspect_region_info_provider().await,
73        }
74    }
75}
76
77impl RegionServer {
78    /// Expose SSTs listed in Manifest as an in-memory table for inspection.
79    pub async fn inspect_sst_manifest_provider(&self) -> Result<Arc<dyn TableProvider>> {
80        let mito = {
81            let guard = self.inner.mito_engine.read().unwrap();
82            guard.as_ref().cloned().context(UnexpectedSnafu {
83                violated: "mito engine not available",
84            })?
85        };
86
87        let entries = mito.all_ssts_from_manifest().await;
88        let schema = ManifestSstEntry::schema().arrow_schema().clone();
89        let batch = ManifestSstEntry::to_record_batch(&entries)
90            .map_err(DataFusionError::from)
91            .context(DataFusionSnafu)?;
92
93        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
94        Ok(Arc::new(table))
95    }
96
97    /// Expose SSTs found in storage as an in-memory table for inspection.
98    pub async fn inspect_sst_storage_provider(&self) -> Result<Arc<dyn TableProvider>> {
99        let mito = {
100            let guard = self.inner.mito_engine.read().unwrap();
101            guard.as_ref().cloned().context(UnexpectedSnafu {
102                violated: "mito engine not available",
103            })?
104        };
105        let entries = mito
106            .all_ssts_from_storage()
107            .try_collect::<Vec<_>>()
108            .await
109            .context(ListStorageSstsSnafu)?;
110        let schema = StorageSstEntry::schema().arrow_schema().clone();
111        let batch = StorageSstEntry::to_record_batch(&entries)
112            .map_err(DataFusionError::from)
113            .context(DataFusionSnafu)?;
114
115        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
116        Ok(Arc::new(table))
117    }
118
119    /// Expose index metadata across the engine as an in-memory table.
120    pub async fn inspect_sst_index_meta_provider(&self) -> Result<Arc<dyn TableProvider>> {
121        let mito = {
122            let guard = self.inner.mito_engine.read().unwrap();
123            guard.as_ref().cloned().context(UnexpectedSnafu {
124                violated: "mito engine not available",
125            })?
126        };
127
128        let entries = mito.all_index_metas().await;
129        let schema = PuffinIndexMetaEntry::schema().arrow_schema().clone();
130        let batch = PuffinIndexMetaEntry::to_record_batch(&entries)
131            .map_err(DataFusionError::from)
132            .context(DataFusionSnafu)?;
133
134        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
135        Ok(Arc::new(table))
136    }
137
138    /// Expose region info across the engine as an in-memory table.
139    pub async fn inspect_region_info_provider(&self) -> Result<Arc<dyn TableProvider>> {
140        let mito = {
141            let guard = self.inner.mito_engine.read().unwrap();
142            guard.as_ref().cloned().context(UnexpectedSnafu {
143                violated: "mito engine not available",
144            })?
145        };
146
147        let entries = mito.all_region_infos().await;
148        let schema = RegionInfoEntry::schema().arrow_schema().clone();
149        let batch = RegionInfoEntry::to_record_batch(&entries)
150            .map_err(DataFusionError::from)
151            .context(DataFusionSnafu)?;
152
153        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
154        Ok(Arc::new(table))
155    }
156}
157
158/// A catalog list that resolves `TableProvider` by table name:
159/// - For reserved internal names, return inspection providers;
160/// - Otherwise, fall back to the Region provider.
161#[derive(Clone, Debug)]
162pub(crate) struct NameAwareCatalogList {
163    catalog: NameAwareCatalogProvider,
164}
165
166impl NameAwareCatalogList {
167    /// Creates the catalog list.
168    pub fn new(server: RegionServer, region_id: RegionId, query_ctx: QueryContextRef) -> Self {
169        let schema_provider = NameAwareSchemaProvider {
170            server,
171            region_id,
172            query_ctx,
173        };
174        let catalog = NameAwareCatalogProvider {
175            schema: schema_provider,
176        };
177        Self { catalog }
178    }
179}
180
181impl CatalogProviderList for NameAwareCatalogList {
182    fn register_catalog(
183        &self,
184        _name: String,
185        _catalog: Arc<dyn CatalogProvider>,
186    ) -> Option<Arc<dyn CatalogProvider>> {
187        None
188    }
189    fn catalog_names(&self) -> Vec<String> {
190        vec![]
191    }
192    fn catalog(&self, _name: &str) -> Option<Arc<dyn CatalogProvider>> {
193        Some(Arc::new(self.catalog.clone()))
194    }
195}
196
197#[derive(Clone, Debug)]
198struct NameAwareCatalogProvider {
199    schema: NameAwareSchemaProvider,
200}
201
202impl CatalogProvider for NameAwareCatalogProvider {
203    fn schema_names(&self) -> Vec<String> {
204        vec![]
205    }
206    fn schema(&self, _name: &str) -> Option<Arc<dyn SchemaProvider>> {
207        Some(Arc::new(self.schema.clone()))
208    }
209}
210
211#[derive(Clone)]
212struct NameAwareSchemaProvider {
213    server: RegionServer,
214    region_id: RegionId,
215    query_ctx: QueryContextRef,
216}
217
218impl std::fmt::Debug for NameAwareSchemaProvider {
219    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
220        write!(f, "NameAwareSchemaProvider")
221    }
222}
223
224#[async_trait::async_trait]
225impl SchemaProvider for NameAwareSchemaProvider {
226    fn table_names(&self) -> Vec<String> {
227        vec![]
228    }
229
230    async fn table(&self, name: &str) -> DfResult<Option<Arc<dyn TableProvider>>> {
231        // Resolve inspect providers by reserved names.
232        if let Some(kind) = InternalTableKind::from_table_name(name) {
233            return kind
234                .table_provider(&self.server)
235                .await
236                .map(Some)
237                .map_err(|e| df_error::DataFusionError::External(Box::new(e)));
238        }
239
240        // Fallback to region provider for any other table name.
241        let provider = self
242            .server
243            .table_provider(self.region_id, Some(self.query_ctx.clone()))
244            .await
245            .map_err(|e| df_error::DataFusionError::External(Box::new(e)))?;
246        Ok(Some(provider))
247    }
248
249    fn table_exist(&self, _name: &str) -> bool {
250        true
251    }
252}
253/// Builds a `NameAwareDataSourceInjector` from a logical plan.
254///
255/// It scans the plan to determine:
256/// - whether a Region `TableSource` is required, and
257/// - which internal inspection sources are referenced.
258pub(crate) struct NameAwareDataSourceInjectorBuilder {
259    /// Whether the plan requires a Region `TableSource`.
260    need_region_provider: bool,
261    /// Internal table kinds referenced by the plan.
262    reserved_table_needed: Vec<InternalTableKind>,
263}
264
265impl NameAwareDataSourceInjectorBuilder {
266    /// Walk the `LogicalPlan` to determine whether a Region source is needed,
267    /// and collect the kinds of internal sources required.
268    pub fn from_plan(plan: &LogicalPlan) -> DfResult<Self> {
269        let mut need_region_provider = false;
270        let mut reserved_table_needed = Vec::new();
271        plan.apply(|node| {
272            if let LogicalPlan::TableScan(ts) = node {
273                let name = ts.table_name.to_string();
274                if let Some(kind) = InternalTableKind::from_table_name(&name) {
275                    if !reserved_table_needed.contains(&kind) {
276                        reserved_table_needed.push(kind);
277                    }
278                } else {
279                    // Any normal table scan implies a Region source is needed.
280                    need_region_provider = true;
281                }
282            }
283            Ok(TreeNodeRecursion::Continue)
284        })?;
285
286        Ok(Self {
287            need_region_provider,
288            reserved_table_needed,
289        })
290    }
291
292    pub async fn build(
293        self,
294        server: &RegionServer,
295        region_id: RegionId,
296        query_ctx: QueryContextRef,
297    ) -> Result<NameAwareDataSourceInjector> {
298        let region = if self.need_region_provider {
299            let provider = server.table_provider(region_id, Some(query_ctx)).await?;
300            Some(provider_as_source(provider))
301        } else {
302            None
303        };
304
305        let mut reserved_sources = HashMap::new();
306        for kind in &self.reserved_table_needed {
307            let provider = kind.table_provider(server).await?;
308            reserved_sources.insert(*kind, provider_as_source(provider));
309        }
310
311        Ok(NameAwareDataSourceInjector {
312            reserved_sources,
313            region_source: region,
314        })
315    }
316}
317
318/// Rewrites `LogicalPlan` to inject proper data sources for `TableScan`.
319/// Uses internal sources for reserved tables; otherwise uses the Region source.
320pub(crate) struct NameAwareDataSourceInjector {
321    /// Sources for reserved internal tables, keyed by kind.
322    reserved_sources: HashMap<InternalTableKind, Arc<dyn TableSource>>,
323    /// Optional Region-level source used for normal tables.
324    region_source: Option<Arc<dyn TableSource>>,
325}
326
327impl TreeNodeRewriter for NameAwareDataSourceInjector {
328    type Node = LogicalPlan;
329
330    fn f_up(&mut self, node: Self::Node) -> DfResult<Transformed<Self::Node>> {
331        Ok(match node {
332            LogicalPlan::TableScan(mut scan) => {
333                let name = scan.table_name.to_string();
334                if let Some(kind) = InternalTableKind::from_table_name(&name)
335                    && let Some(source) = self.reserved_sources.get(&kind)
336                {
337                    // Matched a reserved internal table: rewrite to its dedicated source.
338                    scan.source = source.clone();
339                } else {
340                    let Some(region) = &self.region_source else {
341                        // Region source required but not constructed; this is unexpected.
342                        return Err(datafusion::error::DataFusionError::Plan(
343                            "region provider not available".to_string(),
344                        ));
345                    };
346                    // Normal table: rewrite to the Region source.
347                    scan.source = region.clone();
348                }
349                Transformed::yes(LogicalPlan::TableScan(scan))
350            }
351            _ => Transformed::no(node),
352        })
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use std::sync::Arc;
359
360    use datafusion::catalog::MemTable as DfMemTable;
361    use datafusion_common::tree_node::TreeNode;
362    use datafusion_expr::{LogicalPlanBuilder, table_scan};
363    use datatypes::arrow::array::Int32Array;
364    use datatypes::arrow::datatypes::{DataType, Field, Schema};
365    use datatypes::arrow::record_batch::RecordBatch;
366    use store_api::region_info::RegionInfoEntry;
367
368    use super::*; // bring rewrite() into scope
369
370    fn test_schema() -> Schema {
371        Schema::new(vec![Field::new("a", DataType::Int32, true)])
372    }
373
374    fn empty_mem_table() -> Arc<DfMemTable> {
375        let schema = Arc::new(test_schema());
376        let batch = RecordBatch::try_new(
377            schema.clone(),
378            vec![Arc::new(Int32Array::from(Vec::<i32>::new()))],
379        )
380        .unwrap();
381        Arc::new(DfMemTable::try_new(schema, vec![vec![batch]]).unwrap())
382    }
383
384    #[test]
385    fn test_injector_builder_from_plan_flags() {
386        let schema = test_schema();
387        let reserved = ManifestSstEntry::reserved_table_name_for_inspection();
388        // plan1: reserved table scan only
389        let plan1 = table_scan(Some(reserved), &schema, None)
390            .unwrap()
391            .build()
392            .unwrap();
393        let b1 = NameAwareDataSourceInjectorBuilder::from_plan(&plan1).unwrap();
394        assert!(!b1.need_region_provider);
395        assert_eq!(
396            b1.reserved_table_needed,
397            vec![InternalTableKind::InspectSstManifest]
398        );
399
400        // plan2: normal table scan only
401        let plan2 = table_scan(Some("normal_table"), &schema, None)
402            .unwrap()
403            .build()
404            .unwrap();
405        let b2 = NameAwareDataSourceInjectorBuilder::from_plan(&plan2).unwrap();
406        assert!(b2.need_region_provider);
407        assert!(b2.reserved_table_needed.is_empty());
408
409        // plan3: both reserved and normal (via UNION)
410        let p_res = table_scan(Some(reserved), &schema, None)
411            .unwrap()
412            .build()
413            .unwrap();
414        let p_norm = table_scan(Some("normal_table"), &schema, None)
415            .unwrap()
416            .build()
417            .unwrap();
418        let plan3 = LogicalPlanBuilder::from(p_res)
419            .union(LogicalPlanBuilder::from(p_norm).build().unwrap())
420            .unwrap()
421            .build()
422            .unwrap();
423        let b3 = NameAwareDataSourceInjectorBuilder::from_plan(&plan3).unwrap();
424        assert!(b3.need_region_provider);
425        assert_eq!(
426            b3.reserved_table_needed,
427            vec![InternalTableKind::InspectSstManifest]
428        );
429
430        let region_info = RegionInfoEntry::reserved_table_name_for_inspection();
431        let plan4 = table_scan(Some(region_info), &schema, None)
432            .unwrap()
433            .build()
434            .unwrap();
435        let b4 = NameAwareDataSourceInjectorBuilder::from_plan(&plan4).unwrap();
436        assert!(!b4.need_region_provider);
437        assert_eq!(
438            b4.reserved_table_needed,
439            vec![InternalTableKind::InspectRegionInfo]
440        );
441    }
442
443    #[test]
444    fn test_rewriter_replaces_with_reserved_source() {
445        let schema = test_schema();
446        let table_name = ManifestSstEntry::reserved_table_name_for_inspection();
447        let plan = table_scan(Some(table_name), &schema, None)
448            .unwrap()
449            .build()
450            .unwrap();
451
452        let provider = empty_mem_table();
453        let source = provider_as_source(provider);
454
455        let mut injector = NameAwareDataSourceInjector {
456            reserved_sources: {
457                let mut m = HashMap::new();
458                m.insert(InternalTableKind::InspectSstManifest, source.clone());
459                m
460            },
461            region_source: None,
462        };
463
464        let transformed = plan.rewrite(&mut injector).unwrap();
465        let new_plan = transformed.data;
466
467        if let LogicalPlan::TableScan(scan) = new_plan {
468            // Compare the underlying Arc ptrs to ensure replacement happened
469            let src_ptr = Arc::as_ptr(&scan.source);
470            let want_ptr = Arc::as_ptr(&source);
471            assert!(std::ptr::eq(src_ptr, want_ptr));
472        } else {
473            panic!("expected TableScan after rewrite");
474        }
475    }
476
477    #[test]
478    fn test_rewriter_replaces_with_region_info_reserved_source() {
479        let schema = test_schema();
480        let table_name = RegionInfoEntry::reserved_table_name_for_inspection();
481        let plan = table_scan(Some(table_name), &schema, None)
482            .unwrap()
483            .build()
484            .unwrap();
485
486        let provider = empty_mem_table();
487        let source = provider_as_source(provider);
488
489        let mut injector = NameAwareDataSourceInjector {
490            reserved_sources: {
491                let mut m = HashMap::new();
492                m.insert(InternalTableKind::InspectRegionInfo, source.clone());
493                m
494            },
495            region_source: None,
496        };
497
498        let transformed = plan.rewrite(&mut injector).unwrap();
499        let new_plan = transformed.data;
500
501        if let LogicalPlan::TableScan(scan) = new_plan {
502            let src_ptr = Arc::as_ptr(&scan.source);
503            let want_ptr = Arc::as_ptr(&source);
504            assert!(std::ptr::eq(src_ptr, want_ptr));
505        } else {
506            panic!("expected TableScan after rewrite");
507        }
508    }
509
510    #[test]
511    fn test_rewriter_replaces_with_region_source_for_normal() {
512        let schema = test_schema();
513        let plan = table_scan(Some("normal_table"), &schema, None)
514            .unwrap()
515            .build()
516            .unwrap();
517
518        let provider = empty_mem_table();
519        let region_source = provider_as_source(provider);
520
521        let mut injector = NameAwareDataSourceInjector {
522            reserved_sources: HashMap::new(),
523            region_source: Some(region_source.clone()),
524        };
525
526        let transformed = plan.rewrite(&mut injector).unwrap();
527        let new_plan = transformed.data;
528
529        if let LogicalPlan::TableScan(scan) = new_plan {
530            let src_ptr = Arc::as_ptr(&scan.source);
531            let want_ptr = Arc::as_ptr(&region_source);
532            assert!(std::ptr::eq(src_ptr, want_ptr));
533        } else {
534            panic!("expected TableScan after rewrite");
535        }
536    }
537}
datanode/region_server/catalog.rs

datanode/region_server/
catalog.rs
