surrealdb-core 3.1.2

#![cfg(feature = "kv-indxdb")]

use std::ops::Range;
use std::sync::atomic::{AtomicBool, Ordering};

use indxdb::{Database as Db, Transaction as Tx};
use tokio::sync::RwLock;

use super::api::{BoxFut, KeysResult, ScanLimit, ScanResult};
use super::err::{Error, Result};
use super::{ESTIMATED_BYTES_PER_KEY, ESTIMATED_BYTES_PER_KV, util};
use crate::key::debug::Sprintable;
use crate::kvs::api::Transactable;
use crate::kvs::{Key, Val};

pub struct Datastore {
	db: Db,
}

pub struct Transaction {
	/// Is the transaction complete?
	done: AtomicBool,
	/// Is the transaction writeable?
	write: bool,
	/// The underlying datastore transaction
	inner: RwLock<Tx>,
}

impl Datastore {
	/// Open a new database
	pub async fn new(path: &str) -> Result<Datastore> {
		match indxdb::Database::new(path).await {
			Ok(db) => Ok(Datastore {
				db,
			}),
			Err(e) => Err(Error::Datastore(e.to_string())),
		}
	}
	/// Shutdown the database
	pub(crate) async fn shutdown(&self) -> Result<()> {
		// Nothing to do here
		Ok(())
	}
	/// Start a new transaction
	pub async fn transaction(&self, write: bool, _: bool) -> Result<Box<dyn Transactable>> {
		// Create a new transaction
		match self.db.begin(write).await {
			Ok(txn) => Ok(Box::new(Transaction {
				done: AtomicBool::new(false),
				write,
				inner: RwLock::new(txn),
			})),
			Err(e) => Err(Error::from(e)),
		}
	}
}

