|
| 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 | + |
| 15 | +use async_trait::async_trait; |
| 16 | +use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; |
| 17 | +use common_procedure::{Context as ProcedureContext, LockKey, Procedure, Status}; |
| 18 | +use common_telemetry::tracing::info; |
| 19 | +use serde::{Deserialize, Serialize}; |
| 20 | +use snafu::{ensure, ResultExt}; |
| 21 | +use strum::AsRefStr; |
| 22 | + |
| 23 | +use super::utils::handle_retry_error; |
| 24 | +use crate::cache_invalidator::Context; |
| 25 | +use crate::ddl::DdlContext; |
| 26 | +use crate::error::{Result, SchemaNotFoundSnafu}; |
| 27 | +use crate::instruction::CacheIdent; |
| 28 | +use crate::key::schema_name::{SchemaName, SchemaNameKey, SchemaNameValue}; |
| 29 | +use crate::key::DeserializedValueWithBytes; |
| 30 | +use crate::lock_key::{CatalogLock, SchemaLock}; |
| 31 | +use crate::rpc::ddl::UnsetDatabaseOption::{self}; |
| 32 | +use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, SetDatabaseOption}; |
| 33 | +use crate::ClusterId; |
| 34 | + |
| 35 | +pub struct AlterDatabaseProcedure { |
| 36 | + pub context: DdlContext, |
| 37 | + pub data: AlterDatabaseData, |
| 38 | +} |
| 39 | + |
| 40 | +fn build_new_schema_value( |
| 41 | + mut value: SchemaNameValue, |
| 42 | + alter_kind: &AlterDatabaseKind, |
| 43 | +) -> Result<SchemaNameValue> { |
| 44 | + match alter_kind { |
| 45 | + AlterDatabaseKind::SetDatabaseOptions(options) => { |
| 46 | + for option in options.0.iter() { |
| 47 | + match option { |
| 48 | + SetDatabaseOption::Ttl(ttl) => { |
| 49 | + if ttl.is_zero() { |
| 50 | + value.ttl = None; |
| 51 | + } else { |
| 52 | + value.ttl = Some(*ttl); |
| 53 | + } |
| 54 | + } |
| 55 | + } |
| 56 | + } |
| 57 | + } |
| 58 | + AlterDatabaseKind::UnsetDatabaseOptions(keys) => { |
| 59 | + for key in keys.0.iter() { |
| 60 | + match key { |
| 61 | + UnsetDatabaseOption::Ttl => value.ttl = None, |
| 62 | + } |
| 63 | + } |
| 64 | + } |
| 65 | + } |
| 66 | + Ok(value) |
| 67 | +} |
| 68 | + |
| 69 | +impl AlterDatabaseProcedure { |
| 70 | + pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterDatabase"; |
| 71 | + |
| 72 | + pub fn new( |
| 73 | + cluster_id: ClusterId, |
| 74 | + task: AlterDatabaseTask, |
| 75 | + context: DdlContext, |
| 76 | + ) -> Result<Self> { |
| 77 | + Ok(Self { |
| 78 | + context, |
| 79 | + data: AlterDatabaseData::new(task, cluster_id)?, |
| 80 | + }) |
| 81 | + } |
| 82 | + |
| 83 | + pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult<Self> { |
| 84 | + let data = serde_json::from_str(json).context(FromJsonSnafu)?; |
| 85 | + |
| 86 | + Ok(Self { context, data }) |
| 87 | + } |
| 88 | + |
| 89 | + pub async fn on_prepare(&mut self) -> Result<Status> { |
| 90 | + let value = self |
| 91 | + .context |
| 92 | + .table_metadata_manager |
| 93 | + .schema_manager() |
| 94 | + .get(SchemaNameKey::new(self.data.catalog(), self.data.schema())) |
| 95 | + .await?; |
| 96 | + |
| 97 | + ensure!( |
| 98 | + value.is_some(), |
| 99 | + SchemaNotFoundSnafu { |
| 100 | + table_schema: self.data.schema(), |
| 101 | + } |
| 102 | + ); |
| 103 | + |
| 104 | + self.data.schema_value = value; |
| 105 | + self.data.state = AlterDatabaseState::UpdateMetadata; |
| 106 | + |
| 107 | + Ok(Status::executing(true)) |
| 108 | + } |
| 109 | + |
| 110 | + pub async fn on_update_metadata(&mut self) -> Result<Status> { |
| 111 | + let schema_name = SchemaNameKey::new(self.data.catalog(), self.data.schema()); |
| 112 | + |
| 113 | + // Safety: schema_value is not None. |
| 114 | + let current_schema_value = self.data.schema_value.as_ref().unwrap(); |
| 115 | + |
| 116 | + let new_schema_value = build_new_schema_value( |
| 117 | + current_schema_value.get_inner_ref().clone(), |
| 118 | + &self.data.kind, |
| 119 | + )?; |
| 120 | + |
| 121 | + self.context |
| 122 | + .table_metadata_manager |
| 123 | + .schema_manager() |
| 124 | + .update(schema_name, current_schema_value, &new_schema_value) |
| 125 | + .await?; |
| 126 | + |
| 127 | + info!("Updated database metadata for schema {schema_name}"); |
| 128 | + self.data.state = AlterDatabaseState::InvalidateSchemaCache; |
| 129 | + Ok(Status::executing(true)) |
| 130 | + } |
| 131 | + |
| 132 | + pub async fn on_invalidate_schema_cache(&mut self) -> Result<Status> { |
| 133 | + let cache_invalidator = &self.context.cache_invalidator; |
| 134 | + cache_invalidator |
| 135 | + .invalidate( |
| 136 | + &Context::default(), |
| 137 | + &[CacheIdent::SchemaName(SchemaName { |
| 138 | + catalog_name: self.data.catalog().to_string(), |
| 139 | + schema_name: self.data.schema().to_string(), |
| 140 | + })], |
| 141 | + ) |
| 142 | + .await?; |
| 143 | + |
| 144 | + Ok(Status::done()) |
| 145 | + } |
| 146 | +} |
| 147 | + |
| 148 | +#[async_trait] |
| 149 | +impl Procedure for AlterDatabaseProcedure { |
| 150 | + fn type_name(&self) -> &str { |
| 151 | + Self::TYPE_NAME |
| 152 | + } |
| 153 | + |
| 154 | + async fn execute(&mut self, _ctx: &ProcedureContext) -> ProcedureResult<Status> { |
| 155 | + match self.data.state { |
| 156 | + AlterDatabaseState::Prepare => self.on_prepare().await, |
| 157 | + AlterDatabaseState::UpdateMetadata => self.on_update_metadata().await, |
| 158 | + AlterDatabaseState::InvalidateSchemaCache => self.on_invalidate_schema_cache().await, |
| 159 | + } |
| 160 | + .map_err(handle_retry_error) |
| 161 | + } |
| 162 | + |
| 163 | + fn dump(&self) -> ProcedureResult<String> { |
| 164 | + serde_json::to_string(&self.data).context(ToJsonSnafu) |
| 165 | + } |
| 166 | + |
| 167 | + fn lock_key(&self) -> LockKey { |
| 168 | + let catalog = self.data.catalog(); |
| 169 | + let schema = self.data.schema(); |
| 170 | + |
| 171 | + let lock_key = vec![ |
| 172 | + CatalogLock::Read(catalog).into(), |
| 173 | + SchemaLock::write(catalog, schema).into(), |
| 174 | + ]; |
| 175 | + |
| 176 | + LockKey::new(lock_key) |
| 177 | + } |
| 178 | +} |
| 179 | + |
| 180 | +#[derive(Debug, Serialize, Deserialize, AsRefStr)] |
| 181 | +enum AlterDatabaseState { |
| 182 | + Prepare, |
| 183 | + UpdateMetadata, |
| 184 | + InvalidateSchemaCache, |
| 185 | +} |
| 186 | + |
| 187 | +/// The data of alter database procedure. |
| 188 | +#[derive(Debug, Serialize, Deserialize)] |
| 189 | +pub struct AlterDatabaseData { |
| 190 | + cluster_id: ClusterId, |
| 191 | + state: AlterDatabaseState, |
| 192 | + kind: AlterDatabaseKind, |
| 193 | + catalog_name: String, |
| 194 | + schema_name: String, |
| 195 | + schema_value: Option<DeserializedValueWithBytes<SchemaNameValue>>, |
| 196 | +} |
| 197 | + |
| 198 | +impl AlterDatabaseData { |
| 199 | + pub fn new(task: AlterDatabaseTask, cluster_id: ClusterId) -> Result<Self> { |
| 200 | + Ok(Self { |
| 201 | + cluster_id, |
| 202 | + state: AlterDatabaseState::Prepare, |
| 203 | + kind: AlterDatabaseKind::try_from(task.alter_expr.kind.unwrap())?, |
| 204 | + catalog_name: task.alter_expr.catalog_name, |
| 205 | + schema_name: task.alter_expr.schema_name, |
| 206 | + schema_value: None, |
| 207 | + }) |
| 208 | + } |
| 209 | + |
| 210 | + pub fn catalog(&self) -> &str { |
| 211 | + &self.catalog_name |
| 212 | + } |
| 213 | + |
| 214 | + pub fn schema(&self) -> &str { |
| 215 | + &self.schema_name |
| 216 | + } |
| 217 | +} |
| 218 | + |
| 219 | +#[cfg(test)] |
| 220 | +mod tests { |
| 221 | + use std::time::Duration; |
| 222 | + |
| 223 | + use crate::ddl::alter_database::build_new_schema_value; |
| 224 | + use crate::key::schema_name::SchemaNameValue; |
| 225 | + use crate::rpc::ddl::{ |
| 226 | + AlterDatabaseKind, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, |
| 227 | + UnsetDatabaseOptions, |
| 228 | + }; |
| 229 | + |
| 230 | + #[test] |
| 231 | + fn test_build_new_schema_value() { |
| 232 | + let set_ttl = AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(vec![ |
| 233 | + SetDatabaseOption::Ttl(Duration::from_secs(10)), |
| 234 | + ])); |
| 235 | + let current_schema_value = SchemaNameValue::default(); |
| 236 | + let new_schema_value = |
| 237 | + build_new_schema_value(current_schema_value.clone(), &set_ttl).unwrap(); |
| 238 | + assert_eq!(new_schema_value.ttl, Some(Duration::from_secs(10))); |
| 239 | + |
| 240 | + let unset_ttl_alter_kind = |
| 241 | + AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(vec![ |
| 242 | + UnsetDatabaseOption::Ttl, |
| 243 | + ])); |
| 244 | + let new_schema_value = |
| 245 | + build_new_schema_value(current_schema_value, &unset_ttl_alter_kind).unwrap(); |
| 246 | + assert_eq!(new_schema_value.ttl, None); |
| 247 | + } |
| 248 | +} |
0 commit comments