Skip to main content

common_meta/ddl/
table_meta.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 common_telemetry::{debug, info};
19use snafu::ensure;
20use store_api::storage::{RegionNumber, TableId};
21
22use crate::ddl::TableMetadata;
23use crate::ddl::allocator::region_routes::RegionRoutesAllocatorRef;
24use crate::ddl::allocator::resource_id::ResourceIdAllocatorRef;
25use crate::ddl::allocator::wal_options::WalOptionsAllocatorRef;
26use crate::error::{Result, UnsupportedSnafu};
27use crate::key::table_route::PhysicalTableRouteValue;
28use crate::peer::{NoopPeerAllocator, PeerAllocContext, PeerAllocatorRef};
29use crate::rpc::ddl::CreateTableTask;
30
31pub type TableMetadataAllocatorRef = Arc<TableMetadataAllocator>;
32
33#[derive(Clone)]
34pub struct TableMetadataAllocator {
35    table_id_allocator: ResourceIdAllocatorRef,
36    wal_options_allocator: WalOptionsAllocatorRef,
37    region_routes_allocator: RegionRoutesAllocatorRef,
38}
39
40impl TableMetadataAllocator {
41    pub fn new(
42        table_id_allocator: ResourceIdAllocatorRef,
43        wal_options_allocator: WalOptionsAllocatorRef,
44    ) -> Self {
45        Self::with_peer_allocator(
46            table_id_allocator,
47            wal_options_allocator,
48            Arc::new(NoopPeerAllocator),
49        )
50    }
51
52    pub fn with_peer_allocator(
53        table_id_allocator: ResourceIdAllocatorRef,
54        wal_options_allocator: WalOptionsAllocatorRef,
55        peer_allocator: PeerAllocatorRef,
56    ) -> Self {
57        Self {
58            table_id_allocator,
59            wal_options_allocator,
60            region_routes_allocator: Arc::new(peer_allocator) as _,
61        }
62    }
63
64    pub(crate) async fn allocate_table_id(
65        &self,
66        table_id: &Option<api::v1::TableId>,
67    ) -> Result<TableId> {
68        let table_id = if let Some(table_id) = table_id {
69            let table_id = table_id.id;
70
71            ensure!(
72                !self
73                    .table_id_allocator
74                    .min_max()
75                    .await
76                    .contains(&(table_id as u64)),
77                UnsupportedSnafu {
78                    operation: format!(
79                        "create table by id {} that is reserved in this node",
80                        table_id
81                    )
82                }
83            );
84
85            info!(
86                "Received explicitly allocated table id {}, will use it directly.",
87                table_id
88            );
89
90            table_id
91        } else {
92            self.table_id_allocator.next().await? as TableId
93        };
94        Ok(table_id)
95    }
96
97    async fn create_wal_options(
98        &self,
99        region_numbers: &[RegionNumber],
100        skip_wal: bool,
101    ) -> Result<HashMap<RegionNumber, String>> {
102        self.wal_options_allocator
103            .allocate(region_numbers, skip_wal)
104            .await
105    }
106
107    async fn create_table_route(
108        &self,
109        table_id: TableId,
110        partition_exprs: &[&str],
111        alloc_context: &PeerAllocContext,
112    ) -> Result<PhysicalTableRouteValue> {
113        let region_number_and_partition_exprs = partition_exprs
114            .iter()
115            .enumerate()
116            .map(|(i, partition)| (i as u32, *partition))
117            .collect::<Vec<_>>();
118        let region_routes = self
119            .region_routes_allocator
120            .allocate(table_id, &region_number_and_partition_exprs, alloc_context)
121            .await?;
122
123        Ok(PhysicalTableRouteValue::new(region_routes))
124    }
125
126    /// Create VIEW metadata
127    pub async fn create_view(&self, table_id: &Option<api::v1::TableId>) -> Result<TableMetadata> {
128        let table_id = self.allocate_table_id(table_id).await?;
129
130        Ok(TableMetadata {
131            table_id,
132            ..Default::default()
133        })
134    }
135
136    pub async fn create(&self, task: &CreateTableTask) -> Result<TableMetadata> {
137        self.create_with_context(task, &PeerAllocContext::default())
138            .await
139    }
140
141    pub async fn create_with_context(
142        &self,
143        task: &CreateTableTask,
144        alloc_context: &PeerAllocContext,
145    ) -> Result<TableMetadata> {
146        let table_id = self.allocate_table_id(&task.create_table.table_id).await?;
147        let partition_exprs = task
148            .partitions
149            .iter()
150            .map(|p| p.expression.as_str())
151            .collect::<Vec<_>>();
152        let table_route = self
153            .create_table_route(table_id, &partition_exprs, alloc_context)
154            .await?;
155        let region_numbers = table_route
156            .region_routes
157            .iter()
158            .map(|route| route.region.id.region_number())
159            .collect::<Vec<_>>();
160        let region_wal_options = self
161            .create_wal_options(&region_numbers, task.table_info.meta.options.skip_wal)
162            .await?;
163
164        debug!(
165            "Allocated region wal options {:?} for table {}",
166            region_wal_options, table_id
167        );
168
169        Ok(TableMetadata {
170            table_id,
171            table_route,
172            region_wal_options,
173        })
174    }
175
176    /// Returns the table id allocator.
177    pub fn table_id_allocator(&self) -> ResourceIdAllocatorRef {
178        self.table_id_allocator.clone()
179    }
180
181    /// Returns the wal options allocator.
182    pub fn wal_options_allocator(&self) -> WalOptionsAllocatorRef {
183        self.wal_options_allocator.clone()
184    }
185
186    /// Returns the region routes allocator.
187    pub fn region_routes_allocator(&self) -> RegionRoutesAllocatorRef {
188        self.region_routes_allocator.clone()
189    }
190}