common_meta/ddl/create_table/
template.rs1use std::collections::HashMap;
16
17use api::v1::column_def::try_as_column_def;
18use api::v1::meta::Partition;
19use api::v1::region::{CreateRequest, RegionColumnDef};
20use api::v1::{ColumnDef, CreateTableExpr, SemanticType};
21use common_telemetry::warn;
22use snafu::{OptionExt, ResultExt};
23use store_api::metric_engine_consts::LOGICAL_TABLE_METADATA_KEY;
24use store_api::storage::{RegionId, RegionNumber};
25use table::metadata::{TableId, TableInfo};
26
27use crate::error::{self, Result};
28use crate::reconciliation::utils::build_column_metadata_from_table_info;
29use crate::wal_provider::prepare_wal_options;
30
31pub fn build_template_from_raw_table_info(table_info: &TableInfo) -> Result<CreateRequest> {
39 let primary_key_indices = &table_info.meta.primary_key_indices;
40 let column_defs = table_info
41 .meta
42 .schema
43 .column_schemas()
44 .iter()
45 .enumerate()
46 .map(|(i, c)| {
47 let is_primary_key = primary_key_indices.contains(&i);
48 let column_def = try_as_column_def(c, is_primary_key)
49 .context(error::ConvertColumnDefSnafu { column: &c.name })?;
50 Ok(RegionColumnDef {
51 column_def: Some(column_def),
52 column_id: i as u32,
55 })
56 })
57 .collect::<Result<Vec<_>>>()?;
58
59 let options = HashMap::from(&table_info.meta.options);
60 let template = CreateRequest {
61 region_id: 0,
62 engine: table_info.meta.engine.clone(),
63 column_defs,
64 primary_key: table_info
65 .meta
66 .primary_key_indices
67 .iter()
68 .map(|i| *i as u32)
69 .collect(),
70 path: String::new(),
71 options,
72 partition: None,
73 };
74
75 Ok(template)
76}
77
78pub fn build_template_from_raw_table_info_for_physical_table(
85 table_info: &TableInfo,
86) -> Result<CreateRequest> {
87 let name_to_ids = table_info
88 .name_to_ids()
89 .context(error::MissingColumnIdsSnafu)?;
90 let column_metadatas = build_column_metadata_from_table_info(
91 table_info.meta.schema.column_schemas(),
92 &table_info.meta.primary_key_indices,
93 &name_to_ids,
94 )?;
95 let primary_key_ids = column_metadatas
96 .iter()
97 .filter(|c| c.semantic_type == SemanticType::Tag)
98 .map(|c| c.column_id)
99 .collect::<Vec<_>>();
100 let column_defs = column_metadatas
101 .iter()
102 .map(|c| {
103 let column_def =
104 try_as_column_def(&c.column_schema, c.semantic_type == SemanticType::Tag).context(
105 error::ConvertColumnDefSnafu {
106 column: &c.column_schema.name,
107 },
108 )?;
109 let region_column_def = RegionColumnDef {
110 column_def: Some(column_def),
111 column_id: c.column_id,
112 };
113
114 Ok(region_column_def)
115 })
116 .collect::<Result<Vec<_>>>()?;
117
118 let options = HashMap::from(&table_info.meta.options);
119 let template = CreateRequest {
120 region_id: 0,
121 engine: table_info.meta.engine.clone(),
122 column_defs,
123 primary_key: primary_key_ids,
124 path: String::new(),
125 options,
126 partition: None,
127 };
128
129 Ok(template)
130}
131
132pub(crate) fn build_template(create_table_expr: &CreateTableExpr) -> Result<CreateRequest> {
133 let column_defs = create_table_expr
134 .column_defs
135 .iter()
136 .enumerate()
137 .map(|(i, c)| {
138 let semantic_type = if create_table_expr.time_index == c.name {
139 SemanticType::Timestamp
140 } else if create_table_expr.primary_keys.contains(&c.name) {
141 SemanticType::Tag
142 } else {
143 SemanticType::Field
144 };
145
146 RegionColumnDef {
147 column_def: Some(ColumnDef {
148 name: c.name.clone(),
149 data_type: c.data_type,
150 is_nullable: c.is_nullable,
151 default_constraint: c.default_constraint.clone(),
152 semantic_type: semantic_type as i32,
153 comment: String::new(),
154 datatype_extension: c.datatype_extension.clone(),
155 options: c.options.clone(),
156 }),
157 column_id: i as u32,
158 }
159 })
160 .collect::<Vec<_>>();
161
162 let primary_key = create_table_expr
163 .primary_keys
164 .iter()
165 .map(|key| {
166 column_defs
167 .iter()
168 .find_map(|c| {
169 c.column_def.as_ref().and_then(|x| {
170 if &x.name == key {
171 Some(c.column_id)
172 } else {
173 None
174 }
175 })
176 })
177 .context(error::PrimaryKeyNotFoundSnafu { key })
178 })
179 .collect::<Result<_>>()?;
180
181 let template = CreateRequest {
182 region_id: 0,
183 engine: create_table_expr.engine.clone(),
184 column_defs,
185 primary_key,
186 path: String::new(),
187 options: create_table_expr.table_options.clone(),
188 partition: None,
189 };
190
191 Ok(template)
192}
193
194pub struct CreateRequestBuilder {
196 template: CreateRequest,
197 physical_table_id: Option<TableId>,
199}
200
201impl CreateRequestBuilder {
202 pub fn new(template: CreateRequest, physical_table_id: Option<TableId>) -> Self {
203 Self {
204 template,
205 physical_table_id,
206 }
207 }
208
209 pub fn template(&self) -> &CreateRequest {
210 &self.template
211 }
212
213 pub fn build_one(
214 &self,
215 region_id: RegionId,
216 storage_path: String,
217 region_wal_options: &HashMap<RegionNumber, String>,
218 partition_exprs: &HashMap<RegionNumber, String>,
219 ) -> CreateRequest {
220 let mut request = self.template.clone();
221
222 request.region_id = region_id.as_u64();
223 request.path = storage_path;
224 prepare_wal_options(&mut request.options, region_id, region_wal_options);
226 request.partition = Some(prepare_partition_expr(region_id, partition_exprs));
227
228 if let Some(physical_table_id) = self.physical_table_id {
229 let physical_region_id = RegionId::new(physical_table_id, region_id.region_number());
233
234 request.options.insert(
235 LOGICAL_TABLE_METADATA_KEY.to_string(),
236 physical_region_id.as_u64().to_string(),
237 );
238 }
239
240 request
241 }
242}
243
244fn prepare_partition_expr(
245 region_id: RegionId,
246 partition_exprs: &HashMap<RegionNumber, String>,
247) -> Partition {
248 let expr = partition_exprs.get(®ion_id.region_number()).cloned();
249 if expr.is_none() {
250 warn!("region {} has no partition expr", region_id);
251 }
252
253 Partition {
254 expression: expr.unwrap_or_default(),
255 ..Default::default()
256 }
257}
258
259#[cfg(test)]
260mod tests {
261 use std::collections::HashMap;
262
263 use store_api::storage::{RegionId, RegionNumber};
264
265 use super::*;
266
267 #[test]
268 fn test_build_one_sets_partition_expr_per_region() {
269 let template = CreateRequest {
271 region_id: 0,
272 engine: "mito".to_string(),
273 column_defs: vec![],
274 primary_key: vec![],
275 path: String::new(),
276 options: Default::default(),
277 partition: None,
278 };
279 let builder = CreateRequestBuilder::new(template, None);
280
281 let mut partition_exprs: HashMap<RegionNumber, String> = HashMap::new();
282 let expr_a =
283 r#"{"Expr":{"lhs":{"Column":"a"},"op":"Eq","rhs":{"Value":{"UInt32":1}}}}"#.to_string();
284 partition_exprs.insert(0, expr_a.clone());
285
286 let r0 = builder.build_one(
287 RegionId::new(42, 0),
288 "/p".to_string(),
289 &Default::default(),
290 &partition_exprs,
291 );
292 assert_eq!(r0.partition.as_ref().unwrap().expression, expr_a);
293 }
294}