impl Transactable for Transaction {
	fn kind(&self) -> &'static str {
		"indxdb"
	}

	/// Check if closed
	fn closed(&self) -> bool {
		self.done.load(Ordering::Relaxed)
	}

	/// Check if writeable
	fn writeable(&self) -> bool {
		self.write
	}

	/// Cancel a transaction
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self))]
	fn cancel(&self) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Atomically mark transaction as done and check if it was already closed
			if self.done.swap(true, Ordering::AcqRel) {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Cancel this transaction
			inner.cancel().await?;
			// Continue
			Ok(())
		})
	}

	/// Commit a transaction
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self))]
	fn commit(&self) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Atomically mark transaction as done and check if it was already closed
			if self.done.swap(true, Ordering::AcqRel) {
				return Err(Error::TransactionFinished);
			}
			// Check to see if transaction is writable
			if !self.writeable() {
				return Err(Error::TransactionReadonly);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Cancel this transaction
			inner.commit().await?;
			// Continue
			Ok(())
		})
	}

	/// Check if a key exists
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn exists(&self, key: Key, version: Option<u64>) -> BoxFut<'_, Result<bool>> {
		Box::pin(async move {
			// IndxDB does not support versioned queries.
			if version.is_some() {
				return Err(Error::UnsupportedVersionedQueries);
			}
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let inner = self.inner.read().await;
			// Check the key
			let res = inner.exists(key).await?;
			// Return result
			Ok(res)
		})
	}

	/// Fetch a key from the database
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn get(&self, key: Key, version: Option<u64>) -> BoxFut<'_, Result<Option<Val>>> {
		Box::pin(async move {
			// IndxDB does not support versioned queries.
			if version.is_some() {
				return Err(Error::UnsupportedVersionedQueries);
			}
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let inner = self.inner.read().await;
			// Get the key
			let res = inner.get(key).await?;
			// Return result
			Ok(res)
		})
	}

	/// Insert or update a key in the database
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn set(&self, key: Key, val: Val) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Check to see if transaction is writable
			if !self.writeable() {
				return Err(Error::TransactionReadonly);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Set the key
			inner.set(key, val).await?;
			// Return result
			Ok(())
		})
	}

	/// Insert a key if it doesn't exist in the database
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn put(&self, key: Key, val: Val) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Check to see if transaction is writable
			if !self.writeable() {
				return Err(Error::TransactionReadonly);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Set the key
			inner.put(key, val).await?;
			// Return result
			Ok(())
		})
	}

	/// Insert a key if the current value matches a condition
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn putc(&self, key: Key, val: Val, chk: Option<Val>) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Check to see if transaction is writable
			if !self.writeable() {
				return Err(Error::TransactionReadonly);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Set the key
			inner.putc(key, val, chk.map(Into::into)).await?;
			// Return result
			Ok(())
		})
	}

	/// Delete a key
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn del(&self, key: Key) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Check to see if transaction is writable
			if !self.writeable() {
				return Err(Error::TransactionReadonly);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Remove the key
			let res = inner.del(key).await?;
			// Return result
			Ok(res)
		})
	}

	/// Delete a key if the current value matches a condition
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(key = key.sprint()))]
	fn delc(&self, key: Key, chk: Option<Val>) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Check to see if transaction is writable
			if !self.writeable() {
				return Err(Error::TransactionReadonly);
			}
			// Load the inner transaction
			let mut inner = self.inner.write().await;
			// Remove the key
			let res = inner.delc(key, chk.map(Into::into)).await?;
			// Return result
			Ok(res)
		})
	}

	/// Retrieve a range of keys
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(rng = rng.sprint()))]
	fn keys(
		&self,
		rng: Range<Key>,
		limit: ScanLimit,
		skip: u32,
		version: Option<u64>,
	) -> BoxFut<'_, Result<KeysResult>> {
		Box::pin(async move {
			// IndxDB does not support versioned queries.
			if version.is_some() {
				return Err(Error::UnsupportedVersionedQueries);
			}
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let inner = self.inner.read().await;
			// Extract the limit count, adding skip to fetch enough entries
			let count = match limit {
				ScanLimit::Count(c) => c.saturating_add(skip),
				ScanLimit::Bytes(b) => (b / ESTIMATED_BYTES_PER_KEY).max(1).saturating_add(skip),
				ScanLimit::BytesOrCount(_, c) => c.saturating_add(skip),
			};
			// Scan the keys
			let res = inner.keys(rng, count).await?;
			// Consume the results
			Ok(consume_keys(&mut res.into_iter(), limit, skip))
		})
	}

	/// Retrieve a range of keys, in reverse
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(rng = rng.sprint()))]
	fn keysr(
		&self,
		rng: Range<Key>,
		limit: ScanLimit,
		skip: u32,
		version: Option<u64>,
	) -> BoxFut<'_, Result<KeysResult>> {
		Box::pin(async move {
			// IndxDB does not support versioned queries.
			if version.is_some() {
				return Err(Error::UnsupportedVersionedQueries);
			}
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let inner = self.inner.read().await;
			// Extract the limit count, adding skip to fetch enough entries
			let count = match limit {
				ScanLimit::Count(c) => c.saturating_add(skip),
				ScanLimit::Bytes(b) => (b / ESTIMATED_BYTES_PER_KEY).max(1).saturating_add(skip),
				ScanLimit::BytesOrCount(_, c) => c.saturating_add(skip),
			};
			// Scan the keys
			let res = inner.keysr(rng, count).await?;
			// Consume the results
			Ok(consume_keys(&mut res.into_iter(), limit, skip))
		})
	}

	/// Retrieve a range of key-value pairs
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(rng = rng.sprint()))]
	fn scan(
		&self,
		rng: Range<Key>,
		limit: ScanLimit,
		skip: u32,
		version: Option<u64>,
	) -> BoxFut<'_, Result<ScanResult>> {
		Box::pin(async move {
			// IndxDB does not support versioned queries.
			if version.is_some() {
				return Err(Error::UnsupportedVersionedQueries);
			}
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let inner = self.inner.read().await;
			// Skip entries using keys-only scan to avoid fetching values
			let rng = if skip > 0 {
				let skipped = inner.keys(rng.clone(), skip).await?;
				match skipped.last() {
					Some(last) => {
						let mut start = last.clone();
						util::advance_key(&mut start);
						start..rng.end
					}
					// Fewer entries than skip -- nothing to return
					None => return Ok(ScanResult::default()),
				}
			} else {
				rng
			};
			// Extract the limit count
			let count = match limit {
				ScanLimit::Count(c) => c,
				ScanLimit::Bytes(b) => (b / ESTIMATED_BYTES_PER_KV).max(1),
				ScanLimit::BytesOrCount(_, c) => c,
			};
			// Scan the keys
			let res = inner.scan(rng, count).await?;
			// Consume the results
			Ok(consume_vals(&mut res.into_iter(), limit))
		})
	}

	/// Retrieve a range of key-value pairs, in reverse
	#[instrument(level = "trace", target = "surrealdb::core::kvs::api", skip(self), fields(rng = rng.sprint()))]
	fn scanr(
		&self,
		rng: Range<Key>,
		limit: ScanLimit,
		skip: u32,
		version: Option<u64>,
	) -> BoxFut<'_, Result<ScanResult>> {
		Box::pin(async move {
			// IndxDB does not support versioned queries.
			if version.is_some() {
				return Err(Error::UnsupportedVersionedQueries);
			}
			// Check to see if transaction is closed
			if self.closed() {
				return Err(Error::TransactionFinished);
			}
			// Load the inner transaction
			let inner = self.inner.read().await;
			// Skip entries using keys-only scan to avoid fetching values
			let rng = if skip > 0 {
				let skipped = inner.keysr(rng.clone(), skip).await?;
				match skipped.last() {
					Some(last) => {
						let end = last.clone();
						rng.start..end
					}
					// Fewer entries than skip -- nothing to return
					None => return Ok(ScanResult::default()),
				}
			} else {
				rng
			};
			// Extract the limit count
			let count = match limit {
				ScanLimit::Count(c) => c,
				ScanLimit::Bytes(b) => (b / ESTIMATED_BYTES_PER_KV).max(1),
				ScanLimit::BytesOrCount(_, c) => c,
			};
			// Scan the keys in reverse
			let res = inner.scanr(rng, count).await?;
			// Consume the results
			Ok(consume_vals(&mut res.into_iter(), limit))
		})
	}

	/// Set a new save point on the transaction.
	fn new_save_point(&self) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			self.inner.write().await.set_savepoint().await?;
			Ok(())
		})
	}

	/// Rollback to the last save point.
	fn rollback_to_save_point(&self) -> BoxFut<'_, Result<()>> {
		Box::pin(async move {
			self.inner.write().await.rollback_to_savepoint().await?;
			Ok(())
		})
	}

	/// Release the last save point.
	fn release_last_save_point(&self) -> BoxFut<'_, Result<()>> {
		Box::pin(async move { Ok(()) })
	}
}

