1use api::v1::SemanticType;
16use common_telemetry::{debug, info, warn};
17use datatypes::schema::{SkippingIndexOptions, SkippingIndexType};
18use mito2::engine::MitoEngine;
19use snafu::ResultExt;
20use store_api::metadata::ColumnMetadata;
21use store_api::region_engine::RegionEngine;
22use store_api::region_request::{
23 AddColumn, AffectedRows, AlterKind, RegionAlterRequest, RegionPutRequest, RegionRequest,
24};
25use store_api::storage::consts::ReservedColumnId;
26use store_api::storage::{ConcreteDataType, RegionId};
27
28use crate::engine::IndexOptions;
29use crate::error::{
30 ColumnTypeMismatchSnafu, ForbiddenPhysicalAlterSnafu, MitoReadOperationSnafu,
31 MitoWriteOperationSnafu, Result, SetSkippingIndexOptionSnafu,
32};
33use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_DDL_DURATION, PHYSICAL_COLUMN_COUNT};
34use crate::utils;
35
36pub struct DataRegion {
42 mito: MitoEngine,
43}
44
45impl DataRegion {
46 pub fn new(mito: MitoEngine) -> Self {
47 Self { mito }
48 }
49
50 pub async fn add_columns(
62 &self,
63 region_id: RegionId,
64 columns: Vec<ColumnMetadata>,
65 index_options: IndexOptions,
66 ) -> Result<()> {
67 if columns.is_empty() {
69 return Ok(());
70 }
71
72 let region_id = utils::to_data_region_id(region_id);
73
74 let num_columns = columns.len();
75 let request = self
76 .assemble_alter_request(region_id, columns, index_options)
77 .await?;
78
79 let _timer = MITO_DDL_DURATION.start_timer();
80
81 let _ = self
82 .mito
83 .handle_request(region_id, request)
84 .await
85 .context(MitoWriteOperationSnafu)?;
86
87 PHYSICAL_COLUMN_COUNT.add(num_columns as _);
88
89 Ok(())
90 }
91
92 async fn assemble_alter_request(
95 &self,
96 region_id: RegionId,
97 columns: Vec<ColumnMetadata>,
98 index_options: IndexOptions,
99 ) -> Result<RegionRequest> {
100 let region_metadata = self
102 .mito
103 .get_metadata(region_id)
104 .await
105 .context(MitoReadOperationSnafu)?;
106
107 let new_column_id_start = 1 + region_metadata
109 .column_metadatas
110 .iter()
111 .filter_map(|c| {
112 if ReservedColumnId::is_reserved(c.column_id) {
113 None
114 } else {
115 Some(c.column_id)
116 }
117 })
118 .max()
119 .unwrap_or(0);
120
121 let new_columns = columns
123 .into_iter()
124 .enumerate()
125 .map(|(delta, mut c)| {
126 if c.semantic_type == SemanticType::Tag {
127 if !c.column_schema.data_type.is_string() {
128 return ColumnTypeMismatchSnafu {
129 expect: ConcreteDataType::string_datatype(),
130 actual: c.column_schema.data_type.clone(),
131 }
132 .fail();
133 }
134 } else {
135 warn!(
136 "Column {} in region {region_id} is not a tag",
137 c.column_schema.name
138 );
139 };
140
141 c.column_id = new_column_id_start + delta as u32;
142 c.column_schema.set_nullable();
143 match index_options {
144 IndexOptions::None => {}
145 IndexOptions::Inverted => {
146 c.column_schema.set_inverted_index(true);
147 }
148 IndexOptions::Skipping {
149 granularity,
150 false_positive_rate,
151 } => {
152 c.column_schema
153 .set_skipping_options(
154 &SkippingIndexOptions::new(
155 granularity,
156 false_positive_rate,
157 SkippingIndexType::BloomFilter,
158 )
159 .context(SetSkippingIndexOptionSnafu)?,
160 )
161 .context(SetSkippingIndexOptionSnafu)?;
162 }
163 }
164
165 Ok(AddColumn {
166 column_metadata: c.clone(),
167 location: None,
168 })
169 })
170 .collect::<Result<_>>()?;
171
172 debug!("Adding (Column id assigned) columns {new_columns:?} to region {region_id:?}");
173 let alter_request = RegionRequest::Alter(RegionAlterRequest {
175 kind: AlterKind::AddColumns {
176 columns: new_columns,
177 },
178 });
179
180 Ok(alter_request)
181 }
182
183 pub async fn write_data(
184 &self,
185 region_id: RegionId,
186 request: RegionPutRequest,
187 ) -> Result<AffectedRows> {
188 let region_id = utils::to_data_region_id(region_id);
189 self.mito
190 .handle_request(region_id, RegionRequest::Put(request))
191 .await
192 .context(MitoWriteOperationSnafu)
193 .map(|result| result.affected_rows)
194 }
195
196 pub async fn physical_columns(
197 &self,
198 physical_region_id: RegionId,
199 ) -> Result<Vec<ColumnMetadata>> {
200 let data_region_id = utils::to_data_region_id(physical_region_id);
201 let metadata = self
202 .mito
203 .get_metadata(data_region_id)
204 .await
205 .context(MitoReadOperationSnafu)?;
206 Ok(metadata.column_metadatas.clone())
207 }
208
209 pub async fn alter_region_options(
210 &self,
211 region_id: RegionId,
212 request: RegionAlterRequest,
213 ) -> Result<AffectedRows> {
214 match request.kind {
215 AlterKind::SetRegionOptions { options: _ }
216 | AlterKind::UnsetRegionOptions { keys: _ }
217 | AlterKind::SetIndexes { options: _ }
218 | AlterKind::UnsetIndexes { options: _ }
219 | AlterKind::SyncColumns {
220 column_metadatas: _,
221 } => {
222 let region_id = utils::to_data_region_id(region_id);
223 self.mito
224 .handle_request(region_id, RegionRequest::Alter(request))
225 .await
226 .context(MitoWriteOperationSnafu)
227 .map(|result| result.affected_rows)
228 }
229 _ => {
230 info!(
231 "Metric region received alter request {request:?} on physical region {region_id:?}"
232 );
233 FORBIDDEN_OPERATION_COUNT.inc();
234
235 ForbiddenPhysicalAlterSnafu.fail()
236 }
237 }
238 }
239}
240
241#[cfg(test)]
242mod test {
243 use datatypes::prelude::ConcreteDataType;
244 use datatypes::schema::ColumnSchema;
245
246 use super::*;
247 use crate::test_util::TestEnv;
248
249 #[tokio::test]
250 async fn test_add_columns() {
251 let env = TestEnv::new().await;
252 env.init_metric_region().await;
253
254 let current_version = env
255 .mito()
256 .get_metadata(utils::to_data_region_id(env.default_physical_region_id()))
257 .await
258 .unwrap()
259 .schema_version;
260 assert_eq!(current_version, 1);
262
263 let new_columns = vec![
264 ColumnMetadata {
265 column_id: 0,
266 semantic_type: SemanticType::Tag,
267 column_schema: ColumnSchema::new(
268 "tag2",
269 ConcreteDataType::string_datatype(),
270 false,
271 ),
272 },
273 ColumnMetadata {
274 column_id: 0,
275 semantic_type: SemanticType::Tag,
276 column_schema: ColumnSchema::new(
277 "tag3",
278 ConcreteDataType::string_datatype(),
279 false,
280 ),
281 },
282 ];
283 env.data_region()
284 .add_columns(
285 env.default_physical_region_id(),
286 new_columns,
287 IndexOptions::Inverted,
288 )
289 .await
290 .unwrap();
291
292 let new_metadata = env
293 .mito()
294 .get_metadata(utils::to_data_region_id(env.default_physical_region_id()))
295 .await
296 .unwrap();
297 let column_names = new_metadata
298 .column_metadatas
299 .iter()
300 .map(|c| &c.column_schema.name)
301 .collect::<Vec<_>>();
302 let expected = vec![
303 "greptime_timestamp",
304 "greptime_value",
305 "__table_id",
306 "__tsid",
307 "job",
308 "tag2",
309 "tag3",
310 ];
311 assert_eq!(column_names, expected);
312 }
313
314 #[tokio::test]
316 async fn test_add_invalid_column() {
317 let env = TestEnv::new().await;
318 env.init_metric_region().await;
319
320 let new_columns = vec![ColumnMetadata {
321 column_id: 0,
322 semantic_type: SemanticType::Tag,
323 column_schema: ColumnSchema::new("tag2", ConcreteDataType::int64_datatype(), false),
324 }];
325 let result = env
326 .data_region()
327 .add_columns(
328 env.default_physical_region_id(),
329 new_columns,
330 IndexOptions::Inverted,
331 )
332 .await;
333 assert!(result.is_err());
334 }
335}