common_meta/kv_backend/rds/
mysql.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::marker::PhantomData;
16use std::sync::Arc;
17
18use common_telemetry::{debug, info, warn};
19use lazy_static::lazy_static;
20use regex::Regex;
21use snafu::{OptionExt, ResultExt};
22use sqlx::mysql::MySqlRow;
23use sqlx::pool::Pool;
24use sqlx::{AssertSqlSafe, MySql, MySqlPool, Row, Transaction as MySqlTransaction};
25use strum::AsRefStr;
26
27use crate::error::{
28    CreateMySqlPoolSnafu, MySqlExecutionSnafu, MySqlTransactionSnafu, Result, UnexpectedSnafu,
29};
30use crate::kv_backend::KvBackendRef;
31use crate::kv_backend::rds::{
32    Executor, ExecutorFactory, ExecutorImpl, KvQueryExecutor, RDS_STORE_OP_BATCH_DELETE,
33    RDS_STORE_OP_BATCH_GET, RDS_STORE_OP_BATCH_PUT, RDS_STORE_OP_RANGE_DELETE,
34    RDS_STORE_OP_RANGE_QUERY, RDS_STORE_TXN_RETRY_COUNT, RdsStore, Transaction,
35    ensure_rustls_crypto_provider_installed,
36};
37use crate::rpc::KeyValue;
38use crate::rpc::store::{
39    BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest,
40    BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, RangeRequest, RangeResponse,
41};
42
43const MYSQL_STORE_NAME: &str = "mysql_store";
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46enum ValueBlobType {
47    Blob,
48    MediumBlob,
49    LongBlob,
50}
51
52lazy_static! {
53    static ref VALUE_COLUMN_BLOB_TYPE_RE: Regex =
54        Regex::new(r#"(?i)(?:\(|,)\s*[`"]?v[`"]?\s+(longblob|mediumblob|blob)\b"#).unwrap();
55}
56
57type MySqlClient = Arc<Pool<MySql>>;
58pub struct MySqlTxnClient(MySqlTransaction<'static, MySql>);
59
60fn key_value_from_row(row: MySqlRow) -> KeyValue {
61    // Safety: key and value are the first two columns in the row
62    KeyValue {
63        key: row.get_unchecked(0),
64        value: row.get_unchecked(1),
65    }
66}
67
68const EMPTY: &[u8] = &[0];
69
70/// Type of range template.
71#[derive(Debug, Clone, Copy, AsRefStr)]
72enum RangeTemplateType {
73    Point,
74    Range,
75    Full,
76    LeftBounded,
77    Prefix,
78}
79
80/// Builds params for the given range template type.
81impl RangeTemplateType {
82    /// Builds the parameters for the given range template type.
83    /// You can check out the conventions at [RangeRequest]
84    fn build_params(&self, mut key: Vec<u8>, range_end: Vec<u8>) -> Vec<Vec<u8>> {
85        match self {
86            RangeTemplateType::Point => vec![key],
87            RangeTemplateType::Range => vec![key, range_end],
88            RangeTemplateType::Full => vec![],
89            RangeTemplateType::LeftBounded => vec![key],
90            RangeTemplateType::Prefix => {
91                key.push(b'%');
92                vec![key]
93            }
94        }
95    }
96}
97
98/// Templates for range request.
99#[derive(Debug, Clone)]
100struct RangeTemplate {
101    point: String,
102    range: String,
103    full: String,
104    left_bounded: String,
105    prefix: String,
106}
107
108impl RangeTemplate {
109    /// Gets the template for the given type.
110    fn get(&self, typ: RangeTemplateType) -> &str {
111        match typ {
112            RangeTemplateType::Point => &self.point,
113            RangeTemplateType::Range => &self.range,
114            RangeTemplateType::Full => &self.full,
115            RangeTemplateType::LeftBounded => &self.left_bounded,
116            RangeTemplateType::Prefix => &self.prefix,
117        }
118    }
119
120    /// Adds limit to the template.
121    fn with_limit(template: &str, limit: i64) -> String {
122        if limit == 0 {
123            return format!("{};", template);
124        }
125        format!("{} LIMIT {};", template, limit)
126    }
127}
128
129fn is_prefix_range(start: &[u8], end: &[u8]) -> bool {
130    if start.len() != end.len() {
131        return false;
132    }
133    let l = start.len();
134    let same_prefix = start[0..l - 1] == end[0..l - 1];
135    if let (Some(rhs), Some(lhs)) = (start.last(), end.last()) {
136        return same_prefix && (*rhs + 1) == *lhs;
137    }
138    false
139}
140
141/// Determine the template type for range request.
142fn range_template(key: &[u8], range_end: &[u8]) -> RangeTemplateType {
143    match (key, range_end) {
144        (_, &[]) => RangeTemplateType::Point,
145        (EMPTY, EMPTY) => RangeTemplateType::Full,
146        (_, EMPTY) => RangeTemplateType::LeftBounded,
147        (start, end) => {
148            if is_prefix_range(start, end) {
149                RangeTemplateType::Prefix
150            } else {
151                RangeTemplateType::Range
152            }
153        }
154    }
155}
156
157/// Generate in placeholders for MySQL.
158fn mysql_generate_in_placeholders(from: usize, to: usize) -> Vec<String> {
159    (from..=to).map(|_| "?".to_string()).collect()
160}
161
162/// Factory for building sql templates.
163struct MySqlTemplateFactory<'a> {
164    table_name: &'a str,
165}
166
167impl<'a> MySqlTemplateFactory<'a> {
168    /// Creates a new [`SqlTemplateFactory`] with the given table name.
169    fn new(table_name: &'a str) -> Self {
170        Self { table_name }
171    }
172
173    /// Builds the template set for the given table name.
174    fn build(&self) -> MySqlTemplateSet {
175        let table_name = self.table_name;
176        // Some of queries don't end with `;`, because we need to add `LIMIT` clause.
177        MySqlTemplateSet {
178            table_name: table_name.to_string(),
179            create_table_statement: format!(
180                // Cannot be more than 3072 bytes in PRIMARY KEY
181                "CREATE TABLE IF NOT EXISTS `{table_name}`(k VARBINARY(3072) PRIMARY KEY, v MEDIUMBLOB);",
182            ),
183            show_create_table_statement: format!("SHOW CREATE TABLE `{table_name}`"),
184            alter_value_column_statement: format!(
185                "ALTER TABLE `{table_name}` MODIFY COLUMN v MEDIUMBLOB;"
186            ),
187            range_template: RangeTemplate {
188                point: format!("SELECT k, v FROM `{table_name}` WHERE k = ?"),
189                range: format!("SELECT k, v FROM `{table_name}` WHERE k >= ? AND k < ? ORDER BY k"),
190                full: format!("SELECT k, v FROM `{table_name}` ORDER BY k"),
191                left_bounded: format!("SELECT k, v FROM `{table_name}` WHERE k >= ? ORDER BY k"),
192                prefix: format!("SELECT k, v FROM `{table_name}` WHERE k LIKE ? ORDER BY k"),
193            },
194            delete_template: RangeTemplate {
195                point: format!("DELETE FROM `{table_name}` WHERE k = ?;"),
196                range: format!("DELETE FROM `{table_name}` WHERE k >= ? AND k < ?;"),
197                full: format!("DELETE FROM `{table_name}`"),
198                left_bounded: format!("DELETE FROM `{table_name}` WHERE k >= ?;"),
199                prefix: format!("DELETE FROM `{table_name}` WHERE k LIKE ?;"),
200            },
201        }
202    }
203}
204
205/// Templates for the given table name.
206#[derive(Debug, Clone)]
207pub struct MySqlTemplateSet {
208    table_name: String,
209    create_table_statement: String,
210    show_create_table_statement: String,
211    alter_value_column_statement: String,
212    range_template: RangeTemplate,
213    delete_template: RangeTemplate,
214}
215
216impl MySqlTemplateSet {
217    /// Generates the sql for batch get.
218    fn generate_batch_get_query(&self, key_len: usize) -> String {
219        let table_name = &self.table_name;
220        let in_clause = mysql_generate_in_placeholders(1, key_len).join(", ");
221        format!(
222            "SELECT k, v FROM `{table_name}` WHERE k in ({});",
223            in_clause
224        )
225    }
226
227    /// Generates the sql for batch delete.
228    fn generate_batch_delete_query(&self, key_len: usize) -> String {
229        let table_name = &self.table_name;
230        let in_clause = mysql_generate_in_placeholders(1, key_len).join(", ");
231        format!("DELETE FROM `{table_name}` WHERE k in ({});", in_clause)
232    }
233
234    /// Generates the sql for batch upsert.
235    /// For MySQL, it also generates a select query to get the previous values.
236    fn generate_batch_upsert_query(&self, kv_len: usize) -> (String, String) {
237        let table_name = &self.table_name;
238        let in_placeholders: Vec<String> = (1..=kv_len).map(|_| "?".to_string()).collect();
239        let in_clause = in_placeholders.join(", ");
240        let mut values_placeholders = Vec::new();
241        for _ in 0..kv_len {
242            values_placeholders.push("(?, ?)".to_string());
243        }
244        let values_clause = values_placeholders.join(", ");
245
246        (
247            format!(r#"SELECT k, v FROM `{table_name}` WHERE k IN ({in_clause})"#,),
248            format!(
249                r#"INSERT INTO `{table_name}` (k, v) VALUES {values_clause} ON DUPLICATE KEY UPDATE v = VALUES(v);"#,
250            ),
251        )
252    }
253}
254
255#[async_trait::async_trait]
256impl Executor for MySqlClient {
257    type Transaction<'a>
258        = MySqlTxnClient
259    where
260        Self: 'a;
261
262    fn name() -> &'static str {
263        "MySql"
264    }
265
266    async fn query(&mut self, raw_query: &str, params: &[&Vec<u8>]) -> Result<Vec<KeyValue>> {
267        // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data
268        // is passed via bind parameters.
269        let query = sqlx::query(AssertSqlSafe(raw_query));
270        let query = params.iter().fold(query, |query, param| query.bind(param));
271        let rows = query
272            .fetch_all(&**self)
273            .await
274            .context(MySqlExecutionSnafu { sql: raw_query })?;
275        Ok(rows.into_iter().map(key_value_from_row).collect())
276    }
277
278    async fn execute(&mut self, raw_query: &str, params: &[&Vec<u8>]) -> Result<()> {
279        // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data
280        // is passed via bind parameters.
281        let query = sqlx::query(AssertSqlSafe(raw_query));
282        let query = params.iter().fold(query, |query, param| query.bind(param));
283        query
284            .execute(&**self)
285            .await
286            .context(MySqlExecutionSnafu { sql: raw_query })?;
287        Ok(())
288    }
289
290    async fn txn_executor<'a>(&'a mut self) -> Result<Self::Transaction<'a>> {
291        // sqlx has no isolation level support for now, so we have to set it manually.
292        // TODO(CookiePie): Waiting for https://github.com/launchbadge/sqlx/pull/3614 and remove this.
293        sqlx::query("SET SESSION TRANSACTION ISOLATION LEVEL SERIALIZABLE")
294            .execute(&**self)
295            .await
296            .context(MySqlExecutionSnafu {
297                sql: "SET SESSION TRANSACTION ISOLATION LEVEL SERIALIZABLE",
298            })?;
299        let txn = self
300            .begin()
301            .await
302            .context(MySqlExecutionSnafu { sql: "begin" })?;
303        Ok(MySqlTxnClient(txn))
304    }
305}
306
307#[async_trait::async_trait]
308impl Transaction<'_> for MySqlTxnClient {
309    async fn query(&mut self, raw_query: &str, params: &[&Vec<u8>]) -> Result<Vec<KeyValue>> {
310        // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data
311        // is passed via bind parameters.
312        let query = sqlx::query(AssertSqlSafe(raw_query));
313        let query = params.iter().fold(query, |query, param| query.bind(param));
314        // As said in https://docs.rs/sqlx/latest/sqlx/trait.Executor.html, we need a `&mut *transaction`. Weird.
315        let rows = query
316            .fetch_all(&mut *(self.0))
317            .await
318            .context(MySqlExecutionSnafu { sql: raw_query })?;
319        Ok(rows.into_iter().map(key_value_from_row).collect())
320    }
321
322    async fn execute(&mut self, raw_query: &str, params: &[&Vec<u8>]) -> Result<()> {
323        // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data
324        // is passed via bind parameters.
325        let query = sqlx::query(AssertSqlSafe(raw_query));
326        let query = params.iter().fold(query, |query, param| query.bind(param));
327        // As said in https://docs.rs/sqlx/latest/sqlx/trait.Executor.html, we need a `&mut *transaction`. Weird.
328        query
329            .execute(&mut *(self.0))
330            .await
331            .context(MySqlExecutionSnafu { sql: raw_query })?;
332        Ok(())
333    }
334
335    /// Caution: sqlx will stuck on the query if two transactions conflict with each other.
336    /// Don't know if it's a feature or it depends on the database. Be careful.
337    async fn commit(self) -> Result<()> {
338        self.0.commit().await.context(MySqlTransactionSnafu {
339            operation: "commit",
340        })?;
341        Ok(())
342    }
343}
344
345pub struct MySqlExecutorFactory {
346    pool: Arc<Pool<MySql>>,
347}
348
349#[async_trait::async_trait]
350impl ExecutorFactory<MySqlClient> for MySqlExecutorFactory {
351    async fn default_executor(&self) -> Result<MySqlClient> {
352        Ok(self.pool.clone())
353    }
354
355    async fn txn_executor<'a>(
356        &self,
357        default_executor: &'a mut MySqlClient,
358    ) -> Result<MySqlTxnClient> {
359        default_executor.txn_executor().await
360    }
361}
362
363/// A MySQL-backed key-value store.
364/// It uses [sqlx::Pool<MySql>] as the connection pool for [RdsStore].
365pub type MySqlStore = RdsStore<MySqlClient, MySqlExecutorFactory, MySqlTemplateSet>;
366
367#[async_trait::async_trait]
368impl KvQueryExecutor<MySqlClient> for MySqlStore {
369    async fn range_with_query_executor(
370        &self,
371        query_executor: &mut ExecutorImpl<'_, MySqlClient>,
372        req: RangeRequest,
373    ) -> Result<RangeResponse> {
374        let template_type = range_template(&req.key, &req.range_end);
375        let template = self.sql_template_set.range_template.get(template_type);
376        let params = template_type.build_params(req.key, req.range_end);
377        let params_ref = params.iter().collect::<Vec<_>>();
378        // Always add 1 to limit to check if there is more data
379        let query =
380            RangeTemplate::with_limit(template, if req.limit == 0 { 0 } else { req.limit + 1 });
381        let limit = req.limit as usize;
382        debug!("query: {:?}, params: {:?}", query, params);
383        let mut kvs = crate::record_rds_sql_execute_elapsed!(
384            query_executor.query(&query, &params_ref).await,
385            MYSQL_STORE_NAME,
386            RDS_STORE_OP_RANGE_QUERY,
387            template_type.as_ref()
388        )?;
389        if req.keys_only {
390            kvs.iter_mut().for_each(|kv| kv.value = vec![]);
391        }
392        // If limit is 0, we always return all data
393        if limit == 0 || kvs.len() <= limit {
394            return Ok(RangeResponse { kvs, more: false });
395        }
396        // If limit is greater than the number of rows, we remove the last row and set more to true
397        let removed = kvs.pop();
398        debug_assert!(removed.is_some());
399        Ok(RangeResponse { kvs, more: true })
400    }
401
402    async fn batch_put_with_query_executor(
403        &self,
404        query_executor: &mut ExecutorImpl<'_, MySqlClient>,
405        req: BatchPutRequest,
406    ) -> Result<BatchPutResponse> {
407        let mut in_params = Vec::with_capacity(req.kvs.len() * 3);
408        let mut values_params = Vec::with_capacity(req.kvs.len() * 2);
409
410        for kv in &req.kvs {
411            let processed_key = &kv.key;
412            in_params.push(processed_key);
413
414            let processed_value = &kv.value;
415            values_params.push(processed_key);
416            values_params.push(processed_value);
417        }
418        let in_params = in_params.iter().map(|x| x as _).collect::<Vec<_>>();
419        let values_params = values_params.iter().map(|x| x as _).collect::<Vec<_>>();
420        let (select, update) = self
421            .sql_template_set
422            .generate_batch_upsert_query(req.kvs.len());
423
424        // Fast path: if we don't need previous kvs, we can just upsert the keys.
425        if !req.prev_kv {
426            crate::record_rds_sql_execute_elapsed!(
427                query_executor.execute(&update, &values_params).await,
428                MYSQL_STORE_NAME,
429                RDS_STORE_OP_BATCH_PUT,
430                ""
431            )?;
432            return Ok(BatchPutResponse::default());
433        }
434        // Should use transaction to ensure atomicity.
435        if let ExecutorImpl::Default(query_executor) = query_executor {
436            let txn = query_executor.txn_executor().await?;
437            let mut txn = ExecutorImpl::Txn(txn);
438            let res = self.batch_put_with_query_executor(&mut txn, req).await;
439            txn.commit().await?;
440            return res;
441        }
442        let prev_kvs = crate::record_rds_sql_execute_elapsed!(
443            query_executor.query(&select, &in_params).await,
444            MYSQL_STORE_NAME,
445            RDS_STORE_OP_BATCH_PUT,
446            ""
447        )?;
448        query_executor.execute(&update, &values_params).await?;
449        Ok(BatchPutResponse { prev_kvs })
450    }
451
452    async fn batch_get_with_query_executor(
453        &self,
454        query_executor: &mut ExecutorImpl<'_, MySqlClient>,
455        req: BatchGetRequest,
456    ) -> Result<BatchGetResponse> {
457        if req.keys.is_empty() {
458            return Ok(BatchGetResponse { kvs: vec![] });
459        }
460        let query = self
461            .sql_template_set
462            .generate_batch_get_query(req.keys.len());
463        let params = req.keys.iter().map(|x| x as _).collect::<Vec<_>>();
464        let kvs = crate::record_rds_sql_execute_elapsed!(
465            query_executor.query(&query, &params).await,
466            MYSQL_STORE_NAME,
467            RDS_STORE_OP_BATCH_GET,
468            ""
469        )?;
470        Ok(BatchGetResponse { kvs })
471    }
472
473    async fn delete_range_with_query_executor(
474        &self,
475        query_executor: &mut ExecutorImpl<'_, MySqlClient>,
476        req: DeleteRangeRequest,
477    ) -> Result<DeleteRangeResponse> {
478        // Since we need to know the number of deleted keys, we have no fast path here.
479        // Should use transaction to ensure atomicity.
480        if let ExecutorImpl::Default(query_executor) = query_executor {
481            let txn = query_executor.txn_executor().await?;
482            let mut txn = ExecutorImpl::Txn(txn);
483            let res = self.delete_range_with_query_executor(&mut txn, req).await;
484            txn.commit().await?;
485            return res;
486        }
487        let range_get_req = RangeRequest {
488            key: req.key.clone(),
489            range_end: req.range_end.clone(),
490            limit: 0,
491            keys_only: false,
492        };
493        let prev_kvs = self
494            .range_with_query_executor(query_executor, range_get_req)
495            .await?
496            .kvs;
497        let template_type = range_template(&req.key, &req.range_end);
498        let template = self.sql_template_set.delete_template.get(template_type);
499        let params = template_type.build_params(req.key, req.range_end);
500        let params_ref = params.iter().map(|x| x as _).collect::<Vec<_>>();
501        crate::record_rds_sql_execute_elapsed!(
502            query_executor.execute(template, &params_ref).await,
503            MYSQL_STORE_NAME,
504            RDS_STORE_OP_RANGE_DELETE,
505            template_type.as_ref()
506        )?;
507        let mut resp = DeleteRangeResponse::new(prev_kvs.len() as i64);
508        if req.prev_kv {
509            resp.with_prev_kvs(prev_kvs);
510        }
511        Ok(resp)
512    }
513
514    async fn batch_delete_with_query_executor(
515        &self,
516        query_executor: &mut ExecutorImpl<'_, MySqlClient>,
517        req: BatchDeleteRequest,
518    ) -> Result<BatchDeleteResponse> {
519        if req.keys.is_empty() {
520            return Ok(BatchDeleteResponse::default());
521        }
522        let query = self
523            .sql_template_set
524            .generate_batch_delete_query(req.keys.len());
525        let params = req.keys.iter().map(|x| x as _).collect::<Vec<_>>();
526        // Fast path: if we don't need previous kvs, we can just delete the keys.
527        if !req.prev_kv {
528            crate::record_rds_sql_execute_elapsed!(
529                query_executor.execute(&query, &params).await,
530                MYSQL_STORE_NAME,
531                RDS_STORE_OP_BATCH_DELETE,
532                ""
533            )?;
534            return Ok(BatchDeleteResponse::default());
535        }
536        // Should use transaction to ensure atomicity.
537        if let ExecutorImpl::Default(query_executor) = query_executor {
538            let txn = query_executor.txn_executor().await?;
539            let mut txn = ExecutorImpl::Txn(txn);
540            let res = self.batch_delete_with_query_executor(&mut txn, req).await;
541            txn.commit().await?;
542            return res;
543        }
544        // Should get previous kvs first
545        let batch_get_req = BatchGetRequest {
546            keys: req.keys.clone(),
547        };
548        let prev_kvs = self
549            .batch_get_with_query_executor(query_executor, batch_get_req)
550            .await?
551            .kvs;
552        // Pure `DELETE` has no return value, so we need to use `execute` instead of `query`.
553        crate::record_rds_sql_execute_elapsed!(
554            query_executor.execute(&query, &params).await,
555            MYSQL_STORE_NAME,
556            RDS_STORE_OP_BATCH_DELETE,
557            ""
558        )?;
559        if req.prev_kv {
560            Ok(BatchDeleteResponse { prev_kvs })
561        } else {
562            Ok(BatchDeleteResponse::default())
563        }
564    }
565}
566
567impl MySqlStore {
568    /// Reads the current table definition for best-effort schema upgrades.
569    async fn fetch_create_table_sql(
570        pool: &Pool<MySql>,
571        sql_template_set: &MySqlTemplateSet,
572    ) -> Result<Option<String>> {
573        let row = sqlx::query(AssertSqlSafe(
574            sql_template_set.show_create_table_statement.as_str(),
575        ))
576        .fetch_optional(pool)
577        .await
578        .with_context(|_| MySqlExecutionSnafu {
579            sql: sql_template_set.show_create_table_statement.clone(),
580        })?;
581        Ok(row.map(|row| row.get(1)))
582    }
583
584    /// Parses the blob type of the `v` column from `SHOW CREATE TABLE` output.
585    fn parse_value_column_blob_type(create_table_sql: &str) -> Option<ValueBlobType> {
586        // `SHOW CREATE TABLE` returns MySQL-specific DDL. A minimal parser keeps the
587        // upgrade check small and avoids introducing a SQL parser just for one column.
588        let captures = VALUE_COLUMN_BLOB_TYPE_RE.captures(create_table_sql)?;
589        match captures.get(1)?.as_str().to_ascii_lowercase().as_str() {
590            "blob" => Some(ValueBlobType::Blob),
591            "mediumblob" => Some(ValueBlobType::MediumBlob),
592            "longblob" => Some(ValueBlobType::LongBlob),
593            _ => None,
594        }
595    }
596
597    /// Upgrades the metadata value column to `MEDIUMBLOB` when an old table still uses `BLOB`.
598    async fn maybe_upgrade_value_column_to_mediumblob(
599        pool: &Pool<MySql>,
600        sql_template_set: &MySqlTemplateSet,
601    ) -> Result<()> {
602        let table_name = &sql_template_set.table_name;
603        let create_table_sql = Self::fetch_create_table_sql(pool, sql_template_set)
604            .await?
605            .context(UnexpectedSnafu {
606                err_msg: format!("Failed to fetch CREATE TABLE SQL for `{table_name}`"),
607            })?;
608
609        match Self::parse_value_column_blob_type(&create_table_sql) {
610            Some(ValueBlobType::Blob) => {
611                sqlx::query(AssertSqlSafe(
612                    sql_template_set.alter_value_column_statement.as_str(),
613                ))
614                .execute(pool)
615                .await
616                .with_context(|_| MySqlExecutionSnafu {
617                    sql: sql_template_set.alter_value_column_statement.clone(),
618                })?;
619                info!("Upgraded MySQL metadata value column to MEDIUMBLOB for `{table_name}`");
620            }
621            Some(ValueBlobType::MediumBlob | ValueBlobType::LongBlob) => {
622                debug!("MySQL metadata value column for `{table_name}` is already compatible");
623            }
624            None => {
625                warn!(
626                    "Failed to determine MySQL metadata value column type from table definition for `{table_name}`, skip automatic MEDIUMBLOB upgrade"
627                );
628            }
629        }
630
631        Ok(())
632    }
633
634    /// Create [MySqlStore] impl of [KvBackendRef] from url.
635    pub async fn with_url(url: &str, table_name: &str, max_txn_ops: usize) -> Result<KvBackendRef> {
636        ensure_rustls_crypto_provider_installed()?;
637        let pool = MySqlPool::connect(url)
638            .await
639            .context(CreateMySqlPoolSnafu)?;
640        Self::with_mysql_pool(pool, table_name, max_txn_ops).await
641    }
642
643    /// Create [MySqlStore] impl of [KvBackendRef] from [sqlx::Pool<MySql>].
644    pub async fn with_mysql_pool(
645        pool: Pool<MySql>,
646        table_name: &str,
647        max_txn_ops: usize,
648    ) -> Result<KvBackendRef> {
649        // This step ensures the mysql metadata backend is ready to use.
650        // We check if greptime_metakv table exists, and we will create a new table
651        // if it does not exist.
652        let sql_template_set = MySqlTemplateFactory::new(table_name).build();
653        sqlx::query(AssertSqlSafe(
654            sql_template_set.create_table_statement.as_str(),
655        ))
656        .execute(&pool)
657        .await
658        .with_context(|_| MySqlExecutionSnafu {
659            sql: sql_template_set.create_table_statement.clone(),
660        })?;
661        Self::maybe_upgrade_value_column_to_mediumblob(&pool, &sql_template_set).await?;
662        Ok(Arc::new(MySqlStore {
663            max_txn_ops,
664            sql_template_set,
665            txn_retry_count: RDS_STORE_TXN_RETRY_COUNT,
666            executor_factory: MySqlExecutorFactory {
667                pool: Arc::new(pool),
668            },
669            _phantom: PhantomData,
670        }))
671    }
672}
673
674#[cfg(test)]
675mod tests {
676    use common_telemetry::init_default_ut_logging;
677    use sqlx::mysql::{MySqlConnectOptions, MySqlSslMode};
678    use uuid::Uuid;
679
680    use super::*;
681    use crate::kv_backend::test::{
682        prepare_kv_with_prefix, test_kv_batch_delete_with_prefix, test_kv_batch_get_with_prefix,
683        test_kv_compare_and_put_with_prefix, test_kv_delete_range_with_prefix,
684        test_kv_put_with_prefix, test_kv_range_2_with_prefix, test_kv_range_with_prefix,
685        test_simple_kv_range, test_txn_compare_equal, test_txn_compare_greater,
686        test_txn_compare_less, test_txn_compare_not_equal, test_txn_one_compare_op,
687        text_txn_multi_compare_op, unprepare_kv,
688    };
689    use crate::maybe_skip_mysql_integration_test;
690    use crate::rpc::store::{PutRequest, RangeRequest};
691    use crate::test_util::test_certs_dir;
692
693    fn new_test_table_name(prefix: &str) -> String {
694        let uuid = Uuid::new_v4().simple().to_string();
695        let max_prefix_len = 63usize.saturating_sub(uuid.len() + 1);
696        let prefix = &prefix[..prefix.len().min(max_prefix_len)];
697        format!("{prefix}_{uuid}")
698    }
699
700    async fn mysql_pool() -> Option<MySqlPool> {
701        init_default_ut_logging();
702        let endpoints = std::env::var("GT_MYSQL_ENDPOINTS").unwrap_or_default();
703        if endpoints.is_empty() {
704            return None;
705        }
706        ensure_rustls_crypto_provider_installed().unwrap();
707        Some(MySqlPool::connect(&endpoints).await.unwrap())
708    }
709
710    async fn show_create_table(pool: &MySqlPool, table_name: &str) -> String {
711        let sql = format!("SHOW CREATE TABLE `{table_name}`");
712        let row = sqlx::query(AssertSqlSafe(sql.as_str()))
713            .fetch_one(pool)
714            .await
715            .unwrap();
716        row.get::<String, _>(1)
717    }
718
719    async fn create_legacy_blob_table(pool: &MySqlPool, table_name: &str) {
720        let sql = format!(
721            "CREATE TABLE IF NOT EXISTS `{table_name}`(k VARBINARY(3072) PRIMARY KEY, v BLOB);"
722        );
723        sqlx::query(AssertSqlSafe(sql.as_str()))
724            .execute(pool)
725            .await
726            .unwrap();
727    }
728
729    async fn drop_table(pool: &MySqlPool, table_name: &str) {
730        let sql = format!("DROP TABLE IF EXISTS `{table_name}`;");
731        sqlx::query(AssertSqlSafe(sql.as_str()))
732            .execute(pool)
733            .await
734            .unwrap();
735    }
736
737    async fn build_mysql_kv_backend(table_name: &str) -> Option<MySqlStore> {
738        let pool = mysql_pool().await?;
739        let sql_templates = MySqlTemplateFactory::new(table_name).build();
740        sqlx::query(AssertSqlSafe(sql_templates.create_table_statement.as_str()))
741            .execute(&pool)
742            .await
743            .unwrap();
744        Some(MySqlStore {
745            max_txn_ops: 128,
746            sql_template_set: sql_templates,
747            txn_retry_count: RDS_STORE_TXN_RETRY_COUNT,
748            executor_factory: MySqlExecutorFactory {
749                pool: Arc::new(pool),
750            },
751            _phantom: PhantomData,
752        })
753    }
754
755    #[test]
756    fn test_parse_value_column_blob_type() {
757        let sql = r#"CREATE TABLE `greptime_metakv` (
758  `k` varbinary(3072) NOT NULL,
759  `v` MEDIUMBLOB,
760  PRIMARY KEY (`k`)
761) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4"#;
762        assert_eq!(
763            Some(ValueBlobType::MediumBlob),
764            MySqlStore::parse_value_column_blob_type(sql)
765        );
766
767        let sql = r#"CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, `v` blob, PRIMARY KEY (`k`))"#;
768        assert_eq!(
769            Some(ValueBlobType::Blob),
770            MySqlStore::parse_value_column_blob_type(sql)
771        );
772
773        let sql = r#"CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, `v` longblob, PRIMARY KEY (`k`))"#;
774        assert_eq!(
775            Some(ValueBlobType::LongBlob),
776            MySqlStore::parse_value_column_blob_type(sql)
777        );
778
779        let sql = "CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, `v`    BLOB NOT NULL, PRIMARY KEY (`k`))";
780        assert_eq!(
781            Some(ValueBlobType::Blob),
782            MySqlStore::parse_value_column_blob_type(sql)
783        );
784
785        let sql = "CREATE TABLE `greptime_metakv` (\n  `k` varbinary(3072) NOT NULL,\n  \"v\" MediumBlob,\n  PRIMARY KEY (`k`)\n)";
786        assert_eq!(
787            Some(ValueBlobType::MediumBlob),
788            MySqlStore::parse_value_column_blob_type(sql)
789        );
790
791        let sql = "CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, vv blob, `v` longblob, PRIMARY KEY (`k`))";
792        assert_eq!(
793            Some(ValueBlobType::LongBlob),
794            MySqlStore::parse_value_column_blob_type(sql)
795        );
796
797        let sql = "CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, value blob, PRIMARY KEY (`k`))";
798        assert_eq!(None, MySqlStore::parse_value_column_blob_type(sql));
799
800        let sql = "CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, `v` varchar(255), PRIMARY KEY (`k`))";
801        assert_eq!(None, MySqlStore::parse_value_column_blob_type(sql));
802
803        let sql =
804            "CREATE TABLE `greptime_metakv` (`k` varbinary(3072) NOT NULL, PRIMARY KEY (`k`))";
805        assert_eq!(None, MySqlStore::parse_value_column_blob_type(sql));
806    }
807
808    #[tokio::test]
809    async fn test_mysql_new_metadata_table_uses_mediumblob() {
810        maybe_skip_mysql_integration_test!();
811        let pool = mysql_pool().await.unwrap();
812        let table_name = new_test_table_name("test_mysql_mediumblob_schema");
813
814        MySqlStore::with_mysql_pool(pool.clone(), &table_name, 128)
815            .await
816            .unwrap();
817
818        let create_table = show_create_table(&pool, &table_name).await;
819        assert!(create_table.to_ascii_uppercase().contains("MEDIUMBLOB"));
820
821        drop_table(&pool, &table_name).await;
822    }
823
824    #[tokio::test]
825    async fn test_mysql_legacy_blob_metadata_table_is_upgraded() {
826        maybe_skip_mysql_integration_test!();
827        let pool = mysql_pool().await.unwrap();
828        let table_name = new_test_table_name("test_mysql_legacy_blob_upgrade");
829
830        create_legacy_blob_table(&pool, &table_name).await;
831        MySqlStore::with_mysql_pool(pool.clone(), &table_name, 128)
832            .await
833            .unwrap();
834
835        let create_table = show_create_table(&pool, &table_name).await;
836        assert!(create_table.to_ascii_uppercase().contains("MEDIUMBLOB"));
837
838        drop_table(&pool, &table_name).await;
839    }
840
841    #[tokio::test]
842    async fn test_mysql_metadata_table_stores_large_values() {
843        maybe_skip_mysql_integration_test!();
844        let pool = mysql_pool().await.unwrap();
845        let table_name = new_test_table_name("test_mysql_large_metadata_value");
846        let kv_backend = MySqlStore::with_mysql_pool(pool.clone(), &table_name, 128)
847            .await
848            .unwrap();
849        let key = b"large-value".to_vec();
850        let value = vec![b'x'; 70 * 1024];
851
852        kv_backend
853            .put(
854                PutRequest::new()
855                    .with_key(key.clone())
856                    .with_value(value.clone()),
857            )
858            .await
859            .unwrap();
860        let response = kv_backend
861            .range(RangeRequest::new().with_key(key.clone()))
862            .await
863            .unwrap();
864
865        assert_eq!(1, response.kvs.len());
866        assert_eq!(key, response.kvs[0].key);
867        assert_eq!(value, response.kvs[0].value);
868
869        drop_table(&pool, &table_name).await;
870    }
871
872    #[tokio::test]
873    async fn test_mysql_upgraded_metadata_table_stores_large_values() {
874        maybe_skip_mysql_integration_test!();
875        let pool = mysql_pool().await.unwrap();
876        let table_name = new_test_table_name("test_mysql_upgraded_large_metadata_value");
877
878        create_legacy_blob_table(&pool, &table_name).await;
879        let kv_backend = MySqlStore::with_mysql_pool(pool.clone(), &table_name, 128)
880            .await
881            .unwrap();
882        let key = b"large-value".to_vec();
883        let value = vec![b'y'; 70 * 1024];
884
885        kv_backend
886            .put(
887                PutRequest::new()
888                    .with_key(key.clone())
889                    .with_value(value.clone()),
890            )
891            .await
892            .unwrap();
893        let response = kv_backend
894            .range(RangeRequest::new().with_key(key.clone()))
895            .await
896            .unwrap();
897
898        assert_eq!(1, response.kvs.len());
899        assert_eq!(key, response.kvs[0].key);
900        assert_eq!(value, response.kvs[0].value);
901
902        drop_table(&pool, &table_name).await;
903    }
904
905    #[tokio::test]
906    async fn test_mysql_put() {
907        maybe_skip_mysql_integration_test!();
908        let kv_backend = build_mysql_kv_backend("put-test").await.unwrap();
909        let prefix = b"put/";
910        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
911        test_kv_put_with_prefix(&kv_backend, prefix.to_vec()).await;
912        unprepare_kv(&kv_backend, prefix).await;
913    }
914
915    #[tokio::test]
916    async fn test_mysql_range() {
917        maybe_skip_mysql_integration_test!();
918        let kv_backend = build_mysql_kv_backend("range-test").await.unwrap();
919        let prefix = b"range/";
920        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
921        test_kv_range_with_prefix(&kv_backend, prefix.to_vec()).await;
922        unprepare_kv(&kv_backend, prefix).await;
923    }
924
925    #[tokio::test]
926    async fn test_mysql_range_2() {
927        maybe_skip_mysql_integration_test!();
928        let kv_backend = build_mysql_kv_backend("range2-test").await.unwrap();
929        let prefix = b"range2/";
930        test_kv_range_2_with_prefix(&kv_backend, prefix.to_vec()).await;
931        unprepare_kv(&kv_backend, prefix).await;
932    }
933
934    #[tokio::test]
935    async fn test_mysql_all_range() {
936        maybe_skip_mysql_integration_test!();
937        let kv_backend = build_mysql_kv_backend("simple_range-test").await.unwrap();
938        let prefix = b"";
939        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
940        test_simple_kv_range(&kv_backend).await;
941        unprepare_kv(&kv_backend, prefix).await;
942    }
943
944    #[tokio::test]
945    async fn test_mysql_batch_get() {
946        maybe_skip_mysql_integration_test!();
947        let kv_backend = build_mysql_kv_backend("batch_get-test").await.unwrap();
948        let prefix = b"batch_get/";
949        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
950        test_kv_batch_get_with_prefix(&kv_backend, prefix.to_vec()).await;
951        unprepare_kv(&kv_backend, prefix).await;
952    }
953
954    #[tokio::test]
955    async fn test_mysql_batch_delete() {
956        maybe_skip_mysql_integration_test!();
957        let kv_backend = build_mysql_kv_backend("batch_delete-test").await.unwrap();
958        let prefix = b"batch_delete/";
959        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
960        test_kv_delete_range_with_prefix(&kv_backend, prefix.to_vec()).await;
961        unprepare_kv(&kv_backend, prefix).await;
962    }
963
964    #[tokio::test]
965    async fn test_mysql_batch_delete_with_prefix() {
966        maybe_skip_mysql_integration_test!();
967        let kv_backend = build_mysql_kv_backend("batch_delete_with_prefix-test")
968            .await
969            .unwrap();
970        let prefix = b"batch_delete/";
971        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
972        test_kv_batch_delete_with_prefix(&kv_backend, prefix.to_vec()).await;
973        unprepare_kv(&kv_backend, prefix).await;
974    }
975
976    #[tokio::test]
977    async fn test_mysql_delete_range() {
978        maybe_skip_mysql_integration_test!();
979        let kv_backend = build_mysql_kv_backend("delete_range-test").await.unwrap();
980        let prefix = b"delete_range/";
981        prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
982        test_kv_delete_range_with_prefix(&kv_backend, prefix.to_vec()).await;
983        unprepare_kv(&kv_backend, prefix).await;
984    }
985
986    #[tokio::test]
987    async fn test_mysql_compare_and_put() {
988        maybe_skip_mysql_integration_test!();
989        let kv_backend = build_mysql_kv_backend("compare_and_put-test")
990            .await
991            .unwrap();
992        let prefix = b"compare_and_put/";
993        let kv_backend = Arc::new(kv_backend);
994        test_kv_compare_and_put_with_prefix(kv_backend.clone(), prefix.to_vec()).await;
995    }
996
997    #[tokio::test]
998    async fn test_mysql_txn() {
999        maybe_skip_mysql_integration_test!();
1000        let kv_backend = build_mysql_kv_backend("txn-test").await.unwrap();
1001        test_txn_one_compare_op(&kv_backend).await;
1002        text_txn_multi_compare_op(&kv_backend).await;
1003        test_txn_compare_equal(&kv_backend).await;
1004        test_txn_compare_greater(&kv_backend).await;
1005        test_txn_compare_less(&kv_backend).await;
1006        test_txn_compare_not_equal(&kv_backend).await;
1007    }
1008
1009    #[tokio::test]
1010    async fn test_mysql_with_tls() {
1011        common_telemetry::init_default_ut_logging();
1012        maybe_skip_mysql_integration_test!();
1013        ensure_rustls_crypto_provider_installed().unwrap();
1014        let endpoint = std::env::var("GT_MYSQL_ENDPOINTS").unwrap();
1015
1016        let opts = endpoint
1017            .parse::<MySqlConnectOptions>()
1018            .unwrap()
1019            .ssl_mode(MySqlSslMode::Required);
1020        let pool = MySqlPool::connect_with(opts).await.unwrap();
1021        sqlx::query("SELECT 1").execute(&pool).await.unwrap();
1022    }
1023
1024    #[tokio::test]
1025    async fn test_mysql_with_mtls() {
1026        common_telemetry::init_default_ut_logging();
1027        maybe_skip_mysql_integration_test!();
1028        ensure_rustls_crypto_provider_installed().unwrap();
1029        let endpoint = std::env::var("GT_MYSQL_ENDPOINTS").unwrap();
1030        let certs_dir = test_certs_dir();
1031
1032        let opts = endpoint
1033            .parse::<MySqlConnectOptions>()
1034            .unwrap()
1035            .ssl_mode(MySqlSslMode::Required)
1036            .ssl_client_cert(certs_dir.join("client.crt").to_string_lossy().to_string())
1037            .ssl_client_key(certs_dir.join("client.key").to_string_lossy().to_string());
1038        let pool = MySqlPool::connect_with(opts).await.unwrap();
1039        sqlx::query("SELECT 1").execute(&pool).await.unwrap();
1040    }
1041
1042    #[tokio::test]
1043    async fn test_mysql_with_tls_verify_ca() {
1044        common_telemetry::init_default_ut_logging();
1045        maybe_skip_mysql_integration_test!();
1046        ensure_rustls_crypto_provider_installed().unwrap();
1047        let endpoint = std::env::var("GT_MYSQL_ENDPOINTS").unwrap();
1048        let certs_dir = test_certs_dir();
1049
1050        let opts = endpoint
1051            .parse::<MySqlConnectOptions>()
1052            .unwrap()
1053            .ssl_mode(MySqlSslMode::VerifyCa)
1054            .ssl_ca(certs_dir.join("root.crt").to_string_lossy().to_string())
1055            .ssl_client_cert(certs_dir.join("client.crt").to_string_lossy().to_string())
1056            .ssl_client_key(certs_dir.join("client.key").to_string_lossy().to_string());
1057        let pool = MySqlPool::connect_with(opts).await.unwrap();
1058        sqlx::query("SELECT 1").execute(&pool).await.unwrap();
1059    }
1060
1061    #[tokio::test]
1062    async fn test_mysql_with_tls_verify_ident() {
1063        common_telemetry::init_default_ut_logging();
1064        maybe_skip_mysql_integration_test!();
1065        ensure_rustls_crypto_provider_installed().unwrap();
1066        let endpoint = std::env::var("GT_MYSQL_ENDPOINTS").unwrap();
1067        let certs_dir = test_certs_dir();
1068
1069        let opts = endpoint
1070            .parse::<MySqlConnectOptions>()
1071            .unwrap()
1072            .ssl_mode(MySqlSslMode::VerifyIdentity)
1073            .ssl_ca(certs_dir.join("root.crt").to_string_lossy().to_string())
1074            .ssl_client_cert(certs_dir.join("client.crt").to_string_lossy().to_string())
1075            .ssl_client_key(certs_dir.join("client.key").to_string_lossy().to_string());
1076        let pool = MySqlPool::connect_with(opts).await.unwrap();
1077        sqlx::query("SELECT 1").execute(&pool).await.unwrap();
1078    }
1079}
common_meta/kv_backend/rds/mysql.rs

common_meta/kv_backend/rds/
mysql.rs
