Skip to main content

linera_views/backends/
scylla_db.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Implements [`crate::store::KeyValueStore`] for the ScyllaDB database.
5//!
6//! The current connection is done via a Session and a corresponding primary key called
7//! "namespace". The maximum number of concurrent queries is controlled by
8//! `max_concurrent_queries`.
9
10use std::{
11    collections::{BTreeSet, HashMap},
12    ops::Deref,
13    sync::{
14        atomic::{AtomicI64, Ordering},
15        Arc,
16    },
17    time::{SystemTime, UNIX_EPOCH},
18};
19
20use async_lock::{Semaphore, SemaphoreGuard};
21use futures::{future::join_all, StreamExt as _};
22use linera_base::{ensure, util::future::FutureSyncExt as _};
23use scylla::{
24    client::{
25        execution_profile::{ExecutionProfile, ExecutionProfileHandle},
26        session::Session,
27        session_builder::SessionBuilder,
28    },
29    deserialize::{DeserializationError, TypeCheckError},
30    errors::{
31        DbError, ExecutionError, IntoRowsResultError, NewSessionError, NextPageError, NextRowError,
32        PagerExecutionError, PrepareError, RequestAttemptError, RequestError, RowsError,
33    },
34    policies::{
35        load_balancing::{DefaultPolicy, LoadBalancingPolicy},
36        retry::DefaultRetryPolicy,
37    },
38    response::PagingState,
39    statement::{batch::BatchType, prepared::PreparedStatement, Consistency},
40    value::CqlValue,
41};
42use serde::{Deserialize, Serialize};
43use thiserror::Error;
44
45#[cfg(with_metrics)]
46use crate::metering::MeteredDatabase;
47#[cfg(with_testing)]
48use crate::store::TestKeyValueDatabase;
49use crate::{
50    batch::{SimpleUnorderedBatch, UnorderedBatch},
51    common::{get_uleb128_size, get_upper_bound_option},
52    journaling::{JournalingError, JournalingKeyValueDatabase},
53    lru_caching::{LruCachingConfig, LruCachingDatabase},
54    store::{
55        DirectWritableKeyValueStore, KeyValueDatabase, KeyValueStoreError, ReadableKeyValueStore,
56        WithError,
57    },
58    value_splitting::{ValueSplittingDatabase, ValueSplittingError},
59};
60
61/// Fundamental constant in ScyllaDB: The maximum size of a multi keys query
62/// The limit is in reality 100. But we need one entry for the root key.
63const MAX_MULTI_KEYS: usize = 100 - 1;
64
65/// The maximal size of an operation on ScyllaDB seems to be 16 MiB
66/// https://www.scylladb.com/2019/03/27/best-practices-for-scylla-applications/
67/// "There is a hard limit at 16 MiB, and nothing bigger than that can arrive at once
68///  at the database at any particular time"
69/// So, we set up the maximal size of 16 MiB - 10 KiB for the values and 10 KiB for the keys
70/// We also arbitrarily decrease the size by 4000 bytes because an amount of size is
71/// taken internally by the database.
72const RAW_MAX_VALUE_SIZE: usize = 16 * 1024 * 1024 - 10 * 1024 - 4000;
73const MAX_KEY_SIZE: usize = 10 * 1024;
74const MAX_BATCH_TOTAL_SIZE: usize = RAW_MAX_VALUE_SIZE + MAX_KEY_SIZE;
75
76/// The `RAW_MAX_VALUE_SIZE` is the maximum size on the ScyllaDB storage.
77/// However, the value being written can also be the serialization of a `SimpleUnorderedBatch`
78/// Therefore the actual `MAX_VALUE_SIZE` is lower.
79/// At the maximum the key size is 1024 bytes (see below) and we pack just one entry.
80/// So if the key has 1024 bytes this gets us the inequality
81/// `1 + 1 + 1 + serialized_size(MAX_KEY_SIZE)? + serialized_size(x)? <= RAW_MAX_VALUE_SIZE`.
82/// and so this simplifies to `1 + 1 + 1 + (2 + 10240) + (4 + x) <= RAW_MAX_VALUE_SIZE`
83/// Note on the above formula:
84/// * We write 4 because `get_uleb128_size(RAW_MAX_VALUE_SIZE) = 4)`
85/// * We write `1 + 1 + 1`  because the `UnorderedBatch` has three entries.
86///
87/// This gets us to a maximal value of 16752727.
88const VISIBLE_MAX_VALUE_SIZE: usize = RAW_MAX_VALUE_SIZE
89    - MAX_KEY_SIZE
90    - get_uleb128_size(RAW_MAX_VALUE_SIZE)
91    - get_uleb128_size(MAX_KEY_SIZE)
92    - 3;
93
94/// The constant 14000 is an empirical constant that was found to be necessary
95/// to make the ScyllaDB system work. We have not been able to find this or
96/// a similar constant in the source code or the documentation.
97/// An experimental approach gets us that 14796 is the latest value that is
98/// correct.
99const MAX_BATCH_SIZE: usize = 5000;
100
101/// The keyspace to use for the ScyllaDB database.
102const KEYSPACE: &str = "kv";
103
104/// The client for ScyllaDB:
105/// * The session allows to pass queries
106/// * The namespace that is being assigned to the database
107/// * The prepared queries used for implementing the features of `KeyValueStore`.
108struct ScyllaDbClient {
109    session: Session,
110    namespace: String,
111    read_value: PreparedStatement,
112    read_writetime: PreparedStatement,
113    contains_key: PreparedStatement,
114    write_batch_delete_prefix_unbounded: PreparedStatement,
115    write_batch_delete_prefix_bounded: PreparedStatement,
116    write_batch_deletion: PreparedStatement,
117    write_batch_insertion: PreparedStatement,
118    // Variants carrying an explicit `USING TIMESTAMP ?` marker, used by the
119    // single-batch exclusive-mode write path (`write_batch_exclusive`).
120    write_batch_delete_prefix_unbounded_ts: PreparedStatement,
121    write_batch_delete_prefix_bounded_ts: PreparedStatement,
122    write_batch_deletion_ts: PreparedStatement,
123    write_batch_insertion_ts: PreparedStatement,
124    find_keys_by_prefix_unbounded: PreparedStatement,
125    find_keys_by_prefix_bounded: PreparedStatement,
126    find_key_values_by_prefix_unbounded: PreparedStatement,
127    find_key_values_by_prefix_bounded: PreparedStatement,
128    multi_key_values: papaya::HashMap<usize, PreparedStatement>,
129    multi_keys: papaya::HashMap<usize, PreparedStatement>,
130}
131
132impl ScyllaDbClient {
133    async fn new(session: Session, namespace: &str) -> Result<Self, ScyllaDbStoreInternalError> {
134        let namespace = namespace.to_string();
135        let read_value = session
136            .prepare(format!(
137                "SELECT v FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k = ?"
138            ))
139            .await?;
140
141        let read_writetime = session
142            .prepare(format!(
143                "SELECT WRITETIME(v) FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k = ?"
144            ))
145            .await?;
146
147        let contains_key = session
148            .prepare(format!(
149                "SELECT root_key FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k = ?"
150            ))
151            .await?;
152
153        let write_batch_delete_prefix_unbounded = session
154            .prepare(format!(
155                "DELETE FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k >= ?"
156            ))
157            .await?;
158
159        let write_batch_delete_prefix_bounded = session
160            .prepare(format!(
161                "DELETE FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k >= ? AND k < ?"
162            ))
163            .await?;
164
165        let write_batch_deletion = session
166            .prepare(format!(
167                "DELETE FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k = ?"
168            ))
169            .await?;
170
171        let write_batch_insertion = session
172            .prepare(format!(
173                "INSERT INTO {KEYSPACE}.\"{namespace}\" (root_key, k, v) VALUES (?, ?, ?)"
174            ))
175            .await?;
176
177        // Timestamped variants used by the single-batch exclusive-mode path. The
178        // explicit `USING TIMESTAMP ?` lets prefix-deletions (`T`) and the
179        // insertions/deletions (`T + 1`) share one atomic batch without the range
180        // tombstone shadowing the inserts.
181        let write_batch_delete_prefix_unbounded_ts = session
182            .prepare(format!(
183                "DELETE FROM {KEYSPACE}.\"{namespace}\" USING TIMESTAMP ? WHERE root_key = ? AND k >= ?"
184            ))
185            .await?;
186
187        let write_batch_delete_prefix_bounded_ts = session
188            .prepare(format!(
189                "DELETE FROM {KEYSPACE}.\"{namespace}\" USING TIMESTAMP ? \
190                 WHERE root_key = ? AND k >= ? AND k < ?"
191            ))
192            .await?;
193
194        let write_batch_deletion_ts = session
195            .prepare(format!(
196                "DELETE FROM {KEYSPACE}.\"{namespace}\" USING TIMESTAMP ? WHERE root_key = ? AND k = ?"
197            ))
198            .await?;
199
200        let write_batch_insertion_ts = session
201            .prepare(format!(
202                "INSERT INTO {KEYSPACE}.\"{namespace}\" (root_key, k, v) VALUES (?, ?, ?) \
203                 USING TIMESTAMP ?"
204            ))
205            .await?;
206
207        let find_keys_by_prefix_unbounded = session
208            .prepare(format!(
209                "SELECT k FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k >= ?"
210            ))
211            .await?;
212
213        let find_keys_by_prefix_bounded = session
214            .prepare(format!(
215                "SELECT k FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k >= ? AND k < ?"
216            ))
217            .await?;
218
219        let find_key_values_by_prefix_unbounded = session
220            .prepare(format!(
221                "SELECT k,v FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k >= ?"
222            ))
223            .await?;
224
225        let find_key_values_by_prefix_bounded = session
226            .prepare(format!(
227                "SELECT k,v FROM {KEYSPACE}.\"{namespace}\" WHERE root_key = ? AND k >= ? AND k < ?"
228            ))
229            .await?;
230
231        Ok(Self {
232            session,
233            namespace,
234            read_value,
235            read_writetime,
236            contains_key,
237            write_batch_delete_prefix_unbounded,
238            write_batch_delete_prefix_bounded,
239            write_batch_deletion,
240            write_batch_insertion,
241            write_batch_delete_prefix_unbounded_ts,
242            write_batch_delete_prefix_bounded_ts,
243            write_batch_deletion_ts,
244            write_batch_insertion_ts,
245            find_keys_by_prefix_unbounded,
246            find_keys_by_prefix_bounded,
247            find_key_values_by_prefix_unbounded,
248            find_key_values_by_prefix_bounded,
249            multi_key_values: papaya::HashMap::new(),
250            multi_keys: papaya::HashMap::new(),
251        })
252    }
253
254    fn build_default_policy() -> Arc<dyn LoadBalancingPolicy> {
255        DefaultPolicy::builder().token_aware(true).build()
256    }
257
258    fn build_default_execution_profile_handle(
259        policy: Arc<dyn LoadBalancingPolicy>,
260    ) -> ExecutionProfileHandle {
261        let default_profile = ExecutionProfile::builder()
262            .load_balancing_policy(policy)
263            .retry_policy(Arc::new(DefaultRetryPolicy::new()))
264            .consistency(Consistency::LocalQuorum)
265            .build();
266        default_profile.into_handle()
267    }
268
269    async fn build_default_session(uri: &str) -> Result<Session, ScyllaDbStoreInternalError> {
270        // This explicitly sets a lot of default parameters for clarity and for making future changes
271        // easier.
272        SessionBuilder::new()
273            .known_node(uri)
274            .default_execution_profile_handle(Self::build_default_execution_profile_handle(
275                Self::build_default_policy(),
276            ))
277            .build()
278            .boxed_sync()
279            .await
280            .map_err(Into::into)
281    }
282
283    async fn get_multi_key_values_statement(
284        &self,
285        num_markers: usize,
286    ) -> Result<PreparedStatement, ScyllaDbStoreInternalError> {
287        if let Some(prepared_statement) = self.multi_key_values.pin().get(&num_markers) {
288            return Ok(prepared_statement.clone());
289        }
290        let markers = std::iter::repeat_n("?", num_markers)
291            .collect::<Vec<_>>()
292            .join(",");
293        let prepared_statement = self
294            .session
295            .prepare(format!(
296                "SELECT k,v FROM {}.\"{}\" WHERE root_key = ? AND k IN ({})",
297                KEYSPACE, self.namespace, markers
298            ))
299            .await?;
300        self.multi_key_values
301            .pin()
302            .insert(num_markers, prepared_statement.clone());
303        Ok(prepared_statement)
304    }
305
306    async fn get_multi_keys_statement(
307        &self,
308        num_markers: usize,
309    ) -> Result<PreparedStatement, ScyllaDbStoreInternalError> {
310        if let Some(prepared_statement) = self.multi_keys.pin().get(&num_markers) {
311            return Ok(prepared_statement.clone());
312        };
313        let markers = std::iter::repeat_n("?", num_markers)
314            .collect::<Vec<_>>()
315            .join(",");
316        let prepared_statement = self
317            .session
318            .prepare(format!(
319                "SELECT k FROM {}.\"{}\" WHERE root_key = ? AND k IN ({})",
320                KEYSPACE, self.namespace, markers
321            ))
322            .await?;
323        self.multi_keys
324            .pin()
325            .insert(num_markers, prepared_statement.clone());
326        Ok(prepared_statement)
327    }
328
329    fn check_key_size(key: &[u8]) -> Result<(), ScyllaDbStoreInternalError> {
330        ensure!(
331            key.len() <= MAX_KEY_SIZE,
332            ScyllaDbStoreInternalError::KeyTooLong
333        );
334        Ok(())
335    }
336
337    fn check_value_size(value: &[u8]) -> Result<(), ScyllaDbStoreInternalError> {
338        ensure!(
339            value.len() <= RAW_MAX_VALUE_SIZE,
340            ScyllaDbStoreInternalError::ValueTooLong
341        );
342        Ok(())
343    }
344
345    /// Validates a key supplied by a caller's batch. Besides the size limit, the
346    /// key must be non-empty: the empty (zero-length) key is `WRITETIME_SENTINEL_KEY`,
347    /// reserved for the per-store timestamp sentinel that exclusive mode writes
348    /// internally. Prefix scans now deliberately hide that key, so any caller
349    /// content stored there would be silently invisible to reads.
350    fn check_batch_key(key: &[u8]) -> Result<(), ScyllaDbStoreInternalError> {
351        Self::check_key_size(key)?;
352        ensure!(!key.is_empty(), ScyllaDbStoreInternalError::ZeroLengthKey);
353        Ok(())
354    }
355
356    fn check_batch_len(batch: &UnorderedBatch) -> Result<(), ScyllaDbStoreInternalError> {
357        ensure!(
358            batch.len() <= MAX_BATCH_SIZE,
359            ScyllaDbStoreInternalError::BatchTooLong
360        );
361        Ok(())
362    }
363
364    async fn read_value_internal(
365        &self,
366        root_key: &[u8],
367        key: Vec<u8>,
368    ) -> Result<Option<Vec<u8>>, ScyllaDbStoreInternalError> {
369        Self::check_key_size(&key)?;
370        let session = &self.session;
371        // Read the value of a key
372        let values = (root_key.to_vec(), key);
373
374        let (result, _) = session
375            .execute_single_page(&self.read_value, &values, PagingState::start())
376            .await
377            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
378        let rows = result.into_rows_result()?;
379        let mut rows = rows.rows::<(Vec<u8>,)>()?;
380        Ok(match rows.next() {
381            Some(row) => Some(row?.0),
382            None => None,
383        })
384    }
385
386    fn get_occurrences_map(
387        keys: Vec<Vec<u8>>,
388    ) -> Result<HashMap<Vec<u8>, Vec<usize>>, ScyllaDbStoreInternalError> {
389        let mut map = HashMap::<Vec<u8>, Vec<usize>>::new();
390        for (i_key, key) in keys.into_iter().enumerate() {
391            Self::check_key_size(&key)?;
392            map.entry(key).or_default().push(i_key);
393        }
394        Ok(map)
395    }
396
397    async fn read_multi_values_internal(
398        &self,
399        root_key: &[u8],
400        keys: Vec<Vec<u8>>,
401    ) -> Result<Vec<Option<Vec<u8>>>, ScyllaDbStoreInternalError> {
402        let mut values = vec![None; keys.len()];
403        let map = Self::get_occurrences_map(keys)?;
404        let statement = self.get_multi_key_values_statement(map.len()).await?;
405        let mut inputs = vec![root_key.to_vec()];
406        inputs.extend(map.keys().cloned());
407        let mut rows = Box::pin(self.session.execute_iter(statement, &inputs))
408            .await?
409            .rows_stream::<(Vec<u8>, Vec<u8>)>()?;
410
411        while let Some(row) = rows.next().await {
412            let (key, value) = row?;
413            if let Some((&last, rest)) = map[&key].split_last() {
414                for position in rest {
415                    values[*position] = Some(value.clone());
416                }
417                values[last] = Some(value);
418            }
419        }
420        Ok(values)
421    }
422
423    async fn contains_keys_internal(
424        &self,
425        root_key: &[u8],
426        keys: Vec<Vec<u8>>,
427    ) -> Result<Vec<bool>, ScyllaDbStoreInternalError> {
428        let mut values = vec![false; keys.len()];
429        let map = Self::get_occurrences_map(keys)?;
430        let statement = self.get_multi_keys_statement(map.len()).await?;
431        let mut inputs = vec![root_key.to_vec()];
432        inputs.extend(map.keys().cloned());
433        let mut rows = Box::pin(self.session.execute_iter(statement, &inputs))
434            .await?
435            .rows_stream::<(Vec<u8>,)>()?;
436
437        while let Some(row) = rows.next().await {
438            let (key,) = row?;
439            for i_key in &map[&key] {
440                values[*i_key] = true;
441            }
442        }
443
444        Ok(values)
445    }
446
447    async fn contains_key_internal(
448        &self,
449        root_key: &[u8],
450        key: Vec<u8>,
451    ) -> Result<bool, ScyllaDbStoreInternalError> {
452        Self::check_key_size(&key)?;
453        let session = &self.session;
454        // Read the value of a key
455        let values = (root_key.to_vec(), key);
456
457        let (result, _) = session
458            .execute_single_page(&self.contains_key, &values, PagingState::start())
459            .await
460            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
461        let rows = result.into_rows_result()?;
462        let mut rows = rows.rows::<(Vec<u8>,)>()?;
463        Ok(rows.next().is_some())
464    }
465
466    /// Reads the write-time of a single row in microseconds since Unix epoch,
467    /// returning `None` if the row does not exist or carries no live value.
468    async fn read_writetime_internal(
469        &self,
470        root_key: &[u8],
471        key: Vec<u8>,
472    ) -> Result<Option<i64>, ScyllaDbStoreInternalError> {
473        Self::check_key_size(&key)?;
474        let session = &self.session;
475        let values = (root_key.to_vec(), key);
476        let (result, _) = session
477            .execute_single_page(&self.read_writetime, &values, PagingState::start())
478            .await
479            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
480        let rows = result.into_rows_result()?;
481        let mut rows = rows.rows::<(Option<i64>,)>()?;
482        Ok(match rows.next() {
483            Some(row) => row?.0,
484            None => None,
485        })
486    }
487
488    /// Issues an unlogged batch that contains only prefix-delete statements,
489    /// letting the coordinator assign the write timestamp.
490    async fn write_batch_prefix_deletes(
491        &self,
492        root_key: &[u8],
493        key_prefix_deletions: Vec<Vec<u8>>,
494    ) -> Result<(), ScyllaDbStoreInternalError> {
495        if key_prefix_deletions.is_empty() {
496            return Ok(());
497        }
498        let session = &self.session;
499        let mut batch_query = scylla::statement::batch::Batch::new(BatchType::Unlogged);
500        let mut batch_values = Vec::new();
501        let q_unbounded = &self.write_batch_delete_prefix_unbounded;
502        let q_bounded = &self.write_batch_delete_prefix_bounded;
503        for key_prefix in key_prefix_deletions {
504            Self::check_key_size(&key_prefix)?;
505            match get_upper_bound_option(&key_prefix) {
506                None => {
507                    batch_values.push(vec![root_key.to_vec(), key_prefix]);
508                    batch_query.append_statement(q_unbounded.clone());
509                }
510                Some(upper_bound) => {
511                    batch_values.push(vec![root_key.to_vec(), key_prefix, upper_bound]);
512                    batch_query.append_statement(q_bounded.clone());
513                }
514            }
515        }
516        session
517            .batch(&batch_query, batch_values)
518            .await
519            .map_err(ScyllaDbStoreInternalError::WriteBatchExecutionError)?;
520        Ok(())
521    }
522
523    /// Issues an unlogged batch containing the single-key deletions and the
524    /// insertions, letting the coordinator assign the write timestamp.
525    async fn write_simple_batch(
526        &self,
527        root_key: &[u8],
528        batch: SimpleUnorderedBatch,
529    ) -> Result<(), ScyllaDbStoreInternalError> {
530        if batch.deletions.is_empty() && batch.insertions.is_empty() {
531            return Ok(());
532        }
533        let session = &self.session;
534        let mut batch_query = scylla::statement::batch::Batch::new(BatchType::Unlogged);
535        let mut batch_values = Vec::new();
536        let q_deletion = &self.write_batch_deletion;
537        for key in batch.deletions {
538            Self::check_batch_key(&key)?;
539            batch_values.push(vec![root_key.to_vec(), key]);
540            batch_query.append_statement(q_deletion.clone());
541        }
542        let q_insertion = &self.write_batch_insertion;
543        for (key, value) in batch.insertions {
544            Self::check_batch_key(&key)?;
545            Self::check_value_size(&value)?;
546            batch_values.push(vec![root_key.to_vec(), key, value]);
547            batch_query.append_statement(q_insertion.clone());
548        }
549        session
550            .batch(&batch_query, batch_values)
551            .await
552            .map_err(ScyllaDbStoreInternalError::WriteBatchExecutionError)?;
553        Ok(())
554    }
555
556    /// Issues the whole write as a single atomic unlogged batch, used in
557    /// exclusive mode. Every statement carries an explicit `USING TIMESTAMP`:
558    /// the prefix-deletions use `t`, while the single-key deletions, the
559    /// insertions, and the sentinel write use `t + 1`. The higher timestamp on
560    /// the data ensures a range tombstone never shadows an insertion belonging
561    /// to the same logical batch (at equal timestamps, dead cells win over live
562    /// cells). Because the intended ordering is fixed by these timestamps rather
563    /// than by send order, the prefix-deletions and the data can — and must —
564    /// share one batch, preserving the atomicity that `write_batch` callers rely
565    /// on. The sentinel write at `WRITETIME_SENTINEL_KEY` lets a future process
566    /// recover this store's timestamp floor (see `ensure_ts_seeded`).
567    async fn write_batch_exclusive(
568        &self,
569        root_key: &[u8],
570        batch: UnorderedBatch,
571        t: i64,
572    ) -> Result<(), ScyllaDbStoreInternalError> {
573        let UnorderedBatch {
574            key_prefix_deletions,
575            simple_unordered_batch:
576                SimpleUnorderedBatch {
577                    deletions,
578                    insertions,
579                },
580        } = batch;
581        let session = &self.session;
582        let mut batch_query = scylla::statement::batch::Batch::new(BatchType::Unlogged);
583        let mut batch_values = Vec::new();
584
585        // Prefix-deletions at timestamp `t`.
586        for key_prefix in key_prefix_deletions {
587            Self::check_key_size(&key_prefix)?;
588            match get_upper_bound_option(&key_prefix) {
589                None => {
590                    batch_values.push(vec![
591                        CqlValue::BigInt(t),
592                        CqlValue::Blob(root_key.to_vec()),
593                        CqlValue::Blob(key_prefix),
594                    ]);
595                    batch_query
596                        .append_statement(self.write_batch_delete_prefix_unbounded_ts.clone());
597                }
598                Some(upper_bound) => {
599                    batch_values.push(vec![
600                        CqlValue::BigInt(t),
601                        CqlValue::Blob(root_key.to_vec()),
602                        CqlValue::Blob(key_prefix),
603                        CqlValue::Blob(upper_bound),
604                    ]);
605                    batch_query.append_statement(self.write_batch_delete_prefix_bounded_ts.clone());
606                }
607            }
608        }
609
610        // Single-key deletions, insertions, and the sentinel at timestamp `t + 1`.
611        let t_data = t + 1;
612        for key in deletions {
613            Self::check_batch_key(&key)?;
614            batch_values.push(vec![
615                CqlValue::BigInt(t_data),
616                CqlValue::Blob(root_key.to_vec()),
617                CqlValue::Blob(key),
618            ]);
619            batch_query.append_statement(self.write_batch_deletion_ts.clone());
620        }
621        for (key, value) in insertions {
622            Self::check_batch_key(&key)?;
623            Self::check_value_size(&value)?;
624            batch_values.push(vec![
625                CqlValue::Blob(root_key.to_vec()),
626                CqlValue::Blob(key),
627                CqlValue::Blob(value),
628                CqlValue::BigInt(t_data),
629            ]);
630            batch_query.append_statement(self.write_batch_insertion_ts.clone());
631        }
632        batch_values.push(vec![
633            CqlValue::Blob(root_key.to_vec()),
634            CqlValue::Blob(WRITETIME_SENTINEL_KEY.to_vec()),
635            CqlValue::Blob(Vec::new()),
636            CqlValue::BigInt(t_data),
637        ]);
638        batch_query.append_statement(self.write_batch_insertion_ts.clone());
639
640        session
641            .batch(&batch_query, batch_values)
642            .await
643            .map_err(ScyllaDbStoreInternalError::WriteBatchExecutionError)?;
644        Ok(())
645    }
646
647    async fn find_keys_by_prefix_internal(
648        &self,
649        root_key: &[u8],
650        key_prefix: Vec<u8>,
651    ) -> Result<Vec<Vec<u8>>, ScyllaDbStoreInternalError> {
652        Self::check_key_size(&key_prefix)?;
653        let session = &self.session;
654        // Read the value of a key
655        let len = key_prefix.len();
656        let query_unbounded = &self.find_keys_by_prefix_unbounded;
657        let query_bounded = &self.find_keys_by_prefix_bounded;
658        let rows = match get_upper_bound_option(&key_prefix) {
659            None => {
660                let values = (root_key.to_vec(), key_prefix.clone());
661                Box::pin(session.execute_iter(query_unbounded.clone(), values)).await?
662            }
663            Some(upper_bound) => {
664                let values = (root_key.to_vec(), key_prefix.clone(), upper_bound);
665                Box::pin(session.execute_iter(query_bounded.clone(), values)).await?
666            }
667        };
668        let mut rows = rows.rows_stream::<(Vec<u8>,)>()?;
669        let mut keys = Vec::new();
670        while let Some(row) = rows.next().await {
671            let (key,) = row?;
672            // Skip the reserved timestamp sentinel (exclusive mode writes it at the
673            // empty clustering key). It is an internal implementation detail and must
674            // not surface to callers; it can only match an empty-prefix scan.
675            if key == WRITETIME_SENTINEL_KEY {
676                continue;
677            }
678            let short_key = key[len..].to_vec();
679            keys.push(short_key);
680        }
681        Ok(keys)
682    }
683
684    async fn find_key_values_by_prefix_internal(
685        &self,
686        root_key: &[u8],
687        key_prefix: Vec<u8>,
688    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, ScyllaDbStoreInternalError> {
689        Self::check_key_size(&key_prefix)?;
690        let session = &self.session;
691        // Read the value of a key
692        let len = key_prefix.len();
693        let query_unbounded = &self.find_key_values_by_prefix_unbounded;
694        let query_bounded = &self.find_key_values_by_prefix_bounded;
695        let rows = match get_upper_bound_option(&key_prefix) {
696            None => {
697                let values = (root_key.to_vec(), key_prefix.clone());
698                Box::pin(session.execute_iter(query_unbounded.clone(), values)).await?
699            }
700            Some(upper_bound) => {
701                let values = (root_key.to_vec(), key_prefix.clone(), upper_bound);
702                Box::pin(session.execute_iter(query_bounded.clone(), values)).await?
703            }
704        };
705        let mut rows = rows.rows_stream::<(Vec<u8>, Vec<u8>)>()?;
706        let mut key_values = Vec::new();
707        while let Some(row) = rows.next().await {
708            let (key, value) = row?;
709            // Skip the reserved timestamp sentinel; see `find_keys_by_prefix_internal`.
710            if key == WRITETIME_SENTINEL_KEY {
711                continue;
712            }
713            let short_key = key[len..].to_vec();
714            key_values.push((short_key, value));
715        }
716        Ok(key_values)
717    }
718}
719
720/// The client itself and the keeping of the count of active connections.
721#[derive(Clone)]
722pub struct ScyllaDbStoreInternal {
723    store: Arc<ScyllaDbClient>,
724    semaphore: Option<Arc<Semaphore>>,
725    root_key: Vec<u8>,
726    /// Whether this store was opened with `open_exclusive`. When true, `write_batch`
727    /// resolves in-batch prefix/insert collisions via per-statement `USING TIMESTAMP`;
728    /// when false, it splits the batch into two sequential sub-batches with
729    /// server-side timestamps to preserve ordering across writers.
730    is_exclusive: bool,
731    /// Per-partition timestamp floor for exclusive-mode `USING TIMESTAMP` writes.
732    /// Value 0 means unseeded; populated lazily on first write by reading
733    /// `WRITETIME` of a sentinel row. Each batch reserves 2 µs (T and T+1).
734    ts_floor: Arc<AtomicI64>,
735}
736
737/// Database-level connection to ScyllaDB for managing namespaces and partitions.
738#[derive(Clone)]
739pub struct ScyllaDbDatabaseInternal {
740    store: Arc<ScyllaDbClient>,
741    semaphore: Option<Arc<Semaphore>>,
742}
743
744impl WithError for ScyllaDbDatabaseInternal {
745    type Error = ScyllaDbStoreInternalError;
746}
747
748/// The error type for [`ScyllaDbStoreInternal`]
749#[derive(Error, Debug)]
750pub enum ScyllaDbStoreInternalError {
751    /// BCS serialization error.
752    #[error(transparent)]
753    BcsError(#[from] bcs::Error),
754
755    /// A deserialization error
756    #[error(transparent)]
757    DeserializationError(#[from] DeserializationError),
758
759    /// A row error
760    #[error(transparent)]
761    RowsError(#[from] RowsError),
762
763    /// A conversion error in the accessed data
764    #[error(transparent)]
765    IntoRowsResultError(#[from] IntoRowsResultError),
766
767    /// A type check error
768    #[error(transparent)]
769    TypeCheckError(#[from] TypeCheckError),
770
771    /// A pager execution error
772    #[error(transparent)]
773    PagerExecutionError(#[from] PagerExecutionError),
774
775    /// A prepare error
776    #[error(transparent)]
777    PrepareError(#[from] PrepareError),
778
779    /// An execution error during a query (except write-batch).
780    #[error(transparent)]
781    ExecutionError(ExecutionError),
782
783    /// An execution error during a write-batch operation.
784    #[error(transparent)]
785    WriteBatchExecutionError(ExecutionError),
786
787    /// A session creation error
788    #[error(transparent)]
789    NewSessionError(#[from] NewSessionError),
790
791    /// A next row error in ScyllaDB
792    #[error(transparent)]
793    NextRowError(#[from] NextRowError),
794
795    /// Namespace contains forbidden characters
796    #[error("Namespace contains forbidden characters")]
797    InvalidNamespace,
798
799    /// The key must have at most `MAX_KEY_SIZE` bytes
800    #[error("The key must have at most MAX_KEY_SIZE")]
801    KeyTooLong,
802
803    /// The value must have at most `RAW_MAX_VALUE_SIZE` bytes
804    #[error("The value must have at most RAW_MAX_VALUE_SIZE")]
805    ValueTooLong,
806
807    /// The batch is too long to be written
808    #[error("The batch is too long to be written")]
809    BatchTooLong,
810
811    /// Keys have to be of nonzero length (the empty key is reserved for the
812    /// timestamp sentinel).
813    #[error("The key must be of nonzero length")]
814    ZeroLengthKey,
815}
816
817impl KeyValueStoreError for ScyllaDbStoreInternalError {
818    const BACKEND: &'static str = "scylla_db";
819
820    fn must_reload_view(&self) -> bool {
821        // Errors (notably timeouts) during a `write_batch` may leave the view in a
822        // undetermined state where the batch may or may not have happened.
823        matches!(self, Self::WriteBatchExecutionError(_))
824    }
825}
826
827impl WithError for ScyllaDbStoreInternal {
828    type Error = ScyllaDbStoreInternalError;
829}
830
831impl ReadableKeyValueStore for ScyllaDbStoreInternal {
832    const MAX_KEY_SIZE: usize = MAX_KEY_SIZE;
833
834    fn root_key(&self) -> Result<Vec<u8>, ScyllaDbStoreInternalError> {
835        Ok(self.root_key[1..].to_vec())
836    }
837
838    async fn read_value_bytes(
839        &self,
840        key: &[u8],
841    ) -> Result<Option<Vec<u8>>, ScyllaDbStoreInternalError> {
842        let store = self.store.deref();
843        let _guard = self.acquire().await;
844        Box::pin(store.read_value_internal(&self.root_key, key.to_vec())).await
845    }
846
847    async fn contains_key(&self, key: &[u8]) -> Result<bool, ScyllaDbStoreInternalError> {
848        let store = self.store.deref();
849        let _guard = self.acquire().await;
850        Box::pin(store.contains_key_internal(&self.root_key, key.to_vec())).await
851    }
852
853    async fn contains_keys(
854        &self,
855        keys: &[Vec<u8>],
856    ) -> Result<Vec<bool>, ScyllaDbStoreInternalError> {
857        if keys.is_empty() {
858            return Ok(Vec::new());
859        }
860        let store = self.store.deref();
861        let _guard = self.acquire().await;
862        let handles = keys
863            .chunks(MAX_MULTI_KEYS)
864            .map(|keys| store.contains_keys_internal(&self.root_key, keys.to_vec()));
865        let results: Vec<_> = join_all(handles)
866            .await
867            .into_iter()
868            .collect::<Result<_, _>>()?;
869        Ok(results.into_iter().flatten().collect())
870    }
871
872    async fn read_multi_values_bytes(
873        &self,
874        keys: &[Vec<u8>],
875    ) -> Result<Vec<Option<Vec<u8>>>, ScyllaDbStoreInternalError> {
876        if keys.is_empty() {
877            return Ok(Vec::new());
878        }
879        let store = self.store.deref();
880        let _guard = self.acquire().await;
881        let handles = keys
882            .chunks(MAX_MULTI_KEYS)
883            .map(|keys| store.read_multi_values_internal(&self.root_key, keys.to_vec()));
884        let results: Vec<_> = join_all(handles)
885            .await
886            .into_iter()
887            .collect::<Result<_, _>>()?;
888        Ok(results.into_iter().flatten().collect())
889    }
890
891    async fn find_keys_by_prefix(
892        &self,
893        key_prefix: &[u8],
894    ) -> Result<Vec<Vec<u8>>, ScyllaDbStoreInternalError> {
895        let store = self.store.deref();
896        let _guard = self.acquire().await;
897        Box::pin(store.find_keys_by_prefix_internal(&self.root_key, key_prefix.to_vec())).await
898    }
899
900    async fn find_key_values_by_prefix(
901        &self,
902        key_prefix: &[u8],
903    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, ScyllaDbStoreInternalError> {
904        let store = self.store.deref();
905        let _guard = self.acquire().await;
906        Box::pin(store.find_key_values_by_prefix_internal(&self.root_key, key_prefix.to_vec()))
907            .await
908    }
909}
910
911impl DirectWritableKeyValueStore for ScyllaDbStoreInternal {
912    const MAX_BATCH_SIZE: usize = MAX_BATCH_SIZE;
913    const MAX_BATCH_TOTAL_SIZE: usize = MAX_BATCH_TOTAL_SIZE;
914    const MAX_VALUE_SIZE: usize = VISIBLE_MAX_VALUE_SIZE;
915
916    // ScyllaDB cannot take a `crate::batch::Batch` directly. Indeed, if a delete is
917    // followed by a write, then the delete takes priority. See the sentence "The first
918    // tie-breaking rule when two cells have the same write timestamp is that dead cells
919    // win over live cells" from
920    // https://github.com/scylladb/scylladb/blob/master/docs/dev/timestamp-conflict-resolution.md
921    //
922    // We therefore order the prefix-deletions strictly before the insertions:
923    //   * In exclusive mode we own the timestamps, so we issue a single atomic CQL
924    //     batch with explicit per-statement `USING TIMESTAMP` (`T` for the
925    //     prefix-deletions, `T + 1` for the data). See `write_batch_exclusive`.
926    //   * In shared mode the coordinator owns the timestamps, so we split the write
927    //     into two sequential CQL batches.
928    type Batch = UnorderedBatch;
929
930    async fn write_batch(&self, batch: Self::Batch) -> Result<(), ScyllaDbStoreInternalError> {
931        let store = self.store.deref();
932        let _guard = self.acquire().await;
933        ScyllaDbClient::check_batch_len(&batch)?;
934        if self.is_exclusive {
935            // A single atomic batch; ordering is pinned by the explicit timestamps.
936            let t = self.next_batch_ts().await?;
937            store.write_batch_exclusive(&self.root_key, batch, t).await
938        } else {
939            store
940                .write_batch_prefix_deletes(&self.root_key, batch.key_prefix_deletions)
941                .await?;
942            store
943                .write_simple_batch(&self.root_key, batch.simple_unordered_batch)
944                .await?;
945            Ok(())
946        }
947    }
948}
949
950impl ScyllaDbStoreInternal {
951    /// Seeds the per-store timestamp floor on first write in exclusive mode.
952    /// Reads `WRITETIME` of this chain's row in the reserved sentinel
953    /// partition (written by every prior exclusive batch). Falls back to the
954    /// current wall clock if the row does not yet exist. Idempotent — only
955    /// the first caller wins the compare-exchange.
956    async fn ensure_ts_seeded(&self) -> Result<(), ScyllaDbStoreInternalError> {
957        if self.ts_floor.load(Ordering::Relaxed) > 0 {
958            return Ok(());
959        }
960        let writetime = self
961            .store
962            .read_writetime_internal(&self.root_key, WRITETIME_SENTINEL_KEY.to_vec())
963            .await?
964            .unwrap_or(0);
965        let now_us = SystemTime::now()
966            .duration_since(UNIX_EPOCH)
967            .ok()
968            .and_then(|d| i64::try_from(d.as_micros()).ok())
969            .unwrap_or(0);
970        // `writetime` is the last batch's `T + 1`, i.e. the highest timestamp it
971        // consumed; that is exactly what `ts_floor` tracks, so seed it directly.
972        let seed = now_us.max(writetime);
973        if self
974            .ts_floor
975            .compare_exchange(0, seed, Ordering::Relaxed, Ordering::Relaxed)
976            .is_err()
977        {
978            // Another caller seeded first; their value wins.
979        }
980        Ok(())
981    }
982
983    /// Returns the base timestamp `T` for the next batch in exclusive mode.
984    /// The batch may also use `T + 1`; the generator advances by 2 per call,
985    /// preserving monotonicity across batches in this process.
986    async fn next_batch_ts(&self) -> Result<i64, ScyllaDbStoreInternalError> {
987        self.ensure_ts_seeded().await?;
988        loop {
989            let prev = self.ts_floor.load(Ordering::Relaxed);
990            let now_us = SystemTime::now()
991                .duration_since(UNIX_EPOCH)
992                .ok()
993                .and_then(|d| i64::try_from(d.as_micros()).ok())
994                .unwrap_or(prev);
995            let next = std::cmp::max(now_us, prev + 1);
996            // The batch uses `next` (`T`) and `next + 1` (`T + 1`); store the latter
997            // so the following batch starts strictly above both.
998            if self
999                .ts_floor
1000                .compare_exchange_weak(prev, next + 1, Ordering::Relaxed, Ordering::Relaxed)
1001                .is_ok()
1002            {
1003                return Ok(next);
1004            }
1005        }
1006    }
1007}
1008
1009// ScyllaDB requires that the keys are non-empty.
1010fn get_big_root_key(root_key: &[u8]) -> Vec<u8> {
1011    let mut big_key = vec![0];
1012    big_key.extend(root_key);
1013    big_key
1014}
1015
1016/// Reserved clustering key inside each chain's partition that holds the
1017/// timestamp sentinel used to seed the per-store client timestamp generator
1018/// in exclusive mode. The empty clustering key is unused by any caller:
1019/// views always write keys prefixed with a tag byte (>= `MIN_VIEW_TAG`),
1020/// and the journaling layer writes 6-byte keys starting with `[0, ...]`.
1021const WRITETIME_SENTINEL_KEY: &[u8] = &[];
1022
1023/// The type for building a new ScyllaDB Key Value Store
1024#[derive(Clone, Debug, Deserialize, Serialize)]
1025pub struct ScyllaDbStoreInternalConfig {
1026    /// The URL to which the requests have to be sent
1027    pub uri: String,
1028    /// Maximum number of concurrent database queries allowed for this client.
1029    pub max_concurrent_queries: Option<usize>,
1030    /// The replication factor.
1031    pub replication_factor: u32,
1032}
1033
1034impl KeyValueDatabase for ScyllaDbDatabaseInternal {
1035    type Config = ScyllaDbStoreInternalConfig;
1036    type Store = ScyllaDbStoreInternal;
1037
1038    fn get_name() -> String {
1039        "scylladb internal".to_string()
1040    }
1041
1042    async fn connect(
1043        config: &Self::Config,
1044        namespace: &str,
1045    ) -> Result<Self, ScyllaDbStoreInternalError> {
1046        Self::check_namespace(namespace)?;
1047        let session = ScyllaDbClient::build_default_session(&config.uri).await?;
1048        let store = ScyllaDbClient::new(session, namespace).await?;
1049        let store = Arc::new(store);
1050        let semaphore = config
1051            .max_concurrent_queries
1052            .map(|n| Arc::new(Semaphore::new(n)));
1053        Ok(Self { store, semaphore })
1054    }
1055
1056    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, ScyllaDbStoreInternalError> {
1057        let store = self.store.clone();
1058        let semaphore = self.semaphore.clone();
1059        let root_key = get_big_root_key(root_key);
1060        Ok(ScyllaDbStoreInternal {
1061            store,
1062            semaphore,
1063            root_key,
1064            is_exclusive: false,
1065            ts_floor: Arc::new(AtomicI64::new(0)),
1066        })
1067    }
1068
1069    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, ScyllaDbStoreInternalError> {
1070        let store = self.store.clone();
1071        let semaphore = self.semaphore.clone();
1072        let root_key = get_big_root_key(root_key);
1073        Ok(ScyllaDbStoreInternal {
1074            store,
1075            semaphore,
1076            root_key,
1077            is_exclusive: true,
1078            ts_floor: Arc::new(AtomicI64::new(0)),
1079        })
1080    }
1081
1082    async fn list_all(config: &Self::Config) -> Result<Vec<String>, ScyllaDbStoreInternalError> {
1083        let session = ScyllaDbClient::build_default_session(&config.uri).await?;
1084        let statement = session
1085            .prepare(format!("DESCRIBE KEYSPACE {KEYSPACE}"))
1086            .await?;
1087        let result = Box::pin(session.execute_iter(statement, &[])).await;
1088        let miss_msg = format!("'{KEYSPACE}' not found in keyspaces");
1089        let result = match result {
1090            Ok(result) => result,
1091            Err(error) => {
1092                let invalid_or_keyspace_not_found = match &error {
1093                    PagerExecutionError::NextPageError(NextPageError::RequestFailure(
1094                        RequestError::LastAttemptError(RequestAttemptError::DbError(db_error, msg)),
1095                    )) => *db_error == DbError::Invalid && msg.as_str() == miss_msg,
1096                    _ => false,
1097                };
1098                if invalid_or_keyspace_not_found {
1099                    return Ok(Vec::new());
1100                } else {
1101                    return Err(ScyllaDbStoreInternalError::PagerExecutionError(error));
1102                }
1103            }
1104        };
1105        let mut namespaces = Vec::new();
1106        let mut rows_stream = result.rows_stream::<(String, String, String, String)>()?;
1107        while let Some(row) = rows_stream.next().await {
1108            let (_, object_kind, name, _) = row?;
1109            if object_kind == "table" {
1110                namespaces.push(name);
1111            }
1112        }
1113        Ok(namespaces)
1114    }
1115
1116    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, ScyllaDbStoreInternalError> {
1117        let statement = self
1118            .store
1119            .session
1120            .prepare(format!(
1121                "SELECT root_key FROM {}.\"{}\" ALLOW FILTERING",
1122                KEYSPACE, self.store.namespace
1123            ))
1124            .await?;
1125
1126        // Execute the query
1127        let rows = Box::pin(self.store.session.execute_iter(statement, &[])).await?;
1128        let mut rows = rows.rows_stream::<(Vec<u8>,)>()?;
1129        let mut root_keys = BTreeSet::new();
1130        while let Some(row) = rows.next().await {
1131            let (root_key,) = row?;
1132            let root_key = root_key[1..].to_vec();
1133            root_keys.insert(root_key);
1134        }
1135        Ok(root_keys.into_iter().collect::<Vec<_>>())
1136    }
1137
1138    async fn delete_all(store_config: &Self::Config) -> Result<(), ScyllaDbStoreInternalError> {
1139        let session = ScyllaDbClient::build_default_session(&store_config.uri).await?;
1140        let statement = session
1141            .prepare(format!("DROP KEYSPACE IF EXISTS {KEYSPACE}"))
1142            .await?;
1143
1144        session
1145            .execute_single_page(&statement, &[], PagingState::start())
1146            .await
1147            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
1148        Ok(())
1149    }
1150
1151    async fn exists(
1152        config: &Self::Config,
1153        namespace: &str,
1154    ) -> Result<bool, ScyllaDbStoreInternalError> {
1155        Self::check_namespace(namespace)?;
1156        let session = ScyllaDbClient::build_default_session(&config.uri).await?;
1157
1158        // We check the way the test can fail. It can fail in different ways.
1159        let result = session
1160            .prepare(format!(
1161                "SELECT root_key FROM {KEYSPACE}.\"{namespace}\" LIMIT 1 ALLOW FILTERING"
1162            ))
1163            .await;
1164
1165        // The missing table translates into a very specific error that we matched
1166        let miss_msg1 = format!("unconfigured table {namespace}");
1167        let miss_msg1 = miss_msg1.as_str();
1168        let miss_msg2 = "Undefined name root_key in selection clause";
1169        let miss_msg3 = format!("Keyspace {KEYSPACE} does not exist");
1170        let Err(error) = result else {
1171            // If OK, then the table exists
1172            return Ok(true);
1173        };
1174        let missing_table = match &error {
1175            PrepareError::AllAttemptsFailed {
1176                first_attempt: RequestAttemptError::DbError(db_error, msg),
1177            } => {
1178                if *db_error != DbError::Invalid {
1179                    false
1180                } else {
1181                    msg.as_str() == miss_msg1
1182                        || msg.as_str() == miss_msg2
1183                        || msg.as_str() == miss_msg3
1184                }
1185            }
1186            _ => false,
1187        };
1188        if missing_table {
1189            Ok(false)
1190        } else {
1191            Err(ScyllaDbStoreInternalError::PrepareError(error))
1192        }
1193    }
1194
1195    async fn create(
1196        config: &Self::Config,
1197        namespace: &str,
1198    ) -> Result<(), ScyllaDbStoreInternalError> {
1199        Self::check_namespace(namespace)?;
1200        let session = ScyllaDbClient::build_default_session(&config.uri).await?;
1201
1202        // Create a keyspace if it doesn't exist
1203        let statement = session
1204            .prepare(format!(
1205                "CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{ \
1206                    'class' : 'NetworkTopologyStrategy', \
1207                    'replication_factor' : {} \
1208                }}",
1209                KEYSPACE, config.replication_factor
1210            ))
1211            .await?;
1212        session
1213            .execute_single_page(&statement, &[], PagingState::start())
1214            .await
1215            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
1216
1217        // This explicitly sets a lot of default parameters for clarity and for making future
1218        // changes easier.
1219        let statement = session
1220            .prepare(format!(
1221                "CREATE TABLE {KEYSPACE}.\"{namespace}\" (\
1222                    root_key blob, \
1223                    k blob, \
1224                    v blob, \
1225                    PRIMARY KEY (root_key, k) \
1226                ) \
1227                WITH compaction = {{ \
1228                    'class'          : 'LeveledCompactionStrategy', \
1229                    'sstable_size_in_mb' : 160 \
1230                }} \
1231                AND compression = {{ \
1232                    'sstable_compression': 'LZ4Compressor', \
1233                    'chunk_length_in_kb':'4' \
1234                }} \
1235                AND caching = {{ \
1236                    'enabled': 'true' \
1237                }} \
1238                AND gc_grace_seconds = 0 \
1239                AND tombstone_gc = {{'mode': 'immediate'}}"
1240            ))
1241            .await?;
1242        session
1243            .execute_single_page(&statement, &[], PagingState::start())
1244            .await
1245            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
1246        Ok(())
1247    }
1248
1249    async fn delete(
1250        config: &Self::Config,
1251        namespace: &str,
1252    ) -> Result<(), ScyllaDbStoreInternalError> {
1253        Self::check_namespace(namespace)?;
1254        let session = ScyllaDbClient::build_default_session(&config.uri).await?;
1255        let statement = session
1256            .prepare(format!("DROP TABLE IF EXISTS {KEYSPACE}.\"{namespace}\";"))
1257            .await?;
1258        session
1259            .execute_single_page(&statement, &[], PagingState::start())
1260            .await
1261            .map_err(ScyllaDbStoreInternalError::ExecutionError)?;
1262        Ok(())
1263    }
1264}
1265
1266impl ScyllaDbStoreInternal {
1267    /// Obtains the semaphore lock on the database if needed.
1268    async fn acquire(&self) -> Option<SemaphoreGuard<'_>> {
1269        match &self.semaphore {
1270            None => None,
1271            Some(count) => Some(count.acquire().await),
1272        }
1273    }
1274}
1275
1276impl ScyllaDbDatabaseInternal {
1277    fn check_namespace(namespace: &str) -> Result<(), ScyllaDbStoreInternalError> {
1278        if !namespace.is_empty()
1279            && namespace.len() <= 48
1280            && namespace
1281                .chars()
1282                .all(|c| c.is_ascii_alphanumeric() || c == '_')
1283        {
1284            return Ok(());
1285        }
1286        Err(ScyllaDbStoreInternalError::InvalidNamespace)
1287    }
1288}
1289
1290#[cfg(with_testing)]
1291impl TestKeyValueDatabase for JournalingKeyValueDatabase<ScyllaDbDatabaseInternal> {
1292    async fn new_test_config(
1293    ) -> Result<ScyllaDbStoreInternalConfig, JournalingError<ScyllaDbStoreInternalError>> {
1294        // TODO(#4114): Read the port from an environment variable.
1295        let uri = "localhost:9042".to_string();
1296        Ok(ScyllaDbStoreInternalConfig {
1297            uri,
1298            max_concurrent_queries: Some(10),
1299            replication_factor: 1,
1300        })
1301    }
1302}
1303
1304/// The `ScyllaDbDatabase` composed type with metrics
1305#[cfg(with_metrics)]
1306pub type ScyllaDbDatabase = MeteredDatabase<
1307    LruCachingDatabase<
1308        MeteredDatabase<
1309            ValueSplittingDatabase<
1310                MeteredDatabase<JournalingKeyValueDatabase<ScyllaDbDatabaseInternal>>,
1311            >,
1312        >,
1313    >,
1314>;
1315
1316/// The `ScyllaDbDatabase` composed type
1317#[cfg(not(with_metrics))]
1318pub type ScyllaDbDatabase = LruCachingDatabase<
1319    ValueSplittingDatabase<JournalingKeyValueDatabase<ScyllaDbDatabaseInternal>>,
1320>;
1321
1322/// The `ScyllaDbStoreConfig` input type
1323pub type ScyllaDbStoreConfig = LruCachingConfig<ScyllaDbStoreInternalConfig>;
1324
1325/// The combined error type for the `ScyllaDbDatabase`.
1326pub type ScyllaDbStoreError = ValueSplittingError<JournalingError<ScyllaDbStoreInternalError>>;