From 77e162c5f9017df616b683aa2236ba50218698e8 Mon Sep 17 00:00:00 2001
From: "Meggle (Sebastian Bathke)" <sebastian.bathke@camunda.com>
Date: Tue, 2 Apr 2024 07:46:08 +0200
Subject: [PATCH 1/2] fix: allow insertion of deleted key within single
 transaction

---
 .../test/engine/db/InMemoryDbTransaction.java |  3 ++-
 .../db/InMemoryZeebeDbTransactionTest.java    | 20 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java b/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java
index 3891846fd..e0816e002 100644
--- a/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java
+++ b/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java
@@ -108,6 +108,7 @@ public InMemoryDbIterator newIterator() {
   @Override
   public boolean contains(final FullyQualifiedKey fullyQualifiedKey) {
     final Bytes keyBytes = fullyQualifiedKey.getKeyBytes();
-    return transactionCache.containsKey(keyBytes) || database.containsKey(keyBytes);
+    return !deletedKeys.contains(keyBytes)
+        && (transactionCache.containsKey(keyBytes) || database.containsKey(keyBytes));
   }
 }
diff --git a/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java b/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java
index 3d15101d7..ef9a3ec6b 100644
--- a/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java
+++ b/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java
@@ -440,6 +440,26 @@ void shouldWriteAndDeleteSameKeyValuePairInTransaction() {
     assertThat(oneColumnFamily.exists(oneKey)).isFalse();
   }
 
+  @Test
+  void shouldAllowDeleteAndInsertInTransaction() throws Exception {
+    // given
+    oneKey.wrapLong(1);
+    oneValue.wrapLong(-1L);
+    twoValue.wrapLong(-2L);
+    oneColumnFamily.insert(oneKey, oneValue);
+    transactionContext.getCurrentTransaction().commit();
+
+    // when
+    transactionContext.runInTransaction(
+        () -> {
+          oneColumnFamily.deleteExisting(oneKey);
+          oneColumnFamily.insert(oneKey, twoValue);
+        });
+
+    // then
+    assertThat(oneColumnFamily.get(oneKey).getValue()).isEqualTo(twoValue.getValue());
+  }
+
   @Test
   void shouldNotCommitOnError() {
     // given

From bf9e1eb5c34f2d11387443253ba60dfdf03e11ed Mon Sep 17 00:00:00 2001
From: "Meggle (Sebastian Bathke)" <sebastian.bathke@camunda.com>
Date: Tue, 2 Apr 2024 07:46:08 +0200
Subject: [PATCH 2/2] fix: prevent getting by deleted key within transaction

---
 .../test/engine/db/InMemoryDbTransaction.java |  4 ++++
 .../db/InMemoryZeebeDbTransactionTest.java    | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java b/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java
index e0816e002..467b2d8e9 100644
--- a/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java
+++ b/engine/src/main/java/io/camunda/zeebe/process/test/engine/db/InMemoryDbTransaction.java
@@ -73,6 +73,10 @@ public void put(final FullyQualifiedKey fullyQualifiedKey, final DbValue value)
   public byte[] get(final FullyQualifiedKey fullyQualifiedKey) {
     final Bytes key = fullyQualifiedKey.getKeyBytes();
 
+    if (deletedKeys.contains(key)) {
+      return null;
+    }
+
     final Bytes valueInCache = transactionCache.get(key);
 
     if (valueInCache != null) {
diff --git a/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java b/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java
index ef9a3ec6b..4220267b7 100644
--- a/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java
+++ b/engine/src/test/java/io/camunda/zeebe/process/test/engine/db/InMemoryZeebeDbTransactionTest.java
@@ -460,6 +460,29 @@ void shouldAllowDeleteAndInsertInTransaction() throws Exception {
     assertThat(oneColumnFamily.get(oneKey).getValue()).isEqualTo(twoValue.getValue());
   }
 
+  @Test
+  void shouldNotGetByKeyIfDeletedInTransaction() throws Exception {
+    // given
+    oneKey.wrapLong(1);
+    oneValue.wrapLong(-1L);
+    twoValue.wrapLong(-2L);
+    oneColumnFamily.insert(oneKey, oneValue);
+    transactionContext.getCurrentTransaction().commit();
+
+    // when
+    transactionContext.runInTransaction(
+        () -> {
+          oneColumnFamily.deleteExisting(oneKey);
+
+          if (oneColumnFamily.get(oneKey) != null) {
+            fail("Should not be able to get deleted key.");
+          }
+        });
+
+    // then
+    assertThat(oneColumnFamily.get(oneKey)).isNull();
+  }
+
   @Test
   void shouldNotCommitOnError() {
     // given