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
15mod inspect;
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use datafusion::catalog::{
21    CatalogProvider, CatalogProviderList, MemTable, SchemaProvider, TableProvider,
22};
23use datafusion::datasource::provider_as_source;
24use datafusion::error as df_error;
25use datafusion::error::Result as DfResult;
26use datafusion_common::DataFusionError;
27use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
28use datafusion_expr::{LogicalPlan, TableSource};
29use futures::TryStreamExt;
30use session::context::QueryContextRef;
31use snafu::{OptionExt, ResultExt};
32use store_api::region_info::RegionInfoEntry;
33use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry};
34use store_api::storage::RegionId;
35
36use crate::error::{DataFusionSnafu, ListStorageSstsSnafu, Result, UnexpectedSnafu};
37use crate::region_server::RegionServer;
38use crate::region_server::catalog::inspect::{InspectTable, ManifestSstSource};
39
40/// Reserved internal table kinds used.
41/// These are recognized by reserved table names and mapped to providers.
42#[allow(clippy::enum_variant_names)]
43#[derive(Clone, Debug, PartialEq, Eq, Hash, Copy)]
44enum InternalTableKind {
45    InspectSstManifest,
46    InspectSstStorage,
47    InspectSstIndexMeta,
48    InspectRegionInfo,
49}
50
51impl InternalTableKind {
52    /// Determine if the name is a reserved internal table (case-insensitive).
53    pub fn from_table_name(name: &str) -> Option<Self> {
54        if name.eq_ignore_ascii_case(ManifestSstEntry::reserved_table_name_for_inspection()) {
55            return Some(Self::InspectSstManifest);
56        }
57        if name.eq_ignore_ascii_case(StorageSstEntry::reserved_table_name_for_inspection()) {
58            return Some(Self::InspectSstStorage);
59        }
60        if name.eq_ignore_ascii_case(PuffinIndexMetaEntry::reserved_table_name_for_inspection()) {
61            return Some(Self::InspectSstIndexMeta);
62        }
63        if name.eq_ignore_ascii_case(RegionInfoEntry::reserved_table_name_for_inspection()) {
64            return Some(Self::InspectRegionInfo);
65        }
66        None
67    }
68
69    /// Return the `TableProvider` for the internal table.
70    pub async fn table_provider(&self, server: &RegionServer) -> Result<Arc<dyn TableProvider>> {
71        match self {
72            Self::InspectSstManifest => server.inspect_sst_manifest_provider().await,
73            Self::InspectSstStorage => server.inspect_sst_storage_provider().await,
74            Self::InspectSstIndexMeta => server.inspect_sst_index_meta_provider().await,
75            Self::InspectRegionInfo => server.inspect_region_info_provider().await,
76        }
77    }
78}
79
80impl RegionServer {
81    /// Exposes manifest SSTs through a lazy table with metadata predicate pruning.
82    pub async fn inspect_sst_manifest_provider(&self) -> Result<Arc<dyn TableProvider>> {
83        let mito = {
84            let guard = self.inner.mito_engine.read().unwrap();
85            guard.as_ref().cloned().context(UnexpectedSnafu {
86                violated: "mito engine not available",
87            })?
88        };
89
90        Ok(Arc::new(InspectTable::new(Arc::new(
91            ManifestSstSource::new(mito),
92        ))))
93    }
94
95    /// Expose SSTs found in storage as an in-memory table for inspection.
96    pub async fn inspect_sst_storage_provider(&self) -> Result<Arc<dyn TableProvider>> {
97        let mito = {
98            let guard = self.inner.mito_engine.read().unwrap();
99            guard.as_ref().cloned().context(UnexpectedSnafu {
100                violated: "mito engine not available",
101            })?
102        };
103        let entries = mito
104            .all_ssts_from_storage()
105            .try_collect::<Vec<_>>()
106            .await
107            .context(ListStorageSstsSnafu)?;
108        let schema = StorageSstEntry::schema().arrow_schema().clone();
109        let batch = StorageSstEntry::to_record_batch(&entries)
110            .map_err(DataFusionError::from)
111            .context(DataFusionSnafu)?;
112
113        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
114        Ok(Arc::new(table))
115    }
116
117    /// Expose index metadata across the engine as an in-memory table.
118    pub async fn inspect_sst_index_meta_provider(&self) -> Result<Arc<dyn TableProvider>> {
119        let mito = {
120            let guard = self.inner.mito_engine.read().unwrap();
121            guard.as_ref().cloned().context(UnexpectedSnafu {
122                violated: "mito engine not available",
123            })?
124        };
125
126        let entries = mito.all_index_metas().await;
127        let schema = PuffinIndexMetaEntry::schema().arrow_schema().clone();
128        let batch = PuffinIndexMetaEntry::to_record_batch(&entries)
129            .map_err(DataFusionError::from)
130            .context(DataFusionSnafu)?;
131
132        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
133        Ok(Arc::new(table))
134    }
135
136    /// Expose region info across the engine as an in-memory table.
137    pub async fn inspect_region_info_provider(&self) -> Result<Arc<dyn TableProvider>> {
138        let mito = {
139            let guard = self.inner.mito_engine.read().unwrap();
140            guard.as_ref().cloned().context(UnexpectedSnafu {
141                violated: "mito engine not available",
142            })?
143        };
144
145        let entries = mito.all_region_infos().await;
146        let schema = RegionInfoEntry::schema().arrow_schema().clone();
147        let batch = RegionInfoEntry::to_record_batch(&entries)
148            .map_err(DataFusionError::from)
149            .context(DataFusionSnafu)?;
150
151        let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?;
152        Ok(Arc::new(table))
153    }
154}
155
156/// A catalog list that resolves `TableProvider` by table name:
157/// - For reserved internal names, return inspection providers;
158/// - Otherwise, fall back to the Region provider.
159#[derive(Clone, Debug)]
160pub(crate) struct NameAwareCatalogList {
161    catalog: NameAwareCatalogProvider,
162}
163
164impl NameAwareCatalogList {
165    /// Creates the catalog list.
166    pub fn new(server: RegionServer, region_id: RegionId, query_ctx: QueryContextRef) -> Self {
167        let schema_provider = NameAwareSchemaProvider {
168            server,
169            region_id,
170            query_ctx,
171        };
172        let catalog = NameAwareCatalogProvider {
173            schema: schema_provider,
174        };
175        Self { catalog }
176    }
177}
178
179impl CatalogProviderList for NameAwareCatalogList {
180    fn register_catalog(
181        &self,
182        _name: String,
183        _catalog: Arc<dyn CatalogProvider>,
184    ) -> Option<Arc<dyn CatalogProvider>> {
185        None
186    }
187    fn catalog_names(&self) -> Vec<String> {
188        vec![]
189    }
190    fn catalog(&self, _name: &str) -> Option<Arc<dyn CatalogProvider>> {
191        Some(Arc::new(self.catalog.clone()))
192    }
193}
194
195#[derive(Clone, Debug)]
196struct NameAwareCatalogProvider {
197    schema: NameAwareSchemaProvider,
198}
199
200impl CatalogProvider for NameAwareCatalogProvider {
201    fn schema_names(&self) -> Vec<String> {
202        vec![]
203    }
204    fn schema(&self, _name: &str) -> Option<Arc<dyn SchemaProvider>> {
205        Some(Arc::new(self.schema.clone()))
206    }
207}
208
209#[derive(Clone)]
210struct NameAwareSchemaProvider {
211    server: RegionServer,
212    region_id: RegionId,
213    query_ctx: QueryContextRef,
214}
215
216impl std::fmt::Debug for NameAwareSchemaProvider {
217    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
218        write!(f, "NameAwareSchemaProvider")
219    }
220}
221
222#[async_trait::async_trait]
223impl SchemaProvider for NameAwareSchemaProvider {
224    fn table_names(&self) -> Vec<String> {
225        vec![]
226    }
227
228    async fn table(&self, name: &str) -> DfResult<Option<Arc<dyn TableProvider>>> {
229        // Resolve inspect providers by reserved names.
230        if let Some(kind) = InternalTableKind::from_table_name(name) {
231            return kind
232                .table_provider(&self.server)
233                .await
234                .map(Some)
235                .map_err(|e| df_error::DataFusionError::External(Box::new(e)));
236        }
237
238        // Fallback to region provider for any other table name.
239        let provider = self
240            .server
241            .table_provider(self.region_id, Some(self.query_ctx.clone()))
242            .await
243            .map_err(|e| df_error::DataFusionError::External(Box::new(e)))?;
244        Ok(Some(provider))
245    }
246
247    fn table_exist(&self, _name: &str) -> bool {
248        true
249    }
250}
251/// Builds a `NameAwareDataSourceInjector` from a logical plan.
252///
253/// It scans the plan to determine:
254/// - whether a Region `TableSource` is required, and
255/// - which internal inspection sources are referenced.
256pub(crate) struct NameAwareDataSourceInjectorBuilder {
257    /// Whether the plan requires a Region `TableSource`.
258    need_region_provider: bool,
259    /// Internal table kinds referenced by the plan.
260    reserved_table_needed: Vec<InternalTableKind>,
261}
262
263impl NameAwareDataSourceInjectorBuilder {
264    /// Walk the `LogicalPlan` to determine whether a Region source is needed,
265    /// and collect the kinds of internal sources required.
266    pub fn from_plan(plan: &LogicalPlan) -> DfResult<Self> {
267        let mut need_region_provider = false;
268        let mut reserved_table_needed = Vec::new();
269        plan.apply(|node| {
270            if let LogicalPlan::TableScan(ts) = node {
271                let name = ts.table_name.to_string();
272                if let Some(kind) = InternalTableKind::from_table_name(&name) {
273                    if !reserved_table_needed.contains(&kind) {
274                        reserved_table_needed.push(kind);
275                    }
276                } else {
277                    // Any normal table scan implies a Region source is needed.
278                    need_region_provider = true;
279                }
280            }
281            Ok(TreeNodeRecursion::Continue)
282        })?;
283
284        Ok(Self {
285            need_region_provider,
286            reserved_table_needed,
287        })
288    }
289
290    pub async fn build(
291        self,
292        server: &RegionServer,
293        region_id: RegionId,
294        query_ctx: QueryContextRef,
295    ) -> Result<NameAwareDataSourceInjector> {
296        let region = if self.need_region_provider {
297            let provider = server.table_provider(region_id, Some(query_ctx)).await?;
298            Some(provider_as_source(provider))
299        } else {
300            None
301        };
302
303        let mut reserved_sources = HashMap::new();
304        for kind in &self.reserved_table_needed {
305            let provider = kind.table_provider(server).await?;
306            reserved_sources.insert(*kind, provider_as_source(provider));
307        }
308
309        Ok(NameAwareDataSourceInjector {
310            reserved_sources,
311            region_source: region,
312        })
313    }
314}
315
316/// Rewrites `LogicalPlan` to inject proper data sources for `TableScan`.
317/// Uses internal sources for reserved tables; otherwise uses the Region source.
318pub(crate) struct NameAwareDataSourceInjector {
319    /// Sources for reserved internal tables, keyed by kind.
320    reserved_sources: HashMap<InternalTableKind, Arc<dyn TableSource>>,
321    /// Optional Region-level source used for normal tables.
322    region_source: Option<Arc<dyn TableSource>>,
323}
324
325impl TreeNodeRewriter for NameAwareDataSourceInjector {
326    type Node = LogicalPlan;
327
328    fn f_up(&mut self, node: Self::Node) -> DfResult<Transformed<Self::Node>> {
329        Ok(match node {
330            LogicalPlan::TableScan(mut scan) => {
331                let name = scan.table_name.to_string();
332                if let Some(kind) = InternalTableKind::from_table_name(&name)
333                    && let Some(source) = self.reserved_sources.get(&kind)
334                {
335                    // Matched a reserved internal table: rewrite to its dedicated source.
336                    scan.source = source.clone();
337                } else {
338                    let Some(region) = &self.region_source else {
339                        // Region source required but not constructed; this is unexpected.
340                        return Err(datafusion::error::DataFusionError::Plan(
341                            "region provider not available".to_string(),
342                        ));
343                    };
344                    // Normal table: rewrite to the Region source.
345                    scan.source = region.clone();
346                }
347                Transformed::yes(LogicalPlan::TableScan(scan))
348            }
349            _ => Transformed::no(node),
350        })
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    mod manifest;
357
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
