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

datanode/region_server/
catalog.rs
