common_meta/
kv_backend.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::any::Any;
16use std::sync::Arc;
17
18use async_trait::async_trait;
19use common_error::ext::ErrorExt;
20pub use txn::TxnService;
21
22use crate::error::Error;
23use crate::kv_backend::txn::{Txn, TxnOpResponse};
24use crate::rpc::KeyValue;
25use crate::rpc::store::{
26    BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest,
27    BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest,
28    DeleteRangeResponse, PutRequest, PutResponse, RangeRequest, RangeResponse,
29};
30
31pub mod chroot;
32pub mod etcd;
33pub mod memory;
34#[cfg(any(feature = "mysql_kvbackend", feature = "pg_kvbackend"))]
35pub mod rds;
36pub mod read_only;
37pub mod test;
38#[cfg(any(test, feature = "testing"))]
39pub mod test_util;
40pub mod txn;
41pub mod util;
42pub type KvBackendRef<E = Error> = Arc<dyn KvBackend<Error = E> + Send + Sync>;
43
44/// A [KvBackend] holding resources that must be released explicitly on shutdown, because
45/// components sharing it may keep it alive until the process exits.
46///
47/// Wrappers of a closable backend must forward [ClosableKvBackend::close] to it.
48#[async_trait]
49pub trait ClosableKvBackend: KvBackend<Error = Error> {
50    /// Releases the backend's resources; any later use fails.
51    async fn close(&self) -> Result<(), Error>;
52}
53
54pub type ClosableKvBackendRef = Arc<dyn ClosableKvBackend + Send + Sync>;
55
56#[cfg(any(feature = "pg_kvbackend", feature = "mysql_kvbackend"))]
57// The default meta table name, default is "greptime_metakv".
58pub const DEFAULT_META_TABLE_NAME: &str = "greptime_metakv";
59
60#[cfg(any(feature = "pg_kvbackend", feature = "mysql_kvbackend"))]
61// The default lock id for election, default is 1.
62pub const DEFAULT_META_ELECTION_LOCK_ID: u64 = 1;
63
64#[async_trait]
65pub trait KvBackend: TxnService
66where
67    Self::Error: ErrorExt,
68{
69    fn name(&self) -> &str;
70
71    fn as_any(&self) -> &dyn Any;
72
73    async fn range(&self, req: RangeRequest) -> Result<RangeResponse, Self::Error>;
74
75    async fn put(&self, req: PutRequest) -> Result<PutResponse, Self::Error>;
76
77    async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse, Self::Error>;
78
79    async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse, Self::Error>;
80
81    async fn delete_range(
82        &self,
83        req: DeleteRangeRequest,
84    ) -> Result<DeleteRangeResponse, Self::Error>;
85
86    async fn batch_delete(
87        &self,
88        req: BatchDeleteRequest,
89    ) -> Result<BatchDeleteResponse, Self::Error>;
90
91    // The following methods are implemented based on the above methods,
92    // and a higher-level interface is provided for to simplify usage.
93
94    async fn get(&self, key: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
95        let req = RangeRequest::new().with_key(key.to_vec());
96        let mut resp = self.range(req).await?;
97        Ok(if resp.kvs.is_empty() {
98            None
99        } else {
100            Some(resp.kvs.remove(0))
101        })
102    }
103
104    /// CAS: Compares the value at the key with the given value, and if they are
105    /// equal, puts the new value at the key.
106    async fn compare_and_put(
107        &self,
108        req: CompareAndPutRequest,
109    ) -> Result<CompareAndPutResponse, Self::Error> {
110        let CompareAndPutRequest { key, expect, value } = req;
111        let txn = if expect.is_empty() {
112            Txn::put_if_not_exists(key, value)
113        } else {
114            Txn::compare_and_put(key, expect, value)
115        };
116        let txn_res = self.txn(txn).await?;
117
118        let success = txn_res.succeeded;
119        // The response is guaranteed to have at most one element.
120        let op_res = txn_res.responses.into_iter().next();
121        let prev_kv = match op_res {
122            Some(TxnOpResponse::ResponsePut(res)) => res.prev_kv,
123            Some(TxnOpResponse::ResponseGet(res)) => res.kvs.into_iter().next(),
124            Some(TxnOpResponse::ResponseDelete(res)) => res.prev_kvs.into_iter().next(),
125            None => None,
126        };
127
128        Ok(CompareAndPutResponse { success, prev_kv })
129    }
130
131    /// Puts a value at a key. If `if_not_exists` is `true`, the operation
132    /// ensures the key does not exist before applying the PUT operation.
133    /// Otherwise, it simply applies the PUT operation without checking for
134    /// the key's existence.
135    async fn put_conditionally(
136        &self,
137        key: Vec<u8>,
138        value: Vec<u8>,
139        if_not_exists: bool,
140    ) -> Result<bool, Self::Error> {
141        let success = if if_not_exists {
142            let req = CompareAndPutRequest::new()
143                .with_key(key)
144                .with_expect(vec![])
145                .with_value(value);
146            let res = self.compare_and_put(req).await?;
147            res.success
148        } else {
149            let req = PutRequest::new().with_key(key).with_value(value);
150            self.put(req).await?;
151            true
152        };
153
154        Ok(success)
155    }
156
157    /// Check if the key exists, not returning the value.
158    /// If the value is large, this method is more efficient than `get`.
159    async fn exists(&self, key: &[u8]) -> Result<bool, Self::Error> {
160        let req = RangeRequest::new().with_key(key.to_vec()).with_keys_only();
161        let resp = self.range(req).await?;
162        Ok(!resp.kvs.is_empty())
163    }
164
165    /// Returns previous key-value pair if `prev_kv` is `true`.
166    async fn delete(&self, key: &[u8], prev_kv: bool) -> Result<Option<KeyValue>, Self::Error> {
167        let mut req = DeleteRangeRequest::new().with_key(key.to_vec());
168        if prev_kv {
169            req = req.with_prev_kv();
170        }
171
172        let resp = self.delete_range(req).await?;
173
174        if prev_kv {
175            Ok(resp.prev_kvs.into_iter().next())
176        } else {
177            Ok(None)
178        }
179    }
180}
181
182pub trait ResettableKvBackend: KvBackend
183where
184    Self::Error: ErrorExt,
185{
186    fn reset(&self);
187
188    /// Upcast as `KvBackendRef`. Since https://github.com/rust-lang/rust/issues/65991 is not yet stable.
189    fn as_kv_backend_ref(self: Arc<Self>) -> KvBackendRef<Self::Error>;
190}
191
192pub type ResettableKvBackendRef<E = Error> = Arc<dyn ResettableKvBackend<Error = E> + Send + Sync>;
common_meta/kv_backend.rs

common_meta/
kv_backend.rs