// Consume and iterate over keys
fn consume_keys<I: Iterator<Item = Key>>(iter: &mut I, limit: ScanLimit, skip: u32) -> KeysResult {
	// Skip entries from the pre-fetched iterator
	for _ in 0..skip {
		if iter.next().is_none() {
			return KeysResult::default();
		}
	}
	let mut key_bytes = 0u64;
	let keys = match limit {
		ScanLimit::Count(c) => {
			// Create the result set
			let mut res = Vec::with_capacity(c.min(4096) as usize);
			// Check that we don't exceed the count limit
			while res.len() < c as usize {
				// Check the key
				if let Some(k) = iter.next() {
					key_bytes += k.len() as u64;
					res.push(k);
				} else {
					break;
				}
			}
			res
		}
		ScanLimit::Bytes(b) => {
			// Create the result set
			let mut res = Vec::with_capacity((b / ESTIMATED_BYTES_PER_KEY).min(4096) as usize);
			// Check that we don't exceed the byte limit
			while key_bytes < b as u64 {
				// Check the key
				if let Some(k) = iter.next() {
					key_bytes += k.len() as u64;
					res.push(k);
				} else {
					break;
				}
			}
			res
		}
		ScanLimit::BytesOrCount(b, c) => {
			// Create the result set
			let mut res = Vec::with_capacity(c.min(4096) as usize);
			// Check that we don't exceed the count limit AND the byte limit
			while res.len() < c as usize && key_bytes < b as u64 {
				// Check the key
				if let Some(k) = iter.next() {
					key_bytes += k.len() as u64;
					res.push(k);
				} else {
					break;
				}
			}
			res
		}
	};
	KeysResult {
		keys,
		key_bytes,
	}
}

// Consume and iterate over keys and values
fn consume_vals<I: Iterator<Item = (Key, Val)>>(iter: &mut I, limit: ScanLimit) -> ScanResult {
	// Track the cumulative key/value bytes for the metric. The byte-bounded
	// limit branches still rely on `bytes_fetched` (key + value bytes) to
	// decide when to stop, so the two counters are kept separate.
	let mut key_bytes = 0u64;
	let mut value_bytes = 0u64;
	let values = match limit {
		ScanLimit::Count(c) => {
			// Create the result set
			let mut res = Vec::with_capacity(c.min(4096) as usize);
			// Check that we don't exceed the count limit
			while res.len() < c as usize {
				// Check the key and value
				if let Some((k, v)) = iter.next() {
					key_bytes += k.len() as u64;
					value_bytes += v.len() as u64;
					res.push((k, v));
				} else {
					break;
				}
			}
			res
		}
		ScanLimit::Bytes(b) => {
			// Create the result set
			let mut res = Vec::with_capacity((b / ESTIMATED_BYTES_PER_KV).min(4096) as usize);
			// Count the bytes fetched
			let mut bytes_fetched = 0u64;
			// Check that we don't exceed the byte limit
			while bytes_fetched < b as u64 {
				// Check the key and value
				if let Some((k, v)) = iter.next() {
					let key_len = k.len() as u64;
					let value_len = v.len() as u64;
					bytes_fetched += key_len + value_len;
					key_bytes += key_len;
					value_bytes += value_len;
					res.push((k, v));
				} else {
					break;
				}
			}
			res
		}
		ScanLimit::BytesOrCount(b, c) => {
			// Create the result set
			let mut res = Vec::with_capacity(c.min(4096) as usize);
			// Count the bytes fetched
			let mut bytes_fetched = 0u64;
			// Check that we don't exceed the count limit AND the byte limit
			while res.len() < c as usize && bytes_fetched < b as u64 {
				// Check the key and value
				if let Some((k, v)) = iter.next() {
					let key_len = k.len() as u64;
					let value_len = v.len() as u64;
					bytes_fetched += key_len + value_len;
					key_bytes += key_len;
					value_bytes += value_len;
					res.push((k, v));
				} else {
					break;
				}
			}
			res
		}
	};
	ScanResult {
		values,
		key_bytes,
		value_bytes,
	}
}