duckdb 1.3.4 → 1.4.0

package/src/duckdb/src/storage/temporary_file_manager.cpp

@@ -6,6 +6,7 @@
 #include "duckdb/storage/buffer/temporary_file_information.hpp"
 #include "duckdb/storage/standard_buffer_manager.hpp"
 #include "duckdb/main/database.hpp"
+#include "duckdb/common/encryption_functions.hpp"
 #include "zstd.h"
 
 namespace duckdb {
@@ -69,6 +70,16 @@ TemporaryFileIdentifier::TemporaryFileIdentifier(TemporaryBufferSize size_p, idx
     : size(size_p), file_index(file_index_p) {
 }
 
+TemporaryFileIdentifier::TemporaryFileIdentifier(DatabaseInstance &db, TemporaryBufferSize size_p, idx_t file_index_p,
+                                                 bool encrypted_p)
+    : size(size_p), file_index(file_index_p), encrypted(encrypted_p) {
+
+	if (encrypted) {
+		// generate a random encryption key ID and corresponding key
+		EncryptionEngine::AddTempKeyToCache(db);
+	}
+}
+
 bool TemporaryFileIdentifier::IsValid() const {
 	return size != TemporaryBufferSize::INVALID && file_index.IsValid();
 }
@@ -76,12 +87,13 @@ bool TemporaryFileIdentifier::IsValid() const {
 TemporaryFileIndex::TemporaryFileIndex() {
 }
 
-TemporaryFileIndex::TemporaryFileIndex(TemporaryFileIdentifier identifier_p, idx_t block_index_p)
-    : identifier(identifier_p), block_index(block_index_p) {
+TemporaryFileIndex::TemporaryFileIndex(TemporaryFileIdentifier identifier_p, idx_t block_index_p,
+                                       idx_t block_header_size_p)
+    : identifier(identifier_p), block_index(block_index_p), block_header_size(block_header_size_p) {
 }
 
 bool TemporaryFileIndex::IsValid() const {
-	return identifier.IsValid() && block_index.IsValid();
+	return identifier.IsValid() && block_index.IsValid() && block_header_size.IsValid();
 }
 
 //===--------------------------------------------------------------------===//
@@ -179,10 +191,13 @@ TemporaryFileHandle::TemporaryFileHandle(TemporaryFileManager &manager, Temporar
       path(manager.CreateTemporaryFileName(identifier)), index_manager(manager) {
 }
 
+TemporaryFileHandle::~TemporaryFileHandle() {
+}
+
 TemporaryFileHandle::TemporaryFileLock::TemporaryFileLock(mutex &mutex) : lock(mutex) {
 }
 
-TemporaryFileIndex TemporaryFileHandle::TryGetBlockIndex() {
+TemporaryFileIndex TemporaryFileHandle::TryGetBlockIndex(idx_t block_header_size) {
 	TemporaryFileLock lock(file_lock);
 	if (index_manager.GetMaxIndex() >= max_allowed_index && index_manager.HasFreeBlocks()) {
 		// file is at capacity
@@ -192,26 +207,51 @@ TemporaryFileIndex TemporaryFileHandle::TryGetBlockIndex() {
 	CreateFileIfNotExists(lock);
 	// fetch a new block index to write to
 	auto block_index = index_manager.GetNewBlockIndex(identifier.size);
-	return TemporaryFileIndex(identifier, block_index);
+	return TemporaryFileIndex(identifier, block_index, block_header_size);
 }
 
-unique_ptr<FileBuffer> TemporaryFileHandle::ReadTemporaryBuffer(idx_t block_index,
+unique_ptr<FileBuffer> TemporaryFileHandle::ReadTemporaryBuffer(QueryContext context,
+                                                                const TemporaryFileIndex &index_in_file,
                                                                 unique_ptr<FileBuffer> reusable_buffer) const {
 	auto &buffer_manager = BufferManager::GetBufferManager(db);
-	if (identifier.size == TemporaryBufferSize::DEFAULT) {
-		return StandardBufferManager::ReadTemporaryBufferInternal(
-		    buffer_manager, *handle, GetPositionInFile(block_index), buffer_manager.GetBlockSize(),
-		    std::move(reusable_buffer));
+	auto block_index = index_in_file.block_index.GetIndex();
+	auto block_header_size = index_in_file.block_header_size.GetIndex();
+
+	auto buffer = buffer_manager.ConstructManagedBuffer(buffer_manager.GetBlockAllocSize() - block_header_size,
+	                                                    block_header_size, std::move(reusable_buffer));
+	AllocatedData compressed_buffer;
+	data_ptr_t read_buffer;
+	idx_t read_size;
+	bool is_uncompressed = identifier.size == TemporaryBufferSize::DEFAULT;
+	if (is_uncompressed) {
+		// not compressed - read directly into buffer
+		read_buffer = buffer->InternalBuffer();
+		read_size = buffer->AllocSize();
+	} else {
+		// compressed - read into separate compression buffer first
+		compressed_buffer = Allocator::Get(db).Allocate(TemporaryBufferSizeToSize(identifier.size));
+		read_buffer = compressed_buffer.get();
+		read_size = compressed_buffer.GetSize();
+	}
+
+	idx_t read_position = GetPositionInFile(block_index);
+	if (IsEncrypted()) {
+		uint8_t encryption_metadata[DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE];
+		//! Read nonce and tag.
+		handle->Read(context, encryption_metadata, DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE, read_position);
+		//! Read the encrypted compressed buffer.
+		handle->Read(context, read_buffer, read_size, read_position + DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE);
+		//! Decrypt the compressed buffer.
+		EncryptionEngine::DecryptTemporaryBuffer(db, read_buffer, read_size, encryption_metadata);
+	} else {
+		handle->Read(context, read_buffer, read_size, read_position);
 	}
 
-	// Read compressed buffer
-	auto compressed_buffer = Allocator::Get(db).Allocate(TemporaryBufferSizeToSize(identifier.size));
-	handle->Read(compressed_buffer.get(), compressed_buffer.GetSize(), GetPositionInFile(block_index));
+	if (is_uncompressed) {
+		return buffer;
+	}
 
 	// Decompress into buffer
-	auto buffer = buffer_manager.ConstructManagedBuffer(
-	    buffer_manager.GetBlockSize(), buffer_manager.GetTemporaryBlockHeaderSize(), std::move(reusable_buffer));
-
 	const auto compressed_size = Load<idx_t>(compressed_buffer.get());
 	D_ASSERT(!duckdb_zstd::ZSTD_isError(compressed_size));
 	const auto decompressed_size = duckdb_zstd::ZSTD_decompress(
@@ -227,11 +267,28 @@ void TemporaryFileHandle::WriteTemporaryBuffer(FileBuffer &buffer, const idx_t b
                                                AllocatedData &compressed_buffer) const {
 	// We group DEFAULT_BLOCK_ALLOC_SIZE blocks into the same file.
 	D_ASSERT(buffer.AllocSize() == BufferManager::GetBufferManager(db).GetBlockAllocSize());
-	if (identifier.size == TemporaryBufferSize::DEFAULT) {
-		buffer.Write(*handle, GetPositionInFile(block_index));
+	bool is_uncompressed = identifier.size == TemporaryBufferSize::DEFAULT;
+	data_ptr_t write_buffer;
+	idx_t write_size;
+	if (is_uncompressed) {
+		// write the uncompressed buffer to the file directly
+		write_buffer = buffer.InternalBuffer();
+		write_size = buffer.AllocSize();
+	} else {
+		// write the compressed buffer to the file
+		write_buffer = compressed_buffer.get();
+		write_size = TemporaryBufferSizeToSize(identifier.size);
+	}
+	idx_t write_position = GetPositionInFile(block_index);
+	if (IsEncrypted()) {
+		uint8_t encryption_metadata[DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE];
+		EncryptionEngine::EncryptTemporaryBuffer(db, write_buffer, write_size, encryption_metadata);
+
+		handle->Write(QueryContext(), encryption_metadata, DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE, write_position);
+		handle->Write(QueryContext(), write_buffer, write_size, write_position + DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE);
 	} else {
-		handle->Write(compressed_buffer.get(), TemporaryBufferSizeToSize(identifier.size),
-		              GetPositionInFile(block_index));
+		// write file directly
+		handle->Write(QueryContext(), write_buffer, write_size, write_position);
 	}
 }
 
@@ -255,6 +312,10 @@ bool TemporaryFileHandle::DeleteIfEmpty() {
 	return true;
 }
 
+bool TemporaryFileHandle::IsEncrypted() const {
+	return identifier.encrypted;
+}
+
 TemporaryFileInformation TemporaryFileHandle::GetTemporaryFile() {
 	TemporaryFileLock lock(file_lock);
 	TemporaryFileInformation info;
@@ -289,7 +350,12 @@ void TemporaryFileHandle::RemoveTempBlockIndex(TemporaryFileLock &, idx_t index)
 }
 
 idx_t TemporaryFileHandle::GetPositionInFile(const idx_t index) const {
-	return index * static_cast<idx_t>(identifier.size);
+	idx_t metadata_size = 0;
+	if (IsEncrypted()) {
+		// if the file is encrypted, we need to account for the header size
+		metadata_size = DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE;
+	}
+	return index * (static_cast<idx_t>(identifier.size) + metadata_size);
 }
 
 //===--------------------------------------------------------------------===//
@@ -431,10 +497,11 @@ TemporaryFileManager::~TemporaryFileManager() {
 TemporaryFileManager::TemporaryFileManagerLock::TemporaryFileManagerLock(mutex &mutex) : lock(mutex) {
 }
 
-void TemporaryFileManager::WriteTemporaryBuffer(block_id_t block_id, FileBuffer &buffer) {
+idx_t TemporaryFileManager::WriteTemporaryBuffer(block_id_t block_id, FileBuffer &buffer) {
 	// We group DEFAULT_BLOCK_ALLOC_SIZE blocks into the same file.
 	D_ASSERT(buffer.AllocSize() == BufferManager::GetBufferManager(db).GetBlockAllocSize());
 
+	auto header_size = buffer.GetHeaderSize();
 	const auto adaptivity_idx = TaskScheduler::GetEstimatedCPUId() % COMPRESSION_ADAPTIVITIES;
 	auto &compression_adaptivity = compression_adaptivities[adaptivity_idx];
 
@@ -449,7 +516,7 @@ void TemporaryFileManager::WriteTemporaryBuffer(block_id_t block_id, FileBuffer
 		// first check if we can write to an open existing file
 		for (auto &entry : files.GetMapForSize(compression_result.size)) {
 			auto &temp_file = entry.second;
-			index = temp_file->TryGetBlockIndex();
+			index = temp_file->TryGetBlockIndex(header_size);
 			if (index.IsValid()) {
 				handle = entry.second.get();
 				break;
@@ -458,9 +525,10 @@ void TemporaryFileManager::WriteTemporaryBuffer(block_id_t block_id, FileBuffer
 		if (!handle) {
 			// no existing handle to write to; we need to create & open a new file
 			auto &size = compression_result.size;
-			const TemporaryFileIdentifier identifier(size, index_managers[size].GetNewBlockIndex(size));
+			const TemporaryFileIdentifier identifier(db, size, index_managers[size].GetNewBlockIndex(size),
+			                                         IsEncrypted());
 			auto &new_file = files.CreateFile(identifier);
-			index = new_file.TryGetBlockIndex();
+			index = new_file.TryGetBlockIndex(header_size);
 			handle = &new_file;
 		}
 		D_ASSERT(used_blocks.find(block_id) == used_blocks.end());
@@ -472,6 +540,7 @@ void TemporaryFileManager::WriteTemporaryBuffer(block_id_t block_id, FileBuffer
 	handle->WriteTemporaryBuffer(buffer, index.block_index.GetIndex(), compressed_buffer);
 
 	compression_adaptivity.Update(compression_result.level, time_before_ns);
+	return static_cast<idx_t>(compression_result.size);
 }
 
 TemporaryFileManager::CompressionResult
@@ -572,7 +641,11 @@ void TemporaryFileManager::DecreaseSizeOnDisk(idx_t bytes) {
 	size_on_disk -= bytes;
 }
 
-unique_ptr<FileBuffer> TemporaryFileManager::ReadTemporaryBuffer(block_id_t id,
+bool TemporaryFileManager::IsEncrypted() const {
+	return db.config.options.temp_file_encryption;
+}
+
+unique_ptr<FileBuffer> TemporaryFileManager::ReadTemporaryBuffer(QueryContext context, block_id_t id,
                                                                  unique_ptr<FileBuffer> reusable_buffer) {
 	TemporaryFileIndex index;
 	optional_ptr<TemporaryFileHandle> handle;
@@ -582,7 +655,8 @@ unique_ptr<FileBuffer> TemporaryFileManager::ReadTemporaryBuffer(block_id_t id,
 		handle = GetFileHandle(lock, index.identifier);
 	}
 
-	auto buffer = handle->ReadTemporaryBuffer(index.block_index.GetIndex(), std::move(reusable_buffer));
+	// before the reusable buffer is given,
+	auto buffer = handle->ReadTemporaryBuffer(context, index, std::move(reusable_buffer));
 	{
 		// remove the block (and potentially erase the temp file)
 		TemporaryFileManagerLock lock(manager_lock);
@@ -591,11 +665,12 @@ unique_ptr<FileBuffer> TemporaryFileManager::ReadTemporaryBuffer(block_id_t id,
 	return buffer;
 }
 
-void TemporaryFileManager::DeleteTemporaryBuffer(block_id_t id) {
+idx_t TemporaryFileManager::DeleteTemporaryBuffer(block_id_t id) {
 	TemporaryFileManagerLock lock(manager_lock);
 	auto index = GetTempBlockIndex(lock, id);
 	auto handle = GetFileHandle(lock, index.identifier);
 	EraseUsedBlock(lock, id, *handle, index);
+	return static_cast<idx_t>(index.identifier.size);
 }
 
 vector<TemporaryFileInformation> TemporaryFileManager::GetTemporaryFiles() {

package/src/duckdb/src/storage/wal_replay.cpp

@@ -4,8 +4,10 @@
 #include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp"
 #include "duckdb/catalog/catalog_entry/type_catalog_entry.hpp"
 #include "duckdb/catalog/catalog_entry/view_catalog_entry.hpp"
+#include "duckdb/catalog/duck_catalog.hpp"
 #include "duckdb/common/checksum.hpp"
-#include "duckdb/common/printer.hpp"
+#include "duckdb/common/encryption_functions.hpp"
+#include "duckdb/common/encryption_key_manager.hpp"
 #include "duckdb/common/serializer/binary_deserializer.hpp"
 #include "duckdb/common/serializer/buffered_file_reader.hpp"
 #include "duckdb/common/serializer/memory_stream.hpp"
@@ -24,11 +26,12 @@
 #include "duckdb/planner/binder.hpp"
 #include "duckdb/planner/expression_binder/index_binder.hpp"
 #include "duckdb/planner/parsed_data/bound_create_table_info.hpp"
+#include "duckdb/storage/single_file_block_manager.hpp"
 #include "duckdb/storage/storage_manager.hpp"
+#include "duckdb/storage/table/column_data.hpp"
 #include "duckdb/storage/table/delete_state.hpp"
 #include "duckdb/storage/write_ahead_log.hpp"
 #include "duckdb/transaction/meta_transaction.hpp"
-#include "duckdb/storage/table/column_data.hpp"
 
 namespace duckdb {
 
@@ -43,6 +46,7 @@ public:
 	optional_ptr<TableCatalogEntry> current_table;
 	MetaBlockPointer checkpoint_id;
 	idx_t wal_version = 1;
+	optional_idx expected_checkpoint_id;
 
 	struct ReplayIndexInfo {
 		ReplayIndexInfo(TableIndexList &index_list, unique_ptr<Index> index, const string &table_schema,
@@ -77,35 +81,105 @@ public:
 			// old WAL versions do not have checksums
 			return WriteAheadLogDeserializer(state_p, stream, deserialize_only);
 		}
-		if (state_p.wal_version != 2) {
-			throw IOException("Failed to read WAL of version %llu - can only read version 1 and 2",
-			                  state_p.wal_version);
-		}
-		// read the checksum and size
-		auto size = stream.Read<uint64_t>();
-		auto stored_checksum = stream.Read<uint64_t>();
-		auto offset = stream.CurrentOffset();
-		auto file_size = stream.FileSize();
-
-		if (offset + size > file_size) {
-			throw SerializationException(
-			    "Corrupt WAL file: entry size exceeded remaining data in file at byte position %llu "
-			    "(found entry with size %llu bytes, file size %llu bytes)",
-			    offset, size, file_size);
+
+		if (state_p.wal_version == 2) {
+			// read the size and checksum
+			auto size = stream.Read<uint64_t>();
+			auto stored_checksum = stream.Read<uint64_t>();
+			auto offset = stream.CurrentOffset();
+			auto file_size = stream.FileSize();
+
+			if (offset + size > file_size) {
+				throw SerializationException(
+				    "Corrupt WAL file: entry size exceeded remaining data in file at byte position %llu "
+				    "(found entry with size %llu bytes, file size %llu bytes)",
+				    offset, size, file_size);
+			}
+
+			// allocate a buffer and read data into the buffer
+			auto buffer = unique_ptr<data_t[]>(new data_t[size]);
+			stream.ReadData(buffer.get(), size);
+
+			// compute and verify the checksum
+			auto computed_checksum = Checksum(buffer.get(), size);
+			if (stored_checksum != computed_checksum) {
+				throw IOException("Corrupt WAL file: entry at byte position %llu computed checksum %llu does not match "
+				                  "stored checksum %llu",
+				                  offset, computed_checksum, stored_checksum);
+			}
+
+			return WriteAheadLogDeserializer(state_p, std::move(buffer), size, deserialize_only);
 		}
 
-		// allocate a buffer and read data into the buffer
-		auto buffer = unique_ptr<data_t[]>(new data_t[size]);
-		stream.ReadData(buffer.get(), size);
+		if (state_p.wal_version == 3) {
+			auto &database = state_p.db.GetDatabase();
+			//! Version 3 means that the WAL is encrypted
+			//! For encryption, the length field remains plaintext
+			//! After the length field, we store a 12-byte nonce (for GCM)
+			//! After the nonce, we store the checksum, followed by the actual entry
+			//! After the stored entry, we store a 16-byte nonce (for GCM).
+
+			// read the size (this is excluding the nonce, checksum and tag)
+			auto size = stream.Read<uint64_t>();
+			// the ciphertext size is including the checksum
+			const auto ciphertext_size = size + sizeof(uint64_t);
+
+			auto offset = stream.CurrentOffset();
+			auto file_size = stream.FileSize();
+
+			EncryptionNonce nonce;
+			EncryptionTag tag;
+
+			if (offset + nonce.size() + ciphertext_size + tag.size() > file_size) {
+				throw SerializationException(
+				    "Corrupt Encrypted WAL file: entry size exceeded remaining data in file at byte position %llu "
+				    "(found entry with size %llu bytes, file size %llu bytes)",
+				    offset, size, file_size);
+			}
+
+			stream.ReadData(nonce.data(), nonce.size());
+
+			auto &keys = EncryptionKeyManager::Get(state_p.db.GetDatabase());
+			auto &catalog = state_p.db.GetCatalog().Cast<DuckCatalog>();
+			auto derived_key = keys.GetKey(catalog.GetEncryptionKeyId());
+
+			//! initialize the decryption
+			auto encryption_state = database.GetEncryptionUtil()->CreateEncryptionState(
+			    state_p.db.GetStorageManager().GetCipher(), MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH);
+			encryption_state->InitializeDecryption(nonce.data(), nonce.size(), derived_key,
+			                                       MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH);
+
+			//! Allocate a decryption buffer
+			auto buffer = unique_ptr<data_t[]>(new data_t[ciphertext_size]);
+			auto out_buffer = unique_ptr<data_t[]>(new data_t[size]);
+
+			stream.ReadData(buffer.get(), ciphertext_size);
+			encryption_state->Process(buffer.get(), ciphertext_size, buffer.get(), ciphertext_size);
+
+			//! read and verify the stored tag
+			stream.ReadData(tag.data(), tag.size());
+
+			encryption_state->Finalize(buffer.get(), ciphertext_size, tag.data(), tag.size());
+
+			//! read the stored checksum
+			auto stored_checksum = Load<uint64_t>(buffer.get());
 
-		// compute and verify the checksum
-		auto computed_checksum = Checksum(buffer.get(), size);
-		if (stored_checksum != computed_checksum) {
-			throw IOException("Corrupt WAL file: entry at byte position %llu computed checksum %llu does not match "
-			                  "stored checksum %llu",
-			                  offset, computed_checksum, stored_checksum);
+			//! copy the decrypted data to the output buffer
+			memcpy(out_buffer.get(), buffer.get() + sizeof(stored_checksum), size);
+
+			// compute and verify the checksum
+			auto computed_checksum = Checksum(out_buffer.get(), size);
+			if (stored_checksum != computed_checksum) {
+				throw IOException("Corrupt WAL file: entry at byte position %llu computed checksum %llu does not match "
+				                  "stored checksum %llu",
+				                  offset, computed_checksum, stored_checksum);
+			}
+
+			return WriteAheadLogDeserializer(state_p, std::move(out_buffer), size, deserialize_only);
 		}
-		return WriteAheadLogDeserializer(state_p, std::move(buffer), size, deserialize_only);
+
+		throw IOException("Failed to read WAL of version %llu - can only read version 1, 2 and 3 (encrypted)",
+		                  state_p.wal_version);
 	}
 
 	bool ReplayEntry() {
@@ -120,10 +194,12 @@ public:
 		return false;
 	}
 
-	bool DeserializeOnly() {
+	bool DeserializeOnly() const {
 		return deserialize_only;
 	}
 
+	static void ThrowVersionError(idx_t checkpoint_iteration, idx_t expected_checkpoint_iteration);
+
 protected:
 	void ReplayEntry(WALType wal_type);
 
@@ -162,6 +238,9 @@ protected:
 	void ReplayUpdate();
 	void ReplayCheckpoint();
 
+private:
+	void ReplayIndexData(IndexStorageInfo &info);
+
 private:
 	ReplayState &state;
 	AttachedDatabase &db;
@@ -171,6 +250,7 @@ private:
 	MemoryStream stream;
 	BinaryDeserializer deserializer;
 	bool deserialize_only;
+	optional_idx expected_checkpoint_id;
 };
 
 //===--------------------------------------------------------------------===//
@@ -236,6 +316,11 @@ unique_ptr<WriteAheadLog> WriteAheadLog::ReplayInternal(AttachedDatabase &databa
 			return nullptr;
 		}
 	}
+	if (checkpoint_state.expected_checkpoint_id.IsValid()) {
+		// we expected a checkpoint id - but no checkpoint has happened - abort!
+		auto expected_id = checkpoint_state.expected_checkpoint_id.GetIndex();
+		WriteAheadLogDeserializer::ThrowVersionError(expected_id - 1, expected_id);
+	}
 
 	// we need to recover from the WAL: actually set up the replay state
 	ReplayState state(database, *con.context);
@@ -261,7 +346,7 @@ unique_ptr<WriteAheadLog> WriteAheadLog::ReplayInternal(AttachedDatabase &databa
 				}
 				state.replay_index_infos.clear();
 
-				successful_offset = reader.offset;
+				successful_offset = reader.CurrentOffset();
 				// check if the file is exhausted
 				if (reader.Finished()) {
 					// we finished reading the file: break
@@ -379,8 +464,46 @@ void WriteAheadLogDeserializer::ReplayEntry(WALType entry_type) {
 //===--------------------------------------------------------------------===//
 // Replay Version
 //===--------------------------------------------------------------------===//
+void WriteAheadLogDeserializer::ThrowVersionError(idx_t checkpoint_iteration, idx_t expected_checkpoint_iteration) {
+	string relation = checkpoint_iteration < expected_checkpoint_iteration ? "an older" : "a newer";
+	throw IOException("This WAL was created for this database file, but the WAL checkpoint iteration does not "
+	                  "match the database file. "
+	                  "That means the WAL was created for %s version of this database. File checkpoint "
+	                  "iteration: %d, WAL checkpoint iteration: %d",
+	                  relation, expected_checkpoint_iteration, checkpoint_iteration);
+}
+
 void WriteAheadLogDeserializer::ReplayVersion() {
 	state.wal_version = deserializer.ReadProperty<idx_t>(101, "version");
+
+	auto &single_file_block_manager = db.GetStorageManager().GetBlockManager().Cast<SingleFileBlockManager>();
+	data_t db_identifier[MainHeader::DB_IDENTIFIER_LEN];
+	bool is_set = false;
+	deserializer.ReadOptionalList(102, "db_identifier", [&](Deserializer::List &list, idx_t i) {
+		db_identifier[i] = list.ReadElement<uint8_t>();
+		is_set = true;
+	});
+	auto checkpoint_iteration = deserializer.ReadPropertyWithDefault<optional_idx>(103, "checkpoint_iteration");
+	if (!is_set || !checkpoint_iteration.IsValid()) {
+		return;
+	}
+	auto expected_db_identifier = single_file_block_manager.GetDBIdentifier();
+	if (!MainHeader::CompareDBIdentifiers(db_identifier, expected_db_identifier)) {
+		throw IOException("WAL does not match database file.");
+	}
+
+	auto wal_checkpoint_iteration = checkpoint_iteration.GetIndex();
+	auto expected_checkpoint_iteration = single_file_block_manager.GetCheckpointIteration();
+	if (expected_checkpoint_iteration != wal_checkpoint_iteration) {
+		if (wal_checkpoint_iteration + 1 == expected_checkpoint_iteration) {
+			// this iteration is exactly one lower than the expected iteration
+			// this can happen if we aborted AFTER checkpointing the file, but BEFORE truncating the WAL
+			// expect this situation to occur - we will throw an error if it does not later on
+			state.expected_checkpoint_id = expected_checkpoint_iteration;
+			return;
+		}
+		ThrowVersionError(wal_checkpoint_iteration, expected_checkpoint_iteration);
+	}
 }
 
 //===--------------------------------------------------------------------===//
@@ -427,12 +550,10 @@ void ReplayWithoutIndex(ClientContext &context, Catalog &catalog, AlterInfo &inf
 	catalog.Alter(context, info);
 }
 
-void ReplayIndexData(AttachedDatabase &db, BinaryDeserializer &deserializer, IndexStorageInfo &info,
-                     const bool deserialize_only) {
+void WriteAheadLogDeserializer::ReplayIndexData(IndexStorageInfo &info) {
 	D_ASSERT(info.IsValid() && !info.name.empty());
 
-	auto &storage_manager = db.GetStorageManager();
-	auto &single_file_sm = storage_manager.Cast<SingleFileStorageManager>();
+	auto &single_file_sm = db.GetStorageManager().Cast<SingleFileStorageManager>();
 	auto &block_manager = single_file_sm.block_manager;
 	auto &buffer_manager = block_manager->buffer_manager;
 
@@ -452,7 +573,8 @@ void ReplayIndexData(AttachedDatabase &db, BinaryDeserializer &deserializer, Ind
 			// Convert the buffer handle to a persistent block and store the block id.
 			if (!deserialize_only) {
 				auto block_id = block_manager->GetFreeBlockId();
-				block_manager->ConvertToPersistent(block_id, std::move(block_handle), std::move(buffer_handle));
+				block_manager->ConvertToPersistent(QueryContext(context), block_id, std::move(block_handle),
+				                                   std::move(buffer_handle));
 				data_info.block_pointers[j].block_id = block_id;
 			}
 		}
@@ -467,7 +589,7 @@ void WriteAheadLogDeserializer::ReplayAlter() {
 	}
 
 	auto index_storage_info = deserializer.ReadProperty<IndexStorageInfo>(102, "index_storage_info");
-	ReplayIndexData(db, deserializer, index_storage_info, DeserializeOnly());
+	ReplayIndexData(index_storage_info);
 	if (DeserializeOnly()) {
 		return;
 	}
@@ -684,7 +806,7 @@ void WriteAheadLogDeserializer::ReplayCreateIndex() {
 	auto create_info = deserializer.ReadProperty<unique_ptr<CreateInfo>>(101, "index_catalog_entry");
 	auto index_info = deserializer.ReadProperty<IndexStorageInfo>(102, "index_storage_info");
 
-	ReplayIndexData(db, deserializer, index_info, DeserializeOnly());
+	ReplayIndexData(index_info);
 	if (DeserializeOnly()) {
 		return;
 	}