use anyhow::Result;
use storekey::decode_borrow;
use tracing::{debug, warn};
use crate::catalog::providers::{DatabaseProvider, NamespaceProvider};
use crate::catalog::{DatabaseId, NamespaceId, SequenceDefinition};
use crate::key::schema::{DbRoot, SequenceKey};
use crate::key::{KVSubspace, KVValue, Key, KeyRange};
use crate::kvs::TransactionType;
use crate::kvs::ds::Datastore;
const TARGET: &str = "surrealdb::core::kvs::migration";
const LEGACY_TAG: &[u8] = b"*sq";
pub(super) async fn copy_definitions_out_of_the_table_band(ds: &Datastore) -> Result<()> {
let txn = ds.transaction(TransactionType::Read).await?;
let databases = async {
let namespaces = txn.all_ns(None).await?;
let mut databases = Vec::new();
for ns in namespaces.iter() {
for db in txn.all_db(ns.namespace_id, None).await?.iter() {
databases.push((ns.namespace_id, db.database_id));
}
}
Ok::<_, anyhow::Error>(databases)
}
.await;
txn.cancel().await?;
let databases = databases?;
let mut copied = 0usize;
for (ns, db) in databases {
copied += migrate_database(ds, ns, db).await?;
}
if copied > 0 {
debug!(target: TARGET, count = copied, "Copied sequence definitions to the !sd subspace");
}
Ok(())
}
async fn migrate_database(ds: &Datastore, ns: NamespaceId, db: DatabaseId) -> Result<usize> {
let legacy = collect_legacy(ds, ns, db).await?;
for (_, definition) in &legacy {
let txn = ds.transaction(TransactionType::Write).await?;
let copied = txn
.put_key(
&SequenceKey::new(ns, db, std::borrow::Cow::Borrowed(definition.name.as_str())),
definition,
)
.await;
match run!(txn, copied) {
Err(e) if super::already_exists(&e) => continue,
other => other?,
}
}
Ok(legacy.len())
}
const SCAN_BATCH: u32 = 1000;
async fn collect_legacy(
ds: &Datastore,
ns: NamespaceId,
db: DatabaseId,
) -> Result<Vec<(Key<'static>, SequenceDefinition)>> {
let root = DbRoot::new(ns, db);
let mut prefix = root.encode_bound()?.as_ref().to_vec();
prefix.extend_from_slice(LEGACY_TAG);
let region = Key::from(prefix.clone()).prefix_expect();
let region_end = region.end.clone();
let mut next = Some(root.raw(region));
let mut found = Vec::new();
while let Some(range) = next {
let txn = ds.transaction(TransactionType::Read).await?;
let batch = catch!(txn, txn.batch_keys_raw(range, SCAN_BATCH, None).await);
let mut skip_to = None;
let mut names: Vec<(Key<'static>, String)> = Vec::new();
for key in &batch.result {
match classify(key, &prefix) {
Legacy::Definition(name) => names.push((Key::from(key.clone()), name)),
Legacy::TableSubtree(root_len) => {
skip_to = Some(Key::from(key[..root_len].to_vec()).next_neighbour_expect());
break;
}
Legacy::Foreign => {}
}
}
for (key, name) in names {
let value = catch!(txn, txn.get(key.clone(), None).await);
let Some(value) = value else {
continue;
};
match SequenceDefinition::kv_decode_value(&value, ()) {
Ok(definition) => {
if definition.name.as_str() != name {
warn!(
target: TARGET,
key_name = %name,
value_name = %definition.name,
"Skipping a legacy sequence definition whose stored name \
disagrees with its key"
);
continue;
}
found.push((key, definition));
}
Err(e) => warn!(
target: TARGET,
name = %name,
error = %e,
"Skipping a key under the legacy sequence prefix that does not \
decode as a sequence definition"
),
}
}
txn.cancel().await?;
next = match skip_to {
Some(start) if start < region_end => Some(root.raw(KeyRange {
start,
end: region_end.clone(),
})),
Some(_) => None,
None => batch.next.map(|rng| root.raw(rng)),
};
}
Ok(found)
}
#[derive(Debug, PartialEq, Eq)]
enum Legacy {
Definition(String),
TableSubtree(usize),
Foreign,
}
fn classify(key: &[u8], prefix: &[u8]) -> Legacy {
let Some(rest) = key.strip_prefix(prefix) else {
return Legacy::Foreign;
};
let Some(end) = name_end(rest) else {
return Legacy::Foreign;
};
if end == rest.len() {
return match decode_borrow::<String>(rest) {
Ok(name) => Legacy::Definition(name),
Err(_) => Legacy::Foreign,
};
}
Legacy::TableSubtree(prefix.len() + end)
}
fn name_end(rest: &[u8]) -> Option<usize> {
let mut i = 0;
while i < rest.len() {
match rest[i] {
0x01 => i += 2,
0x00 => return Some(i + 1),
_ => i += 1,
}
}
None
}
#[cfg(test)]
fn legacy_name(key: &[u8], prefix: &[u8]) -> Option<String> {
match classify(key, prefix) {
Legacy::Definition(name) => Some(name),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn prefix() -> Vec<u8> {
let mut p =
DbRoot::new(NamespaceId(1), DatabaseId(2)).encode_bound().unwrap().as_ref().to_vec();
p.extend_from_slice(LEGACY_TAG);
p
}
fn legacy_key(name: &str) -> Vec<u8> {
let mut k = prefix();
k.extend_from_slice(&storekey::encode_vec(&name).unwrap());
k
}
#[test]
fn the_legacy_layout_matches_what_3_2_wrote() {
assert_eq!(legacy_key("test"), b"/*\x00\x00\x00\x01*\x00\x00\x00\x02*sqtest\0");
}
#[test]
fn a_definition_key_yields_its_name() {
let prefix = prefix();
assert_eq!(legacy_name(&legacy_key("foo"), &prefix).as_deref(), Some("foo"));
assert_eq!(legacy_name(&legacy_key(""), &prefix).as_deref(), Some(""));
}
#[test]
fn awkward_names_round_trip() {
let prefix = prefix();
for name in [
"foo",
"sq",
"sqfoo",
"*",
"!fd",
"*sq",
"a*b!c",
"\u{0}",
"a\u{0}b",
"\u{1}",
"\u{1}\u{0}\u{1}",
"ключ",
"emoji-🔑",
"a".repeat(512).as_str(),
] {
assert_eq!(
legacy_name(&legacy_key(name), &prefix).as_deref(),
Some(name),
"name {name:?} did not round trip"
);
}
}
#[test]
fn keys_belonging_to_a_table_are_rejected() {
let prefix = prefix();
for suffix in [
&b"!fd"[..], &b"!ix"[..], &b"!ev"[..], &b"*"[..], &b"+"[..], &b"~"[..], &b"&"[..], &b"!lq"[..], &b"\x00"[..], ] {
let mut key = legacy_key("foo");
key.extend_from_slice(suffix);
assert_eq!(
legacy_name(&key, &prefix),
None,
"a table key ending {suffix:?} was mistaken for a definition"
);
}
}
#[test]
fn a_foreign_key_is_rejected() {
let prefix = prefix();
assert_eq!(legacy_name(b"/!nstest\0", &prefix), None);
let other_db =
DbRoot::new(NamespaceId(9), DatabaseId(9)).encode_bound().unwrap().as_ref().to_vec();
assert_eq!(legacy_name(&other_db, &prefix), None);
}
#[test]
fn a_truncated_key_is_rejected() {
let prefix = prefix();
let mut key = prefix.clone();
key.extend_from_slice(b"foo");
assert_eq!(legacy_name(&key, &prefix), None);
}
#[cfg(feature = "kv-mem")]
mod against_a_datastore {
use std::sync::Arc;
use super::*;
use crate::catalog::{DatabaseDefinition, NamespaceDefinition};
use crate::key::schema::SequencePrefix;
use crate::kvs::Datastore;
async fn create_db(ds: &Datastore, ns: NamespaceId, db: DatabaseId, name: &str) {
let txn = ds.transaction(TransactionType::Write).await.unwrap();
txn.put_ns(NamespaceDefinition {
namespace_id: ns,
name: name.into(),
comment: None,
})
.await
.unwrap();
txn.put_db(name, database(ns, db, name)).await.unwrap();
txn.commit().await.unwrap();
}
fn database(ns: NamespaceId, db: DatabaseId, name: &str) -> DatabaseDefinition {
DatabaseDefinition {
namespace_id: ns,
database_id: db,
name: name.into(),
strict: false,
comment: None,
changefeed: None,
}
}
async fn fixture() -> (Arc<Datastore>, NamespaceId, DatabaseId) {
let ds = Datastore::new("memory").await.unwrap();
let (ns, db) = (NamespaceId(1), DatabaseId(1));
create_db(&ds, ns, db, "test").await;
(ds, ns, db)
}
fn definition(name: &str, batch: u32, start: i64) -> SequenceDefinition {
SequenceDefinition {
name: name.into(),
batch,
start,
timeout: None,
}
}
async fn write_legacy(
ds: &Datastore,
ns: NamespaceId,
db: DatabaseId,
d: &SequenceDefinition,
) {
let mut key = DbRoot::new(ns, db).encode_bound().unwrap().as_ref().to_vec();
key.extend_from_slice(LEGACY_TAG);
key.extend_from_slice(&storekey::encode_vec(&d.name.as_str()).unwrap());
let txn = ds.transaction(TransactionType::Write).await.unwrap();
txn.set(Key::from(key), d.kv_encode_value().unwrap()).await.unwrap();
txn.commit().await.unwrap();
}
async fn definitions(
ds: &Datastore,
ns: NamespaceId,
db: DatabaseId,
) -> Vec<SequenceDefinition> {
let txn = ds.transaction(TransactionType::Read).await.unwrap();
let range = SequencePrefix::new(ns, db).range().unwrap();
let found = txn.getr(range, None).await.unwrap();
txn.cancel().await.unwrap();
found.into_iter().map(|(_, d)| d).collect()
}
async fn legacy_bytes(
ds: &Datastore,
ns: NamespaceId,
db: DatabaseId,
) -> Vec<(Vec<u8>, Vec<u8>)> {
let root = DbRoot::new(ns, db);
let mut prefix = root.encode_bound().unwrap().as_ref().to_vec();
prefix.extend_from_slice(LEGACY_TAG);
let range = root.raw(Key::from(prefix).prefix_expect());
let txn = ds.transaction(TransactionType::Read).await.unwrap();
let found = txn.getr_raw(range, None).await.unwrap();
txn.cancel().await.unwrap();
found
}
async fn legacy_keys(ds: &Datastore, ns: NamespaceId, db: DatabaseId) -> Vec<Vec<u8>> {
let root = DbRoot::new(ns, db);
let mut prefix = root.encode_bound().unwrap().as_ref().to_vec();
prefix.extend_from_slice(LEGACY_TAG);
let range = root.raw(Key::from(prefix).prefix_expect());
let txn = ds.transaction(TransactionType::Read).await.unwrap();
let found = txn.getr_raw(range, None).await.unwrap();
txn.cancel().await.unwrap();
found.into_iter().map(|(k, _)| k).collect()
}
#[tokio::test]
async fn definitions_move_with_their_values_intact() {
let (ds, ns, db) = fixture().await;
let a = definition("orders", 50, 7);
let b = definition("invoices", 1, -3);
write_legacy(&ds, ns, db, &a).await;
write_legacy(&ds, ns, db, &b).await;
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
let mut moved = definitions(&ds, ns, db).await;
moved.sort_by(|x, y| x.name.cmp(&y.name));
assert_eq!(moved, vec![b, a]);
assert_eq!(legacy_keys(&ds, ns, db).await.len(), 2);
}
#[tokio::test]
async fn a_colliding_table_is_left_alone() {
let (ds, ns, db) = fixture().await;
let sequence = definition("foo", 10, 0);
write_legacy(&ds, ns, db, &sequence).await;
let mut roots = Vec::new();
for table in ["sq", "sqfoo", "sqzzz"] {
let mut root = DbRoot::new(ns, db).encode_bound().unwrap().as_ref().to_vec();
root.extend_from_slice(b"*");
root.extend_from_slice(&storekey::encode_vec(&table).unwrap());
roots.push(root);
}
let txn = ds.transaction(TransactionType::Write).await.unwrap();
let mut planted = Vec::new();
for root in &roots {
for suffix in
[&b"!fda\0"[..], &b"*\x01record"[..], &b"~edge"[..], &b"+\x00\x00\x00\x01"[..]]
{
let mut key = root.clone();
key.extend_from_slice(suffix);
txn.set(Key::from(key.clone()), b"table data".to_vec()).await.unwrap();
planted.push(key);
}
}
txn.commit().await.unwrap();
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
assert_eq!(definitions(&ds, ns, db).await, vec![sequence]);
let remaining = legacy_keys(&ds, ns, db).await;
for key in &planted {
assert!(
remaining.contains(key),
"the migration deleted a key belonging to a table: {key:?}"
);
}
assert_eq!(remaining.len(), planted.len() + 1);
}
#[tokio::test]
async fn the_legacy_key_survives_unchanged() {
let (ds, ns, db) = fixture().await;
let definition = definition("orders", 50, 7);
write_legacy(&ds, ns, db, &definition).await;
let before = legacy_bytes(&ds, ns, db).await;
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
assert_eq!(definitions(&ds, ns, db).await, vec![definition.clone()]);
assert_eq!(legacy_bytes(&ds, ns, db).await, before);
assert_eq!(SequenceDefinition::kv_decode_value(&before[0].1, ()).unwrap(), definition);
}
#[tokio::test]
async fn running_it_twice_converges() {
let (ds, ns, db) = fixture().await;
write_legacy(&ds, ns, db, &definition("orders", 50, 7)).await;
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
let after_first = definitions(&ds, ns, db).await;
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
assert_eq!(definitions(&ds, ns, db).await, after_first);
assert_eq!(legacy_keys(&ds, ns, db).await.len(), 1);
}
#[tokio::test]
async fn an_existing_definition_is_never_overwritten() {
let (ds, ns, db) = fixture().await;
let stale = definition("orders", 1, 0);
write_legacy(&ds, ns, db, &stale).await;
let newer = definition("orders", 999, 42);
let txn = ds.transaction(TransactionType::Write).await.unwrap();
txn.set_key(&SequenceKey::new(ns, db, std::borrow::Cow::Borrowed("orders")), &newer)
.await
.unwrap();
txn.commit().await.unwrap();
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
assert_eq!(
definitions(&ds, ns, db).await,
vec![newer],
"the migration clobbered a definition written after it read its copy"
);
}
#[tokio::test]
async fn repeated_passes_converge() {
let (ds, ns, db) = fixture().await;
for i in 0..8 {
write_legacy(&ds, ns, db, &definition(&format!("seq{i}"), 10, i)).await;
}
for _ in 0..4 {
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
}
let mut moved = definitions(&ds, ns, db).await;
moved.sort_by(|x, y| x.name.cmp(&y.name));
assert_eq!(moved.len(), 8);
for (i, d) in moved.iter().enumerate() {
assert_eq!(d.name.as_str(), format!("seq{i}"));
assert_eq!(d.start, i as i64);
}
assert_eq!(legacy_keys(&ds, ns, db).await.len(), 8);
}
#[tokio::test]
async fn an_empty_datastore_is_a_no_op() {
let ds = Datastore::new("memory").await.unwrap();
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
}
#[tokio::test]
async fn an_undecodable_value_is_left_in_place() {
let (ds, ns, db) = fixture().await;
let mut key = DbRoot::new(ns, db).encode_bound().unwrap().as_ref().to_vec();
key.extend_from_slice(LEGACY_TAG);
key.extend_from_slice(&storekey::encode_vec(&"mystery").unwrap());
let txn = ds.transaction(TransactionType::Write).await.unwrap();
txn.set(Key::from(key.clone()), vec![0xff, 0xff, 0xff]).await.unwrap();
txn.commit().await.unwrap();
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
assert!(definitions(&ds, ns, db).await.is_empty());
assert_eq!(legacy_keys(&ds, ns, db).await, vec![key]);
}
#[tokio::test]
async fn adversarial_names_survive_the_move() {
let (ds, ns, db) = fixture().await;
let names = ["sq", "sqfoo", "*", "!fd", "a\u{0}b", "\u{1}", "ключ", "🔑"];
for name in names {
write_legacy(&ds, ns, db, &definition(name, 1, 0)).await;
}
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
let moved = definitions(&ds, ns, db).await;
assert_eq!(moved.len(), names.len());
for name in names {
assert!(
moved.iter().any(|d| d.name.as_str() == name),
"sequence {name:?} did not survive the copy"
);
}
assert_eq!(legacy_keys(&ds, ns, db).await.len(), names.len());
}
#[tokio::test]
async fn definitions_are_found_past_the_first_scan_batch() {
let (ds, ns, db) = fixture().await;
write_legacy(&ds, ns, db, &definition("aaa-first", 1, 1)).await;
write_legacy(&ds, ns, db, &definition("zzz-last", 1, 2)).await;
let mut root = DbRoot::new(ns, db).encode_bound().unwrap().as_ref().to_vec();
root.extend_from_slice(b"*");
root.extend_from_slice(&storekey::encode_vec(&"sqbulk").unwrap());
let padding = (SCAN_BATCH as usize) * 2 + 7;
let txn = ds.transaction(TransactionType::Write).await.unwrap();
for i in 0..padding {
let mut key = root.clone();
key.extend_from_slice(b"*");
key.extend_from_slice(&storekey::encode_vec(&format!("{i:08}")).unwrap());
txn.set(Key::from(key), b"record".to_vec()).await.unwrap();
}
txn.commit().await.unwrap();
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
let mut moved = definitions(&ds, ns, db).await;
moved.sort_by(|x, y| x.name.cmp(&y.name));
assert_eq!(
moved.iter().map(|d| d.name.to_string()).collect::<Vec<_>>(),
vec!["aaa-first".to_string(), "zzz-last".to_string()]
);
assert_eq!(legacy_keys(&ds, ns, db).await.len(), padding + 2);
}
#[tokio::test]
async fn all_databases_are_migrated() {
let ds = Datastore::new("memory").await.unwrap();
let places = [
(NamespaceId(1), DatabaseId(1)),
(NamespaceId(2), DatabaseId(1)),
(NamespaceId(2), DatabaseId(2)),
];
create_db(&ds, NamespaceId(1), DatabaseId(1), "one").await;
create_db(&ds, NamespaceId(2), DatabaseId(1), "two").await;
let txn = ds.transaction(TransactionType::Write).await.unwrap();
txn.put_db("two", database(NamespaceId(2), DatabaseId(2), "gamma")).await.unwrap();
txn.commit().await.unwrap();
for (ns, db) in places {
write_legacy(&ds, ns, db, &definition("shared", 5, 1)).await;
}
copy_definitions_out_of_the_table_band(&ds).await.unwrap();
for (ns, db) in places {
assert_eq!(definitions(&ds, ns, db).await.len(), 1, "ns {ns:?} db {db:?}");
assert_eq!(legacy_keys(&ds, ns, db).await.len(), 1, "ns {ns:?} db {db:?}");
}
}
}
}