diff --git a/Cargo.lock b/Cargo.lock index 1db9f7d62..da6719404 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3150,6 +3150,7 @@ dependencies = [ "arrow", "async-trait", "bytes", + "chrono", "dashmap", "datafusion", "datafusion-proto", diff --git a/crates/modelardb_embedded/bindings/python/modelardb/operations.py b/crates/modelardb_embedded/bindings/python/modelardb/operations.py index 7430f7b6e..75a246591 100644 --- a/crates/modelardb_embedded/bindings/python/modelardb/operations.py +++ b/crates/modelardb_embedded/bindings/python/modelardb/operations.py @@ -188,6 +188,10 @@ def __find_library(build: str) -> str: int modelardb_embedded_drop(void* maybe_operations_ptr, bool is_data_folder, char* table_name_ptr); + + int modelardb_embedded_vacuum(void* maybe_operations_ptr, + bool is_data_folder, + char* table_name_ptr); char* modelardb_embedded_error(); """ @@ -724,6 +728,19 @@ def drop(self, table_name: str): ) self.__check_return_code_and_raise_error(return_code) + def vacuum(self, table_name: str): + """Vacuum the table with `table_name`. + + :param table_name: The name of the table to vacuum. + :type table_name: str + :raises ValueError: If incorrect arguments are provided. + """ + table_name_ptr = ffi.new("char[]", bytes(table_name, "UTF-8")) + return_code = self.__library.modelardb_embedded_vacuum( + self.__operations_ptr, self.__is_data_folder, table_name_ptr + ) + self.__check_return_code_and_raise_error(return_code) + def __check_return_code_and_raise_error(self, return_code: int): """Raises an appropriate exception based on the return code. diff --git a/crates/modelardb_embedded/bindings/python/tests/test_operations.py b/crates/modelardb_embedded/bindings/python/tests/test_operations.py index 95e3a8681..1c96982b7 100644 --- a/crates/modelardb_embedded/bindings/python/tests/test_operations.py +++ b/crates/modelardb_embedded/bindings/python/tests/test_operations.py @@ -437,6 +437,37 @@ def test_data_folder_drop_error(self): ) self.assertEqual(error_message, str(context.exception)) + def test_data_folder_vacuum(self): + with TemporaryDirectory() as temp_dir: + os.environ["MODELARDBD_RETENTION_PERIOD_IN_SECONDS"] = "0" + + data_folder = Operations.open_local(temp_dir) + create_tables_in_data_folder(data_folder) + + data_folder.write(TIME_SERIES_TABLE_NAME, time_series_table_data()) + data_folder.truncate(TIME_SERIES_TABLE_NAME) + + # The files should still exist on disk even though they are no longer active. + folder_path = os.path.join(temp_dir, "tables", TIME_SERIES_TABLE_NAME, "field_column=2") + file_count = len(os.listdir(folder_path)) + self.assertEqual(file_count, 1) + + data_folder.vacuum(TIME_SERIES_TABLE_NAME) + + # No files should remain in the column folder. + file_count = len(os.listdir(folder_path)) + self.assertEqual(file_count, 0) + + def test_data_folder_vacuum_error(self): + with TemporaryDirectory() as temp_dir: + data_folder = Operations.open_local(temp_dir) + + with self.assertRaises(RuntimeError) as context: + data_folder.vacuum(MISSING_TABLE_NAME) + + error_message = f"Invalid Argument Error: Table with name '{MISSING_TABLE_NAME}' does not exist." + self.assertEqual(error_message, str(context.exception)) + def create_tables_in_data_folder(data_folder: Operations): table_type = NormalTable(normal_table_schema()) diff --git a/crates/modelardb_embedded/src/capi.rs b/crates/modelardb_embedded/src/capi.rs index a894cc480..604a24a8a 100644 --- a/crates/modelardb_embedded/src/capi.rs +++ b/crates/modelardb_embedded/src/capi.rs @@ -932,6 +932,31 @@ unsafe fn drop( TOKIO_RUNTIME.block_on(modelardb.drop(table_name)) } +/// Vacuums the table with the name in `table_name_ptr` in the [`DataFolder`] or [`Client`] in +/// `maybe_operations_ptr`. Assumes `maybe_operations_ptr` points to a [`DataFolder`] or [`Client`]; +/// and `table_name_ptr` points to a valid C string. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn modelardb_embedded_vacuum( + maybe_operations_ptr: *mut c_void, + is_data_folder: bool, + table_name_ptr: *const c_char, +) -> c_int { + let maybe_unit = unsafe { vacuum(maybe_operations_ptr, is_data_folder, table_name_ptr) }; + set_error_and_return_code(maybe_unit) +} + +/// See documentation for [`modelardb_embedded_vacuum`]. +unsafe fn vacuum( + maybe_operations_ptr: *mut c_void, + is_data_folder: bool, + table_name_ptr: *const c_char, +) -> Result<()> { + let modelardb = unsafe { c_void_to_operations(maybe_operations_ptr, is_data_folder)? }; + let table_name = unsafe { c_char_ptr_to_str(table_name_ptr)? }; + + TOKIO_RUNTIME.block_on(modelardb.vacuum(table_name)) +} + /// Return a read-only [`*const c_char`] with a human-readable representation of the last error the /// current thread encountered. The lifetime of the returned [`*const c_char`] ends when /// [`modelardb_embedded_error()`] is called again. If no errors have occurred, a zero-initialized diff --git a/crates/modelardb_embedded/src/operations/client.rs b/crates/modelardb_embedded/src/operations/client.rs index 81021a85c..711d004bb 100644 --- a/crates/modelardb_embedded/src/operations/client.rs +++ b/crates/modelardb_embedded/src/operations/client.rs @@ -357,4 +357,13 @@ impl Operations for Client { Ok(()) } + + /// Vacuum the table with the name in `table_name`. If the table could not be vacuumed, + /// [`ModelarDbEmbeddedError`] is returned. + async fn vacuum(&mut self, table_name: &str) -> Result<()> { + let ticket = Ticket::new(format!("VACUUM {table_name}")); + self.flight_client.do_get(ticket).await?; + + Ok(()) + } } diff --git a/crates/modelardb_embedded/src/operations/data_folder.rs b/crates/modelardb_embedded/src/operations/data_folder.rs index 22a6a3199..c9fb7dca1 100644 --- a/crates/modelardb_embedded/src/operations/data_folder.rs +++ b/crates/modelardb_embedded/src/operations/data_folder.rs @@ -17,6 +17,7 @@ use std::any::Any; use std::collections::HashMap; +use std::env; use std::fmt::{Debug, Formatter, Result as FmtResult}; use std::path::Path as StdPath; use std::pin::Pin; @@ -851,6 +852,24 @@ impl Operations for DataFolder { Ok(()) } + + /// Vacuum the table with the name in `table_name`. If the table does not exist or the + /// table could not be vacuumed, [`ModelarDbEmbeddedError`] is returned. + async fn vacuum(&mut self, table_name: &str) -> Result<()> { + if self.tables().await?.contains(&table_name.to_owned()) { + let retention_period_in_seconds = env::var("MODELARDBD_RETENTION_PERIOD_IN_SECONDS") + .map_or(60 * 60 * 24 * 7, |value| value.parse().unwrap()); + + self.delta_lake + .vacuum_table(table_name, retention_period_in_seconds) + .await + .map_err(|error| error.into()) + } else { + Err(ModelarDbEmbeddedError::InvalidArgument(format!( + "Table with name '{table_name}' does not exist." + ))) + } + } } /// Sort the `uncompressed_data` from the time series table with `time_series_table_metadata` @@ -922,6 +941,8 @@ fn schemas_are_compatible(source_schema: &Schema, target_schema: &Schema) -> boo mod tests { use super::*; + use std::sync::{LazyLock, Mutex}; + use arrow::array::{Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array}; use arrow::datatypes::{ArrowPrimitiveType, DataType, Field}; use arrow_flight::flight_service_client::FlightServiceClient; @@ -942,6 +963,9 @@ mod tests { const TIME_SERIES_TABLE_WITH_GENERATED_COLUMN_NAME: &str = "time_series_table_with_generated"; const INVALID_COLUMN_NAME: &str = "invalid_column"; + /// Lock used for env::set_var() as it is not guaranteed to be thread-safe. + static SET_VAR_LOCK: LazyLock> = LazyLock::new(|| Mutex::new(())); + #[tokio::test] async fn test_create_normal_table() { let (_temp_dir, data_folder) = create_data_folder_with_normal_table().await; @@ -2540,6 +2564,87 @@ mod tests { ); } + #[tokio::test] + async fn test_vacuum_normal_table() { + // env::set_var is safe to call in a single-threaded program. + unsafe { + let _mutex_guard = SET_VAR_LOCK.lock(); + env::set_var("MODELARDBD_RETENTION_PERIOD_IN_SECONDS", "0"); + } + + let (temp_dir, mut data_folder) = create_data_folder_with_normal_table().await; + + data_folder + .write(NORMAL_TABLE_NAME, normal_table_data()) + .await + .unwrap(); + + data_folder.truncate(NORMAL_TABLE_NAME).await.unwrap(); + + // The files should still exist on disk even though they are no longer active. + let table_path = format!( + "{}/tables/{}", + temp_dir.path().to_str().unwrap(), + NORMAL_TABLE_NAME + ); + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 2); + + data_folder.vacuum(NORMAL_TABLE_NAME).await.unwrap(); + + // Only the _delta_log folder should remain. + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 1); + } + + #[tokio::test] + async fn test_vacuum_time_series_table() { + // env::set_var is safe to call in a single-threaded program. + unsafe { + let _mutex_guard = SET_VAR_LOCK.lock(); + env::set_var("MODELARDBD_RETENTION_PERIOD_IN_SECONDS", "0"); + } + + let (temp_dir, mut data_folder) = create_data_folder_with_time_series_table().await; + + data_folder + .write(TIME_SERIES_TABLE_NAME, time_series_table_data()) + .await + .unwrap(); + + data_folder.truncate(TIME_SERIES_TABLE_NAME).await.unwrap(); + + // The files should still exist on disk even though they are no longer active. + let column_path = format!( + "{}/tables/{}/field_column=3", + temp_dir.path().to_str().unwrap(), + TIME_SERIES_TABLE_NAME + ); + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 1); + + data_folder.vacuum(TIME_SERIES_TABLE_NAME).await.unwrap(); + + // No files should remain in the column folder. + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 0); + } + + #[tokio::test] + async fn test_vacuum_missing_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let mut data_folder = DataFolder::open_local(temp_dir.path()).await.unwrap(); + + let result = data_folder.vacuum(MISSING_TABLE_NAME).await; + + assert_eq!( + result.unwrap_err().to_string(), + format!( + "Invalid Argument Error: Table with name '{MISSING_TABLE_NAME}' does not exist." + ) + ); + } + #[tokio::test] async fn test_move_normal_table_to_normal_table() { let (_temp_dir, mut source) = create_data_folder_with_normal_table().await; diff --git a/crates/modelardb_embedded/src/operations/mod.rs b/crates/modelardb_embedded/src/operations/mod.rs index 1e0f81e0a..650fd8545 100644 --- a/crates/modelardb_embedded/src/operations/mod.rs +++ b/crates/modelardb_embedded/src/operations/mod.rs @@ -107,6 +107,9 @@ pub trait Operations: Sync + Send { /// Drop the table with the name in `table_name`. async fn drop(&mut self, table_name: &str) -> Result<()>; + + /// Vacuum the table with the name in `table_name`. + async fn vacuum(&mut self, table_name: &str) -> Result<()>; } /// Use the time series table metadata in `table_name`, `schema`, `error_bounds`, and `generated_columns` diff --git a/crates/modelardb_manager/src/remote.rs b/crates/modelardb_manager/src/remote.rs index 11d3f08ad..60ae293ce 100644 --- a/crates/modelardb_manager/src/remote.rs +++ b/crates/modelardb_manager/src/remote.rs @@ -21,8 +21,8 @@ use std::error::Error; use std::net::SocketAddr; use std::pin::Pin; use std::result::Result as StdResult; -use std::str; use std::sync::Arc; +use std::{env, str}; use arrow::datatypes::Schema; use arrow::ipc::writer::IpcWriteOptions; @@ -333,6 +333,33 @@ impl FlightServiceHandler { Ok(()) } + + /// Vacuum the table in the remote data folder and at each node controlled by the manager. If + /// the table does not exist or the table cannot be vacuumed in the remote data folder + /// and at each node, return [`Status`]. + async fn vacuum_cluster_table(&self, table_name: &str) -> StdResult<(), Status> { + let retention_period_in_seconds = env::var("MODELARDBD_RETENTION_PERIOD_IN_SECONDS") + .map_or(60 * 60 * 24 * 7, |value| value.parse().unwrap()); + + // Vacuum the table in the remote data folder Delta lake. + self.context + .remote_data_folder + .delta_lake + .vacuum_table(table_name, retention_period_in_seconds) + .await + .map_err(error_to_status_internal)?; + + // Vacuum the table in the nodes controlled by the manager. + self.context + .cluster + .read() + .await + .cluster_do_get(&format!("VACUUM {table_name}"), &self.context.key) + .await + .map_err(error_to_status_internal)?; + + Ok(()) + } } #[tonic::async_trait] @@ -453,8 +480,8 @@ impl FlightService for FlightServiceHandler { } /// Execute a SQL statement provided in UTF-8 and return the schema of the result followed by - /// the result itself. Currently, CREATE TABLE, CREATE TIME SERIES TABLE, TRUNCATE TABLE, and - /// DROP TABLE are supported. + /// the result itself. Currently, CREATE TABLE, CREATE TIME SERIES TABLE, TRUNCATE TABLE, + /// DROP TABLE, and VACUUM are supported. async fn do_get( &self, request: Request, @@ -497,6 +524,23 @@ impl FlightService for FlightServiceHandler { self.drop_cluster_table(&table_name).await?; } } + ModelarDbStatement::Vacuum(mut table_names) => { + // Vacuum all tables if no table names are provided. + if table_names.is_empty() { + table_names = self + .context + .remote_data_folder + .metadata_manager + .table_metadata_manager + .table_names() + .await + .map_err(error_to_status_internal)?; + } + + for table_name in table_names { + self.vacuum_cluster_table(&table_name).await?; + } + } // .. is not used so a compile error is raised if a new ModelarDbStatement is added. ModelarDbStatement::Statement(_) | ModelarDbStatement::IncludeSelect(..) => { return Err(Status::invalid_argument( diff --git a/crates/modelardb_server/src/configuration.rs b/crates/modelardb_server/src/configuration.rs index 7826be7e7..f29f937d6 100644 --- a/crates/modelardb_server/src/configuration.rs +++ b/crates/modelardb_server/src/configuration.rs @@ -45,6 +45,8 @@ pub struct ConfigurationManager { /// The number of seconds between each transfer of data to the remote object store. If [`None`], /// data is not transferred based on time. transfer_time_in_seconds: Option, + /// The number of seconds to retain deleted data in storage before it can be removed by vacuum. + retention_period_in_seconds: usize, /// Number of threads to allocate for converting multivariate time series to univariate /// time series. pub(crate) ingestion_threads: usize, @@ -74,6 +76,9 @@ impl ConfigurationManager { let transfer_time_in_seconds = env::var("MODELARDBD_TRANSFER_TIME_IN_SECONDS") .map_or(None, |value| Some(value.parse().unwrap())); + let retention_period_in_seconds = env::var("MODELARDBD_RETENTION_PERIOD_IN_SECONDS") + .map_or(60 * 60 * 24 * 7, |value| value.parse().unwrap()); + Self { cluster_mode, multivariate_reserved_memory_in_bytes, @@ -81,6 +86,7 @@ impl ConfigurationManager { compressed_reserved_memory_in_bytes, transfer_batch_size_in_bytes, transfer_time_in_seconds, + retention_period_in_seconds, // TODO: Add support for running multiple threads per component. The individual // components in the storage engine have not been validated with multiple threads, e.g., // UncompressedDataManager may have race conditions finishing buffers if multiple @@ -217,6 +223,18 @@ impl ConfigurationManager { Ok(()) } + pub(crate) fn retention_period_in_seconds(&self) -> usize { + self.retention_period_in_seconds + } + + /// Set the new value for the retention period in seconds. + pub(crate) fn set_retention_period_in_seconds( + &mut self, + new_retention_period_in_seconds: usize, + ) { + self.retention_period_in_seconds = new_retention_period_in_seconds; + } + /// Encode the current configuration into a [`Configuration`](protocol::Configuration) /// protobuf message and serialize it. pub(crate) fn encode_and_serialize(&self) -> Vec { @@ -228,6 +246,7 @@ impl ConfigurationManager { compressed_reserved_memory_in_bytes: self.compressed_reserved_memory_in_bytes as u64, transfer_batch_size_in_bytes: self.transfer_batch_size_in_bytes.map(|v| v as u64), transfer_time_in_seconds: self.transfer_time_in_seconds.map(|v| v as u64), + retention_period_in_seconds: self.retention_period_in_seconds as u64, ingestion_threads: self.ingestion_threads as u32, compression_threads: self.compression_threads as u32, writer_threads: self.writer_threads as u32, @@ -268,10 +287,11 @@ mod tests { 512 * 1024 * 1024 ); + let new_value = 1024; configuration_manager .write() .await - .set_multivariate_reserved_memory_in_bytes(1024, storage_engine) + .set_multivariate_reserved_memory_in_bytes(new_value, storage_engine) .await; assert_eq!( @@ -279,7 +299,7 @@ mod tests { .read() .await .multivariate_reserved_memory_in_bytes(), - 1024 + new_value ); } @@ -296,10 +316,11 @@ mod tests { 512 * 1024 * 1024 ); + let new_value = 1024; configuration_manager .write() .await - .set_uncompressed_reserved_memory_in_bytes(1024, storage_engine) + .set_uncompressed_reserved_memory_in_bytes(new_value, storage_engine) .await .unwrap(); @@ -308,7 +329,7 @@ mod tests { .read() .await .uncompressed_reserved_memory_in_bytes(), - 1024 + new_value ); } @@ -325,10 +346,11 @@ mod tests { 512 * 1024 * 1024 ); + let new_value = 1024; configuration_manager .write() .await - .set_compressed_reserved_memory_in_bytes(1024, storage_engine) + .set_compressed_reserved_memory_in_bytes(new_value, storage_engine) .await .unwrap(); @@ -337,7 +359,7 @@ mod tests { .read() .await .compressed_reserved_memory_in_bytes(), - 1024 + new_value ); } @@ -354,10 +376,11 @@ mod tests { Some(64 * 1024 * 1024) ); + let new_value = Some(1024); configuration_manager .write() .await - .set_transfer_batch_size_in_bytes(Some(1024), storage_engine) + .set_transfer_batch_size_in_bytes(new_value, storage_engine) .await .unwrap(); @@ -366,7 +389,7 @@ mod tests { .read() .await .transfer_batch_size_in_bytes(), - Some(1024) + new_value ); } @@ -383,10 +406,11 @@ mod tests { None ); + let new_value = Some(60); configuration_manager .write() .await - .set_transfer_time_in_seconds(Some(60), storage_engine) + .set_transfer_time_in_seconds(new_value, storage_engine) .await .unwrap(); @@ -395,7 +419,35 @@ mod tests { .read() .await .transfer_time_in_seconds(), - Some(60) + new_value + ); + } + + #[tokio::test] + async fn test_set_retention_period_in_seconds() { + let temp_dir = tempfile::tempdir().unwrap(); + let (_, configuration_manager) = create_components(&temp_dir).await; + + assert_eq!( + configuration_manager + .read() + .await + .retention_period_in_seconds(), + 60 * 60 * 24 * 7 + ); + + let new_value = 60; + configuration_manager + .write() + .await + .set_retention_period_in_seconds(new_value); + + assert_eq!( + configuration_manager + .read() + .await + .retention_period_in_seconds(), + new_value ); } diff --git a/crates/modelardb_server/src/context.rs b/crates/modelardb_server/src/context.rs index ac3f36d04..e30675a44 100644 --- a/crates/modelardb_server/src/context.rs +++ b/crates/modelardb_server/src/context.rs @@ -350,6 +350,24 @@ impl Context { Ok(()) } + /// Vacuum the table with `table_name` if it exists. If the table does not exist or if it could + /// not be vacuumed, [`ModelarDbServerError`] is returned. + pub async fn vacuum_table(&self, table_name: &str) -> Result<()> { + let retention_period_in_seconds = self + .configuration_manager + .read() + .await + .retention_period_in_seconds(); + + self.data_folders + .local_data_folder + .delta_lake + .vacuum_table(table_name, retention_period_in_seconds) + .await?; + + Ok(()) + } + /// Lookup the [`TimeSeriesTableMetadata`] of the time series table with name `table_name` if it /// exists. Specifically, the method returns: /// * [`TimeSeriesTableMetadata`] if a time series table with the name `table_name` exists. @@ -426,10 +444,11 @@ impl Context { mod tests { use super::*; - use crate::data_folders::DataFolder; use modelardb_test::table::{self, NORMAL_TABLE_NAME, TIME_SERIES_TABLE_NAME}; use tempfile::TempDir; + use crate::data_folders::DataFolder; + // Tests for Context. #[tokio::test] async fn test_create_tables_from_bytes() { @@ -706,12 +725,7 @@ mod tests { #[tokio::test] async fn test_truncate_normal_table() { let temp_dir = tempfile::tempdir().unwrap(); - let context = create_context(&temp_dir).await; - - context - .create_normal_table(NORMAL_TABLE_NAME, &table::normal_table_schema()) - .await - .unwrap(); + let context = create_context_with_normal_table(&temp_dir).await; let local_data_folder = &context.data_folders.local_data_folder; let mut delta_table = local_data_folder @@ -720,17 +734,6 @@ mod tests { .await .unwrap(); - // Write data to the normal table that should be deleted when the table is truncated. - local_data_folder - .delta_lake - .write_record_batches_to_normal_table( - NORMAL_TABLE_NAME, - vec![table::normal_table_record_batch()], - ) - .await - .unwrap(); - - delta_table.load().await.unwrap(); assert_eq!(delta_table.get_files_count(), 1); context.truncate_table(NORMAL_TABLE_NAME).await.unwrap(); @@ -752,12 +755,7 @@ mod tests { #[tokio::test] async fn test_truncate_time_series_table() { let temp_dir = tempfile::tempdir().unwrap(); - let context = create_context(&temp_dir).await; - - context - .create_time_series_table(&table::time_series_table_metadata()) - .await - .unwrap(); + let context = create_context_with_time_series_table(&temp_dir).await; let local_data_folder = &context.data_folders.local_data_folder; let mut delta_table = local_data_folder @@ -766,18 +764,6 @@ mod tests { .await .unwrap(); - // Write data to the time series table that should be deleted when the table is truncated. - let record_batch = table::compressed_segments_record_batch(); - local_data_folder - .delta_lake - .write_compressed_segments_to_time_series_table( - TIME_SERIES_TABLE_NAME, - vec![record_batch], - ) - .await - .unwrap(); - - delta_table.load().await.unwrap(); assert_eq!(delta_table.get_files_count(), 1); context @@ -812,6 +798,122 @@ mod tests { ); } + #[tokio::test] + async fn test_vacuum_normal_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context_with_normal_table(&temp_dir).await; + + context + .configuration_manager + .write() + .await + .set_retention_period_in_seconds(0); + + context.truncate_table(NORMAL_TABLE_NAME).await.unwrap(); + + // The files should still exist on disk even though they are no longer active. + let table_path = format!( + "{}/tables/{}", + temp_dir.path().to_str().unwrap(), + NORMAL_TABLE_NAME + ); + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 2); + + context.vacuum_table(NORMAL_TABLE_NAME).await.unwrap(); + + // Only the _delta_log folder should remain. + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 1); + } + + /// Create a [`Context`] with a normal table named `NORMAL_TABLE_NAME` and write data to it. + async fn create_context_with_normal_table(temp_dir: &TempDir) -> Arc { + let context = create_context(temp_dir).await; + + context + .create_normal_table(NORMAL_TABLE_NAME, &table::normal_table_schema()) + .await + .unwrap(); + + // Write data to the normal table. + let local_data_folder = &context.data_folders.local_data_folder; + local_data_folder + .delta_lake + .write_record_batches_to_normal_table( + NORMAL_TABLE_NAME, + vec![table::normal_table_record_batch()], + ) + .await + .unwrap(); + + context + } + + #[tokio::test] + async fn test_vacuum_time_series_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context_with_time_series_table(&temp_dir).await; + + context + .configuration_manager + .write() + .await + .set_retention_period_in_seconds(0); + + context + .truncate_table(TIME_SERIES_TABLE_NAME) + .await + .unwrap(); + + // The files should still exist on disk even though they are no longer active. + let column_path = format!( + "{}/tables/{}/field_column=0", + temp_dir.path().to_str().unwrap(), + TIME_SERIES_TABLE_NAME + ); + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 1); + + context.vacuum_table(TIME_SERIES_TABLE_NAME).await.unwrap(); + + // No files should remain in the column folder. + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 0); + } + + /// Create a [`Context`] with a time series table named `TIME_SERIES_TABLE_NAME` and write data + /// to it. + async fn create_context_with_time_series_table(temp_dir: &TempDir) -> Arc { + let context = create_context(temp_dir).await; + + context + .create_time_series_table(&table::time_series_table_metadata()) + .await + .unwrap(); + + // Write data to the time series table. + let local_data_folder = &context.data_folders.local_data_folder; + local_data_folder + .delta_lake + .write_compressed_segments_to_time_series_table( + TIME_SERIES_TABLE_NAME, + vec![table::compressed_segments_record_batch()], + ) + .await + .unwrap(); + + context + } + + #[tokio::test] + async fn test_vacuum_missing_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context(&temp_dir).await; + + assert!(context.vacuum_table(TIME_SERIES_TABLE_NAME).await.is_err()); + } + #[tokio::test] async fn test_time_series_table_metadata_from_default_database_schema() { let temp_dir = tempfile::tempdir().unwrap(); diff --git a/crates/modelardb_server/src/remote.rs b/crates/modelardb_server/src/remote.rs index 57c1cf6fa..fb9074033 100644 --- a/crates/modelardb_server/src/remote.rs +++ b/crates/modelardb_server/src/remote.rs @@ -434,7 +434,7 @@ impl FlightService for FlightServiceHandler { /// Execute a SQL statement provided in UTF-8 and return the schema of the result followed by /// the result itself. Currently, CREATE TABLE, CREATE TIME SERIES TABLE, EXPLAIN, INCLUDE, - /// SELECT, INSERT, TRUNCATE TABLE, and DROP TABLE are supported. + /// SELECT, INSERT, TRUNCATE TABLE, DROP TABLE, and VACUUM are supported. async fn do_get( &self, request: Request, @@ -514,6 +514,25 @@ impl FlightService for FlightServiceHandler { .map_err(error_to_status_invalid_argument)?; } + Ok(empty_record_batch_stream()) + } + ModelarDbStatement::Vacuum(mut table_names) => { + // Vacuum all tables if no table names are provided. + if table_names.is_empty() { + table_names = self + .context + .default_database_schema() + .map_err(error_to_status_internal)? + .table_names(); + }; + + for table_name in table_names { + self.context + .vacuum_table(&table_name) + .await + .map_err(error_to_status_invalid_argument)?; + } + Ok(empty_record_batch_stream()) } } @@ -737,6 +756,13 @@ impl FlightService for FlightServiceHandler { .await .map_err(error_to_status_internal) } + Ok(protocol::update_configuration::Setting::RetentionPeriodInSeconds) => { + let new_value = new_value.ok_or(invalid_null_error)?; + + configuration_manager.set_retention_period_in_seconds(new_value); + + Ok(()) + } _ => Err(Status::unimplemented(format!( "{setting} is not an updatable setting in the server configuration." ))), diff --git a/crates/modelardb_server/tests/integration_test.rs b/crates/modelardb_server/tests/integration_test.rs index b23c0a254..11746b361 100644 --- a/crates/modelardb_server/tests/integration_test.rs +++ b/crates/modelardb_server/tests/integration_test.rs @@ -48,7 +48,6 @@ use tokio::time; use tonic::transport::Channel; use tonic::{Request, Response, Status, Streaming}; -const TABLE_NAME: &str = "table_name"; const HOST: &str = "127.0.0.1"; /// The next port to be used for the server in an integration test. Each test uses a unique port and @@ -255,6 +254,15 @@ impl TestContext { self.client.do_get(ticket).await } + /// Vacuum a table in the server through the `do_get()` method. + async fn vacuum_table( + &mut self, + table_name: &str, + ) -> Result>, Status> { + let ticket = Ticket::new(format!("VACUUM {table_name}")); + self.client.do_get(ticket).await + } + /// Return a [`RecordBatch`] containing a time series with regular or irregular time stamps /// depending on `generate_irregular_timestamps`, generated values with noise depending on /// `multiply_noise_range`, and an optional tag. @@ -464,13 +472,13 @@ async fn test_can_create_normal_table() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::NormalTable) + .create_table(NORMAL_TABLE_NAME, TableType::NormalTable) .await; let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names.len(), 1); - assert_eq!(retrieved_table_names[0], TABLE_NAME); + assert_eq!(retrieved_table_names[0], NORMAL_TABLE_NAME); } #[tokio::test] @@ -478,7 +486,7 @@ async fn test_can_register_normal_table_after_restart() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::NormalTable) + .create_table(NORMAL_TABLE_NAME, TableType::NormalTable) .await; test_context.restart_server().await; @@ -486,7 +494,7 @@ async fn test_can_register_normal_table_after_restart() { let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names.len(), 1); - assert_eq!(retrieved_table_names[0], TABLE_NAME); + assert_eq!(retrieved_table_names[0], NORMAL_TABLE_NAME); } #[tokio::test] @@ -494,13 +502,13 @@ async fn test_can_create_time_series_table() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names.len(), 1); - assert_eq!(retrieved_table_names[0], TABLE_NAME); + assert_eq!(retrieved_table_names[0], TIME_SERIES_TABLE_NAME); } #[tokio::test] @@ -508,7 +516,7 @@ async fn test_can_register_time_series_table_after_restart() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; test_context.restart_server().await; @@ -516,7 +524,7 @@ async fn test_can_register_time_series_table_after_restart() { let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names.len(), 1); - assert_eq!(retrieved_table_names[0], TABLE_NAME); + assert_eq!(retrieved_table_names[0], TIME_SERIES_TABLE_NAME); } #[tokio::test] @@ -571,13 +579,13 @@ async fn test_can_create_register_and_list_multiple_normal_tables_and_time_serie async fn test_can_drop_normal_table() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::NormalTable) + .create_table(NORMAL_TABLE_NAME, TableType::NormalTable) .await; let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); - assert_eq!(retrieved_table_names[0], TABLE_NAME); + assert_eq!(retrieved_table_names[0], NORMAL_TABLE_NAME); - test_context.drop_table(TABLE_NAME).await.unwrap(); + test_context.drop_table(NORMAL_TABLE_NAME).await.unwrap(); let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names.len(), 0); @@ -585,7 +593,7 @@ async fn test_can_drop_normal_table() { // It should be possible to create a normal table, drop it, and then create a new normal table // with the same name. test_context - .create_table(TABLE_NAME, TableType::NormalTable) + .create_table(NORMAL_TABLE_NAME, TableType::NormalTable) .await; } @@ -593,13 +601,16 @@ async fn test_can_drop_normal_table() { async fn test_can_drop_time_series_table() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); - assert_eq!(retrieved_table_names[0], TABLE_NAME); + assert_eq!(retrieved_table_names[0], TIME_SERIES_TABLE_NAME); - test_context.drop_table(TABLE_NAME).await.unwrap(); + test_context + .drop_table(TIME_SERIES_TABLE_NAME) + .await + .unwrap(); let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names.len(), 0); @@ -607,7 +618,7 @@ async fn test_can_drop_time_series_table() { // It should be possible to create a time series table, drop it, and then create a new time // series table with the same name. test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; } @@ -615,7 +626,7 @@ async fn test_can_drop_time_series_table() { async fn test_cannot_drop_missing_table() { let mut test_context = TestContext::new().await; - let result = test_context.drop_table(TABLE_NAME).await; + let result = test_context.drop_table(NORMAL_TABLE_NAME).await; assert!(result.is_err()); } @@ -627,14 +638,18 @@ async fn test_can_truncate_normal_table() { ingest_time_series_and_flush_data( &mut test_context, slice::from_ref(&time_series), + NORMAL_TABLE_NAME, TableType::NormalTable, ) .await; - test_context.truncate_table(TABLE_NAME).await.unwrap(); + test_context + .truncate_table(NORMAL_TABLE_NAME) + .await + .unwrap(); let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {NORMAL_TABLE_NAME}")) .await .unwrap(); @@ -650,14 +665,18 @@ async fn test_can_truncate_time_series_table() { ingest_time_series_and_flush_data( &mut test_context, slice::from_ref(&time_series), + TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable, ) .await; - test_context.truncate_table(TABLE_NAME).await.unwrap(); + test_context + .truncate_table(TIME_SERIES_TABLE_NAME) + .await + .unwrap(); let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); @@ -669,7 +688,100 @@ async fn test_can_truncate_time_series_table() { async fn test_cannot_truncate_missing_table() { let mut test_context = TestContext::new().await; - let result = test_context.truncate_table(TABLE_NAME).await; + let result = test_context.truncate_table(NORMAL_TABLE_NAME).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_can_vacuum_normal_table() { + let mut test_context = TestContext::new().await; + test_context + .update_configuration( + protocol::update_configuration::Setting::RetentionPeriodInSeconds as i32, + Some(0), + ) + .await + .unwrap(); + + let time_series = TestContext::generate_time_series_with_tag(false, None, Some("location")); + ingest_time_series_and_flush_data( + &mut test_context, + slice::from_ref(&time_series), + NORMAL_TABLE_NAME, + TableType::NormalTable, + ) + .await; + + test_context + .truncate_table(NORMAL_TABLE_NAME) + .await + .unwrap(); + + // The files should still exist on disk even though they are no longer active. + let table_path = format!( + "{}/tables/{}", + test_context.temp_dir.path().to_str().unwrap(), + NORMAL_TABLE_NAME + ); + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 2); + + test_context.vacuum_table(NORMAL_TABLE_NAME).await.unwrap(); + + // Only the _delta_log folder should remain. + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 1); +} + +#[tokio::test] +async fn test_can_vacuum_time_series_table() { + let mut test_context = TestContext::new().await; + test_context + .update_configuration( + protocol::update_configuration::Setting::RetentionPeriodInSeconds as i32, + Some(0), + ) + .await + .unwrap(); + + let time_series = TestContext::generate_time_series_with_tag(false, None, Some("location")); + ingest_time_series_and_flush_data( + &mut test_context, + slice::from_ref(&time_series), + TIME_SERIES_TABLE_NAME, + TableType::TimeSeriesTable, + ) + .await; + + test_context + .truncate_table(TIME_SERIES_TABLE_NAME) + .await + .unwrap(); + + // The files should still exist on disk even though they are no longer active. + let column_path = format!( + "{}/tables/{}/field_column=1", + test_context.temp_dir.path().to_str().unwrap(), + TIME_SERIES_TABLE_NAME + ); + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 1); + + test_context + .vacuum_table(TIME_SERIES_TABLE_NAME) + .await + .unwrap(); + + // No files should remain in the column folder. + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 0); +} + +#[tokio::test] +async fn test_cannot_vacuum_missing_table() { + let mut test_context = TestContext::new().await; + + let result = test_context.vacuum_table(NORMAL_TABLE_NAME).await; assert!(result.is_err()); } @@ -678,10 +790,10 @@ async fn test_can_get_schema() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; - let schema = test_context.retrieve_schema(TABLE_NAME).await; + let schema = test_context.retrieve_schema(TIME_SERIES_TABLE_NAME).await; assert_eq!( schema, @@ -737,12 +849,13 @@ async fn test_do_put_can_ingest_time_series_with_tags() { ingest_time_series_and_flush_data( &mut test_context, slice::from_ref(&time_series), + TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable, ) .await; let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); @@ -753,12 +866,12 @@ async fn test_do_put_can_ingest_time_series_with_tags() { async fn test_insert_can_ingest_time_series_with_tags() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; let insert_result = test_context .execute_query(format!( - "INSERT INTO {TABLE_NAME} VALUES\ + "INSERT INTO {TIME_SERIES_TABLE_NAME} VALUES\ ('2020-01-01 13:00:00', 1, 2, 3, 4, 5, 'Aalborg'),\ ('2020-01-01 13:00:01', 1, 2, 3, 4, 5, 'Aalborg'),\ ('2020-01-01 13:00:02', 1, 2, 3, 4, 5, 'Aalborg'),\ @@ -771,7 +884,7 @@ async fn test_insert_can_ingest_time_series_with_tags() { test_context.flush_data_to_disk().await; let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); @@ -788,12 +901,13 @@ async fn test_do_put_can_ingest_time_series_without_tags() { ingest_time_series_and_flush_data( &mut test_context, slice::from_ref(&time_series), + TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTableNoTag, ) .await; let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); @@ -804,12 +918,12 @@ async fn test_do_put_can_ingest_time_series_without_tags() { async fn test_insert_can_ingest_time_series_without_tags() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTableNoTag) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTableNoTag) .await; let insert_result = test_context .execute_query(format!( - "INSERT INTO {TABLE_NAME} VALUES\ + "INSERT INTO {TIME_SERIES_TABLE_NAME} VALUES\ ('2020-01-01 13:00:00', 1, 2, 3, 4, 5),\ ('2020-01-01 13:00:01', 1, 2, 3, 4, 5),\ ('2020-01-01 13:00:02', 1, 2, 3, 4, 5),\ @@ -822,7 +936,7 @@ async fn test_insert_can_ingest_time_series_without_tags() { test_context.flush_data_to_disk().await; let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); @@ -839,13 +953,16 @@ async fn test_do_put_can_ingest_time_series_with_generated_field() { ingest_time_series_and_flush_data( &mut test_context, slice::from_ref(&time_series), + TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTableAsField, ) .await; // The optimizer is allowed to add SortedJoinExec between SortedJoinExec and GeneratedAsExec. let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME} ORDER BY timestamp")) + .execute_query(format!( + "SELECT * FROM {TIME_SERIES_TABLE_NAME} ORDER BY timestamp" + )) .await .unwrap(); @@ -860,12 +977,12 @@ async fn test_do_put_can_ingest_time_series_with_generated_field() { async fn test_insert_can_ingest_time_series_with_generated_field() { let mut test_context = TestContext::new().await; test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTableAsField) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTableAsField) .await; let insert_result = test_context .execute_query(format!( - "INSERT INTO {TABLE_NAME} VALUES\ + "INSERT INTO {TIME_SERIES_TABLE_NAME} VALUES\ ('2020-01-01 13:00:00', 1, 2, 3, 4),\ ('2020-01-01 13:00:01', 1, 2, 3, 4),\ ('2020-01-01 13:00:02', 1, 2, 3, 4),\ @@ -878,7 +995,7 @@ async fn test_insert_can_ingest_time_series_with_generated_field() { test_context.flush_data_to_disk().await; let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); @@ -897,12 +1014,17 @@ async fn test_do_put_can_ingest_multiple_time_series_with_different_tags() { TestContext::generate_time_series_with_tag(false, None, Some("tag_two")); let time_series = &[time_series_with_tag_one, time_series_with_tag_two]; - ingest_time_series_and_flush_data(&mut test_context, time_series, TableType::TimeSeriesTable) - .await; + ingest_time_series_and_flush_data( + &mut test_context, + time_series, + TIME_SERIES_TABLE_NAME, + TableType::TimeSeriesTable, + ) + .await; let query_result = test_context .execute_query(format!( - "SELECT * FROM {TABLE_NAME} ORDER BY tag, timestamp" + "SELECT * FROM {TIME_SERIES_TABLE_NAME} ORDER BY tag, timestamp" )) .await .unwrap(); @@ -921,11 +1043,13 @@ async fn test_do_put_can_ingest_multiple_time_series_with_different_tags() { async fn test_cannot_ingest_invalid_time_series() { let mut test_context = TestContext::new().await; let time_series = TestContext::generate_time_series_with_tag(false, None, None); - let flight_data = - TestContext::create_flight_data_from_time_series(TABLE_NAME.to_owned(), &[time_series]); + let flight_data = TestContext::create_flight_data_from_time_series( + TIME_SERIES_TABLE_NAME.to_owned(), + &[time_series], + ); test_context - .create_table(TABLE_NAME, TableType::TimeSeriesTable) + .create_table(TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable) .await; assert!( @@ -938,7 +1062,7 @@ async fn test_cannot_ingest_invalid_time_series() { test_context.flush_data_to_disk().await; let query_result = test_context - .execute_query(format!("SELECT * FROM {TABLE_NAME}")) + .execute_query(format!("SELECT * FROM {TIME_SERIES_TABLE_NAME}")) .await .unwrap(); assert_eq!(query_result.num_rows(), 0); @@ -965,6 +1089,7 @@ async fn execute_and_assert_include_select(address_count: usize) { ingest_time_series_and_flush_data( &mut test_context, &[time_series], + TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable, ) .await; @@ -975,7 +1100,9 @@ async fn execute_and_assert_include_select(address_count: usize) { let address = addresses_separate.join(", "); let query_result = test_context - .execute_query(format!("INCLUDE {address} SELECT * FROM {TABLE_NAME}")) + .execute_query(format!( + "INCLUDE {address} SELECT * FROM {TIME_SERIES_TABLE_NAME}" + )) .await .unwrap(); @@ -984,32 +1111,47 @@ async fn execute_and_assert_include_select(address_count: usize) { #[tokio::test] async fn test_count_from_segments_equals_count_from_data_points() { - assert_ne_query_plans_and_eq_result(format!("SELECT COUNT(field_one) FROM {TABLE_NAME}"), 0.0) - .await; + assert_ne_query_plans_and_eq_result( + format!("SELECT COUNT(field_one) FROM {TIME_SERIES_TABLE_NAME}"), + 0.0, + ) + .await; } #[tokio::test] async fn test_min_from_segments_equals_min_from_data_points() { - assert_ne_query_plans_and_eq_result(format!("SELECT MIN(field_one) FROM {TABLE_NAME}"), 0.0) - .await; + assert_ne_query_plans_and_eq_result( + format!("SELECT MIN(field_one) FROM {TIME_SERIES_TABLE_NAME}"), + 0.0, + ) + .await; } #[tokio::test] async fn test_max_from_segments_equals_max_from_data_points() { - assert_ne_query_plans_and_eq_result(format!("SELECT MAX(field_one) FROM {TABLE_NAME}"), 0.0) - .await; + assert_ne_query_plans_and_eq_result( + format!("SELECT MAX(field_one) FROM {TIME_SERIES_TABLE_NAME}"), + 0.0, + ) + .await; } #[tokio::test] async fn test_sum_from_segments_equals_sum_from_data_points() { - assert_ne_query_plans_and_eq_result(format!("SELECT SUM(field_one) FROM {TABLE_NAME}"), 0.001) - .await; + assert_ne_query_plans_and_eq_result( + format!("SELECT SUM(field_one) FROM {TIME_SERIES_TABLE_NAME}"), + 0.001, + ) + .await; } #[tokio::test] async fn test_avg_from_segments_equals_avg_from_data_points() { - assert_ne_query_plans_and_eq_result(format!("SELECT AVG(field_one) FROM {TABLE_NAME}"), 0.001) - .await; + assert_ne_query_plans_and_eq_result( + format!("SELECT AVG(field_one) FROM {TIME_SERIES_TABLE_NAME}"), + 0.001, + ) + .await; } /// Asserts that the query executed on segments in `segment_query` returns a result within @@ -1030,6 +1172,7 @@ async fn assert_ne_query_plans_and_eq_result(segment_query: String, error_bound: ingest_time_series_and_flush_data( &mut test_context, &[time_series], + TIME_SERIES_TABLE_NAME, TableType::TimeSeriesTable, ) .await; @@ -1090,12 +1233,13 @@ async fn assert_ne_query_plans_and_eq_result(segment_query: String, error_bound: async fn ingest_time_series_and_flush_data( test_context: &mut TestContext, time_series: &[RecordBatch], + table_name: &str, table_type: TableType, ) { let flight_data = - TestContext::create_flight_data_from_time_series(TABLE_NAME.to_owned(), time_series); + TestContext::create_flight_data_from_time_series(table_name.to_owned(), time_series); - test_context.create_table(TABLE_NAME, table_type).await; + test_context.create_table(table_name, table_type).await; test_context .send_time_series_to_server(flight_data) @@ -1129,6 +1273,7 @@ async fn test_can_get_configuration() { Some(64 * 1024 * 1024) ); assert_eq!(configuration.transfer_time_in_seconds, None); + assert_eq!(configuration.retention_period_in_seconds, 60 * 60 * 24 * 7); assert_eq!(configuration.ingestion_threads, 1); assert_eq!(configuration.compression_threads, 1); assert_eq!(configuration.writer_threads, 1); @@ -1170,6 +1315,16 @@ async fn test_can_update_compressed_reserved_memory_in_bytes() { assert_eq!(updated_configuration.compressed_reserved_memory_in_bytes, 1); } +#[tokio::test] +async fn test_can_update_retention_period_in_seconds() { + let updated_configuration = update_and_get_configuration( + protocol::update_configuration::Setting::RetentionPeriodInSeconds as i32, + ) + .await; + + assert_eq!(updated_configuration.retention_period_in_seconds, 1); +} + async fn update_and_get_configuration(setting: i32) -> protocol::Configuration { let mut test_context = TestContext::new().await; test_context @@ -1221,6 +1376,7 @@ async fn test_cannot_update_non_nullable_setting_with_null_value() { protocol::update_configuration::Setting::MultivariateReservedMemoryInBytes as i32, protocol::update_configuration::Setting::UncompressedReservedMemoryInBytes as i32, protocol::update_configuration::Setting::CompressedReservedMemoryInBytes as i32, + protocol::update_configuration::Setting::RetentionPeriodInSeconds as i32, ] { update_configuration_and_assert_error( setting, diff --git a/crates/modelardb_storage/Cargo.toml b/crates/modelardb_storage/Cargo.toml index a868e5e03..55c00c69f 100644 --- a/crates/modelardb_storage/Cargo.toml +++ b/crates/modelardb_storage/Cargo.toml @@ -26,6 +26,7 @@ workspace = true arrow.workspace = true async-trait.workspace = true bytes.workspace = true +chrono = "0.4.41" dashmap.workspace = true datafusion.workspace = true datafusion-proto.workspace = true diff --git a/crates/modelardb_storage/src/delta_lake.rs b/crates/modelardb_storage/src/delta_lake.rs index 585763739..8d8662903 100644 --- a/crates/modelardb_storage/src/delta_lake.rs +++ b/crates/modelardb_storage/src/delta_lake.rs @@ -22,6 +22,7 @@ use std::sync::Arc; use arrow::array::{ArrayRef, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema}; +use chrono::TimeDelta; use dashmap::DashMap; use datafusion::catalog::TableProvider; use datafusion::parquet::file::properties::WriterProperties; @@ -465,6 +466,31 @@ impl DeltaLake { Ok(()) } + /// Vacuum the Delta Lake table with `table_name` by deleting all files that are older than + /// `retention_period_in_seconds` seconds. If the retention period is out of bounds or the + /// files could not be deleted, a [`ModelarDbStorageError`] is returned. + pub async fn vacuum_table( + &self, + table_name: &str, + retention_period_in_seconds: usize, + ) -> Result<()> { + let delta_table_ops = self.delta_ops(table_name).await?; + + let retention_period = TimeDelta::new(retention_period_in_seconds as i64, 0).ok_or( + ModelarDbStorageError::InvalidArgument(format!( + "Retention period of {retention_period_in_seconds} seconds is out of bounds." + )), + )?; + + delta_table_ops + .vacuum() + .with_retention_period(retention_period) + .with_enforce_retention_duration(false) + .await?; + + Ok(()) + } + /// Write `columns` to a metadata Delta Lake table with `table_name`. Returns an updated /// [`DeltaTable`] version if the file was written successfully, otherwise returns /// [`ModelarDbStorageError`]. diff --git a/crates/modelardb_storage/src/parser.rs b/crates/modelardb_storage/src/parser.rs index 2f2a5cc56..c5da7a146 100644 --- a/crates/modelardb_storage/src/parser.rs +++ b/crates/modelardb_storage/src/parser.rs @@ -50,8 +50,9 @@ use sqlparser::tokenizer::{Span, Token}; use crate::error::{ModelarDbStorageError, Result}; -/// A top-level statement (CREATE, INSERT, SELECT, TRUNCATE, DROP, etc.) that have been tokenized, -/// parsed, and for which semantic checks have verified that it is compatible with ModelarDB. +/// A top-level statement (CREATE, INSERT, SELECT, TRUNCATE, DROP, VACUUM etc.) that has been +/// tokenized, parsed, and for which semantic checks have verified that it is compatible with +/// ModelarDB. #[derive(Debug)] pub enum ModelarDbStatement { /// CREATE TABLE. @@ -66,12 +67,14 @@ pub enum ModelarDbStatement { DropTable(Vec), /// TRUNCATE TABLE. TruncateTable(Vec), + /// VACUUM. + Vacuum(Vec), } -/// Tokenizes and parses the SQL statement in `sql` and return its parsed representation in the form -/// of a [`ModelarDbStatement`]. Returns a [`ModelarDbStorageError`] if `sql` is empty, contain +/// Tokenizes and parses the SQL statement in `sql` and returns its parsed representation in the form +/// of a [`ModelarDbStatement`]. Returns a [`ModelarDbStorageError`] if `sql` is empty, contains /// multiple statements, or the statement is unsupported. Currently, CREATE TABLE, CREATE TIME SERIES -/// TABLE, INSERT, EXPLAIN, INCLUDE, SELECT, TRUNCATE TABLE, and DROP TABLE are supported. +/// TABLE, INSERT, EXPLAIN, INCLUDE, SELECT, TRUNCATE TABLE, DROP TABLE, and VACUUM are supported. pub fn tokenize_and_parse_sql_statement(sql_statement: &str) -> Result { let mut statements = Parser::parse_sql(&ModelarDbDialect::new(), sql_statement)?; @@ -128,6 +131,11 @@ pub fn tokenize_and_parse_sql_statement(sql_statement: &str) -> Result Ok(ModelarDbStatement::Vacuum( + variable.into_iter().map(|ident| ident.value).collect(), + )), Statement::Explain { .. } => Ok(ModelarDbStatement::Statement(statement)), Statement::Query(ref boxed_query) => { if let Some(addresses) = extract_include_addresses(boxed_query) { @@ -138,7 +146,7 @@ pub fn tokenize_and_parse_sql_statement(sql_statement: &str) -> Result Ok(ModelarDbStatement::Statement(statement)), _ => Err(ModelarDbStorageError::InvalidArgument( - "Only CREATE, DROP, TRUNCATE, EXPLAIN, INCLUDE, SELECT, and INSERT are supported." + "Only CREATE, DROP, TRUNCATE, EXPLAIN, INCLUDE, SELECT, INSERT, and VACUUM are supported." .to_owned(), )), } @@ -164,7 +172,8 @@ pub fn tokenize_and_parse_sql_expression( } /// SQL dialect that extends `sqlparsers's` [`GenericDialect`] with support for parsing CREATE TIME -/// SERIES TABLE table_name DDL statements and INCLUDE 'address'[, 'address']+ DQL statements. +/// SERIES TABLE table_name DDL statements, INCLUDE 'address'\[, 'address'\]+ DQL statements, and +/// VACUUM \[table_name\[, table_name\]+\] statements. #[derive(Debug)] struct ModelarDbDialect { /// Dialect to use for identifying identifiers. @@ -476,6 +485,49 @@ impl ModelarDbDialect { Ok(statement) } + + /// Return [`true`] if the token stream starts with VACUUM, otherwise [`false`] is returned. + /// The method does not consume tokens. + fn next_token_is_vacuum(&self, parser: &Parser) -> bool { + // VACUUM. + if let Token::Word(word) = parser.peek_nth_token(0).token { + word.keyword == Keyword::VACUUM + } else { + false + } + } + + /// Parse VACUUM \[table_name\[, table_name\]+\] to a [`Statement::ShowVariable`] with the + /// table names in the `variable` field. Note that [`Statement::ShowVariable`] is used since + /// [`Statement`] does not have a `Vacuum` variant. A [`ParserError`] is returned if VACUUM is + /// not the first word or the table names cannot be extracted. + fn parse_vacuum(&self, parser: &mut Parser) -> StdResult { + // VACUUM. + parser.expect_keyword(Keyword::VACUUM)?; + + let mut table_names = vec![]; + + if Token::EOF != parser.peek_nth_token(0).token { + loop { + match self.parse_word_value(parser) { + Ok(table_name) => { + table_names.push(Ident::new(table_name)); + if Token::Comma == parser.peek_nth_token(0).token { + parser.next_token(); + } else { + break; + }; + } + Err(error) => return Err(error), + } + } + } + + // Return Statement::ShowVariable as a substitute for Vacuum. + Ok(Statement::ShowVariable { + variable: table_names, + }) + } } /// Create a [`Setting`] with `key`, `quote_style`, and `value`. @@ -504,15 +556,18 @@ impl Dialect for ModelarDbDialect { /// Check if the next tokens are CREATE TIME SERIES TABLE, if so, attempt to parse the token stream /// as a CREATE TIME SERIES TABLE DDL statement. If not, check if the next token is INCLUDE, if so, - /// attempt to parse the token stream as an INCLUDE 'address'[, 'address']+ DQL statement. If - /// both checks fail, [`None`] is returned so sqlparser uses its parsing methods for all other - /// statements. If parsing succeeds, a [`Statement`] is returned, and if not, a [`ParserError`] - /// is returned. + /// attempt to parse the token stream as an INCLUDE 'address'\[, 'address'\]+ DQL statement. + /// If not, check if the next token is VACUUM, if so, attempt to parse the token stream as a + /// VACUUM \[table_name\[, table_name\]+\] statement. If all checks fail, [`None`] is returned + /// so [`sqlparser`] uses its parsing methods for all other statements. If parsing succeeds, a + /// [`Statement`] is returned, and if not, a [`ParserError`] is returned. fn parse_statement(&self, parser: &mut Parser) -> Option> { if self.next_tokens_are_create_time_series_table(parser) { Some(self.parse_create_time_series_table(parser)) } else if self.next_token_is_include(parser) { Some(self.parse_include_query(parser)) + } else if self.next_token_is_vacuum(parser) { + Some(self.parse_vacuum(parser)) } else { None } @@ -1603,4 +1658,52 @@ mod tests { fn test_tokenize_and_parse_include_zero_addresses_select() { assert!(tokenize_and_parse_sql_statement("INCLUDE SELECT * FROM table_name",).is_err()); } + + #[test] + fn test_tokenize_and_parse_vacuum_all_tables() { + let table_names = parse_vacuum_and_extract_table_names("VACUUM"); + + assert!(table_names.is_empty()); + } + + #[test] + fn test_tokenize_and_parse_vacuum_single_table() { + let table_names = parse_vacuum_and_extract_table_names("VACUUM table_name"); + + assert_eq!(table_names, vec!["table_name".to_owned()]); + } + + #[test] + fn test_tokenize_and_parse_vacuum_multiple_tables() { + let table_names = parse_vacuum_and_extract_table_names("VACUUM table_name_1, table_name_2"); + + assert_eq!( + table_names, + vec!["table_name_1".to_owned(), "table_name_2".to_owned()] + ); + } + + fn parse_vacuum_and_extract_table_names(sql_statement: &str) -> Vec { + let modelardb_statement = tokenize_and_parse_sql_statement(sql_statement).unwrap(); + + match modelardb_statement { + ModelarDbStatement::Vacuum(table_names) => table_names, + _ => panic!("Expected ModelarDbStatement::Vacuum."), + } + } + + #[test] + fn test_tokenize_and_parse_vacuum_trailing_comma() { + assert!(tokenize_and_parse_sql_statement("VACUUM table_name,").is_err()); + } + + #[test] + fn test_tokenize_and_parse_vacuum_leading_comma() { + assert!(tokenize_and_parse_sql_statement("VACUUM ,table_name").is_err()); + } + + #[test] + fn test_tokenize_and_parse_vacuum_quoted_table_name() { + assert!(tokenize_and_parse_sql_statement("VACUUM 'table_name'").is_err()); + } } diff --git a/crates/modelardb_types/src/flight/protocol.proto b/crates/modelardb_types/src/flight/protocol.proto index d54ba68e6..77e1524a5 100644 --- a/crates/modelardb_types/src/flight/protocol.proto +++ b/crates/modelardb_types/src/flight/protocol.proto @@ -107,14 +107,17 @@ message Configuration { // The number of seconds between each transfer of data to the remote object store. optional uint64 transfer_time_in_seconds = 5; + // The number of seconds to retain deleted data in storage before it can be removed by vacuum. + uint64 retention_period_in_seconds = 6; + // Number of threads to allocate for converting multivariate time series to univariate time series. - uint32 ingestion_threads = 6; + uint32 ingestion_threads = 7; // Number of threads to allocate for compressing univariate time series to segments. - uint32 compression_threads = 7; + uint32 compression_threads = 8; // Number of threads to allocate for writing segments to a local and/or remote data folder. - uint32 writer_threads = 8; + uint32 writer_threads = 9; } // Request to update the configuration of a ModelarDB node. @@ -125,6 +128,7 @@ message UpdateConfiguration { COMPRESSED_RESERVED_MEMORY_IN_BYTES = 2; TRANSFER_BATCH_SIZE_IN_BYTES = 3; TRANSFER_TIME_IN_SECONDS = 4; + RETENTION_PERIOD_IN_SECONDS = 5; } // Setting to update in the configuration.