Add e2e test

This reverts commit 5f2af9050a.

.github/workflows/diff.yaml

@@ -112,7 +112,7 @@ jobs:
           source /opt/toolchain-v4/activate
 
           # Restrict clang-tidy results only to the modified parts
-          git diff -U0 master... -- src ':!*.hpp' | ./tools/github/clang-tidy/clang-tidy-diff.py -p 1 -j $THREADS -path build | tee ./build/clang_tidy_output.txt
+          git diff -U0 master... -- src | ./tools/github/clang-tidy/clang-tidy-diff.py -p 1 -j $THREADS -path build | tee ./build/clang_tidy_output.txt
 
           # Fail if any warning is reported
           ! cat ./build/clang_tidy_output.txt | ./tools/github/clang-tidy/grep_error_lines.sh > /dev/null

libs/.gitignore

@@ -4,5 +4,4 @@
 !cleanup.sh
 !CMakeLists.txt
 !__main.cpp
-!jemalloc.cmake
 !pulsar.patch

libs/jemalloc.cmake

@@ -1,55 +0,0 @@
-set(JEMALLOC_DIR "${LIB_DIR}/jemalloc")
-
-set(JEMALLOC_SRCS
-    ${JEMALLOC_DIR}/src/arena.c
-    ${JEMALLOC_DIR}/src/background_thread.c
-    ${JEMALLOC_DIR}/src/base.c
-    ${JEMALLOC_DIR}/src/bin.c
-    ${JEMALLOC_DIR}/src/bitmap.c
-    ${JEMALLOC_DIR}/src/ckh.c
-    ${JEMALLOC_DIR}/src/ctl.c
-    ${JEMALLOC_DIR}/src/div.c
-    ${JEMALLOC_DIR}/src/extent.c
-    ${JEMALLOC_DIR}/src/extent_dss.c
-    ${JEMALLOC_DIR}/src/extent_mmap.c
-    ${JEMALLOC_DIR}/src/hash.c
-    ${JEMALLOC_DIR}/src/hook.c
-    ${JEMALLOC_DIR}/src/jemalloc.c
-    ${JEMALLOC_DIR}/src/large.c
-    ${JEMALLOC_DIR}/src/log.c
-    ${JEMALLOC_DIR}/src/malloc_io.c
-    ${JEMALLOC_DIR}/src/mutex.c
-    ${JEMALLOC_DIR}/src/mutex_pool.c
-    ${JEMALLOC_DIR}/src/nstime.c
-    ${JEMALLOC_DIR}/src/pages.c
-    ${JEMALLOC_DIR}/src/prng.c
-    ${JEMALLOC_DIR}/src/prof.c
-    ${JEMALLOC_DIR}/src/rtree.c
-    ${JEMALLOC_DIR}/src/sc.c
-    ${JEMALLOC_DIR}/src/stats.c
-    ${JEMALLOC_DIR}/src/sz.c
-    ${JEMALLOC_DIR}/src/tcache.c
-    ${JEMALLOC_DIR}/src/test_hooks.c
-    ${JEMALLOC_DIR}/src/ticker.c
-    ${JEMALLOC_DIR}/src/tsd.c
-    ${JEMALLOC_DIR}/src/witness.c
-    ${JEMALLOC_DIR}/src/safety_check.c
-)
-
-add_library(jemalloc ${JEMALLOC_SRCS})
-target_include_directories(jemalloc PUBLIC "${JEMALLOC_DIR}/include")
-
-find_package(Threads REQUIRED)
-target_link_libraries(jemalloc PUBLIC Threads::Threads)
-
-target_compile_definitions(jemalloc PRIVATE -DJEMALLOC_NO_PRIVATE_NAMESPACE)
-
-if (CMAKE_BUILD_TYPE STREQUAL "DEBUG")
-    target_compile_definitions(jemalloc PRIVATE -DJEMALLOC_DEBUG=1 -DJEMALLOC_PROF=1)
-endif()
-
-target_compile_options(jemalloc PRIVATE -Wno-redundant-decls)
-# for RTLD_NEXT
-target_compile_definitions(jemalloc PRIVATE _GNU_SOURCE)
-
-set_property(TARGET jemalloc APPEND PROPERTY INTERFACE_COMPILE_DEFINITIONS USE_JEMALLOC=1)

src/communication/v2/session.hpp

@@ -105,9 +105,16 @@ class WebsocketSession : public std::enable_shared_from_this<WebsocketSession<TS
     boost::asio::socket_base::keep_alive option(true);
 
     // Set a decorator to change the Server of the handshake
-    ws_.set_option(boost::beast::websocket::stream_base::decorator([](boost::beast::websocket::response_type &res) {
+    ws_.set_option(boost::beast::websocket::stream_base::decorator([&req](boost::beast::websocket::response_type &res) {
       res.set(boost::beast::http::field::server, std::string("Memgraph Bolt WS"));
-      res.set(boost::beast::http::field::sec_websocket_protocol, "binary");
+
+      // We need to do this to support WASM clients, which explicitly send this flag
+      // in their upgrade request
+      // Neo4j client breaks when this flag is sent
+      if (const auto secondary_protocol = req.base().find(boost::beast::http::field::sec_websocket_protocol);
+          secondary_protocol != res.base().end() && secondary_protocol->value() == "binary") {
+        res.set(boost::beast::http::field::sec_websocket_protocol, "binary");
+      }
     }));
     ws_.binary(true);
 
@@ -162,7 +169,7 @@ class WebsocketSession : public std::enable_shared_from_this<WebsocketSession<TS
         boost::asio::bind_executor(strand_, std::bind_front(&WebsocketSession::OnRead, shared_from_this())));
   }
 
-  void OnRead(const boost::system::error_code &ec, [[maybe_unused]] const size_t bytes_transferred) {
+  void OnRead(const boost::system::error_code &ec, const size_t bytes_transferred) {
     // This indicates that the WebsocketSession was closed
     if (ec == boost::beast::websocket::error::closed) {
       return;
@@ -322,7 +329,8 @@ class Session final : public std::enable_shared_from_this<Session<TSession, TSes
       socket.lowest_layer().non_blocking(false);
     });
     timeout_timer_.expires_at(boost::asio::steady_timer::time_point::max());
-    spdlog::info("Accepted a connection from {}:", service_name_, remote_endpoint_.address(), remote_endpoint_.port());
+    spdlog::info("Accepted a connection from {}: {}:{}", service_name_, remote_endpoint_.address(),
+                 remote_endpoint_.port());
   }
 
   void DoRead() {

src/integrations/constants.hpp

@@ -17,6 +17,7 @@
 namespace memgraph::integrations {
 
 inline constexpr int64_t kDefaultCheckBatchLimit{1};
+inline constexpr int64_t kMinimumStartBatchLimit{1};
 inline constexpr std::chrono::milliseconds kDefaultCheckTimeout{30000};
 inline constexpr std::chrono::milliseconds kMinimumInterval{1};
 inline constexpr int64_t kMinimumSize{1};

src/integrations/kafka/consumer.cpp

@@ -74,6 +74,36 @@ utils::BasicResult<std::string, std::vector<Message>> GetBatch(RdKafka::KafkaCon
 
   return std::move(batch);
 }
+
+void CheckAndDestroyLastAssignmentIfNeeded(RdKafka::KafkaConsumer &consumer, const ConsumerInfo &info,
+                                           std::vector<RdKafka::TopicPartition *> &last_assignment) {
+  if (!last_assignment.empty()) {
+    if (const auto err = consumer.assign(last_assignment); err != RdKafka::ERR_NO_ERROR) {
+      throw ConsumerStartFailedException(info.consumer_name,
+                                         fmt::format("Couldn't restore commited offsets: '{}'", RdKafka::err2str(err)));
+    }
+    RdKafka::TopicPartition::destroy(last_assignment);
+  }
+}
+
+void TryToConsumeBatch(RdKafka::KafkaConsumer &consumer, const ConsumerInfo &info,
+                       const ConsumerFunction &consumer_function, const std::vector<Message> &batch) {
+  consumer_function(batch);
+  std::vector<RdKafka::TopicPartition *> partitions;
+  utils::OnScopeExit clear_partitions([&]() { RdKafka::TopicPartition::destroy(partitions); });
+
+  if (const auto err = consumer.assignment(partitions); err != RdKafka::ERR_NO_ERROR) {
+    throw ConsumerCommitFailedException(
+        info.consumer_name, fmt::format("Couldn't get assignment to commit offsets: {}", RdKafka::err2str(err)));
+  }
+  if (const auto err = consumer.position(partitions); err != RdKafka::ERR_NO_ERROR) {
+    throw ConsumerCommitFailedException(info.consumer_name,
+                                        fmt::format("Couldn't get offsets from librdkafka {}", RdKafka::err2str(err)));
+  }
+  if (const auto err = consumer.commitSync(partitions); err != RdKafka::ERR_NO_ERROR) {
+    throw ConsumerCommitFailedException(info.consumer_name, RdKafka::err2str(err));
+  }
+}
 }  // namespace
 
 Message::Message(std::unique_ptr<RdKafka::Message> &&message) : message_{std::move(message)} {
@@ -221,10 +251,21 @@ void Consumer::Start() {
   StartConsuming();
 }
 
-void Consumer::StartIfStopped() {
-  if (!is_running_) {
-    StartConsuming();
+void Consumer::StartWithLimit(const uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const {
+  if (is_running_) {
+    throw ConsumerRunningException(info_.consumer_name);
   }
+  if (limit_batches < kMinimumStartBatchLimit) {
+    throw ConsumerStartFailedException(
+        info_.consumer_name, fmt::format("Batch limit has to be greater than or equal to {}", kMinimumStartBatchLimit));
+  }
+  if (timeout.value_or(kMinimumInterval) < kMinimumInterval) {
+    throw ConsumerStartFailedException(
+        info_.consumer_name,
+        fmt::format("Timeout has to be greater than or equal to {} milliseconds", kMinimumInterval.count()));
+  }
+
+  StartConsumingWithLimit(limit_batches, timeout);
 }
 
 void Consumer::Stop() {
@@ -244,7 +285,7 @@ void Consumer::StopIfRunning() {
   }
 }
 
-void Consumer::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> limit_batches,
+void Consumer::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> limit_batches,
                      const ConsumerFunction &check_consumer_function) const {
   // NOLINTNEXTLINE (modernize-use-nullptr)
   if (timeout.value_or(kMinimumInterval) < kMinimumInterval) {
@@ -344,13 +385,7 @@ void Consumer::StartConsuming() {
 
   is_running_.store(true);
 
-  if (!last_assignment_.empty()) {
-    if (const auto err = consumer_->assign(last_assignment_); err != RdKafka::ERR_NO_ERROR) {
-      throw ConsumerStartFailedException(info_.consumer_name,
-                                         fmt::format("Couldn't restore commited offsets: '{}'", RdKafka::err2str(err)));
-    }
-    RdKafka::TopicPartition::destroy(last_assignment_);
-  }
+  CheckAndDestroyLastAssignmentIfNeeded(*consumer_, info_, last_assignment_);
 
   thread_ = std::thread([this] {
     static constexpr auto kMaxThreadNameSize = utils::GetMaxThreadNameSize();
@@ -361,33 +396,18 @@ void Consumer::StartConsuming() {
     while (is_running_) {
       auto maybe_batch = GetBatch(*consumer_, info_, is_running_);
       if (maybe_batch.HasError()) {
-        spdlog::warn("Error happened in consumer {} while fetching messages: {}!", info_.consumer_name,
-                     maybe_batch.GetError());
-        break;
+        throw ConsumerReadMessagesFailedException(info_.consumer_name, maybe_batch.GetError());
       }
       const auto &batch = maybe_batch.GetValue();
 
-      if (batch.empty()) continue;
+      if (batch.empty()) {
+        continue;
+      }
 
       spdlog::info("Kafka consumer {} is processing a batch", info_.consumer_name);
 
       try {
-        consumer_function_(batch);
-        std::vector<RdKafka::TopicPartition *> partitions;
-        utils::OnScopeExit clear_partitions([&]() { RdKafka::TopicPartition::destroy(partitions); });
-
-        if (const auto err = consumer_->assignment(partitions); err != RdKafka::ERR_NO_ERROR) {
-          throw ConsumerCheckFailedException(
-              info_.consumer_name, fmt::format("Couldn't get assignment to commit offsets: {}", RdKafka::err2str(err)));
-        }
-        if (const auto err = consumer_->position(partitions); err != RdKafka::ERR_NO_ERROR) {
-          throw ConsumerCheckFailedException(
-              info_.consumer_name, fmt::format("Couldn't get offsets from librdkafka {}", RdKafka::err2str(err)));
-        }
-        if (const auto err = consumer_->commitSync(partitions); err != RdKafka::ERR_NO_ERROR) {
-          spdlog::warn("Committing offset of consumer {} failed: {}", info_.consumer_name, RdKafka::err2str(err));
-          break;
-        }
+        TryToConsumeBatch(*consumer_, info_, consumer_function_, batch);
       } catch (const std::exception &e) {
         spdlog::warn("Error happened in consumer {} while processing a batch: {}!", info_.consumer_name, e.what());
         break;
@@ -398,6 +418,44 @@ void Consumer::StartConsuming() {
   });
 }
 
+void Consumer::StartConsumingWithLimit(uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const {
+  MG_ASSERT(!is_running_, "Cannot start already running consumer!");
+
+  if (is_running_.exchange(true)) {
+    throw ConsumerRunningException(info_.consumer_name);
+  }
+  utils::OnScopeExit restore_is_running([this] { is_running_.store(false); });
+
+  CheckAndDestroyLastAssignmentIfNeeded(*consumer_, info_, last_assignment_);
+
+  const auto timeout_to_use = timeout.value_or(kDefaultCheckTimeout);
+  const auto start = std::chrono::steady_clock::now();
+
+  for (uint64_t batch_count = 0; batch_count < limit_batches;) {
+    const auto now = std::chrono::steady_clock::now();
+    if (now - start >= timeout_to_use) {
+      throw ConsumerStartFailedException(info_.consumer_name, "Timeout reached");
+    }
+
+    const auto maybe_batch = GetBatch(*consumer_, info_, is_running_);
+    if (maybe_batch.HasError()) {
+      throw ConsumerReadMessagesFailedException(info_.consumer_name, maybe_batch.GetError());
+    }
+    const auto &batch = maybe_batch.GetValue();
+
+    if (batch.empty()) {
+      continue;
+    }
+    ++batch_count;
+
+    spdlog::info("Kafka consumer {} is processing a batch", info_.consumer_name);
+
+    TryToConsumeBatch(*consumer_, info_, consumer_function_, batch);
+
+    spdlog::info("Kafka consumer {} finished processing", info_.consumer_name);
+  }
+}
+
 void Consumer::StopConsuming() {
   is_running_.store(false);
   if (thread_.joinable()) thread_.join();

src/integrations/kafka/consumer.hpp

@@ -113,11 +113,19 @@ class Consumer final : public RdKafka::EventCb {
   /// This method will start a new thread which will poll all the topics for messages.
   ///
   /// @throws ConsumerRunningException if the consumer is already running
+  /// @throws ConsumerStartFailedException if the commited offsets cannot be restored
   void Start();
 
-  /// Starts consuming messages if it is not started already.
+  /// Starts consuming messages.
   ///
-  void StartIfStopped();
+  /// This method will start a new thread which will poll all the topics for messages.
+  ///
+  /// @param limit_batches the consumer will only consume the given number of batches.
+  /// @param timeout the maximum duration during which the command should run.
+  ///
+  /// @throws ConsumerRunningException if the consumer is already running
+  /// @throws ConsumerStartFailedException if the commited offsets cannot be restored
+  void StartWithLimit(uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const;
 
   /// Stops consuming messages.
   ///
@@ -136,9 +144,9 @@ class Consumer final : public RdKafka::EventCb {
   ///                      used.
   /// @param check_consumer_function a function to feed the received messages in, only used during this dry-run.
   ///
-  /// @throws ConsumerRunningException if the consumer is alredy running.
+  /// @throws ConsumerRunningException if the consumer is already running.
   /// @throws ConsumerCheckFailedException if check isn't successful.
-  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> limit_batches,
+  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> limit_batches,
              const ConsumerFunction &check_consumer_function) const;
 
   /// Returns true if the consumer is actively consuming messages.
@@ -157,6 +165,7 @@ class Consumer final : public RdKafka::EventCb {
   void event_cb(RdKafka::Event &event) override;
 
   void StartConsuming();
+  void StartConsumingWithLimit(uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const;
 
   void StopConsuming();
 
@@ -178,7 +187,6 @@ class Consumer final : public RdKafka::EventCb {
   ConsumerFunction consumer_function_;
   mutable std::atomic<bool> is_running_{false};
   mutable std::vector<RdKafka::TopicPartition *> last_assignment_;  // Protected by is_running_
-  std::optional<int64_t> limit_batches_{std::nullopt};
   std::unique_ptr<RdKafka::KafkaConsumer, std::function<void(RdKafka::KafkaConsumer *)>> consumer_;
   std::thread thread_;
   ConsumerRebalanceCb cb_;

src/integrations/kafka/exceptions.hpp

@@ -64,4 +64,16 @@ class TopicNotFoundException : public KafkaStreamException {
   TopicNotFoundException(const std::string_view consumer_name, const std::string_view topic_name)
       : KafkaStreamException("Kafka consumer {} cannot find topic {}", consumer_name, topic_name) {}
 };
+
+class ConsumerCommitFailedException : public KafkaStreamException {
+ public:
+  ConsumerCommitFailedException(const std::string_view consumer_name, const std::string_view error)
+      : KafkaStreamException("Committing offset of consumer {} failed: {}", consumer_name, error) {}
+};
+
+class ConsumerReadMessagesFailedException : public KafkaStreamException {
+ public:
+  ConsumerReadMessagesFailedException(const std::string_view consumer_name, const std::string_view error)
+      : KafkaStreamException("Error happened in consumer {} while fetching messages: {}", consumer_name, error) {}
+};
 }  // namespace memgraph::integrations::kafka

src/integrations/pulsar/consumer.cpp

@@ -11,13 +11,14 @@
 
 #include "integrations/pulsar/consumer.hpp"
 
+#include <algorithm>
+#include <chrono>
+#include <thread>
+
 #include <fmt/format.h>
 #include <pulsar/Client.h>
 #include <pulsar/InitialPosition.h>
 
-#include <chrono>
-#include <thread>
-
 #include "integrations/constants.hpp"
 #include "integrations/pulsar/exceptions.hpp"
 #include "utils/concepts.hpp"
@@ -33,6 +34,10 @@ namespace {
 template <typename T>
 concept PulsarConsumer = utils::SameAsAnyOf<T, pulsar_client::Consumer, pulsar_client::Reader>;
 
+template <typename TFunc>
+concept PulsarMessageGetter =
+    std::same_as<const pulsar_client::Message &, std::invoke_result_t<TFunc, const Message &>>;
+
 pulsar_client::Result ConsumeMessage(pulsar_client::Consumer &consumer, pulsar_client::Message &message,
                                      int remaining_timeout_in_ms) {
   return consumer.receive(message, remaining_timeout_in_ms);
@@ -97,6 +102,26 @@ pulsar_client::Client CreateClient(const std::string &service_url) {
   conf.setLogger(new SpdlogLoggerFactory);
   return {service_url, conf};
 }
+
+template <PulsarConsumer TConsumer, PulsarMessageGetter TPulsarMessageGetter>
+void TryToConsumeBatch(TConsumer &consumer, const ConsumerInfo &info, const ConsumerFunction &consumer_function,
+                       pulsar_client::MessageId &last_message_id, const std::vector<Message> &batch,
+                       const TPulsarMessageGetter &message_getter) {
+  consumer_function(batch);
+
+  auto has_message_failed = [&consumer, &info, &last_message_id, &message_getter](const auto &message) {
+    if (const auto result = consumer.acknowledge(message_getter(message)); result != pulsar_client::ResultOk) {
+      spdlog::warn("Acknowledging a message of consumer {} failed: {}", info.consumer_name, result);
+      return true;
+    }
+    last_message_id = message_getter(message).getMessageId();
+    return false;
+  };
+
+  if (std::ranges::any_of(batch, has_message_failed)) {
+    throw ConsumerAcknowledgeMessagesFailedException(info.consumer_name);
+  }
+}
 }  // namespace
 
 Message::Message(pulsar_client::Message &&message) : message_{std::move(message)} {}
@@ -137,6 +162,24 @@ void Consumer::Start() {
   StartConsuming();
 }
 
+void Consumer::StartWithLimit(const uint64_t limit_batches,
+                              const std::optional<std::chrono::milliseconds> timeout) const {
+  if (is_running_) {
+    throw ConsumerRunningException(info_.consumer_name);
+  }
+  if (limit_batches < kMinimumStartBatchLimit) {
+    throw ConsumerStartFailedException(
+        info_.consumer_name, fmt::format("Batch limit has to be greater than or equal to {}", kMinimumStartBatchLimit));
+  }
+  if (timeout.value_or(kMinimumInterval) < kMinimumInterval) {
+    throw ConsumerStartFailedException(
+        info_.consumer_name,
+        fmt::format("Timeout has to be greater than or equal to {} milliseconds", kMinimumInterval.count()));
+  }
+
+  StartConsumingWithLimit(limit_batches, timeout);
+}
+
 void Consumer::Stop() {
   if (!is_running_) {
     throw ConsumerStoppedException(info_.consumer_name);
@@ -154,7 +197,7 @@ void Consumer::StopIfRunning() {
   }
 }
 
-void Consumer::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> limit_batches,
+void Consumer::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> limit_batches,
                      const ConsumerFunction &check_consumer_function) const {
   // NOLINTNEXTLINE (modernize-use-nullptr)
   if (timeout.value_or(kMinimumInterval) < kMinimumInterval) {
@@ -240,9 +283,7 @@ void Consumer::StartConsuming() {
       auto maybe_batch = GetBatch(consumer_, info_, is_running_, last_message_id_);
 
       if (maybe_batch.HasError()) {
-        spdlog::warn("Error happened in consumer {} while fetching messages: {}!", info_.consumer_name,
-                     maybe_batch.GetError());
-        break;
+        throw ConsumerReadMessagesFailedException(info_.consumer_name, maybe_batch.GetError());
       }
 
       const auto &batch = maybe_batch.GetValue();
@@ -254,18 +295,8 @@ void Consumer::StartConsuming() {
       spdlog::info("Pulsar consumer {} is processing a batch", info_.consumer_name);
 
       try {
-        consumer_function_(batch);
-
-        if (std::any_of(batch.begin(), batch.end(), [&](const auto &message) {
-              if (const auto result = consumer_.acknowledge(message.message_); result != pulsar_client::ResultOk) {
-                spdlog::warn("Acknowledging a message of consumer {} failed: {}", info_.consumer_name, result);
-                return true;
-              }
-              last_message_id_ = message.message_.getMessageId();
-              return false;
-            })) {
-          break;
-        }
+        TryToConsumeBatch(consumer_, info_, consumer_function_, last_message_id_, batch,
+                          [&](const Message &message) -> const pulsar_client::Message & { return message.message_; });
       } catch (const std::exception &e) {
         spdlog::warn("Error happened in consumer {} while processing a batch: {}!", info_.consumer_name, e.what());
         break;
@@ -277,6 +308,43 @@ void Consumer::StartConsuming() {
   });
 }
 
+void Consumer::StartConsumingWithLimit(uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const {
+  if (is_running_.exchange(true)) {
+    throw ConsumerRunningException(info_.consumer_name);
+  }
+  utils::OnScopeExit restore_is_running([this] { is_running_.store(false); });
+
+  const auto timeout_to_use = timeout.value_or(kDefaultCheckTimeout);
+  const auto start = std::chrono::steady_clock::now();
+
+  for (uint64_t batch_count = 0; batch_count < limit_batches;) {
+    const auto now = std::chrono::steady_clock::now();
+    if (now - start >= timeout_to_use) {
+      throw ConsumerCheckFailedException(info_.consumer_name, "Timeout reached");
+    }
+
+    const auto maybe_batch = GetBatch(consumer_, info_, is_running_, last_message_id_);
+
+    if (maybe_batch.HasError()) {
+      throw ConsumerReadMessagesFailedException(info_.consumer_name, maybe_batch.GetError());
+    }
+
+    const auto &batch = maybe_batch.GetValue();
+
+    if (batch.empty()) {
+      continue;
+    }
+    ++batch_count;
+
+    spdlog::info("Pulsar consumer {} is processing a batch", info_.consumer_name);
+
+    TryToConsumeBatch(consumer_, info_, consumer_function_, last_message_id_, batch,
+                      [](const Message &message) -> const pulsar_client::Message & { return message.message_; });
+
+    spdlog::info("Pulsar consumer {} finished processing", info_.consumer_name);
+  }
+}
+
 void Consumer::StopConsuming() {
   is_running_.store(false);
   if (thread_.joinable()) {

src/integrations/pulsar/consumer.hpp

@@ -58,25 +58,27 @@ class Consumer final {
 
   bool IsRunning() const;
   void Start();
+  void StartWithLimit(uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const;
   void Stop();
   void StopIfRunning();
 
-  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> limit_batches,
+  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> limit_batches,
              const ConsumerFunction &check_consumer_function) const;
 
   const ConsumerInfo &Info() const;
 
  private:
   void StartConsuming();
+  void StartConsumingWithLimit(uint64_t limit_batches, std::optional<std::chrono::milliseconds> timeout) const;
   void StopConsuming();
 
   ConsumerInfo info_;
   mutable pulsar_client::Client client_;
-  pulsar_client::Consumer consumer_;
+  mutable pulsar_client::Consumer consumer_;
   ConsumerFunction consumer_function_;
 
   mutable std::atomic<bool> is_running_{false};
-  pulsar_client::MessageId last_message_id_{pulsar_client::MessageId::earliest()};
+  mutable pulsar_client::MessageId last_message_id_{pulsar_client::MessageId::earliest()};  // Protected by is_running_
   std::thread thread_;
 };
 }  // namespace memgraph::integrations::pulsar

src/integrations/pulsar/exceptions.hpp

@@ -55,4 +55,16 @@ class TopicNotFoundException : public PulsarStreamException {
   TopicNotFoundException(const std::string &consumer_name, const std::string &topic_name)
       : PulsarStreamException("Pulsar consumer {} cannot find topic {}", consumer_name, topic_name) {}
 };
+
+class ConsumerReadMessagesFailedException : public PulsarStreamException {
+ public:
+  ConsumerReadMessagesFailedException(const std::string_view consumer_name, const std::string_view error)
+      : PulsarStreamException("Error happened in consumer {} while fetching messages: {}", consumer_name, error) {}
+};
+
+class ConsumerAcknowledgeMessagesFailedException : public PulsarStreamException {
+ public:
+  explicit ConsumerAcknowledgeMessagesFailedException(const std::string_view consumer_name)
+      : PulsarStreamException("Acknowledging a message of consumer {} has failed!", consumer_name) {}
+};
 }  // namespace memgraph::integrations::pulsar

src/memory/CMakeLists.txt

@@ -9,4 +9,5 @@ target_link_libraries(mg-memory mg-utils fmt)
 
 if (ENABLE_JEMALLOC)
   target_link_libraries(mg-memory Jemalloc::Jemalloc)
+  target_compile_definitions(mg-memory PRIVATE USE_JEMALLOC=1)
 endif()

src/query/db_accessor.hpp

@@ -324,7 +324,7 @@ class DbAccessor final {
 
   void AdvanceCommand() { accessor_->AdvanceCommand(); }
 
-  utils::BasicResult<storage::CommitError, void> Commit() { return accessor_->Commit(); }
+  utils::BasicResult<storage::ConstraintViolation, void> Commit() { return accessor_->Commit(); }
 
   void Abort() { accessor_->Abort(); }
 

src/query/frontend/ast/ast.lcp

@@ -398,7 +398,7 @@ cpp<#
                ,@(loop for op in
                        '(or-operator xor-operator and-operator addition-operator
                          subtraction-operator multiplication-operator division-operator
-                         mod-operator not-equal-operator equal-operator less-operator
+                         mod-operator not-equal-operator less-operator
                          greater-operator less-equal-operator greater-equal-operator
                          in-list-operator subscript-operator)
                        collecting
@@ -458,6 +458,33 @@ cpp<#
                           (:clone))))))
   (define-unary-operators))
 
+(lcp:define-class equal-operator (binary-operator)
+  ((isNullCheckRequired "bool" :initval "false" :scope :public))
+  (:public
+    #>cpp
+    DEFVISITABLE(ExpressionVisitor<TypedValue>);
+    DEFVISITABLE(ExpressionVisitor<void>);
+
+    bool Accept(HierarchicalTreeVisitor &visitor) override {
+      if (visitor.PreVisit(*this)) {
+        expression1_->Accept(visitor) && expression2_->Accept(visitor);
+      }
+      return visitor.PostVisit(*this);
+    }
+    cpp<#)
+  (:protected
+    #>cpp
+    using BinaryOperator::BinaryOperator;
+    EqualOperator(Expression *expression1, Expression *expression2, bool is_nullcheck_required = false)
+        : BinaryOperator(expression1, expression2), isnullcheckrequired_(is_nullcheck_required) {}
+    cpp<#)
+  (:private
+    #>cpp
+    friend class AstStorage;
+    cpp<#)
+  (:serialize (:slk))
+  (:clone))
+
 (lcp:define-class aggregation (binary-operator)
   ((op "Op" :scope :public)
    (symbol-pos :int32_t :initval -1 :scope :public
@@ -2391,6 +2418,9 @@ cpp<#
     (lcp:define-enum sync-mode
         (sync async)
       (:serialize))
+    (lcp:define-enum replica-state
+        (ready replicating recovery invalid)
+      (:serialize))
     #>cpp
     ReplicationQuery() = default;
 

src/query/frontend/ast/cypher_main_visitor.cpp

@@ -775,6 +775,23 @@ antlrcpp::Any CypherMainVisitor::visitDropStream(MemgraphCypher::DropStreamConte
 antlrcpp::Any CypherMainVisitor::visitStartStream(MemgraphCypher::StartStreamContext *ctx) {
   auto *stream_query = storage_->Create<StreamQuery>();
   stream_query->action_ = StreamQuery::Action::START_STREAM;
+
+  if (ctx->BATCH_LIMIT()) {
+    if (!ctx->batchLimit->numberLiteral() || !ctx->batchLimit->numberLiteral()->integerLiteral()) {
+      throw SemanticException("Batch limit should be an integer literal!");
+    }
+    stream_query->batch_limit_ = ctx->batchLimit->accept(this);
+  }
+  if (ctx->TIMEOUT()) {
+    if (!ctx->timeout->numberLiteral() || !ctx->timeout->numberLiteral()->integerLiteral()) {
+      throw SemanticException("Timeout should be an integer literal!");
+    }
+    if (!ctx->BATCH_LIMIT()) {
+      throw SemanticException("Parameter TIMEOUT can only be defined if BATCH_LIMIT is defined");
+    }
+    stream_query->timeout_ = ctx->timeout->accept(this);
+  }
+
   stream_query->stream_name_ = ctx->streamName()->symbolicName()->accept(this).as<std::string>();
   return stream_query;
 }
@@ -2250,9 +2267,9 @@ antlrcpp::Any CypherMainVisitor::visitCaseExpression(MemgraphCypher::CaseExpress
   Expression *else_expression = ctx->else_expression ? ctx->else_expression->accept(this).as<Expression *>()
                                                      : storage_->Create<PrimitiveLiteral>(TypedValue());
   for (auto *alternative : alternatives) {
-    Expression *condition =
-        test_expression ? storage_->Create<EqualOperator>(test_expression, alternative->when_expression->accept(this))
-                        : alternative->when_expression->accept(this).as<Expression *>();
+    Expression *condition = test_expression ? storage_->Create<EqualOperator>(
+                                                  test_expression, alternative->when_expression->accept(this), true)
+                                            : alternative->when_expression->accept(this).as<Expression *>();
     Expression *then_expression = alternative->then_expression->accept(this);
     else_expression = storage_->Create<IfOperator>(condition, then_expression, else_expression);
   }

src/query/frontend/opencypher/grammar/MemgraphCypher.g4

@@ -351,7 +351,7 @@ pulsarCreateStream : CREATE PULSAR STREAM streamName ( pulsarCreateStreamConfig
 
 dropStream : DROP STREAM streamName ;
 
-startStream : START STREAM streamName ;
+startStream : START STREAM streamName ( BATCH_LIMIT batchLimit=literal ) ? ( TIMEOUT timeout=literal ) ? ;
 
 startAllStreams : START ALL STREAMS ;
 

src/query/interpret/eval.hpp

@@ -81,7 +81,6 @@ class ExpressionEvaluator : public ExpressionVisitor<TypedValue> {
   BINARY_OPERATOR_VISITOR(DivisionOperator, /, /);
   BINARY_OPERATOR_VISITOR(ModOperator, %, %);
   BINARY_OPERATOR_VISITOR(NotEqualOperator, !=, <>);
-  BINARY_OPERATOR_VISITOR(EqualOperator, ==, =);
   BINARY_OPERATOR_VISITOR(LessOperator, <, <);
   BINARY_OPERATOR_VISITOR(GreaterOperator, >, >);
   BINARY_OPERATOR_VISITOR(LessEqualOperator, <=, <=);
@@ -94,6 +93,21 @@ class ExpressionEvaluator : public ExpressionVisitor<TypedValue> {
 #undef BINARY_OPERATOR_VISITOR
 #undef UNARY_OPERATOR_VISITOR
 
+  TypedValue Visit(EqualOperator &op) override {
+    auto val1 = op.expression1_->Accept(*this);
+    auto val2 = op.expression2_->Accept(*this);
+
+    try {
+      if (op.isnullcheckrequired_ && val2.IsNull()) {
+        throw QueryRuntimeException("Use the generic form when checking against NULL.");
+      }
+
+      return val1 == val2;
+    } catch (const TypedValueException &) {
+      throw QueryRuntimeException("Invalid types: {} and {} for '='.", val1.type(), val2.type());
+    }
+  }
+
   TypedValue Visit(AndOperator &op) override {
     auto value1 = op.expression1_->Accept(*this);
     if (value1.IsBool() && !value1.ValueBool()) {

src/query/interpreter.cpp

@@ -231,7 +231,20 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
       if (repl_info.timeout) {
         replica.timeout = *repl_info.timeout;
       }
-
+      switch (repl_info.state) {
+        case storage::replication::ReplicaState::READY:
+          replica.state = ReplicationQuery::ReplicaState::READY;
+          break;
+        case storage::replication::ReplicaState::REPLICATING:
+          replica.state = ReplicationQuery::ReplicaState::REPLICATING;
+          break;
+        case storage::replication::ReplicaState::RECOVERY:
+          replica.state = ReplicationQuery::ReplicaState::RECOVERY;
+          break;
+        case storage::replication::ReplicaState::INVALID:
+          replica.state = ReplicationQuery::ReplicaState::INVALID;
+          break;
+      }
       return replica;
     };
 
@@ -486,6 +499,7 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
                                   fmt::format("Replica {} is registered.", repl_query->replica_name_));
       return callback;
     }
+
     case ReplicationQuery::Action::DROP_REPLICA: {
       const auto &name = repl_query->replica_name_;
       callback.fn = [handler = ReplQueryHandler{interpreter_context->db}, name]() mutable {
@@ -496,8 +510,9 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
                                   fmt::format("Replica {} is dropped.", repl_query->replica_name_));
       return callback;
     }
+
     case ReplicationQuery::Action::SHOW_REPLICAS: {
-      callback.header = {"name", "socket_address", "sync_mode", "timeout"};
+      callback.header = {"name", "socket_address", "sync_mode", "timeout", "state"};
       callback.fn = [handler = ReplQueryHandler{interpreter_context->db}, replica_nfields = callback.header.size()] {
         const auto &replicas = handler.ShowReplicas();
         auto typed_replicas = std::vector<std::vector<TypedValue>>{};
@@ -508,6 +523,7 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
 
           typed_replica.emplace_back(TypedValue(replica.name));
           typed_replica.emplace_back(TypedValue(replica.socket_address));
+
           switch (replica.sync_mode) {
             case ReplicationQuery::SyncMode::SYNC:
               typed_replica.emplace_back(TypedValue("sync"));
@@ -516,12 +532,28 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
               typed_replica.emplace_back(TypedValue("async"));
               break;
           }
+
           if (replica.timeout) {
             typed_replica.emplace_back(TypedValue(*replica.timeout));
           } else {
             typed_replica.emplace_back(TypedValue());
           }
 
+          switch (replica.state) {
+            case ReplicationQuery::ReplicaState::READY:
+              typed_replica.emplace_back(TypedValue("ready"));
+              break;
+            case ReplicationQuery::ReplicaState::REPLICATING:
+              typed_replica.emplace_back(TypedValue("replicating"));
+              break;
+            case ReplicationQuery::ReplicaState::RECOVERY:
+              typed_replica.emplace_back(TypedValue("recovery"));
+              break;
+            case ReplicationQuery::ReplicaState::INVALID:
+              typed_replica.emplace_back(TypedValue("invalid"));
+              break;
+          }
+
           typed_replicas.emplace_back(std::move(typed_replica));
         }
         return typed_replicas;
@@ -655,12 +687,26 @@ Callback HandleStreamQuery(StreamQuery *stream_query, const Parameters &paramete
       return callback;
     }
     case StreamQuery::Action::START_STREAM: {
-      callback.fn = [interpreter_context, stream_name = stream_query->stream_name_]() {
-        interpreter_context->streams.Start(stream_name);
-        return std::vector<std::vector<TypedValue>>{};
-      };
-      notifications->emplace_back(SeverityLevel::INFO, NotificationCode::START_STREAM,
-                                  fmt::format("Started stream {}.", stream_query->stream_name_));
+      const auto batch_limit = GetOptionalValue<int64_t>(stream_query->batch_limit_, evaluator);
+      const auto timeout = GetOptionalValue<std::chrono::milliseconds>(stream_query->timeout_, evaluator);
+
+      if (batch_limit.has_value()) {
+        if (batch_limit.value() < 0) {
+          throw utils::BasicException("Parameter BATCH_LIMIT cannot hold negative value");
+        }
+
+        callback.fn = [interpreter_context, stream_name = stream_query->stream_name_, batch_limit, timeout]() {
+          interpreter_context->streams.StartWithLimit(stream_name, static_cast<uint64_t>(batch_limit.value()), timeout);
+          return std::vector<std::vector<TypedValue>>{};
+        };
+      } else {
+        callback.fn = [interpreter_context, stream_name = stream_query->stream_name_]() {
+          interpreter_context->streams.Start(stream_name);
+          return std::vector<std::vector<TypedValue>>{};
+        };
+        notifications->emplace_back(SeverityLevel::INFO, NotificationCode::START_STREAM,
+                                    fmt::format("Started stream {}.", stream_query->stream_name_));
+      }
       return callback;
     }
     case StreamQuery::Action::START_ALL_STREAMS: {
@@ -729,10 +775,16 @@ Callback HandleStreamQuery(StreamQuery *stream_query, const Parameters &paramete
       return callback;
     }
     case StreamQuery::Action::CHECK_STREAM: {
-      callback.header = {"query", "parameters"};
+      callback.header = {"queries", "raw messages"};
+
+      const auto batch_limit = GetOptionalValue<int64_t>(stream_query->batch_limit_, evaluator);
+      if (batch_limit.has_value() && batch_limit.value() < 0) {
+        throw utils::BasicException("Parameter BATCH_LIMIT cannot hold negative value");
+      }
+
       callback.fn = [interpreter_context, stream_name = stream_query->stream_name_,
                      timeout = GetOptionalValue<std::chrono::milliseconds>(stream_query->timeout_, evaluator),
-                     batch_limit = GetOptionalValue<int64_t>(stream_query->batch_limit_, evaluator)]() mutable {
+                     batch_limit]() mutable {
         return interpreter_context->streams.Check(stream_name, timeout, batch_limit);
       };
       notifications->emplace_back(SeverityLevel::INFO, NotificationCode::CHECK_STREAM,
@@ -2209,39 +2261,25 @@ void RunTriggersIndividually(const utils::SkipList<Trigger> &triggers, Interpret
       continue;
     }
 
-    auto maybe_commit_error = db_accessor.Commit();
-    if (maybe_commit_error.HasError()) {
-      const auto &commit_error = maybe_commit_error.GetError();
-      switch (commit_error.type) {
-        case storage::CommitError::Type::UNABLE_TO_SYNC_REPLICATE: {
-          // TODO(gitbuda): This is tricky because this is an internal
-          // operation. Consider stopping main Memgraph instance here.
-          spdlog::warn("Trigger '{}' failed to commit due to inability to replicate data to SYNC replica",
-                       trigger.Name());
+    auto maybe_constraint_violation = db_accessor.Commit();
+    if (maybe_constraint_violation.HasError()) {
+      const auto &constraint_violation = maybe_constraint_violation.GetError();
+      switch (constraint_violation.type) {
+        case storage::ConstraintViolation::Type::EXISTENCE: {
+          const auto &label_name = db_accessor.LabelToName(constraint_violation.label);
+          MG_ASSERT(constraint_violation.properties.size() == 1U);
+          const auto &property_name = db_accessor.PropertyToName(*constraint_violation.properties.begin());
+          spdlog::warn("Trigger '{}' failed to commit due to existence constraint violation on :{}({})", trigger.Name(),
+                       label_name, property_name);
           break;
         }
-        case storage::CommitError::Type::CONSTRAINT_VIOLATION: {
-          MG_ASSERT(commit_error.maybe_constraint_violation.has_value());
-          const auto &constraint_violation = *commit_error.maybe_constraint_violation;
-          switch (constraint_violation.type) {
-            case storage::ConstraintViolation::Type::EXISTENCE: {
-              const auto &label_name = db_accessor.LabelToName(constraint_violation.label);
-              MG_ASSERT(constraint_violation.properties.size() == 1U);
-              const auto &property_name = db_accessor.PropertyToName(*constraint_violation.properties.begin());
-              spdlog::warn("Trigger '{}' failed to commit due to existence constraint violation on :{}({})",
-                           trigger.Name(), label_name, property_name);
-              break;
-            }
-            case storage::ConstraintViolation::Type::UNIQUE: {
-              const auto &label_name = db_accessor.LabelToName(constraint_violation.label);
-              std::stringstream property_names_stream;
-              utils::PrintIterable(property_names_stream, constraint_violation.properties, ", ",
-                                   [&](auto &stream, const auto &prop) { stream << db_accessor.PropertyToName(prop); });
-              spdlog::warn("Trigger '{}' failed to commit due to unique constraint violation on :{}({})",
-                           trigger.Name(), label_name, property_names_stream.str());
-              break;
-            }
-          }
+        case storage::ConstraintViolation::Type::UNIQUE: {
+          const auto &label_name = db_accessor.LabelToName(constraint_violation.label);
+          std::stringstream property_names_stream;
+          utils::PrintIterable(property_names_stream, constraint_violation.properties, ", ",
+                               [&](auto &stream, const auto &prop) { stream << db_accessor.PropertyToName(prop); });
+          spdlog::warn("Trigger '{}' failed to commit due to unique constraint violation on :{}({})", trigger.Name(),
+                       label_name, property_names_stream.str());
           break;
         }
       }
@@ -2286,40 +2324,28 @@ void Interpreter::Commit() {
     trigger_context_collector_.reset();
   };
 
-  auto maybe_commit_error = db_accessor_->Commit();
-  if (maybe_commit_error.HasError()) {
-    const auto &commit_error = maybe_commit_error.GetError();
-    switch (commit_error.type) {
-      case storage::CommitError::Type::UNABLE_TO_SYNC_REPLICATE: {
+  auto maybe_constraint_violation = db_accessor_->Commit();
+  if (maybe_constraint_violation.HasError()) {
+    const auto &constraint_violation = maybe_constraint_violation.GetError();
+    switch (constraint_violation.type) {
+      case storage::ConstraintViolation::Type::EXISTENCE: {
+        auto label_name = execution_db_accessor_->LabelToName(constraint_violation.label);
+        MG_ASSERT(constraint_violation.properties.size() == 1U);
+        auto property_name = execution_db_accessor_->PropertyToName(*constraint_violation.properties.begin());
         reset_necessary_members();
-        throw QueryException("Unable to commit due to inability to replicate to SYNC replica");
+        throw QueryException("Unable to commit due to existence constraint violation on :{}({})", label_name,
+                             property_name);
         break;
       }
-      case storage::CommitError::Type::CONSTRAINT_VIOLATION: {
-        MG_ASSERT(commit_error.maybe_constraint_violation.has_value());
-        const auto &constraint_violation = *commit_error.maybe_constraint_violation;
-        switch (constraint_violation.type) {
-          case storage::ConstraintViolation::Type::EXISTENCE: {
-            auto label_name = execution_db_accessor_->LabelToName(constraint_violation.label);
-            MG_ASSERT(constraint_violation.properties.size() == 1U);
-            auto property_name = execution_db_accessor_->PropertyToName(*constraint_violation.properties.begin());
-            reset_necessary_members();
-            throw QueryException("Unable to commit due to existence constraint violation on :{}({})", label_name,
-                                 property_name);
-            break;
-          }
-          case storage::ConstraintViolation::Type::UNIQUE: {
-            auto label_name = execution_db_accessor_->LabelToName(constraint_violation.label);
-            std::stringstream property_names_stream;
-            utils::PrintIterable(
-                property_names_stream, constraint_violation.properties, ", ",
-                [this](auto &stream, const auto &prop) { stream << execution_db_accessor_->PropertyToName(prop); });
-            reset_necessary_members();
-            throw QueryException("Unable to commit due to unique constraint violation on :{}({})", label_name,
-                                 property_names_stream.str());
-            break;
-          }
-        }
+      case storage::ConstraintViolation::Type::UNIQUE: {
+        auto label_name = execution_db_accessor_->LabelToName(constraint_violation.label);
+        std::stringstream property_names_stream;
+        utils::PrintIterable(
+            property_names_stream, constraint_violation.properties, ", ",
+            [this](auto &stream, const auto &prop) { stream << execution_db_accessor_->PropertyToName(prop); });
+        reset_necessary_members();
+        throw QueryException("Unable to commit due to unique constraint violation on :{}({})", label_name,
+                             property_names_stream.str());
         break;
       }
     }

src/query/interpreter.hpp

@@ -127,6 +127,7 @@ class ReplicationQueryHandler {
     std::string socket_address;
     ReplicationQuery::SyncMode sync_mode;
     std::optional<double> timeout;
+    ReplicationQuery::ReplicaState state;
   };
 
   /// @throw QueryRuntimeException if an error ocurred.

src/query/stream/common.hpp

@@ -52,10 +52,11 @@ concept Stream = requires(TStream stream) {
   typename TStream::Message;
   TStream{std::string{""}, typename TStream::StreamInfo{}, ConsumerFunction<typename TStream::Message>{}};
   { stream.Start() } -> std::same_as<void>;
+  { stream.StartWithLimit(uint64_t{}, std::optional<std::chrono::milliseconds>{}) } -> std::same_as<void>;
   { stream.Stop() } -> std::same_as<void>;
   { stream.IsRunning() } -> std::same_as<bool>;
   {
-    stream.Check(std::optional<std::chrono::milliseconds>{}, std::optional<int64_t>{},
+    stream.Check(std::optional<std::chrono::milliseconds>{}, std::optional<uint64_t>{},
                  ConsumerFunction<typename TStream::Message>{})
     } -> std::same_as<void>;
   requires std::same_as<std::decay_t<decltype(std::declval<typename TStream::StreamInfo>().common_info)>,

src/query/stream/sources.cpp

@@ -44,10 +44,13 @@ KafkaStream::StreamInfo KafkaStream::Info(std::string transformation_name) const
 }
 
 void KafkaStream::Start() { consumer_->Start(); }
+void KafkaStream::StartWithLimit(uint64_t batch_limit, std::optional<std::chrono::milliseconds> timeout) const {
+  consumer_->StartWithLimit(batch_limit, timeout);
+}
 void KafkaStream::Stop() { consumer_->Stop(); }
 bool KafkaStream::IsRunning() const { return consumer_->IsRunning(); }
 
-void KafkaStream::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> batch_limit,
+void KafkaStream::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> batch_limit,
                         const ConsumerFunction<integrations::kafka::Message> &consumer_function) const {
   consumer_->Check(timeout, batch_limit, consumer_function);
 }
@@ -106,10 +109,12 @@ PulsarStream::StreamInfo PulsarStream::Info(std::string transformation_name) con
 }
 
 void PulsarStream::Start() { consumer_->Start(); }
+void PulsarStream::StartWithLimit(uint64_t batch_limit, std::optional<std::chrono::milliseconds> timeout) const {
+  consumer_->StartWithLimit(batch_limit, timeout);
+}
 void PulsarStream::Stop() { consumer_->Stop(); }
 bool PulsarStream::IsRunning() const { return consumer_->IsRunning(); }
-
-void PulsarStream::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> batch_limit,
+void PulsarStream::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> batch_limit,
                          const ConsumerFunction<Message> &consumer_function) const {
   consumer_->Check(timeout, batch_limit, consumer_function);
 }

src/query/stream/sources.hpp

@@ -36,10 +36,11 @@ struct KafkaStream {
   StreamInfo Info(std::string transformation_name) const;
 
   void Start();
+  void StartWithLimit(uint64_t batch_limit, std::optional<std::chrono::milliseconds> timeout) const;
   void Stop();
   bool IsRunning() const;
 
-  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> batch_limit,
+  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> batch_limit,
              const ConsumerFunction<Message> &consumer_function) const;
 
   utils::BasicResult<std::string> SetStreamOffset(int64_t offset);
@@ -71,10 +72,11 @@ struct PulsarStream {
   StreamInfo Info(std::string transformation_name) const;
 
   void Start();
+  void StartWithLimit(uint64_t batch_limit, std::optional<std::chrono::milliseconds> timeout) const;
   void Stop();
   bool IsRunning() const;
 
-  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> batch_limit,
+  void Check(std::optional<std::chrono::milliseconds> timeout, std::optional<uint64_t> batch_limit,
              const ConsumerFunction<Message> &consumer_function) const;
 
  private:

src/query/stream/streams.cpp

@@ -43,6 +43,7 @@ extern const Event MessagesConsumed;
 namespace memgraph::query::stream {
 namespace {
 inline constexpr auto kExpectedTransformationResultSize = 2;
+inline constexpr auto kCheckStreamResultSize = 2;
 const utils::pmr::string query_param_name{"query", utils::NewDeleteResource()};
 const utils::pmr::string params_param_name{"parameters", utils::NewDeleteResource()};
 
@@ -455,7 +456,7 @@ void Streams::Create(const std::string &stream_name, typename TStream::StreamInf
 
   try {
     std::visit(
-        [&](auto &&stream_data) {
+        [&](const auto &stream_data) {
           const auto stream_source_ptr = stream_data.stream_source->ReadLock();
           Persist(CreateStatus(stream_name, stream_data.transformation_name, stream_data.owner, *stream_source_ptr));
         },
@@ -574,7 +575,7 @@ void Streams::RestoreStreams() {
         auto it = CreateConsumer<T>(*locked_streams_map, stream_name, std::move(status.info), std::move(status.owner));
         if (status.is_running) {
           std::visit(
-              [&](auto &&stream_data) {
+              [&](const auto &stream_data) {
                 auto stream_source_ptr = stream_data.stream_source->Lock();
                 stream_source_ptr->Start();
               },
@@ -616,7 +617,7 @@ void Streams::Drop(const std::string &stream_name) {
   // function can be executing with the consumer, nothing else.
   // By acquiring the write lock here for the consumer, we make sure there is
   // no running Test function for this consumer, therefore it can be erased.
-  std::visit([&](auto &&stream_data) { stream_data.stream_source->Lock(); }, it->second);
+  std::visit([&](const auto &stream_data) { stream_data.stream_source->Lock(); }, it->second);
 
   locked_streams->erase(it);
   if (!storage_.Delete(stream_name)) {
@@ -631,7 +632,7 @@ void Streams::Start(const std::string &stream_name) {
   auto it = GetStream(*locked_streams, stream_name);
 
   std::visit(
-      [&, this](auto &&stream_data) {
+      [&, this](const auto &stream_data) {
         auto stream_source_ptr = stream_data.stream_source->Lock();
         stream_source_ptr->Start();
         Persist(CreateStatus(stream_name, stream_data.transformation_name, stream_data.owner, *stream_source_ptr));
@@ -639,12 +640,27 @@ void Streams::Start(const std::string &stream_name) {
       it->second);
 }
 
+void Streams::StartWithLimit(const std::string &stream_name, uint64_t batch_limit,
+                             std::optional<std::chrono::milliseconds> timeout) const {
+  std::optional locked_streams{streams_.ReadLock()};
+  auto it = GetStream(**locked_streams, stream_name);
+
+  std::visit(
+      [&](const auto &stream_data) {
+        const auto locked_stream_source = stream_data.stream_source->ReadLock();
+        locked_streams.reset();
+
+        locked_stream_source->StartWithLimit(batch_limit, timeout);
+      },
+      it->second);
+}
+
 void Streams::Stop(const std::string &stream_name) {
   auto locked_streams = streams_.Lock();
   auto it = GetStream(*locked_streams, stream_name);
 
   std::visit(
-      [&, this](auto &&stream_data) {
+      [&, this](const auto &stream_data) {
         auto stream_source_ptr = stream_data.stream_source->Lock();
         stream_source_ptr->Stop();
 
@@ -656,7 +672,7 @@ void Streams::Stop(const std::string &stream_name) {
 void Streams::StartAll() {
   for (auto locked_streams = streams_.Lock(); auto &[stream_name, stream_data] : *locked_streams) {
     std::visit(
-        [&stream_name = stream_name, this](auto &&stream_data) {
+        [&stream_name = stream_name, this](const auto &stream_data) {
           auto locked_stream_source = stream_data.stream_source->Lock();
           if (!locked_stream_source->IsRunning()) {
             locked_stream_source->Start();
@@ -671,7 +687,7 @@ void Streams::StartAll() {
 void Streams::StopAll() {
   for (auto locked_streams = streams_.Lock(); auto &[stream_name, stream_data] : *locked_streams) {
     std::visit(
-        [&stream_name = stream_name, this](auto &&stream_data) {
+        [&stream_name = stream_name, this](const auto &stream_data) {
           auto locked_stream_source = stream_data.stream_source->Lock();
           if (locked_stream_source->IsRunning()) {
             locked_stream_source->Stop();
@@ -688,7 +704,7 @@ std::vector<StreamStatus<>> Streams::GetStreamInfo() const {
   {
     for (auto locked_streams = streams_.ReadLock(); const auto &[stream_name, stream_data] : *locked_streams) {
       std::visit(
-          [&, &stream_name = stream_name](auto &&stream_data) {
+          [&, &stream_name = stream_name](const auto &stream_data) {
             auto locked_stream_source = stream_data.stream_source->ReadLock();
             auto info = locked_stream_source->Info(stream_data.transformation_name);
             result.emplace_back(StreamStatus<>{stream_name, StreamType(*locked_stream_source),
@@ -702,12 +718,12 @@ std::vector<StreamStatus<>> Streams::GetStreamInfo() const {
 }
 
 TransformationResult Streams::Check(const std::string &stream_name, std::optional<std::chrono::milliseconds> timeout,
-                                    std::optional<int64_t> batch_limit) const {
+                                    std::optional<uint64_t> batch_limit) const {
   std::optional locked_streams{streams_.ReadLock()};
   auto it = GetStream(**locked_streams, stream_name);
 
   return std::visit(
-      [&](auto &&stream_data) {
+      [&](const auto &stream_data) {
         // This depends on the fact that Drop will first acquire a write lock to the consumer, and erase it only after
         // that
         const auto locked_stream_source = stream_data.stream_source->ReadLock();
@@ -724,15 +740,27 @@ TransformationResult Streams::Check(const std::string &stream_name, std::optiona
           auto accessor = interpreter_context->db->Access();
           CallCustomTransformation(transformation_name, messages, result, accessor, *memory_resource, stream_name);
 
-          for (auto &row : result.rows) {
-            auto [query, parameters] = ExtractTransformationResult(row.values, transformation_name, stream_name);
-            std::vector<TypedValue> result_row;
-            result_row.reserve(kExpectedTransformationResultSize);
-            result_row.push_back(std::move(query));
-            result_row.push_back(std::move(parameters));
+          auto result_row = std::vector<TypedValue>();
+          result_row.reserve(kCheckStreamResultSize);
 
-            test_result.push_back(std::move(result_row));
-          }
+          auto queries_and_parameters = std::vector<TypedValue>(result.rows.size());
+          std::transform(
+              result.rows.cbegin(), result.rows.cend(), queries_and_parameters.begin(), [&](const auto &row) {
+                auto [query, parameters] = ExtractTransformationResult(row.values, transformation_name, stream_name);
+
+                return std::map<std::string, TypedValue>{{"query", std::move(query)},
+                                                         {"parameters", std::move(parameters)}};
+              });
+          result_row.emplace_back(std::move(queries_and_parameters));
+
+          auto messages_list = std::vector<TypedValue>(messages.size());
+          std::transform(messages.cbegin(), messages.cend(), messages_list.begin(), [](const auto &message) {
+            return std::string_view(message.Payload().data(), message.Payload().size());
+          });
+
+          result_row.emplace_back(std::move(messages_list));
+
+          test_result.emplace_back(std::move(result_row));
         };
 
         locked_stream_source->Check(timeout, batch_limit, consumer_function);

src/query/stream/streams.hpp

@@ -115,6 +115,17 @@ class Streams final {
   /// @throws ConsumerRunningException if the consumer is already running
   void Start(const std::string &stream_name);
 
+  /// Start consuming from a stream.
+  ///
+  /// @param stream_name name of the stream that needs to be started
+  /// @param batch_limit number of batches we want to consume before stopping
+  /// @param timeout the maximum duration during which the command should run.
+  ///
+  /// @throws StreamsException if the stream doesn't exist
+  /// @throws ConsumerRunningException if the consumer is already running
+  void StartWithLimit(const std::string &stream_name, uint64_t batch_limit,
+                      std::optional<std::chrono::milliseconds> timeout) const;
+
   /// Stop consuming from a stream.
   ///
   /// @param stream_name name of the stream that needs to be stopped
@@ -142,6 +153,7 @@ class Streams final {
   ///
   /// @param stream_name name of the stream we want to test
   /// @param batch_limit number of batches we want to test before stopping
+  /// @param timeout the maximum duration during which the command should run.
   ///
   /// @returns A vector of vectors of TypedValue. Each subvector contains two elements, the query string and the
   /// nullable parameters map.
@@ -151,7 +163,7 @@ class Streams final {
   /// @throws ConsumerCheckFailedException if the transformation function throws any std::exception during processing
   TransformationResult Check(const std::string &stream_name,
                              std::optional<std::chrono::milliseconds> timeout = std::nullopt,
-                             std::optional<int64_t> batch_limit = std::nullopt) const;
+                             std::optional<uint64_t> batch_limit = std::nullopt) const;
 
  private:
   template <Stream TStream>

src/storage/v2/CMakeLists.txt

@@ -1,6 +1,5 @@
 set(storage_v2_src_files
     commit_log.cpp
-    commit_error.cpp
     constraints.cpp
     temporal.cpp
     durability/durability.cpp

src/storage/v2/commit_error.cpp

@@ -1,20 +0,0 @@
-// Copyright 2022 Memgraph Ltd.
-//
-// Use of this software is governed by the Business Source License
-// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
-// License, and you may not use this file except in compliance with the Business Source License.
-//
-// As of the Change Date specified in that file, in accordance with
-// the Business Source License, use of this software will be governed
-// by the Apache License, Version 2.0, included in the file
-// licenses/APL.txt.
-
-#include "storage/v2/commit_error.hpp"
-
-namespace memgraph::storage {
-
-bool operator==(const CommitError &lhs, const CommitError &rhs) {
-  return lhs.type == rhs.type && lhs.maybe_constraint_violation == rhs.maybe_constraint_violation;
-}
-
-}  // namespace memgraph::storage

src/storage/v2/commit_error.hpp

@@ -1,32 +0,0 @@
-// Copyright 2022 Memgraph Ltd.
-//
-// Use of this software is governed by the Business Source License
-// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
-// License, and you may not use this file except in compliance with the Business Source License.
-//
-// As of the Change Date specified in that file, in accordance with
-// the Business Source License, use of this software will be governed
-// by the Apache License, Version 2.0, included in the file
-// licenses/APL.txt.
-
-#pragma once
-
-#include <optional>
-
-#include "storage/v2/constraints.hpp"
-
-namespace memgraph::storage {
-
-struct CommitError {
-  enum class Type {
-    CONSTRAINT_VIOLATION,
-    UNABLE_TO_SYNC_REPLICATE,
-  };
-  Type type;
-
-  std::optional<ConstraintViolation> maybe_constraint_violation;
-};
-
-bool operator==(const CommitError &lhs, const CommitError &rhs);
-
-}  // namespace memgraph::storage

src/storage/v2/replication/replication_client.cpp

@@ -227,18 +227,17 @@ void Storage::ReplicationClient::IfStreamingTransaction(const std::function<void
   }
 }
 
-std::optional<bool> Storage::ReplicationClient::FinalizeTransactionReplication() {
+void Storage::ReplicationClient::FinalizeTransactionReplication() {
   // We can only check the state because it guarantees to be only
   // valid during a single transaction replication (if the assumption
   // that this and other transaction replication functions can only be
   // called from a one thread stands)
   if (replica_state_ != replication::ReplicaState::REPLICATING) {
-    return std::nullopt;
+    return;
   }
 
   if (mode_ == replication::ReplicationMode::ASYNC) {
-    thread_pool_.AddTask([this] { [[maybe_unused]] auto finalized = this->FinalizeTransactionReplicationInternal(); });
-    return true;
+    thread_pool_.AddTask([this] { this->FinalizeTransactionReplicationInternal(); });
   } else if (timeout_) {
     MG_ASSERT(mode_ == replication::ReplicationMode::SYNC, "Only SYNC replica can have a timeout.");
     MG_ASSERT(timeout_dispatcher_, "Timeout thread is missing");
@@ -246,7 +245,7 @@ std::optional<bool> Storage::ReplicationClient::FinalizeTransactionReplication()
 
     timeout_dispatcher_->active = true;
     thread_pool_.AddTask([&, this] {
-      [[maybe_unused]] auto finalized = this->FinalizeTransactionReplicationInternal();
+      this->FinalizeTransactionReplicationInternal();
       std::unique_lock main_guard(timeout_dispatcher_->main_lock);
       // TimerThread can finish waiting for timeout
       timeout_dispatcher_->active = false;
@@ -274,13 +273,12 @@ std::optional<bool> Storage::ReplicationClient::FinalizeTransactionReplication()
       // and acces the `active` variable`
       thread_pool_.AddTask([this] { timeout_dispatcher_.reset(); });
     }
-    return true;
   } else {
-    return FinalizeTransactionReplicationInternal();
+    FinalizeTransactionReplicationInternal();
   }
 }
 
-bool Storage::ReplicationClient::FinalizeTransactionReplicationInternal() {
+void Storage::ReplicationClient::FinalizeTransactionReplicationInternal() {
   MG_ASSERT(replica_stream_, "Missing stream for transaction deltas");
   try {
     auto response = replica_stream_->Finalize();
@@ -291,7 +289,6 @@ bool Storage::ReplicationClient::FinalizeTransactionReplicationInternal() {
       thread_pool_.AddTask([&, this] { this->RecoverReplica(response.current_commit_timestamp); });
     } else {
       replica_state_.store(replication::ReplicaState::READY);
-      return true;
     }
   } catch (const rpc::RpcFailedException &) {
     replica_stream_.reset();
@@ -301,7 +298,6 @@ bool Storage::ReplicationClient::FinalizeTransactionReplicationInternal() {
     }
     HandleRpcFailure();
   }
-  return false;
 }
 
 void Storage::ReplicationClient::RecoverReplica(uint64_t replica_commit) {

src/storage/v2/replication/replication_client.hpp

@@ -103,10 +103,7 @@ class Storage::ReplicationClient {
   // StartTransactionReplication, stream is created.
   void IfStreamingTransaction(const std::function<void(ReplicaStream &handler)> &callback);
 
-  // Return none  -> OK
-  // Return true  -> OK
-  // Return false -> FAIL
-  std::optional<bool> FinalizeTransactionReplication();
+  void FinalizeTransactionReplication();
 
   // Transfer the snapshot file.
   // @param path Path of the snapshot file.
@@ -128,7 +125,7 @@ class Storage::ReplicationClient {
   const auto &Endpoint() const { return rpc_client_->Endpoint(); }
 
  private:
-  [[nodiscard]] bool FinalizeTransactionReplicationInternal();
+  void FinalizeTransactionReplicationInternal();
 
   void RecoverReplica(uint64_t replica_commit);
 

src/storage/v2/storage.cpp

@@ -40,7 +40,6 @@
 #include "utils/uuid.hpp"
 
 /// REPLICATION ///
-#include "storage/v2/commit_error.hpp"
 #include "storage/v2/replication/replication_client.hpp"
 #include "storage/v2/replication/replication_server.hpp"
 #include "storage/v2/replication/rpc.hpp"
@@ -307,6 +306,13 @@ Storage::Storage(Config config)
       uuid_(utils::GenerateUUID()),
       epoch_id_(utils::GenerateUUID()),
       global_locker_(file_retainer_.AddLocker()) {
+  if (config_.durability.snapshot_wal_mode == Config::Durability::SnapshotWalMode::DISABLED &&
+      replication_role_ == ReplicationRole::MAIN) {
+    spdlog::warn(
+        "The instance has the MAIN replication role, but durability logs and snapshots are disabled. Please consider "
+        "enabling durability by using --storage-snapshot-interval-sec and --storage-wal-enabled flags because "
+        "without write-ahead logs this instance is not replicating any data.");
+  }
   if (config_.durability.snapshot_wal_mode != Config::Durability::SnapshotWalMode::DISABLED ||
       config_.durability.snapshot_on_exit || config_.durability.recover_on_startup) {
     // Create the directory initially to crash the database in case of
@@ -814,7 +820,7 @@ EdgeTypeId Storage::Accessor::NameToEdgeType(const std::string_view &name) { ret
 
 void Storage::Accessor::AdvanceCommand() { ++transaction_.command_id; }
 
-utils::BasicResult<CommitError, void> Storage::Accessor::Commit(
+utils::BasicResult<ConstraintViolation, void> Storage::Accessor::Commit(
     const std::optional<uint64_t> desired_commit_timestamp) {
   MG_ASSERT(is_transaction_active_, "The transaction is already terminated!");
   MG_ASSERT(!transaction_.must_abort, "The transaction can't be committed!");
@@ -837,8 +843,7 @@ utils::BasicResult<CommitError, void> Storage::Accessor::Commit(
       auto validation_result = ValidateExistenceConstraints(*prev.vertex, storage_->constraints_);
       if (validation_result) {
         Abort();
-        return CommitError{.type = CommitError::Type::CONSTRAINT_VIOLATION,
-                           .maybe_constraint_violation = *validation_result};
+        return *validation_result;
       }
     }
 
@@ -850,29 +855,6 @@ utils::BasicResult<CommitError, void> Storage::Accessor::Commit(
     // Save these so we can mark them used in the commit log.
     uint64_t start_timestamp = transaction_.start_timestamp;
 
-    // Aboring here and after the locked block obviously has an issue if
-    // replica goes down and back up during the execution of the locked block of
-    // code.
-    // Not enough, before the actual commit check all SYNC replicas for availability.
-    bool unable_to_sync_replicate = false;
-    const auto check_replicas = [&]() {
-      storage_->replication_clients_.WithLock([&](auto &clients) {
-        for (auto &client : clients) {
-          // Exclusively SYNC replicas have to be available to commit the transaction.
-          if (client->Mode() == replication::ReplicationMode::SYNC && !client->Timeout().has_value() &&
-              client->State() == replication::ReplicaState::INVALID) {
-            unable_to_sync_replicate = true;
-          }
-        }
-      });
-    };
-
-    check_replicas();
-    if (unable_to_sync_replicate) {
-      Abort();
-      return CommitError{.type = CommitError::Type::UNABLE_TO_SYNC_REPLICATE};
-    }
-
     {
       std::unique_lock<utils::SpinLock> engine_guard(storage_->engine_lock_);
       commit_timestamp_.emplace(storage_->CommitTimestamp(desired_commit_timestamp));
@@ -918,8 +900,6 @@ utils::BasicResult<CommitError, void> Storage::Accessor::Commit(
         // Replica can log only the write transaction received from Main
         // so the Wal files are consistent
         if (storage_->replication_role_ == ReplicationRole::MAIN || desired_commit_timestamp.has_value()) {
-          // TODO(gitbuda): Possible to abort data operation because in this context there is an abort operation.
-          // TODO(gitbuda): If AppendToWal returns false, we can exit this block and Abort.
           storage_->AppendToWal(transaction_, *commit_timestamp_);
         }
 
@@ -942,27 +922,13 @@ utils::BasicResult<CommitError, void> Storage::Accessor::Commit(
           engine_guard.unlock();
         });
 
-        // NOTE: This will finish/commit the transaction.
         storage_->commit_log_->MarkFinished(start_timestamp);
-        // TODO(gitbuda): Maybe the solution here is to 1. mark 2. check all
-        // the replicas 3. unmark if required... it's not possible to unmark
-        // easily because mark contains mark + update_latest_active which is
-        // not reversable operation
       }
     }
 
-    // TODO(gitbuda): This doesn't have any effect because (it will only if the
-    // constraints are also violated).
-    check_replicas();
-    if (unable_to_sync_replicate) {
-      Abort();
-      return CommitError{.type = CommitError::Type::UNABLE_TO_SYNC_REPLICATE};
-    }
-
     if (unique_constraint_violation) {
       Abort();
-      return CommitError{.type = CommitError::Type::CONSTRAINT_VIOLATION,
-                         .maybe_constraint_violation = *unique_constraint_violation};
+      return *unique_constraint_violation;
     }
   }
   is_transaction_active_ = false;
@@ -1165,18 +1131,6 @@ EdgeTypeId Storage::NameToEdgeType(const std::string_view &name) {
   return EdgeTypeId::FromUint(name_id_mapper_.NameToId(name));
 }
 
-// TODO(gitbuda): Hard to abort global operations in SYNC replication mode
-// because there is no an abort op for that yet, one idea is to just apply
-// reverse operation, e.g., CreateIndex <-> DropIndex.
-//
-// Another idea is to double check that all replicas have relevant data prior
-// to calling MarkFinished. That approach would work in both data replication
-// and global operation cases.
-//
-// EDGE CASE 1: What if the first SYNC replica is alive, receives the delta
-// object, while the second SYNC replica is dead? (replication clients are
-// stored in a vector and accessed one by one)
-
 bool Storage::CreateIndex(LabelId label, const std::optional<uint64_t> desired_commit_timestamp) {
   std::unique_lock<utils::RWLock> storage_guard(main_lock_);
   if (!indices_.label_index.CreateIndex(label, vertices_.access())) return false;
@@ -1625,9 +1579,6 @@ void Storage::FinalizeWalFile() {
   }
 }
 
-// TODO(gitbuda): Hard to abort data operation in SYNC replication mode because:
-//   * Just calling Abort inside AppendToWal for some reason causes infinite loop.
-
 void Storage::AppendToWal(const Transaction &transaction, uint64_t final_commit_timestamp) {
   if (!InitializeWalFile()) return;
   // Traverse deltas and append them to the WAL file.
@@ -1799,15 +1750,10 @@ void Storage::AppendToWal(const Transaction &transaction, uint64_t final_commit_
   FinalizeWalFile();
 
   replication_clients_.WithLock([&](auto &clients) {
-    bool all_sync_replicas_ok = true;
     for (auto &client : clients) {
-      // TODO(gitbuda): SEMI-SYNC should be exculded from here.
-      if (client->Mode() == replication::ReplicationMode::SYNC)
       client->IfStreamingTransaction([&](auto &stream) { stream.AppendTransactionEnd(final_commit_timestamp); });
-      // TODO(gitbuda): FinalizeTransactionReplication should also indicate that eveything went well for SYNC replicas.
       client->FinalizeTransactionReplication();
     }
-    return all_sync_replicas_ok;
   });
 }
 
@@ -1940,13 +1886,22 @@ utils::BasicResult<Storage::RegisterReplicaError> Storage::RegisterReplica(
   MG_ASSERT(replication_role_.load() == ReplicationRole::MAIN, "Only main instance can register a replica!");
 
   const bool name_exists = replication_clients_.WithLock([&](auto &clients) {
-    return std::any_of(clients.begin(), clients.end(), [&](auto &client) { return client->Name() == name; });
+    return std::any_of(clients.begin(), clients.end(), [&name](const auto &client) { return client->Name() == name; });
   });
 
   if (name_exists) {
     return RegisterReplicaError::NAME_EXISTS;
   }
 
+  const auto end_point_exists = replication_clients_.WithLock([&endpoint](auto &clients) {
+    return std::any_of(clients.begin(), clients.end(),
+                       [&endpoint](const auto &client) { return client->Endpoint() == endpoint; });
+  });
+
+  if (end_point_exists) {
+    return RegisterReplicaError::END_POINT_EXISTS;
+  }
+
   MG_ASSERT(replication_mode == replication::ReplicationMode::SYNC || !config.timeout,
             "Only SYNC mode can have a timeout set");
 
@@ -1959,10 +1914,15 @@ utils::BasicResult<Storage::RegisterReplicaError> Storage::RegisterReplica(
     // Another thread could have added a client with same name while
     // we were connecting to this client.
     if (std::any_of(clients.begin(), clients.end(),
-                    [&](auto &other_client) { return client->Name() == other_client->Name(); })) {
+                    [&](const auto &other_client) { return client->Name() == other_client->Name(); })) {
       return RegisterReplicaError::NAME_EXISTS;
     }
 
+    if (std::any_of(clients.begin(), clients.end(),
+                    [&client](const auto &other_client) { return client->Endpoint() == other_client->Endpoint(); })) {
+      return RegisterReplicaError::END_POINT_EXISTS;
+    }
+
     clients.push_back(std::move(client));
     return {};
   });

src/storage/v2/storage.hpp

@@ -43,7 +43,6 @@
 
 /// REPLICATION ///
 #include "rpc/server.hpp"
-#include "storage/v2/commit_error.hpp"
 #include "storage/v2/replication/config.hpp"
 #include "storage/v2/replication/enums.hpp"
 #include "storage/v2/replication/rpc.hpp"
@@ -309,12 +308,11 @@ class Storage final {
 
     void AdvanceCommand();
 
-    /// Commit returns `CommitError` if the changes made by this transaction
-    /// violate an existence, unique constraint or data could NOT be replicated
-    /// to SYNC replica. In that case the transaction is automatically aborted.
-    /// Otherwise, void is returned.
+    /// Commit returns `ConstraintViolation` if the changes made by this
+    /// transaction violate an existence or unique constraint. In that case the
+    /// transaction is automatically aborted. Otherwise, void is returned.
     /// @throw std::bad_alloc
-    utils::BasicResult<CommitError, void> Commit(std::optional<uint64_t> desired_commit_timestamp = {});
+    utils::BasicResult<ConstraintViolation, void> Commit(std::optional<uint64_t> desired_commit_timestamp = {});
 
     /// @throw std::bad_alloc
     void Abort();
@@ -413,7 +411,7 @@ class Storage final {
 
   bool SetMainReplicationRole();
 
-  enum class RegisterReplicaError : uint8_t { NAME_EXISTS, CONNECTION_FAILED };
+  enum class RegisterReplicaError : uint8_t { NAME_EXISTS, END_POINT_EXISTS, CONNECTION_FAILED };
 
   /// @pre The instance should have a MAIN role
   /// @pre Timeout can only be set for SYNC replication

tests/e2e/CMakeLists.txt

@@ -1,25 +1,32 @@
 # Set up C++ functions for e2e tests
 function(add_query_module target_name src)
-  add_library(${target_name} SHARED ${src})
-  SET_TARGET_PROPERTIES(${target_name} PROPERTIES PREFIX "")
-  target_include_directories(${target_name} PRIVATE ${CMAKE_SOURCE_DIR}/include)
+        add_library(${target_name} SHARED ${src})
+        SET_TARGET_PROPERTIES(${target_name} PROPERTIES PREFIX "")
+        target_include_directories(${target_name} PRIVATE ${CMAKE_SOURCE_DIR}/include)
 endfunction()
 
-
 function(copy_e2e_python_files TARGET_PREFIX FILE_NAME)
-add_custom_target(memgraph__e2e__${TARGET_PREFIX}__${FILE_NAME} ALL
-        COMMAND ${CMAKE_COMMAND} -E copy
+        add_custom_target(memgraph__e2e__${TARGET_PREFIX}__${FILE_NAME} ALL
+                COMMAND ${CMAKE_COMMAND} -E copy
                 ${CMAKE_CURRENT_SOURCE_DIR}/${FILE_NAME}
                 ${CMAKE_CURRENT_BINARY_DIR}/${FILE_NAME}
-        DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${FILE_NAME})
+                DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${FILE_NAME})
+endfunction()
+
+function(copy_e2e_python_files_from_parent_folder TARGET_PREFIX EXTRA_PATH FILE_NAME)
+        add_custom_target(memgraph__e2e__${TARGET_PREFIX}__${FILE_NAME} ALL
+                COMMAND ${CMAKE_COMMAND} -E copy
+                ${CMAKE_CURRENT_SOURCE_DIR}/${EXTRA_PATH}/${FILE_NAME}
+                ${CMAKE_CURRENT_BINARY_DIR}/${FILE_NAME}
+                DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${EXTRA_PATH}/${FILE_NAME})
 endfunction()
 
 function(copy_e2e_cpp_files TARGET_PREFIX FILE_NAME)
-add_custom_target(memgraph__e2e__${TARGET_PREFIX}__${FILE_NAME} ALL
-        COMMAND ${CMAKE_COMMAND} -E copy
+        add_custom_target(memgraph__e2e__${TARGET_PREFIX}__${FILE_NAME} ALL
+                COMMAND ${CMAKE_COMMAND} -E copy
                 ${CMAKE_CURRENT_SOURCE_DIR}/${FILE_NAME}
                 ${CMAKE_CURRENT_BINARY_DIR}/${FILE_NAME}
-        DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${FILE_NAME})
+                DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${FILE_NAME})
 endfunction()
 
 add_subdirectory(server)
@@ -33,6 +40,7 @@ add_subdirectory(write_procedures)
 add_subdirectory(magic_functions)
 add_subdirectory(module_file_manager)
 add_subdirectory(monitoring_server)
+add_subdirectory(util_e2e)
 
 copy_e2e_python_files(pytest_runner pytest_runner.sh "")
 file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/memgraph-selfsigned.crt DESTINATION ${CMAKE_CURRENT_BINARY_DIR})

tests/e2e/interactive_mg_runner.py

@@ -0,0 +1,219 @@
+# Copyright 2022 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+# TODO(gitbuda): Add action to print the context/cluster.
+# TODO(gitbuda): Add action to print logs of each Memgraph instance.
+# TODO(gitbuda): Polish naming within script.
+# TODO(gitbuda): Consider moving this somewhere higher in the project or even put inside GQLAlchmey.
+
+# The idea here is to implement simple interactive runner of Memgraph instances because:
+#   * it should be possible to manually create new test cases first
+#     by just running this script and executing command manually from e.g. mgconsole,
+#     running single instance of Memgraph is easy but running multiple instances and
+#     controlling them is not that easy
+#   * it should be easy to create new operational test without huge knowledge overhead
+#     by e.g. calling `process_actions` from any e2e Python test, the test will contain the
+#     string with all actions and should run test code in a different thread.
+#
+# NOTE: The intention here is not to provide infrastructure to write data
+# correctness tests or any heavy workload, the intention is to being able to
+# easily test e2e "operational" cases, simple cluster setup and basic Memgraph
+# operational queries. For any type of data correctness tests Jepsen or similar
+# approaches have to be employed.
+# NOTE: The instance description / context should be compatible with tests/e2e/runner.py
+
+import atexit
+import logging
+import os
+import subprocess
+from argparse import ArgumentParser
+from pathlib import Path
+import time
+import sys
+from inspect import signature
+
+import yaml
+from memgraph import MemgraphInstanceRunner
+
+SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+PROJECT_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", ".."))
+BUILD_DIR = os.path.join(PROJECT_DIR, "build")
+MEMGRAPH_BINARY = os.path.join(BUILD_DIR, "memgraph")
+
+# Cluster description, injectable as the context.
+# If the script argument is not provided, the following will be used as a default.
+MEMGRAPH_INSTANCES_DESCRIPTION = {
+    "replica1": {
+        "args": ["--bolt-port", "7688", "--log-level=TRACE"],
+        "log_file": "replica1.log",
+        "setup_queries": ["SET REPLICATION ROLE TO REPLICA WITH PORT 10001;"],
+    },
+    "replica2": {
+        "args": ["--bolt-port", "7689", "--log-level=TRACE"],
+        "log_file": "replica2.log",
+        "setup_queries": ["SET REPLICATION ROLE TO REPLICA WITH PORT 10002;"],
+    },
+    "main": {
+        "args": ["--bolt-port", "7687", "--log-level=TRACE"],
+        "log_file": "main.log",
+        "setup_queries": [
+            "REGISTER REPLICA replica1 SYNC TO '127.0.0.1:10001'",
+            "REGISTER REPLICA replica2 SYNC WITH TIMEOUT 1 TO '127.0.0.1:10002'",
+        ],
+    },
+}
+MEMGRAPH_INSTANCES = {}
+ACTIONS = {
+    "info": lambda context: info(context),
+    "stop": lambda context, name: stop(context, name),
+    "start": lambda context, name: start(context, name),
+    "sleep": lambda context, delta: time.sleep(float(delta)),
+    "exit": lambda context: sys.exit(1),
+    "quit": lambda context: sys.exit(1),
+}
+
+log = logging.getLogger("memgraph.tests.e2e")
+
+
+def load_args():
+    parser = ArgumentParser()
+    parser.add_argument("--actions", required=False, help="What actions to run", default="")
+    parser.add_argument(
+        "--context-yaml",
+        required=False,
+        help="YAML file with the cluster description",
+        default="",
+    )
+    return parser.parse_args()
+
+
+def _start_instance(name, args, log_file, queries, use_ssl, procdir):
+    mg_instance = MemgraphInstanceRunner(MEMGRAPH_BINARY, use_ssl)
+    MEMGRAPH_INSTANCES[name] = mg_instance
+    log_file_path = os.path.join(BUILD_DIR, "logs", log_file)
+    binary_args = args + ["--log-file", log_file_path]
+
+    if len(procdir) != 0:
+        binary_args.append("--query-modules-directory=" + procdir)
+
+    mg_instance.start(args=binary_args)
+    for query in queries:
+        mg_instance.query(query)
+
+    return mg_instance
+
+
+def stop_all():
+    for mg_instance in MEMGRAPH_INSTANCES.values():
+        mg_instance.stop()
+
+
+def stop_instance(context, name):
+    for key, _ in context.items():
+        if key != name:
+            continue
+        MEMGRAPH_INSTANCES[name].stop()
+
+
+def stop(context, name):
+    if name != "all":
+        stop_instance(context, name)
+        return
+
+    stop_all()
+
+
+@atexit.register
+def cleanup():
+    stop_all()
+
+
+def start_instance(context, name, procdir):
+    mg_instances = {}
+
+    for key, value in context.items():
+        if key != name:
+            continue
+        args = value["args"]
+        log_file = value["log_file"]
+        queries = []
+        if "setup_queries" in value:
+            queries = value["setup_queries"]
+        use_ssl = False
+        if "ssl" in value:
+            use_ssl = bool(value["ssl"])
+            value.pop("ssl")
+
+        instance = _start_instance(name, args, log_file, queries, use_ssl, procdir)
+        mg_instances[name] = instance
+
+    assert len(mg_instances) == 1
+
+    return mg_instances
+
+
+def start_all(context, procdir=""):
+    mg_instances = {}
+    for key, _ in context.items():
+        mg_instances.update(start_instance(context, key, procdir))
+
+    return mg_instances
+
+
+def start(context, name, procdir=""):
+    if name != "all":
+        return start_instance(context, name, procdir)
+
+    return start_all(context)
+
+
+def info(context):
+    print("{:<15s}{:>6s}".format("NAME", "STATUS"))
+    for name, _ in context.items():
+        if name not in MEMGRAPH_INSTANCES:
+            continue
+        instance = MEMGRAPH_INSTANCES[name]
+        print("{:<15s}{:>6s}".format(name, "UP" if instance.is_running() else "DOWN"))
+
+
+def process_actions(context, actions):
+    actions = actions.split(" ")
+    actions.reverse()
+    while len(actions) > 0:
+        name = actions.pop()
+        action = ACTIONS[name]
+        args_no = len(signature(action).parameters) - 1
+        assert (
+            args_no >= 0
+        ), "Wrong action definition, each action has to accept at least 1 argument which is the context."
+        assert args_no <= 1, "Actions with more than one user argument are not yet supported"
+        if args_no == 0:
+            action(context)
+        if args_no == 1:
+            action(context, actions.pop())
+
+
+if __name__ == "__main__":
+    args = load_args()
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(asctime)s %(name)s] %(message)s")
+
+    if args.context_yaml == "":
+        context = MEMGRAPH_INSTANCES_DESCRIPTION
+    else:
+        with open(args.context_yaml, "r") as f:
+            context = yaml.load(f, Loader=yaml.FullLoader)
+    if args.actions != "":
+        process_actions(context, args.actions)
+        sys.exit(0)
+
+    while True:
+        choice = input("ACTION>")
+        process_actions(context, choice)

tests/e2e/memgraph.py

@@ -106,3 +106,10 @@ class MemgraphInstanceRunner:
         self.proc_mg.terminate()
         code = self.proc_mg.wait()
         assert code == 0, "The Memgraph process exited with non-zero!"
+
+    def kill(self):
+        if not self.is_running():
+            return
+        self.proc_mg.kill()
+        code = self.proc_mg.wait()
+        assert code == -9, "The killed Memgraph process exited with non-nine!"

tests/e2e/replication/CMakeLists.txt

@@ -9,3 +9,6 @@ target_link_libraries(memgraph__e2e__replication__read_write_benchmark gflags js
 copy_e2e_python_files(replication_show common.py)
 copy_e2e_python_files(replication_show conftest.py)
 copy_e2e_python_files(replication_show show.py)
+copy_e2e_python_files(replication_show show_while_creating_invalid_state.py)
+copy_e2e_python_files_from_parent_folder(replication_show ".." memgraph.py)
+copy_e2e_python_files_from_parent_folder(replication_show ".." interactive_mg_runner.py)

tests/e2e/replication/show.py

@@ -12,6 +12,7 @@
 import sys
 
 import pytest
+
 from common import execute_and_fetch_all
 
 
@@ -26,21 +27,5 @@ def test_show_replication_role(port, role, connection):
     assert data[0][0] == role
 
 
-def test_show_replicas(connection):
-    cursor = connection(7687, "main").cursor()
-    actual_data = set(execute_and_fetch_all(cursor, "SHOW REPLICAS;"))
-
-    expected_column_names = {"name", "socket_address", "sync_mode", "timeout"}
-    actual_column_names = {x.name for x in cursor.description}
-    assert expected_column_names == actual_column_names
-
-    expected_data = {
-        ("replica_1", "127.0.0.1:10001", "sync", 0),
-        ("replica_2", "127.0.0.1:10002", "sync", 1.0),
-        ("replica_3", "127.0.0.1:10003", "async", None),
-    }
-    assert expected_data == actual_data
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-rA"]))

tests/e2e/replication/show_while_creating_invalid_state.py

@@ -0,0 +1,121 @@
+# Copyright 2022 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+import sys
+
+import atexit
+import os
+import pytest
+import time
+
+from common import execute_and_fetch_all
+import interactive_mg_runner
+
+interactive_mg_runner.SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+interactive_mg_runner.PROJECT_DIR = os.path.normpath(
+    os.path.join(interactive_mg_runner.SCRIPT_DIR, "..", "..", "..", "..")
+)
+interactive_mg_runner.BUILD_DIR = os.path.normpath(os.path.join(interactive_mg_runner.PROJECT_DIR, "build"))
+interactive_mg_runner.MEMGRAPH_BINARY = os.path.normpath(os.path.join(interactive_mg_runner.BUILD_DIR, "memgraph"))
+
+MEMGRAPH_INSTANCES_DESCRIPTION = {
+    "replica_1": {
+        "args": ["--bolt-port", "7688", "--log-level=TRACE"],
+        "log_file": "replica1.log",
+        "setup_queries": ["SET REPLICATION ROLE TO REPLICA WITH PORT 10001;"],
+    },
+    "replica_2": {
+        "args": ["--bolt-port", "7689", "--log-level=TRACE"],
+        "log_file": "replica2.log",
+        "setup_queries": ["SET REPLICATION ROLE TO REPLICA WITH PORT 10002;"],
+    },
+    "replica_3": {
+        "args": ["--bolt-port", "7690", "--log-level=TRACE"],
+        "log_file": "replica3.log",
+        "setup_queries": ["SET REPLICATION ROLE TO REPLICA WITH PORT 10003;"],
+    },
+    "replica_4": {
+        "args": ["--bolt-port", "7691", "--log-level=TRACE"],
+        "log_file": "replica4.log",
+        "setup_queries": ["SET REPLICATION ROLE TO REPLICA WITH PORT 10004;"],
+    },
+    "main": {
+        "args": ["--bolt-port", "7687", "--log-level=TRACE"],
+        "log_file": "main.log",
+        "setup_queries": [
+            "REGISTER REPLICA replica_1 SYNC WITH TIMEOUT 0 TO '127.0.0.1:10001';",
+            "REGISTER REPLICA replica_2 SYNC WITH TIMEOUT 1 TO '127.0.0.1:10002';",
+            "REGISTER REPLICA replica_3 ASYNC TO '127.0.0.1:10003';",
+            "REGISTER REPLICA replica_4 ASYNC TO '127.0.0.1:10004';",
+        ],
+    },
+}
+
+
+def test_show_replicas(connection):
+    # Goal of this test is to check the SHOW REPLICAS command.
+    # 0/ We start all replicas manually: we want to be able to kill them ourselves without relying on external tooling to kill processes.
+    # 1/ We check that all replicas have the correct state: they should all be ready.
+    # 2/ We drop one replica. It should not appear anymore in the SHOW REPLICAS command.
+    # 3/ We kill another replica. It should become invalid in the SHOW REPLICAS command.
+
+    # 0/
+
+    atexit.register(
+        interactive_mg_runner.stop_all
+    )  # Needed in case the test fails due to an assert. One still want the instances to be stoped.
+    mg_instances = interactive_mg_runner.start_all(MEMGRAPH_INSTANCES_DESCRIPTION)
+
+    cursor = connection(7687, "main").cursor()
+
+    # 1/
+    actual_data = set(execute_and_fetch_all(cursor, "SHOW REPLICAS;"))
+    EXPECTED_COLUMN_NAMES = {"name", "socket_address", "sync_mode", "timeout", "state"}
+
+    actual_column_names = {x.name for x in cursor.description}
+    assert EXPECTED_COLUMN_NAMES == actual_column_names
+
+    expected_data = {
+        ("replica_1", "127.0.0.1:10001", "sync", 0, "ready"),
+        ("replica_2", "127.0.0.1:10002", "sync", 1.0, "ready"),
+        ("replica_3", "127.0.0.1:10003", "async", None, "ready"),
+        ("replica_4", "127.0.0.1:10004", "async", None, "ready"),
+    }
+    assert expected_data == actual_data
+
+    # 2/
+    execute_and_fetch_all(cursor, "DROP REPLICA replica_2")
+    actual_data = set(execute_and_fetch_all(cursor, "SHOW REPLICAS;"))
+    expected_data = {
+        ("replica_1", "127.0.0.1:10001", "sync", 0, "ready"),
+        ("replica_3", "127.0.0.1:10003", "async", None, "ready"),
+        ("replica_4", "127.0.0.1:10004", "async", None, "ready"),
+    }
+    assert expected_data == actual_data
+
+    # 3/
+    mg_instances["replica_1"].kill()
+    mg_instances["replica_3"].kill()
+    mg_instances["replica_4"].stop()
+
+    # We leave some time for the main to realise the replicas are down.
+    time.sleep(2)
+    actual_data = set(execute_and_fetch_all(cursor, "SHOW REPLICAS;"))
+    expected_data = {
+        ("replica_1", "127.0.0.1:10001", "sync", 0, "invalid"),
+        ("replica_3", "127.0.0.1:10003", "async", None, "invalid"),
+        ("replica_4", "127.0.0.1:10004", "async", None, "invalid"),
+    }
+    assert expected_data == actual_data
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-rA"]))

tests/e2e/replication/workloads.yaml

@@ -74,3 +74,7 @@ workloads:
           "REGISTER REPLICA replica_3 ASYNC TO '127.0.0.1:10003'"
         ]
         validation_queries: []
+
+  - name: "Show while creating invalid state"
+    binary: "tests/e2e/pytest_runner.sh"
+    args: ["replication/show_while_creating_invalid_state.py"]

tests/e2e/runner.py

@@ -18,12 +18,11 @@ from pathlib import Path
 
 import yaml
 
-from memgraph import MemgraphInstanceRunner
+import interactive_mg_runner
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 PROJECT_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", ".."))
 BUILD_DIR = os.path.join(PROJECT_DIR, "build")
-MEMGRAPH_BINARY = os.path.join(BUILD_DIR, "memgraph")
 
 log = logging.getLogger("memgraph.tests.e2e")
 
@@ -58,30 +57,22 @@ def run(args):
             for mg_instance in mg_instances.values():
                 mg_instance.stop()
 
-        for name, config in workload["cluster"].items():
-            use_ssl = False
-            if "ssl" in config:
-                use_ssl = bool(config["ssl"])
-                config.pop("ssl")
-            mg_instance = MemgraphInstanceRunner(MEMGRAPH_BINARY, use_ssl)
-            mg_instances[name] = mg_instance
-            log_file_path = os.path.join(BUILD_DIR, "logs", config["log_file"])
-            binary_args = config["args"] + ["--log-file", log_file_path]
+        if "cluster" in workload:
+            procdir = ""
             if "proc" in workload:
-                procdir = "--query-modules-directory=" + os.path.join(BUILD_DIR, workload["proc"])
-                binary_args.append(procdir)
-            mg_instance.start(args=binary_args)
-            for query in config.get("setup_queries", []):
-                mg_instance.query(query)
+                procdir = os.path.join(BUILD_DIR, workload["proc"])
+            mg_instances = interactive_mg_runner.start_all(workload["cluster"], procdir)
+
         # Test.
         mg_test_binary = os.path.join(BUILD_DIR, workload["binary"])
         subprocess.run([mg_test_binary] + workload["args"], check=True, stderr=subprocess.STDOUT)
         # Validation.
-        for name, config in workload["cluster"].items():
-            for validation in config.get("validation_queries", []):
-                mg_instance = mg_instances[name]
-                data = mg_instance.query(validation["query"])[0][0]
-                assert data == validation["expected"]
+        if "cluster" in workload:
+            for name, config in workload["cluster"].items():
+                for validation in config.get("validation_queries", []):
+                    mg_instance = mg_instances[name]
+                    data = mg_instance.query(validation["query"])[0][0]
+                    assert data == validation["expected"]
         cleanup()
         log.info("%s PASSED.", workload_name)
 

tests/e2e/streams/common.py

@@ -10,9 +10,10 @@
 # licenses/APL.txt.
 
 import mgclient
+import pytest
 import time
 
-from multiprocessing import Process, Value
+from multiprocessing import Manager, Process, Value
 
 # These are the indices of the different values in the result of SHOW STREAM
 # query
@@ -26,8 +27,10 @@ IS_RUNNING = 6
 
 # These are the indices of the query and parameters in the result of CHECK
 # STREAM query
-QUERY = 0
-PARAMS = 1
+QUERIES = 0
+RAWMESSAGES = 1
+PARAMETERS_LITERAL = "parameters"
+QUERY_LITERAL = "query"
 
 SIMPLE_MSG = b"message"
 
@@ -45,13 +48,13 @@ def connect(**kwargs):
 
 def timed_wait(fun):
     start_time = time.time()
-    seconds = 10
+    SECONDS = 10
 
     while True:
         current_time = time.time()
         elapsed_time = current_time - start_time
 
-        if elapsed_time > seconds:
+        if elapsed_time > SECONDS:
             return False
 
         if fun():
@@ -62,13 +65,13 @@ def timed_wait(fun):
 
 def check_one_result_row(cursor, query):
     start_time = time.time()
-    seconds = 10
+    SECONDS = 10
 
     while True:
         current_time = time.time()
         elapsed_time = current_time - start_time
 
-        if elapsed_time > seconds:
+        if elapsed_time > SECONDS:
             return False
 
         cursor.execute(query)
@@ -81,12 +84,10 @@ def check_one_result_row(cursor, query):
 
 
 def check_vertex_exists_with_properties(cursor, properties):
-    properties_string = ', '.join([f'{k}: {v}' for k, v in properties.items()])
+    properties_string = ", ".join([f"{k}: {v}" for k, v in properties.items()])
     assert check_one_result_row(
         cursor,
-        "MATCH (n: MESSAGE {"
-        f"{properties_string}"
-        "}) RETURN n",
+        f"MATCH (n: MESSAGE {{{properties_string}}}) RETURN n",
     )
 
 
@@ -112,6 +113,13 @@ def start_stream(cursor, stream_name):
     assert get_is_running(cursor, stream_name)
 
 
+def start_stream_with_limit(cursor, stream_name, batch_limit, timeout=None):
+    if timeout is not None:
+        execute_and_fetch_all(cursor, f"START STREAM {stream_name} BATCH_LIMIT {batch_limit} TIMEOUT {timeout} ")
+    else:
+        execute_and_fetch_all(cursor, f"START STREAM {stream_name} BATCH_LIMIT {batch_limit}")
+
+
 def stop_stream(cursor, stream_name):
     execute_and_fetch_all(cursor, f"STOP STREAM {stream_name}")
 
@@ -129,28 +137,27 @@ def validate_info(actual_stream_info, expected_stream_info):
     for info, expected_info in zip(actual_stream_info, expected_stream_info):
         assert info == expected_info
 
+
 def check_stream_info(cursor, stream_name, expected_stream_info):
     stream_info = get_stream_info(cursor, stream_name)
     validate_info(stream_info, expected_stream_info)
 
+
 def kafka_check_vertex_exists_with_topic_and_payload(cursor, topic, payload_bytes):
-    decoded_payload = payload_bytes.decode('utf-8')
-    check_vertex_exists_with_properties(
-            cursor, {'topic': f'"{topic}"', 'payload': f'"{decoded_payload}"'})
+    decoded_payload = payload_bytes.decode("utf-8")
+    check_vertex_exists_with_properties(cursor, {"topic": f'"{topic}"', "payload": f'"{decoded_payload}"'})
 
 
-PULSAR_SERVICE_URL = 'pulsar://127.0.0.1:6650'
+PULSAR_SERVICE_URL = "pulsar://127.0.0.1:6650"
+
 
 def pulsar_default_namespace_topic(topic):
-    return f'persistent://public/default/{topic}'
+    return f"persistent://public/default/{topic}"
 
 
 def test_start_and_stop_during_check(
-        operation,
-        connection,
-        stream_creator,
-        message_sender,
-        already_stopped_error):
+    operation, connection, stream_creator, message_sender, already_stopped_error, batchSize
+):
     # This test is quite complex. The goal is to call START/STOP queries
     # while a CHECK query is waiting for its result. Because the Global
     # Interpreter Lock, running queries on multiple threads is not useful,
@@ -161,11 +168,9 @@ def test_start_and_stop_during_check(
     # synchronize between the different processes. Each value represents a
     # specific phase of the execution of the processes.
     assert operation in ["START", "STOP"]
+    assert batchSize == 1
     cursor = connection.cursor()
-    execute_and_fetch_all(
-        cursor,
-        stream_creator('test_stream')
-    )
+    execute_and_fetch_all(cursor, stream_creator("test_stream"))
 
     check_counter = Value("i", 0)
     check_result_len = Value("i", 0)
@@ -185,7 +190,9 @@ def test_start_and_stop_during_check(
         result = execute_and_fetch_all(cursor, "CHECK STREAM test_stream")
         result_len.value = len(result)
         counter.value = CHECK_AFTER_FETCHALL
-        if len(result) > 0 and "payload: 'message'" in result[0][QUERY]:
+        if (
+            len(result) > 0 and "payload: 'message'" in result[0][QUERIES][0][QUERY_LITERAL]
+        ):  # The 0 is only correct because batchSize is 1
             counter.value = CHECK_CORRECT_RESULT
         else:
             counter.value = CHECK_INCORRECT_RESULT
@@ -213,12 +220,8 @@ def test_start_and_stop_during_check(
         except Exception:
             counter.value = OP_UNEXPECTED_EXCEPTION
 
-    check_stream_proc = Process(
-        target=call_check, daemon=True, args=(check_counter, check_result_len)
-    )
-    operation_proc = Process(
-        target=call_operation, daemon=True, args=(operation_counter,)
-    )
+    check_stream_proc = Process(target=call_check, daemon=True, args=(check_counter, check_result_len))
+    operation_proc = Process(target=call_operation, daemon=True, args=(operation_counter,))
 
     try:
         check_stream_proc.start()
@@ -227,9 +230,7 @@ def test_start_and_stop_during_check(
 
         assert timed_wait(lambda: check_counter.value == CHECK_BEFORE_EXECUTE)
         assert timed_wait(lambda: get_is_running(cursor, "test_stream"))
-        assert check_counter.value == CHECK_BEFORE_EXECUTE, (
-            "SHOW STREAMS " "was blocked until the end of CHECK STREAM"
-        )
+        assert check_counter.value == CHECK_BEFORE_EXECUTE, "SHOW STREAMS " "was blocked until the end of CHECK STREAM"
         operation_proc.start()
         assert timed_wait(lambda: operation_counter.value == OP_BEFORE_EXECUTE)
 
@@ -255,31 +256,393 @@ def test_start_and_stop_during_check(
         if operation_proc.is_alive():
             operation_proc.terminate()
 
+
 def test_start_checked_stream_after_timeout(connection, stream_creator):
     cursor = connection.cursor()
-    execute_and_fetch_all(
-        cursor,
-        stream_creator('test_stream')
-    )
+    execute_and_fetch_all(cursor, stream_creator("test_stream"))
 
-    timeout_ms = 2000
+    TIMEOUT_IN_MS = 2000
+    TIMEOUT_IN_SECONDS = TIMEOUT_IN_MS / 1000
 
     def call_check():
-        execute_and_fetch_all(
-            connect().cursor(),
-            f"CHECK STREAM test_stream TIMEOUT {timeout_ms}")
+        execute_and_fetch_all(connect().cursor(), f"CHECK STREAM test_stream TIMEOUT {TIMEOUT_IN_MS}")
 
     check_stream_proc = Process(target=call_check, daemon=True)
 
     start = time.time()
     check_stream_proc.start()
-    assert timed_wait(
-        lambda: get_is_running(
-            cursor, "test_stream"))
+    assert timed_wait(lambda: get_is_running(cursor, "test_stream"))
     start_stream(cursor, "test_stream")
     end = time.time()
 
-    assert (end - start) < 1.3 * \
-        timeout_ms, "The START STREAM was blocked too long"
+    assert (end - start) < 1.3 * TIMEOUT_IN_SECONDS, "The START STREAM was blocked too long"
     assert get_is_running(cursor, "test_stream")
     stop_stream(cursor, "test_stream")
+
+
+def test_check_stream_same_number_of_queries_than_messages(connection, stream_creator, message_sender):
+    BATCH_SIZE = 2
+    BATCH_LIMIT = 3
+    STREAM_NAME = "test_stream"
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME, BATCH_SIZE))
+    time.sleep(2)
+
+    test_results = Manager().Namespace()
+
+    def check_stream(stream_name, batch_limit):
+        connection = connect()
+        cursor = connection.cursor()
+        test_results.value = execute_and_fetch_all(cursor, f"CHECK STREAM {stream_name} BATCH_LIMIT {batch_limit} ")
+
+    check_stream_proc = Process(target=check_stream, args=(STREAM_NAME, BATCH_LIMIT))
+    check_stream_proc.start()
+    time.sleep(2)
+
+    MESSAGES = [b"01", b"02", b"03", b"04", b"05", b"06"]
+    for message in MESSAGES:
+        message_sender(message)
+
+    check_stream_proc.join()
+
+    # # Transformation does not do any filtering and simply create queries as "Messages: {contentOfMessage}". Queries should be like:
+    # # -Batch 1: [{parameters: {"value": "Parameter: 01"}, query: "Message: 01"},
+    # #            {parameters: {"value": "Parameter: 02"}, query: "Message: 02"}]
+    # # -Batch 2: [{parameters: {"value": "Parameter: 03"}, query: "Message: 03"},
+    # #            {parameters: {"value": "Parameter: 04"}, query: "Message: 04"}]
+    # # -Batch 3: [{parameters: {"value": "Parameter: 05"}, query: "Message: 05"},
+    # #            {parameters: {"value": "Parameter: 06"}, query: "Message: 06"}]
+
+    assert len(test_results.value) == BATCH_LIMIT
+
+    expected_queries_and_raw_messages_1 = (
+        [  # queries
+            {PARAMETERS_LITERAL: {"value": "Parameter: 01"}, QUERY_LITERAL: "Message: 01"},
+            {PARAMETERS_LITERAL: {"value": "Parameter: 02"}, QUERY_LITERAL: "Message: 02"},
+        ],
+        ["01", "02"],  # raw message
+    )
+
+    expected_queries_and_raw_messages_2 = (
+        [  # queries
+            {PARAMETERS_LITERAL: {"value": "Parameter: 03"}, QUERY_LITERAL: "Message: 03"},
+            {PARAMETERS_LITERAL: {"value": "Parameter: 04"}, QUERY_LITERAL: "Message: 04"},
+        ],
+        ["03", "04"],  # raw message
+    )
+
+    expected_queries_and_raw_messages_3 = (
+        [  # queries
+            {PARAMETERS_LITERAL: {"value": "Parameter: 05"}, QUERY_LITERAL: "Message: 05"},
+            {PARAMETERS_LITERAL: {"value": "Parameter: 06"}, QUERY_LITERAL: "Message: 06"},
+        ],
+        ["05", "06"],  # raw message
+    )
+
+    assert expected_queries_and_raw_messages_1 == test_results.value[0]
+    assert expected_queries_and_raw_messages_2 == test_results.value[1]
+    assert expected_queries_and_raw_messages_3 == test_results.value[2]
+
+
+def test_check_stream_different_number_of_queries_than_messages(connection, stream_creator, message_sender):
+    BATCH_SIZE = 2
+    BATCH_LIMIT = 3
+    STREAM_NAME = "test_stream"
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME, BATCH_SIZE))
+    time.sleep(2)
+
+    results = Manager().Namespace()
+
+    def check_stream(stream_name, batch_limit):
+        connection = connect()
+        cursor = connection.cursor()
+        results.value = execute_and_fetch_all(cursor, f"CHECK STREAM {stream_name} BATCH_LIMIT {batch_limit} ")
+
+    check_stream_proc = Process(target=check_stream, args=(STREAM_NAME, BATCH_LIMIT))
+    check_stream_proc.start()
+    time.sleep(2)
+
+    MESSAGES = [b"a_01", b"a_02", b"03", b"04", b"b_05", b"06"]
+    for message in MESSAGES:
+        message_sender(message)
+
+    check_stream_proc.join()
+
+    # Transformation does some filtering: if message contains "a", it is ignored.
+    # Transformation also has special rule to create query if message is "b": it create more queries.
+    #
+    # Queries should be like:
+    # -Batch 1: []
+    # -Batch 2: [{parameters: {"value": "Parameter: 03"}, query: "Message: 03"},
+    #            {parameters: {"value": "Parameter: 04"}, query: "Message: 04"}]
+    # -Batch 3: [{parameters: {"value": "Parameter: 05"}, query: "Message: 05"},
+    #            {parameters: {"value": "Parameter: extra_05"}, query: "Message: extra_05"}
+    #            {parameters: {"value": "Parameter: 06"}, query: "Message: 06"}]
+
+    assert len(results.value) == BATCH_LIMIT
+
+    expected_queries_and_raw_messages_1 = (
+        [],  # queries
+        ["a_01", "a_02"],  # raw message
+    )
+
+    expected_queries_and_raw_messages_2 = (
+        [  # queries
+            {PARAMETERS_LITERAL: {"value": "Parameter: 03"}, QUERY_LITERAL: "Message: 03"},
+            {PARAMETERS_LITERAL: {"value": "Parameter: 04"}, QUERY_LITERAL: "Message: 04"},
+        ],
+        ["03", "04"],  # raw message
+    )
+
+    expected_queries_and_raw_messages_3 = (
+        [  # queries
+            {PARAMETERS_LITERAL: {"value": "Parameter: b_05"}, QUERY_LITERAL: "Message: b_05"},
+            {
+                PARAMETERS_LITERAL: {"value": "Parameter: extra_b_05"},
+                QUERY_LITERAL: "Message: extra_b_05",
+            },
+            {PARAMETERS_LITERAL: {"value": "Parameter: 06"}, QUERY_LITERAL: "Message: 06"},
+        ],
+        ["b_05", "06"],  # raw message
+    )
+
+    assert expected_queries_and_raw_messages_1 == results.value[0]
+    assert expected_queries_and_raw_messages_2 == results.value[1]
+    assert expected_queries_and_raw_messages_3 == results.value[2]
+
+
+def test_start_stream_with_batch_limit(connection, stream_creator, messages_sender):
+    STREAM_NAME = "test"
+    BATCH_LIMIT = 5
+
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME))
+
+    def start_new_stream_with_limit(stream_name, batch_limit):
+        connection = connect()
+        cursor = connection.cursor()
+        start_stream_with_limit(cursor, stream_name, batch_limit)
+
+    thread_stream_running = Process(target=start_new_stream_with_limit, daemon=True, args=(STREAM_NAME, BATCH_LIMIT))
+    thread_stream_running.start()
+
+    time.sleep(2)
+    assert get_is_running(cursor, STREAM_NAME)
+
+    messages_sender(BATCH_LIMIT - 1)
+
+    # We have not sent enough batches to reach the limit. We check that the stream is still correctly running.
+    assert get_is_running(cursor, STREAM_NAME)
+
+    # We send a last message to reach the batch_limit
+    messages_sender(1)
+
+    time.sleep(2)
+
+    # We check that the stream has correctly stoped.
+    assert not get_is_running(cursor, STREAM_NAME)
+
+
+def test_start_stream_with_batch_limit_timeout(connection, stream_creator):
+    # We check that we get the expected exception when trying to run START STREAM while providing TIMEOUT and not BATCH_LIMIT
+    STREAM_NAME = "test"
+
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME))
+
+    with pytest.raises(mgclient.DatabaseError):
+        execute_and_fetch_all(cursor, f"START STREAM {STREAM_NAME} TIMEOUT 3000")
+
+
+def test_start_stream_with_batch_limit_reaching_timeout(connection, stream_creator):
+    # We check that we get the expected exception when running START STREAM while providing TIMEOUT and BATCH_LIMIT
+    STREAM_NAME = "test"
+    BATCH_LIMIT = 5
+    TIMEOUT = 3000
+    TIMEOUT_IN_SECONDS = TIMEOUT / 1000
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME, BATCH_SIZE))
+
+    start_time = time.time()
+
+    with pytest.raises(mgclient.DatabaseError):
+        execute_and_fetch_all(cursor, f"START STREAM {STREAM_NAME} BATCH_LIMIT {BATCH_LIMIT} TIMEOUT {TIMEOUT}")
+
+    end_time = time.time()
+    assert (
+        end_time - start_time
+    ) >= TIMEOUT_IN_SECONDS, "The START STREAM has probably thrown due to something else than timeout!"
+
+
+def test_start_stream_with_batch_limit_while_check_running(
+    connection, stream_creator, message_sender, setup_function=None
+):
+    # 1/ We check we get the correct exception calling START STREAM with BATCH_LIMIT while a CHECK STREAM is already running.
+    # 2/ Afterwards, we terminate the CHECK STREAM and start a START STREAM with BATCH_LIMIT
+    def start_check_stream(stream_name, batch_limit, timeout):
+        connection = connect()
+        cursor = connection.cursor()
+        execute_and_fetch_all(cursor, f"CHECK STREAM {stream_name} BATCH_LIMIT {batch_limit} TIMEOUT {timeout}")
+
+    def start_new_stream_with_limit(stream_name, batch_limit, timeout):
+        connection = connect()
+        cursor = connection.cursor()
+        start_stream_with_limit(cursor, stream_name, batch_limit, timeout=timeout)
+
+    STREAM_NAME = "test_check_and_batch_limit"
+    BATCH_LIMIT = 1
+    TIMEOUT = 10000
+
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME))
+
+    # 0/ Extra setup needed for Kafka to works correctly if Check stream is execute before any messages have been consumed.
+    if setup_function is not None:
+        setup_function(start_check_stream, cursor, STREAM_NAME, BATCH_LIMIT, TIMEOUT)
+
+    # 1/
+    thread_stream_check = Process(target=start_check_stream, daemon=True, args=(STREAM_NAME, BATCH_LIMIT, TIMEOUT))
+    thread_stream_check.start()
+    time.sleep(2)
+    assert get_is_running(cursor, STREAM_NAME)
+
+    with pytest.raises(mgclient.DatabaseError):
+        start_stream_with_limit(cursor, STREAM_NAME, BATCH_LIMIT, timeout=TIMEOUT)
+
+    assert get_is_running(cursor, STREAM_NAME)
+    message_sender(SIMPLE_MSG)
+    thread_stream_check.join()
+
+    assert not get_is_running(cursor, STREAM_NAME)
+
+    # 2/
+    thread_stream_running = Process(
+        target=start_new_stream_with_limit, daemon=True, args=(STREAM_NAME, BATCH_LIMIT + 1, TIMEOUT)
+    )  # Sending BATCH_LIMIT + 1 messages as BATCH_LIMIT messages have already been sent during the CHECK STREAM (and not consumed)
+    thread_stream_running.start()
+    time.sleep(2)
+    assert get_is_running(cursor, STREAM_NAME)
+
+    message_sender(SIMPLE_MSG)
+    time.sleep(2)
+
+    assert not get_is_running(cursor, STREAM_NAME)
+
+
+def test_check_while_stream_with_batch_limit_running(connection, stream_creator, message_sender):
+    # 1/ We check we get the correct exception calling CHECK STREAM while START STREAM with BATCH_LIMIT is already running
+    # 2/ Afterwards, we terminate the START STREAM with BATCH_LIMIT and start a CHECK STREAM
+    def start_new_stream_with_limit(stream_name, batch_limit, timeout):
+        connection = connect()
+        cursor = connection.cursor()
+        start_stream_with_limit(cursor, stream_name, batch_limit, timeout=timeout)
+
+    def start_check_stream(stream_name, batch_limit, timeout):
+        connection = connect()
+        cursor = connection.cursor()
+        execute_and_fetch_all(cursor, f"CHECK STREAM {stream_name} BATCH_LIMIT {batch_limit} TIMEOUT {timeout}")
+
+    STREAM_NAME = "test_batch_limit_and_check"
+    BATCH_LIMIT = 1
+    TIMEOUT = 10000
+    TIMEOUT_IN_SECONDS = TIMEOUT / 1000
+
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME))
+
+    # 1/
+    thread_stream_running = Process(
+        target=start_new_stream_with_limit, daemon=True, args=(STREAM_NAME, BATCH_LIMIT, TIMEOUT)
+    )
+    start_time = time.time()
+    thread_stream_running.start()
+    time.sleep(2)
+    assert get_is_running(cursor, STREAM_NAME)
+
+    with pytest.raises(mgclient.DatabaseError):
+        execute_and_fetch_all(cursor, f"CHECK STREAM {STREAM_NAME} BATCH_LIMIT {BATCH_LIMIT} TIMEOUT {TIMEOUT}")
+
+    end_time = time.time()
+    assert (end_time - start_time) < 0.8 * TIMEOUT, "The CHECK STREAM has probably thrown due to timeout!"
+
+    message_sender(SIMPLE_MSG)
+    time.sleep(2)
+
+    assert not get_is_running(cursor, STREAM_NAME)
+
+    # 2/
+    thread_stream_check = Process(target=start_check_stream, daemon=True, args=(STREAM_NAME, BATCH_LIMIT, TIMEOUT))
+    start_time = time.time()
+    thread_stream_check.start()
+    time.sleep(2)
+    assert get_is_running(cursor, STREAM_NAME)
+
+    message_sender(SIMPLE_MSG)
+    time.sleep(2)
+    end_time = time.time()
+    assert (end_time - start_time) < 0.8 * TIMEOUT_IN_SECONDS, "The CHECK STREAM has probably thrown due to timeout!"
+
+    assert not get_is_running(cursor, STREAM_NAME)
+
+
+def test_start_stream_with_batch_limit_with_invalid_batch_limit(connection, stream_creator):
+    # We check that we get a correct exception when giving a negative batch_limit
+    STREAM_NAME = "test_batch_limit_invalid_batch_limit"
+    TIMEOUT = 10000
+    TIMEOUT_IN_SECONDS = TIMEOUT / 1000
+
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME))
+    time.sleep(2)
+
+    # 1/ checking with batch_limit=-10
+    batch_limit = -10
+    start_time = time.time()
+
+    with pytest.raises(mgclient.DatabaseError):
+        start_stream_with_limit(cursor, STREAM_NAME, batch_limit, timeout=TIMEOUT)
+
+    end_time = time.time()
+    assert (end_time - start_time) < 0.8 * TIMEOUT_IN_SECONDS, "The START STREAM has probably thrown due to timeout!"
+
+    # 2/ checking with batch_limit=0
+    batch_limit = 0
+    start_time = time.time()
+
+    with pytest.raises(mgclient.DatabaseError):
+        start_stream_with_limit(cursor, STREAM_NAME, batch_limit, timeout=TIMEOUT)
+
+    end_time = time.time()
+    assert (end_time - start_time) < 0.8 * TIMEOUT_IN_SECONDS, "The START STREAM has probably thrown due to timeout!"
+
+
+def test_check_stream_with_batch_limit_with_invalid_batch_limit(connection, stream_creator):
+    # We check that we get a correct exception when giving a negative batch_limit
+    STREAM_NAME = "test_batch_limit_invalid_batch_limit"
+    TIMEOUT = 10000
+    TIMEOUT_IN_SECONDS = TIMEOUT / 1000
+
+    cursor = connection.cursor()
+    execute_and_fetch_all(cursor, stream_creator(STREAM_NAME))
+    time.sleep(2)
+
+    # 1/ checking with batch_limit=-10
+    batch_limit = -10
+    start_time = time.time()
+
+    with pytest.raises(mgclient.DatabaseError):
+        execute_and_fetch_all(cursor, f"CHECK STREAM {STREAM_NAME} BATCH_LIMIT {batch_limit} TIMEOUT {TIMEOUT}")
+
+    end_time = time.time()
+    assert (end_time - start_time) < 0.8 * TIMEOUT_IN_SECONDS, "The CHECK STREAM has probably thrown due to timeout!"
+
+    # 2/ checking with batch_limit=0
+    batch_limit = 0
+    start_time = time.time()
+
+    with pytest.raises(mgclient.DatabaseError):
+        execute_and_fetch_all(cursor, f"CHECK STREAM {STREAM_NAME} BATCH_LIMIT {batch_limit} TIMEOUT {TIMEOUT}")
+
+    end_time = time.time()
+    assert (end_time - start_time) < 0.8 * TIMEOUT_IN_SECONDS, "The CHECK STREAM has probably thrown due to timeout!"

tests/e2e/streams/kafka_streams_tests.py

@@ -18,12 +18,10 @@ import time
 from multiprocessing import Process, Value
 import common
 
-TRANSFORMATIONS_TO_CHECK_C = [
-    "empty_transformation"]
+TRANSFORMATIONS_TO_CHECK_C = ["c_transformations.empty_transformation"]
+
+TRANSFORMATIONS_TO_CHECK_PY = ["kafka_transform.simple", "kafka_transform.with_parameters"]
 
-TRANSFORMATIONS_TO_CHECK_PY = [
-    "kafka_transform.simple",
-    "kafka_transform.with_parameters"]
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK_PY)
 def test_simple(kafka_producer, kafka_topics, connection, transformation):
@@ -31,9 +29,7 @@ def test_simple(kafka_producer, kafka_topics, connection, transformation):
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM test "
-        f"TOPICS {','.join(kafka_topics)} "
-        f"TRANSFORM {transformation}",
+        f"CREATE KAFKA STREAM test TOPICS {','.join(kafka_topics)} TRANSFORM {transformation}",
     )
     common.start_stream(cursor, "test")
     time.sleep(5)
@@ -42,16 +38,11 @@ def test_simple(kafka_producer, kafka_topics, connection, transformation):
         kafka_producer.send(topic, common.SIMPLE_MSG).get(timeout=60)
 
     for topic in kafka_topics:
-        common.kafka_check_vertex_exists_with_topic_and_payload(
-            cursor, topic, common.SIMPLE_MSG)
+        common.kafka_check_vertex_exists_with_topic_and_payload(cursor, topic, common.SIMPLE_MSG)
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK_PY)
-def test_separate_consumers(
-        kafka_producer,
-        kafka_topics,
-        connection,
-        transformation):
+def test_separate_consumers(kafka_producer, kafka_topics, connection, transformation):
     assert len(kafka_topics) > 0
     cursor = connection.cursor()
 
@@ -61,9 +52,7 @@ def test_separate_consumers(
         stream_names.append(stream_name)
         common.execute_and_fetch_all(
             cursor,
-            f"CREATE KAFKA STREAM {stream_name} "
-            f"TOPICS {topic} "
-            f"TRANSFORM {transformation}",
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {topic} TRANSFORM {transformation}",
         )
 
     for stream_name in stream_names:
@@ -75,12 +64,10 @@ def test_separate_consumers(
         kafka_producer.send(topic, common.SIMPLE_MSG).get(timeout=60)
 
     for topic in kafka_topics:
-        common.kafka_check_vertex_exists_with_topic_and_payload(
-            cursor, topic, common.SIMPLE_MSG)
+        common.kafka_check_vertex_exists_with_topic_and_payload(cursor, topic, common.SIMPLE_MSG)
 
 
-def test_start_from_last_committed_offset(
-        kafka_producer, kafka_topics, connection):
+def test_start_from_last_committed_offset(kafka_producer, kafka_topics, connection):
     # This test creates a stream, consumes a message to have a committed
     # offset, then destroys the stream. A new message is sent before the
     # stream is recreated and then restarted. This simulates when Memgraph is
@@ -90,16 +77,15 @@ def test_start_from_last_committed_offset(
     assert len(kafka_topics) > 0
     cursor = connection.cursor()
     common.execute_and_fetch_all(
-        cursor, "CREATE KAFKA STREAM test "
-        f"TOPICS {kafka_topics[0]} "
-        "TRANSFORM kafka_transform.simple", )
+        cursor,
+        f"CREATE KAFKA STREAM test TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple",
+    )
     common.start_stream(cursor, "test")
     time.sleep(1)
 
     kafka_producer.send(kafka_topics[0], common.SIMPLE_MSG).get(timeout=60)
 
-    common.kafka_check_vertex_exists_with_topic_and_payload(
-        cursor, kafka_topics[0], common.SIMPLE_MSG)
+    common.kafka_check_vertex_exists_with_topic_and_payload(cursor, kafka_topics[0], common.SIMPLE_MSG)
 
     common.stop_stream(cursor, "test")
     common.drop_stream(cursor, "test")
@@ -111,36 +97,30 @@ def test_start_from_last_committed_offset(
     for message in messages:
         vertices_with_msg = common.execute_and_fetch_all(
             cursor,
-            "MATCH (n: MESSAGE {" f"payload: '{message.decode('utf-8')}'" "}) RETURN n",
+            f"MATCH (n: MESSAGE {{payload: '{message.decode('utf-8')}'}}) RETURN n",
         )
 
         assert len(vertices_with_msg) == 0
 
     common.execute_and_fetch_all(
-        cursor, "CREATE KAFKA STREAM test "
-        f"TOPICS {kafka_topics[0]} "
-        "TRANSFORM kafka_transform.simple", )
+        cursor,
+        f"CREATE KAFKA STREAM test TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple",
+    )
     common.start_stream(cursor, "test")
 
     for message in messages:
-        common.kafka_check_vertex_exists_with_topic_and_payload(
-            cursor, kafka_topics[0], message)
+        common.kafka_check_vertex_exists_with_topic_and_payload(cursor, kafka_topics[0], message)
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK_PY)
-def test_check_stream(
-        kafka_producer,
-        kafka_topics,
-        connection,
-        transformation):
+def test_check_stream(kafka_producer, kafka_topics, connection, transformation):
     assert len(kafka_topics) > 0
+    BATCH_SIZE = 1
+    INDEX_OF_FIRST_BATCH = 0
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM test "
-        f"TOPICS {kafka_topics[0]} "
-        f"TRANSFORM {transformation} "
-        "BATCH_SIZE 1",
+        f"CREATE KAFKA STREAM test TOPICS {kafka_topics[0]} TRANSFORM {transformation} BATCH_SIZE {BATCH_SIZE}",
     )
     common.start_stream(cursor, "test")
     time.sleep(1)
@@ -153,24 +133,28 @@ def test_check_stream(
         kafka_producer.send(kafka_topics[0], message).get(timeout=60)
 
     def check_check_stream(batch_limit):
-        assert (
-            transformation == "kafka_transform.simple"
-            or transformation == "kafka_transform.with_parameters"
-        )
-        test_results = common.execute_and_fetch_all(
-            cursor, f"CHECK STREAM test BATCH_LIMIT {batch_limit}"
-        )
+        assert transformation == "kafka_transform.simple" or transformation == "kafka_transform.with_parameters"
+        test_results = common.execute_and_fetch_all(cursor, f"CHECK STREAM test BATCH_LIMIT {batch_limit}")
         assert len(test_results) == batch_limit
 
         for i in range(batch_limit):
             message_as_str = messages[i].decode("utf-8")
+            assert (
+                BATCH_SIZE == 1
+            )  # If batch size != 1, then the usage of INDEX_OF_FIRST_BATCH must change: the result will have a list of queries (pair<parameters,query>)
+
             if transformation == "kafka_transform.simple":
-                assert f"payload: '{message_as_str}'" in test_results[i][common.QUERY]
-                assert test_results[i][common.PARAMS] is None
+                assert (
+                    f"payload: '{message_as_str}'"
+                    in test_results[i][common.QUERIES][INDEX_OF_FIRST_BATCH][common.QUERY_LITERAL]
+                )
+                assert test_results[i][common.QUERIES][INDEX_OF_FIRST_BATCH][common.PARAMETERS_LITERAL] is None
             else:
-                assert f"payload: $payload" in test_results[i][
-                    common.QUERY] and f"topic: $topic" in test_results[i][common.QUERY]
-                parameters = test_results[i][common.PARAMS]
+                assert (
+                    f"payload: $payload" in test_results[i][common.QUERIES][INDEX_OF_FIRST_BATCH][common.QUERY_LITERAL]
+                    and f"topic: $topic" in test_results[i][common.QUERIES][INDEX_OF_FIRST_BATCH][common.QUERY_LITERAL]
+                )
+                parameters = test_results[i][common.QUERIES][INDEX_OF_FIRST_BATCH][common.PARAMETERS_LITERAL]
                 # this is not a very sofisticated test, but checks if
                 # timestamp has some kind of value
                 assert parameters["timestamp"] > 1000000000000
@@ -183,8 +167,7 @@ def test_check_stream(
     common.start_stream(cursor, "test")
 
     for message in messages:
-        common.kafka_check_vertex_exists_with_topic_and_payload(
-            cursor, kafka_topics[0], message)
+        common.kafka_check_vertex_exists_with_topic_and_payload(cursor, kafka_topics[0], message)
 
 
 def test_show_streams(kafka_producer, kafka_topics, connection):
@@ -192,23 +175,15 @@ def test_show_streams(kafka_producer, kafka_topics, connection):
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM default_values "
-        f"TOPICS {kafka_topics[0]} "
-        f"TRANSFORM kafka_transform.simple "
-        f"BOOTSTRAP_SERVERS 'localhost:9092'",
+        f"CREATE KAFKA STREAM default_values TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BOOTSTRAP_SERVERS 'localhost:9092'",
     )
 
     consumer_group = "my_special_consumer_group"
-    batch_interval = 42
-    batch_size = 3
+    BATCH_INTERVAL = 42
+    BATCH_SIZE = 3
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM complex_values "
-        f"TOPICS {','.join(kafka_topics)} "
-        f"TRANSFORM kafka_transform.with_parameters "
-        f"CONSUMER_GROUP {consumer_group} "
-        f"BATCH_INTERVAL {batch_interval} "
-        f"BATCH_SIZE {batch_size} ",
+        f"CREATE KAFKA STREAM complex_values TOPICS {','.join(kafka_topics)} TRANSFORM kafka_transform.with_parameters CONSUMER_GROUP {consumer_group} BATCH_INTERVAL {BATCH_INTERVAL} BATCH_SIZE {BATCH_SIZE} ",
     )
 
     assert len(common.execute_and_fetch_all(cursor, "SHOW STREAMS")) == 2
@@ -216,13 +191,7 @@ def test_show_streams(kafka_producer, kafka_topics, connection):
     common.check_stream_info(
         cursor,
         "default_values",
-        ("default_values",
-         "kafka",
-         100,
-         1000,
-         "kafka_transform.simple",
-         None,
-         False),
+        ("default_values", "kafka", 100, 1000, "kafka_transform.simple", None, False),
     )
 
     common.check_stream_info(
@@ -231,8 +200,8 @@ def test_show_streams(kafka_producer, kafka_topics, connection):
         (
             "complex_values",
             "kafka",
-            batch_interval,
-            batch_size,
+            BATCH_INTERVAL,
+            BATCH_SIZE,
             "kafka_transform.with_parameters",
             None,
             False,
@@ -241,15 +210,12 @@ def test_show_streams(kafka_producer, kafka_topics, connection):
 
 
 @pytest.mark.parametrize("operation", ["START", "STOP"])
-def test_start_and_stop_during_check(
-        kafka_producer,
-        kafka_topics,
-        connection,
-        operation):
+def test_start_and_stop_during_check(kafka_producer, kafka_topics, connection, operation):
     assert len(kafka_topics) > 1
+    BATCH_SIZE = 1
 
     def stream_creator(stream_name):
-        return f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple"
+        return f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE {BATCH_SIZE}"
 
     def message_sender(msg):
         kafka_producer.send(kafka_topics[0], msg).get(timeout=60)
@@ -259,7 +225,9 @@ def test_start_and_stop_during_check(
         connection,
         stream_creator,
         message_sender,
-        "Kafka consumer test_stream is already stopped")
+        "Kafka consumer test_stream is already stopped",
+        BATCH_SIZE,
+    )
 
 
 def test_check_already_started_stream(kafka_topics, connection):
@@ -268,9 +236,7 @@ def test_check_already_started_stream(kafka_topics, connection):
 
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM started_stream "
-        f"TOPICS {kafka_topics[0]} "
-        f"TRANSFORM kafka_transform.simple",
+        f"CREATE KAFKA STREAM started_stream TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple",
     )
     common.start_stream(cursor, "started_stream")
 
@@ -289,41 +255,29 @@ def test_restart_after_error(kafka_producer, kafka_topics, connection):
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM test_stream "
-        f"TOPICS {kafka_topics[0]} "
-        f"TRANSFORM kafka_transform.query",
+        f"CREATE KAFKA STREAM test_stream TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.query",
     )
 
     common.start_stream(cursor, "test_stream")
     time.sleep(1)
 
     kafka_producer.send(kafka_topics[0], common.SIMPLE_MSG).get(timeout=60)
-    assert common.timed_wait(
-        lambda: not common.get_is_running(
-            cursor, "test_stream"))
+    assert common.timed_wait(lambda: not common.get_is_running(cursor, "test_stream"))
 
     common.start_stream(cursor, "test_stream")
     time.sleep(1)
     kafka_producer.send(kafka_topics[0], b"CREATE (n:VERTEX { id : 42 })")
-    assert common.check_one_result_row(
-        cursor, "MATCH (n:VERTEX { id : 42 }) RETURN n")
+    assert common.check_one_result_row(cursor, "MATCH (n:VERTEX { id : 42 }) RETURN n")
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK_PY)
-def test_bootstrap_server(
-        kafka_producer,
-        kafka_topics,
-        connection,
-        transformation):
+def test_bootstrap_server(kafka_producer, kafka_topics, connection, transformation):
     assert len(kafka_topics) > 0
     cursor = connection.cursor()
-    local = "localhost:9092"
+    LOCAL = "localhost:9092"
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM test "
-        f"TOPICS {','.join(kafka_topics)} "
-        f"TRANSFORM {transformation} "
-        f"BOOTSTRAP_SERVERS '{local}'",
+        f"CREATE KAFKA STREAM test TOPICS {','.join(kafka_topics)} TRANSFORM {transformation} BOOTSTRAP_SERVERS '{LOCAL}'",
     )
     common.start_stream(cursor, "test")
     time.sleep(5)
@@ -332,25 +286,17 @@ def test_bootstrap_server(
         kafka_producer.send(topic, common.SIMPLE_MSG).get(timeout=60)
 
     for topic in kafka_topics:
-        common.kafka_check_vertex_exists_with_topic_and_payload(
-            cursor, topic, common.SIMPLE_MSG)
+        common.kafka_check_vertex_exists_with_topic_and_payload(cursor, topic, common.SIMPLE_MSG)
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK_PY)
-def test_bootstrap_server_empty(
-        kafka_producer,
-        kafka_topics,
-        connection,
-        transformation):
+def test_bootstrap_server_empty(kafka_producer, kafka_topics, connection, transformation):
     assert len(kafka_topics) > 0
     cursor = connection.cursor()
     with pytest.raises(mgclient.DatabaseError):
         common.execute_and_fetch_all(
             cursor,
-            "CREATE KAFKA STREAM test "
-            f"TOPICS {','.join(kafka_topics)} "
-            f"TRANSFORM {transformation} "
-            "BOOTSTRAP_SERVERS ''",
+            f"CREATE KAFKA STREAM test TOPICS {','.join(kafka_topics)} TRANSFORM {transformation} BOOTSTRAP_SERVERS ''",
         )
 
 
@@ -360,10 +306,7 @@ def test_set_offset(kafka_producer, kafka_topics, connection, transformation):
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE KAFKA STREAM test "
-        f"TOPICS {kafka_topics[0]} "
-        f"TRANSFORM {transformation} "
-        "BATCH_SIZE 1",
+        f"CREATE KAFKA STREAM test TOPICS {kafka_topics[0]} TRANSFORM {transformation} BATCH_SIZE 1",
     )
 
     messages = [f"{i} message" for i in range(1, 21)]
@@ -377,27 +320,18 @@ def test_set_offset(kafka_producer, kafka_topics, connection, transformation):
         else:
             assert common.check_one_result_row(
                 cursor,
-                (
-                    f"MATCH (n: MESSAGE {{payload: '{expected_msgs[-1]}'}})"
-                    "RETURN n"
-                ),
+                (f"MATCH (n: MESSAGE {{payload: '{expected_msgs[-1]}'}})" "RETURN n"),
             )
         common.stop_stream(cursor, "test")
-        res = common.execute_and_fetch_all(
-            cursor, "MATCH (n) RETURN n.payload"
-        )
+        res = common.execute_and_fetch_all(cursor, "MATCH (n) RETURN n.payload")
         return res
 
     def execute_set_offset_and_consume(id, expected_msgs):
-        common.execute_and_fetch_all(
-            cursor, f"CALL mg.kafka_set_stream_offset('test', {id})"
-        )
+        common.execute_and_fetch_all(cursor, f"CALL mg.kafka_set_stream_offset('test', {id})")
         return consume(expected_msgs)
 
     with pytest.raises(mgclient.DatabaseError):
-        res = common.execute_and_fetch_all(
-            cursor, "CALL mg.kafka_set_stream_offset('foo', 10)"
-        )
+        res = common.execute_and_fetch_all(cursor, "CALL mg.kafka_set_stream_offset('foo', 10)")
 
     def comparison_check(a, b):
         return a == str(b).strip("'(,)")
@@ -426,40 +360,156 @@ def test_set_offset(kafka_producer, kafka_topics, connection, transformation):
 
 def test_info_procedure(kafka_topics, connection):
     cursor = connection.cursor()
-    stream_name = 'test_stream'
-    configs = {"sasl.username": "michael.scott"}
-    local = "localhost:9092"
-    credentials = {"sasl.password": "S3cr3tP4ssw0rd"}
-    consumer_group = "ConsumerGr"
+    STREAM_NAME = "test_stream"
+    CONFIGS = {"sasl.username": "michael.scott"}
+    LOCAL = "localhost:9092"
+    CREDENTIALS = {"sasl.password": "S3cr3tP4ssw0rd"}
+    CONSUMER_GROUP = "ConsumerGr"
     common.execute_and_fetch_all(
         cursor,
-        f"CREATE KAFKA STREAM {stream_name} "
-        f"TOPICS {','.join(kafka_topics)} "
-        f"TRANSFORM pulsar_transform.simple "
-        f"CONSUMER_GROUP {consumer_group} "
-        f"BOOTSTRAP_SERVERS '{local}' "
-        f"CONFIGS {configs} "
-        f"CREDENTIALS {credentials}"
+        f"CREATE KAFKA STREAM {STREAM_NAME} TOPICS {','.join(kafka_topics)} TRANSFORM kafka_transform.simple CONSUMER_GROUP {CONSUMER_GROUP} BOOTSTRAP_SERVERS '{LOCAL}' CONFIGS {CONFIGS} CREDENTIALS {CREDENTIALS}",
     )
 
-    stream_info = common.execute_and_fetch_all(
-        cursor, f"CALL mg.kafka_stream_info('{stream_name}') YIELD *")
+    stream_info = common.execute_and_fetch_all(cursor, f"CALL mg.kafka_stream_info('{STREAM_NAME}') YIELD *")
 
-    reducted_credentials = {key: "<REDUCTED>" for
-                            key in credentials.keys()}
+    reducted_credentials = {key: "<REDUCTED>" for key in CREDENTIALS.keys()}
 
-    expected_stream_info = [
-        (local, configs, consumer_group, reducted_credentials, kafka_topics)]
+    expected_stream_info = [(LOCAL, CONFIGS, CONSUMER_GROUP, reducted_credentials, kafka_topics)]
     common.validate_info(stream_info, expected_stream_info)
 
-@pytest.mark.parametrize("transformation",TRANSFORMATIONS_TO_CHECK_C)
+
+@pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK_C)
 def test_load_c_transformations(connection, transformation):
     cursor = connection.cursor()
-    query = "CALL mg.transformations() YIELD * WITH name WHERE name STARTS WITH 'c_transformations." + transformation + "' RETURN name" 
-    result = common.execute_and_fetch_all(
-                 cursor, query)
+
+    query = f"CALL mg.transformations() YIELD * WITH name WHERE name STARTS WITH '{transformation}' RETURN name"
+    result = common.execute_and_fetch_all(cursor, query)
     assert len(result) == 1
-    assert result[0][0] == "c_transformations." + transformation
-    
+    assert result[0][0] == transformation
+
+
+def test_check_stream_same_number_of_queries_than_messages(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    TRANSFORMATION = "common_transform.check_stream_no_filtering"
+
+    def stream_creator(stream_name, batch_size):
+        return f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM {TRANSFORMATION} BATCH_INTERVAL 3000 BATCH_SIZE {batch_size}"
+
+    def message_sender(msg):
+        kafka_producer.send(kafka_topics[0], msg).get(timeout=60)
+
+    common.test_check_stream_same_number_of_queries_than_messages(connection, stream_creator, message_sender)
+
+
+def test_check_stream_different_number_of_queries_than_messages(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    TRANSFORMATION = "common_transform.check_stream_with_filtering"
+
+    def stream_creator(stream_name, batch_size):
+        return f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM {TRANSFORMATION} BATCH_INTERVAL 3000  BATCH_SIZE {batch_size}"
+
+    def message_sender(msg):
+        kafka_producer.send(kafka_topics[0], msg).get(timeout=60)
+
+    common.test_check_stream_different_number_of_queries_than_messages(connection, stream_creator, message_sender)
+
+
+def test_start_stream_with_batch_limit(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name):
+        return (
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE 1"
+        )
+
+    def messages_sender(nof_messages):
+        for x in range(nof_messages):
+            kafka_producer.send(kafka_topics[0], common.SIMPLE_MSG).get(timeout=60)
+
+    common.test_start_stream_with_batch_limit(connection, stream_creator, messages_sender)
+
+
+def test_start_stream_with_batch_limit_timeout(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name):
+        return (
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE 1"
+        )
+
+    common.test_start_stream_with_batch_limit_timeout(connection, stream_creator)
+
+
+def test_start_stream_with_batch_limit_reaching_timeout(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name, batch_size):
+        return f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE {batch_size}"
+
+    common.test_start_stream_with_batch_limit_reaching_timeout(connection, stream_creator)
+
+
+def test_start_stream_with_batch_limit_while_check_running(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name):
+        return (
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE 1"
+        )
+
+    def message_sender(message):
+        kafka_producer.send(kafka_topics[0], message).get(timeout=6000)
+
+    def setup_function(start_check_stream, cursor, stream_name, batch_limit, timeout):
+        thread_stream_check = Process(target=start_check_stream, daemon=True, args=(stream_name, batch_limit, timeout))
+        thread_stream_check.start()
+        time.sleep(2)
+        assert common.get_is_running(cursor, stream_name)
+        message_sender(common.SIMPLE_MSG)
+        thread_stream_check.join()
+
+    common.test_start_stream_with_batch_limit_while_check_running(
+        connection, stream_creator, message_sender, setup_function
+    )
+
+
+def test_check_while_stream_with_batch_limit_running(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name):
+        return (
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE 1"
+        )
+
+    def message_sender(message):
+        kafka_producer.send(kafka_topics[0], message).get(timeout=6000)
+
+    common.test_check_while_stream_with_batch_limit_running(connection, stream_creator, message_sender)
+
+
+def test_start_stream_with_batch_limit_with_invalid_batch_limit(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name):
+        return (
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE 1"
+        )
+
+    common.test_start_stream_with_batch_limit_with_invalid_batch_limit(connection, stream_creator)
+
+
+def test_check_stream_with_batch_limit_with_invalid_batch_limit(kafka_producer, kafka_topics, connection):
+    assert len(kafka_topics) > 0
+
+    def stream_creator(stream_name):
+        return (
+            f"CREATE KAFKA STREAM {stream_name} TOPICS {kafka_topics[0]} TRANSFORM kafka_transform.simple BATCH_SIZE 1"
+        )
+
+    common.test_check_stream_with_batch_limit_with_invalid_batch_limit(connection, stream_creator)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-rA"]))

tests/e2e/streams/pulsar_streams_tests.py

@@ -18,17 +18,14 @@ import time
 from multiprocessing import Process, Value
 import common
 
-TRANSFORMATIONS_TO_CHECK = [
-    "pulsar_transform.simple",
-    "pulsar_transform.with_parameters"]
+TRANSFORMATIONS_TO_CHECK = ["pulsar_transform.simple", "pulsar_transform.with_parameters"]
 
 
 def check_vertex_exists_with_topic_and_payload(cursor, topic, payload_byte):
-    decoded_payload = payload_byte.decode('utf-8')
+    decoded_payload = payload_byte.decode("utf-8")
     common.check_vertex_exists_with_properties(
-        cursor, {
-            'topic': f'"{common.pulsar_default_namespace_topic(topic)}"',
-            'payload': f'"{decoded_payload}"'})
+        cursor, {"topic": f'"{common.pulsar_default_namespace_topic(topic)}"', "payload": f'"{decoded_payload}"'}
+    )
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK)
@@ -37,30 +34,23 @@ def test_simple(pulsar_client, pulsar_topics, connection, transformation):
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM test "
-        f"TOPICS '{','.join(pulsar_topics)}' "
-        f"TRANSFORM {transformation}",
+        f"CREATE PULSAR STREAM test TOPICS '{','.join(pulsar_topics)}' TRANSFORM {transformation}",
     )
     common.start_stream(cursor, "test")
     time.sleep(5)
 
     for topic in pulsar_topics:
         producer = pulsar_client.create_producer(
-            common.pulsar_default_namespace_topic(topic),
-            send_timeout_millis=60000)
+            common.pulsar_default_namespace_topic(topic), send_timeout_millis=60000
+        )
         producer.send(common.SIMPLE_MSG)
 
     for topic in pulsar_topics:
-        check_vertex_exists_with_topic_and_payload(
-            cursor, topic, common.SIMPLE_MSG)
+        check_vertex_exists_with_topic_and_payload(cursor, topic, common.SIMPLE_MSG)
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK)
-def test_separate_consumers(
-        pulsar_client,
-        pulsar_topics,
-        connection,
-        transformation):
+def test_separate_consumers(pulsar_client, pulsar_topics, connection, transformation):
     assert len(pulsar_topics) > 0
     cursor = connection.cursor()
 
@@ -70,9 +60,7 @@ def test_separate_consumers(
         stream_names.append(stream_name)
         common.execute_and_fetch_all(
             cursor,
-            f"CREATE PULSAR STREAM {stream_name} "
-            f"TOPICS {topic} "
-            f"TRANSFORM {transformation}",
+            f"CREATE PULSAR STREAM {stream_name} TOPICS {topic} TRANSFORM {transformation}",
         )
 
     for stream_name in stream_names:
@@ -81,13 +69,11 @@ def test_separate_consumers(
     time.sleep(5)
 
     for topic in pulsar_topics:
-        producer = pulsar_client.create_producer(
-            topic, send_timeout_millis=60000)
+        producer = pulsar_client.create_producer(topic, send_timeout_millis=60000)
         producer.send(common.SIMPLE_MSG)
 
     for topic in pulsar_topics:
-        check_vertex_exists_with_topic_and_payload(
-            cursor, topic, common.SIMPLE_MSG)
+        check_vertex_exists_with_topic_and_payload(cursor, topic, common.SIMPLE_MSG)
 
 
 def test_start_from_latest_messages(pulsar_client, pulsar_topics, connection):
@@ -99,118 +85,112 @@ def test_start_from_latest_messages(pulsar_client, pulsar_topics, connection):
     assert len(pulsar_topics) > 0
     cursor = connection.cursor()
     common.execute_and_fetch_all(
-        cursor, "CREATE PULSAR STREAM test "
-        f"TOPICS {pulsar_topics[0]} "
-        "TRANSFORM pulsar_transform.simple", )
+        cursor,
+        f"CREATE PULSAR STREAM test TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple",
+    )
     common.start_stream(cursor, "test")
     time.sleep(1)
 
     def assert_message_not_consumed(message):
         vertices_with_msg = common.execute_and_fetch_all(
             cursor,
-            "MATCH (n: MESSAGE {" f"payload: '{message.decode('utf-8')}'" "}) RETURN n",
+            f"MATCH (n: MESSAGE {{payload: '{message.decode('utf-8')}'}}) RETURN n",
         )
 
         assert len(vertices_with_msg) == 0
 
     producer = pulsar_client.create_producer(
-        common.pulsar_default_namespace_topic(
-            pulsar_topics[0]), send_timeout_millis=60000)
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
     producer.send(common.SIMPLE_MSG)
 
-    check_vertex_exists_with_topic_and_payload(
-        cursor, pulsar_topics[0], common.SIMPLE_MSG)
+    check_vertex_exists_with_topic_and_payload(cursor, pulsar_topics[0], common.SIMPLE_MSG)
 
     common.stop_stream(cursor, "test")
 
-    next_message = b"NEXT"
-    producer.send(next_message)
+    NEXT_MESSAGE = b"NEXT"
+    producer.send(NEXT_MESSAGE)
 
-    assert_message_not_consumed(next_message)
+    assert_message_not_consumed(NEXT_MESSAGE)
 
     common.start_stream(cursor, "test")
-    check_vertex_exists_with_topic_and_payload(
-        cursor, pulsar_topics[0], next_message)
+    check_vertex_exists_with_topic_and_payload(cursor, pulsar_topics[0], NEXT_MESSAGE)
     common.stop_stream(cursor, "test")
 
     common.drop_stream(cursor, "test")
 
-    lost_message = b"LOST"
-    valid_messages = [b"second message", b"third message"]
+    LOST_MESSAGE = b"LOST"
+    VALID_MESSAGES = [b"second message", b"third message"]
 
-    producer.send(lost_message)
+    producer.send(LOST_MESSAGE)
 
-    assert_message_not_consumed(lost_message)
+    assert_message_not_consumed(LOST_MESSAGE)
 
     common.execute_and_fetch_all(
-        cursor, "CREATE PULSAR STREAM test "
-        f"TOPICS {pulsar_topics[0]} "
-        "TRANSFORM pulsar_transform.simple", )
+        cursor,
+        f"CREATE PULSAR STREAM test TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple",
+    )
 
-    for message in valid_messages:
+    for message in VALID_MESSAGES:
         producer.send(message)
         assert_message_not_consumed(message)
 
     common.start_stream(cursor, "test")
 
-    assert_message_not_consumed(lost_message)
+    assert_message_not_consumed(LOST_MESSAGE)
 
-    for message in valid_messages:
-        check_vertex_exists_with_topic_and_payload(
-            cursor, pulsar_topics[0], message)
+    for message in VALID_MESSAGES:
+        check_vertex_exists_with_topic_and_payload(cursor, pulsar_topics[0], message)
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK)
-def test_check_stream(
-        pulsar_client,
-        pulsar_topics,
-        connection,
-        transformation):
+def test_check_stream(pulsar_client, pulsar_topics, connection, transformation):
     assert len(pulsar_topics) > 0
+    BATCH_SIZE = 1
+    INDEX_Of_FIRST_BATCH = 0
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM test "
-        f"TOPICS {pulsar_topics[0]} "
-        f"TRANSFORM {transformation} "
-        "BATCH_SIZE 1",
+        f"CREATE PULSAR STREAM test TOPICS {pulsar_topics[0]} TRANSFORM {transformation} BATCH_SIZE {BATCH_SIZE}",
     )
     common.start_stream(cursor, "test")
     time.sleep(1)
 
     producer = pulsar_client.create_producer(
-        common.pulsar_default_namespace_topic(
-            pulsar_topics[0]), send_timeout_millis=60000)
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
     producer.send(common.SIMPLE_MSG)
-    check_vertex_exists_with_topic_and_payload(
-        cursor, pulsar_topics[0], common.SIMPLE_MSG)
+    check_vertex_exists_with_topic_and_payload(cursor, pulsar_topics[0], common.SIMPLE_MSG)
     common.stop_stream(cursor, "test")
 
-    messages = [b"first message", b"second message", b"third message"]
-    for message in messages:
+    MESSAGES = [b"first message", b"second message", b"third message"]
+    for message in MESSAGES:
         producer.send(message)
 
     def check_check_stream(batch_limit):
-        assert (
-            transformation == "pulsar_transform.simple"
-            or transformation == "pulsar_transform.with_parameters"
-        )
-        test_results = common.execute_and_fetch_all(
-            cursor, f"CHECK STREAM test BATCH_LIMIT {batch_limit}"
-        )
+        assert transformation == "pulsar_transform.simple" or transformation == "pulsar_transform.with_parameters"
+        test_results = common.execute_and_fetch_all(cursor, f"CHECK STREAM test BATCH_LIMIT {batch_limit}")
         assert len(test_results) == batch_limit
 
         for i in range(batch_limit):
-            message_as_str = messages[i].decode("utf-8")
+            message_as_str = MESSAGES[i].decode("utf-8")
+            assert (
+                BATCH_SIZE == 1
+            )  # If batch size != 1, then the usage of INDEX_Of_FIRST_BATCH must change: the result will have a list of queries (pair<parameters,query>)
+
             if transformation == "pulsar_transform.simple":
-                assert f"payload: '{message_as_str}'" in test_results[i][common.QUERY]
-                assert test_results[i][common.PARAMS] is None
+                assert (
+                    f"payload: '{message_as_str}'"
+                    in test_results[i][common.QUERIES][INDEX_Of_FIRST_BATCH][common.QUERY_LITERAL]
+                )
+                assert test_results[i][common.QUERIES][INDEX_Of_FIRST_BATCH][common.PARAMETERS_LITERAL] is None
             else:
-                assert f"payload: $payload" in test_results[i][
-                    common.QUERY] and f"topic: $topic" in test_results[i][common.QUERY]
-                parameters = test_results[i][common.PARAMS]
-                assert parameters["topic"] == common.pulsar_default_namespace_topic(
-                    pulsar_topics[0])
+                assert (
+                    f"payload: $payload" in test_results[i][common.QUERIES][INDEX_Of_FIRST_BATCH][common.QUERY_LITERAL]
+                    and f"topic: $topic" in test_results[i][common.QUERIES][INDEX_Of_FIRST_BATCH][common.QUERY_LITERAL]
+                )
+                parameters = test_results[i][common.QUERIES][INDEX_Of_FIRST_BATCH][common.PARAMETERS_LITERAL]
+                assert parameters["topic"] == common.pulsar_default_namespace_topic(pulsar_topics[0])
                 assert parameters["payload"] == message_as_str
 
     check_check_stream(1)
@@ -218,45 +198,37 @@ def test_check_stream(
     check_check_stream(3)
     common.start_stream(cursor, "test")
 
-    for message in messages:
-        check_vertex_exists_with_topic_and_payload(
-            cursor, pulsar_topics[0], message)
+    for message in MESSAGES:
+        check_vertex_exists_with_topic_and_payload(cursor, pulsar_topics[0], message)
 
 
 def test_info_procedure(pulsar_client, pulsar_topics, connection):
     cursor = connection.cursor()
-    stream_name = 'test_stream'
+    STREAM_NAME = "test_stream"
     common.execute_and_fetch_all(
         cursor,
-        f"CREATE PULSAR STREAM {stream_name} "
-        f"TOPICS {','.join(pulsar_topics)} "
-        f"TRANSFORM pulsar_transform.simple ",
+        f"CREATE PULSAR STREAM {STREAM_NAME} TOPICS {','.join(pulsar_topics)} TRANSFORM pulsar_transform.simple ",
     )
 
-    stream_info = common.execute_and_fetch_all(cursor, f"CALL mg.pulsar_stream_info('{stream_name}') YIELD *")
+    stream_info = common.execute_and_fetch_all(cursor, f"CALL mg.pulsar_stream_info('{STREAM_NAME}') YIELD *")
 
     expected_stream_info = [(common.PULSAR_SERVICE_URL, pulsar_topics)]
     common.validate_info(stream_info, expected_stream_info)
 
+
 def test_show_streams(pulsar_client, pulsar_topics, connection):
     assert len(pulsar_topics) > 1
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM default_values "
-        f"TOPICS {pulsar_topics[0]} "
-        f"TRANSFORM pulsar_transform.simple ",
+        f"CREATE PULSAR STREAM default_values TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple ",
     )
 
-    batch_interval = 42
-    batch_size = 3
+    BATCH_INTERVAL = 42
+    BATCH_SIZE = 3
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM complex_values "
-        f"TOPICS {','.join(pulsar_topics)} "
-        f"TRANSFORM pulsar_transform.with_parameters "
-        f"BATCH_INTERVAL {batch_interval} "
-        f"BATCH_SIZE {batch_size} ",
+        f"CREATE PULSAR STREAM complex_values TOPICS {','.join(pulsar_topics)} TRANSFORM pulsar_transform.with_parameters BATCH_INTERVAL {BATCH_INTERVAL} BATCH_SIZE {BATCH_SIZE} ",
     )
 
     assert len(common.execute_and_fetch_all(cursor, "SHOW STREAMS")) == 2
@@ -264,13 +236,7 @@ def test_show_streams(pulsar_client, pulsar_topics, connection):
     common.check_stream_info(
         cursor,
         "default_values",
-        ("default_values",
-         "pulsar",
-         100,
-         1000,
-         "pulsar_transform.simple",
-         None,
-         False),
+        ("default_values", "pulsar", 100, 1000, "pulsar_transform.simple", None, False),
     )
 
     common.check_stream_info(
@@ -279,8 +245,8 @@ def test_show_streams(pulsar_client, pulsar_topics, connection):
         (
             "complex_values",
             "pulsar",
-            batch_interval,
-            batch_size,
+            BATCH_INTERVAL,
+            BATCH_SIZE,
             "pulsar_transform.with_parameters",
             None,
             False,
@@ -289,19 +255,16 @@ def test_show_streams(pulsar_client, pulsar_topics, connection):
 
 
 @pytest.mark.parametrize("operation", ["START", "STOP"])
-def test_start_and_stop_during_check(
-        pulsar_client,
-        pulsar_topics,
-        connection,
-        operation):
+def test_start_and_stop_during_check(pulsar_client, pulsar_topics, connection, operation):
     assert len(pulsar_topics) > 1
+    BATCH_SIZE = 1
 
     def stream_creator(stream_name):
-        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple"
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE {BATCH_SIZE}"
 
     producer = pulsar_client.create_producer(
-        common.pulsar_default_namespace_topic(
-            pulsar_topics[0]), send_timeout_millis=60000)
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
 
     def message_sender(msg):
         producer.send(msg)
@@ -311,7 +274,9 @@ def test_start_and_stop_during_check(
         connection,
         stream_creator,
         message_sender,
-        "Pulsar consumer test_stream is already stopped")
+        "Pulsar consumer test_stream is already stopped",
+        BATCH_SIZE,
+    )
 
 
 def test_check_already_started_stream(pulsar_topics, connection):
@@ -320,9 +285,7 @@ def test_check_already_started_stream(pulsar_topics, connection):
 
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM started_stream "
-        f"TOPICS {pulsar_topics[0]} "
-        f"TRANSFORM pulsar_transform.simple",
+        f"CREATE PULSAR STREAM started_stream TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple",
     )
     common.start_stream(cursor, "started_stream")
 
@@ -333,6 +296,7 @@ def test_check_already_started_stream(pulsar_topics, connection):
 def test_start_checked_stream_after_timeout(pulsar_topics, connection):
     def stream_creator(stream_name):
         return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple"
+
     common.test_start_checked_stream_after_timeout(connection, stream_creator)
 
 
@@ -340,53 +304,165 @@ def test_restart_after_error(pulsar_client, pulsar_topics, connection):
     cursor = connection.cursor()
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM test_stream "
-        f"TOPICS {pulsar_topics[0]} "
-        f"TRANSFORM pulsar_transform.query",
+        f"CREATE PULSAR STREAM test_stream TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.query",
     )
 
     common.start_stream(cursor, "test_stream")
     time.sleep(1)
 
     producer = pulsar_client.create_producer(
-        common.pulsar_default_namespace_topic(
-            pulsar_topics[0]), send_timeout_millis=60000)
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
     producer.send(common.SIMPLE_MSG)
-    assert common.timed_wait(
-        lambda: not common.get_is_running(
-            cursor, "test_stream"))
+    assert common.timed_wait(lambda: not common.get_is_running(cursor, "test_stream"))
 
     common.start_stream(cursor, "test_stream")
     time.sleep(1)
     producer.send(b"CREATE (n:VERTEX { id : 42 })")
-    assert common.check_one_result_row(
-        cursor, "MATCH (n:VERTEX { id : 42 }) RETURN n")
+    assert common.check_one_result_row(cursor, "MATCH (n:VERTEX { id : 42 }) RETURN n")
 
 
 @pytest.mark.parametrize("transformation", TRANSFORMATIONS_TO_CHECK)
 def test_service_url(pulsar_client, pulsar_topics, connection, transformation):
     assert len(pulsar_topics) > 0
     cursor = connection.cursor()
-    local = "pulsar://127.0.0.1:6650"
+    LOCAL = "pulsar://127.0.0.1:6650"
     common.execute_and_fetch_all(
         cursor,
-        "CREATE PULSAR STREAM test "
-        f"TOPICS {','.join(pulsar_topics)} "
-        f"TRANSFORM {transformation} "
-        f"SERVICE_URL '{local}'",
+        f"CREATE PULSAR STREAM test TOPICS {','.join(pulsar_topics)} TRANSFORM {transformation} SERVICE_URL '{LOCAL}'",
     )
     common.start_stream(cursor, "test")
     time.sleep(5)
 
     for topic in pulsar_topics:
         producer = pulsar_client.create_producer(
-            common.pulsar_default_namespace_topic(topic),
-            send_timeout_millis=60000)
+            common.pulsar_default_namespace_topic(topic), send_timeout_millis=60000
+        )
         producer.send(common.SIMPLE_MSG)
 
     for topic in pulsar_topics:
-        check_vertex_exists_with_topic_and_payload(
-            cursor, topic, common.SIMPLE_MSG)
+        check_vertex_exists_with_topic_and_payload(cursor, topic, common.SIMPLE_MSG)
+
+
+def test_start_stream_with_batch_limit(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 1
+
+    def stream_creator(stream_name):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE 1"
+
+    producer = pulsar_client.create_producer(
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
+
+    def messages_sender(nof_messages):
+        for x in range(nof_messages):
+            producer.send(common.SIMPLE_MSG)
+
+    common.test_start_stream_with_batch_limit(connection, stream_creator, messages_sender)
+
+
+def test_start_stream_with_batch_limit_timeout(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 1
+
+    def stream_creator(stream_name):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE 1"
+
+    common.test_start_stream_with_batch_limit_timeout(connection, stream_creator)
+
+
+def test_start_stream_with_batch_limit_reaching_timeout(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 1
+
+    def stream_creator(stream_name, batch_size):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE {batch_size}"
+
+    common.test_start_stream_with_batch_limit_reaching_timeout(connection, stream_creator)
+
+
+def test_start_stream_with_batch_limit_while_check_running(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 0
+
+    def stream_creator(stream_name):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE 1"
+
+    producer = pulsar_client.create_producer(
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
+
+    def message_sender(message):
+        producer.send(message)
+
+    common.test_start_stream_with_batch_limit_while_check_running(connection, stream_creator, message_sender)
+
+
+def test_check_while_stream_with_batch_limit_running(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 0
+
+    def stream_creator(stream_name):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE 1"
+
+    producer = pulsar_client.create_producer(
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
+
+    def message_sender(message):
+        producer.send(message)
+
+    common.test_check_while_stream_with_batch_limit_running(connection, stream_creator, message_sender)
+
+
+def test_check_stream_same_number_of_queries_than_messages(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 0
+
+    TRANSFORMATION = "common_transform.check_stream_no_filtering"
+
+    def stream_creator(stream_name, batch_size):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM {TRANSFORMATION} BATCH_INTERVAL 3000 BATCH_SIZE  {batch_size} "
+
+    producer = pulsar_client.create_producer(
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
+
+    def message_sender(msg):
+        producer.send(msg)
+
+    common.test_check_stream_same_number_of_queries_than_messages(connection, stream_creator, message_sender)
+
+
+def test_check_stream_different_number_of_queries_than_messages(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 0
+
+    TRANSFORMATION = "common_transform.check_stream_with_filtering"
+
+    def stream_creator(stream_name, batch_size):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM {TRANSFORMATION} BATCH_INTERVAL 3000 BATCH_SIZE  {batch_size} "
+
+    producer = pulsar_client.create_producer(
+        common.pulsar_default_namespace_topic(pulsar_topics[0]), send_timeout_millis=60000
+    )
+
+    def message_sender(msg):
+        producer.send(msg)
+
+    common.test_check_stream_different_number_of_queries_than_messages(connection, stream_creator, message_sender)
+
+
+def test_start_stream_with_batch_limit_with_invalid_batch_limit(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 0
+
+    def stream_creator(stream_name):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE 1"
+
+    common.test_start_stream_with_batch_limit_with_invalid_batch_limit(connection, stream_creator)
+
+
+def test_check_stream_with_batch_limit_with_invalid_batch_limit(pulsar_client, pulsar_topics, connection):
+    assert len(pulsar_topics) > 0
+
+    def stream_creator(stream_name):
+        return f"CREATE PULSAR STREAM {stream_name} TOPICS {pulsar_topics[0]} TRANSFORM pulsar_transform.simple BATCH_SIZE 1"
+
+    common.test_check_stream_with_batch_limit_with_invalid_batch_limit(connection, stream_creator)
 
 
 if __name__ == "__main__":

tests/e2e/streams/transformations/CMakeLists.txt

@@ -1,3 +1,4 @@
 copy_streams_e2e_python_files(kafka_transform.py)
 copy_streams_e2e_python_files(pulsar_transform.py)
+copy_streams_e2e_python_files(common_transform.py)
 add_query_module(c_transformations c_transformations.cpp)

tests/e2e/streams/transformations/common_transform.py

@@ -0,0 +1,57 @@
+# Copyright 2021 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+import mgp
+
+
+@mgp.transformation
+def check_stream_no_filtering(
+    context: mgp.TransCtx, messages: mgp.Messages
+) -> mgp.Record(query=str, parameters=mgp.Map):
+
+    result_queries = []
+
+    for i in range(0, messages.total_messages()):
+        message = messages.message_at(i)
+        payload_as_str = message.payload().decode("utf-8")
+        result_queries.append(
+            mgp.Record(query=f"Message: {payload_as_str}", parameters={"value": f"Parameter: {payload_as_str}"})
+        )
+
+    return result_queries
+
+
+@mgp.transformation
+def check_stream_with_filtering(
+    context: mgp.TransCtx, messages: mgp.Messages
+) -> mgp.Record(query=str, parameters=mgp.Map):
+
+    result_queries = []
+
+    for i in range(0, messages.total_messages()):
+        message = messages.message_at(i)
+        payload_as_str = message.payload().decode("utf-8")
+
+        if "a" in payload_as_str:
+            continue
+
+        result_queries.append(
+            mgp.Record(query=f"Message: {payload_as_str}", parameters={"value": f"Parameter: {payload_as_str}"})
+        )
+
+        if "b" in payload_as_str:
+            result_queries.append(
+                mgp.Record(
+                    query=f"Message: extra_{payload_as_str}", parameters={"value": f"Parameter: extra_{payload_as_str}"}
+                )
+            )
+
+    return result_queries

tests/e2e/streams/transformations/kafka_transform.py

@@ -13,9 +13,7 @@ import mgp
 
 
 @mgp.transformation
-def simple(
-    context: mgp.TransCtx, messages: mgp.Messages
-) -> mgp.Record(query=str, parameters=mgp.Map):
+def simple(context: mgp.TransCtx, messages: mgp.Messages) -> mgp.Record(query=str, parameters=mgp.Map):
 
     result_queries = []
 
@@ -32,15 +30,15 @@ def simple(
                     offset: '{message.offset()}',
                     topic: '{message.topic_name()}'
                 }})""",
-                parameters=None))
+                parameters=None,
+            )
+        )
 
     return result_queries
 
 
 @mgp.transformation
-def with_parameters(
-    context: mgp.TransCtx, messages: mgp.Messages
-) -> mgp.Record(query=str, parameters=mgp.Map):
+def with_parameters(context: mgp.TransCtx, messages: mgp.Messages) -> mgp.Record(query=str, parameters=mgp.Map):
 
     result_queries = []
 
@@ -61,7 +59,10 @@ def with_parameters(
                     "timestamp": message.timestamp(),
                     "payload": payload_as_str,
                     "offset": message.offset(),
-                    "topic": message.topic_name()}))
+                    "topic": message.topic_name(),
+                },
+            )
+        )
 
     return result_queries
 
@@ -76,8 +77,6 @@ def query(
         message = messages.message_at(i)
         assert message.source_type() == mgp.SOURCE_TYPE_KAFKA
         payload_as_str = message.payload().decode("utf-8")
-        result_queries.append(
-            mgp.Record(query=payload_as_str, parameters=None)
-        )
+        result_queries.append(mgp.Record(query=payload_as_str, parameters=None))
 
     return result_queries

tests/e2e/streams/transformations/pulsar_transform.py

@@ -13,9 +13,7 @@ import mgp
 
 
 @mgp.transformation
-def simple(context: mgp.TransCtx,
-           messages: mgp.Messages
-           ) -> mgp.Record(query=str, parameters=mgp.Map):
+def simple(context: mgp.TransCtx, messages: mgp.Messages) -> mgp.Record(query=str, parameters=mgp.Map):
 
     result_queries = []
 
@@ -30,15 +28,15 @@ def simple(context: mgp.TransCtx,
                     payload: '{payload_as_str}',
                     topic: '{message.topic_name()}'
                 }})""",
-                parameters=None))
+                parameters=None,
+            )
+        )
 
     return result_queries
 
 
 @mgp.transformation
-def with_parameters(context: mgp.TransCtx,
-                    messages: mgp.Messages
-                    ) -> mgp.Record(query=str, parameters=mgp.Map):
+def with_parameters(context: mgp.TransCtx, messages: mgp.Messages) -> mgp.Record(query=str, parameters=mgp.Map):
 
     result_queries = []
 
@@ -53,23 +51,21 @@ def with_parameters(context: mgp.TransCtx,
                     payload: $payload,
                     topic: $topic
                 })""",
-                parameters={
-                    "payload": payload_as_str,
-                    "topic": message.topic_name()}))
+                parameters={"payload": payload_as_str, "topic": message.topic_name()},
+            )
+        )
 
     return result_queries
 
 
 @mgp.transformation
-def query(messages: mgp.Messages
-          ) -> mgp.Record(query=str, parameters=mgp.Nullable[mgp.Map]):
+def query(messages: mgp.Messages) -> mgp.Record(query=str, parameters=mgp.Nullable[mgp.Map]):
     result_queries = []
 
     for i in range(0, messages.total_messages()):
         message = messages.message_at(i)
         assert message.source_type() == mgp.SOURCE_TYPE_PULSAR
         payload_as_str = message.payload().decode("utf-8")
-        result_queries.append(mgp.Record(
-            query=payload_as_str, parameters=None))
+        result_queries.append(mgp.Record(query=payload_as_str, parameters=None))
 
     return result_queries

tests/e2e/util_e2e/CMakeLists.txt

@@ -0,0 +1,5 @@
+function(copy_util_e2e_python_files FILE_NAME)
+    copy_e2e_python_files(util_e2e ${FILE_NAME})
+endfunction()
+
+copy_util_e2e_python_files(utility_simple.py)

tests/e2e/util_e2e/utility_simple.py

@@ -0,0 +1,40 @@
+# Copyright 2022 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+import typing
+import mgclient
+import sys
+import pytest
+
+
+def test_does_throw_if_null_is_checked_against_in_generic_case():
+
+    connection = mgclient.connect(host="localhost", port=7687)
+    connection.autocommit = True
+
+    cursor = connection.cursor()
+
+    query = """WITH 2 AS name
+                RETURN CASE name
+                WHEN 3 THEN 'works'
+                WHEN null THEN "doesn't work"
+                ELSE 'something went wrong'
+                END"""
+
+    with pytest.raises(mgclient.DatabaseError) as err:
+        cursor.execute(query)
+
+    error_msg = err.value.args[0]
+    assert error_msg == "Use the generic form when checking against NULL."
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-rA"]))

tests/e2e/util_e2e/workloads.yaml

@@ -0,0 +1,13 @@
+template_cluster: &template_cluster
+  cluster:
+    main:
+      args: ["--bolt-port", "7687", "--log-level=TRACE"]
+      log_file: "util-e2e.log"
+      setup_queries: []
+      validation_queries: []
+
+workloads:
+  - name: "Utility simple"
+    binary: "tests/e2e/pytest_runner.sh"
+    args: ["util_e2e/utility_simple.py"]
+    <<: *template_cluster

tests/unit/cypher_main_visitor.cpp

@@ -786,6 +786,7 @@ TEST_P(CypherMainVisitorTest, CaseSimpleForm) {
   ASSERT_TRUE(condition);
   ast_generator.CheckLiteral(condition->expression1_, 5);
   ast_generator.CheckLiteral(condition->expression2_, 10);
+  ASSERT_TRUE(condition->isnullcheckrequired_);
   ast_generator.CheckLiteral(if_operator->then_expression_, 1);
   ast_generator.CheckLiteral(if_operator->else_expression_, TypedValue());
 }

tests/unit/integrations_kafka_consumer.cpp

@@ -149,7 +149,7 @@ TEST_F(ConsumerTest, BatchInterval) {
   }
 
   consumer->Stop();
-  EXPECT_TRUE(expected_messages_received) << "Some unexpected message have been received";
+  EXPECT_TRUE(expected_messages_received) << "Some unexpected message has been received";
 
   auto check_received_timestamp = [&received_timestamps](size_t index) {
     SCOPED_TRACE("Checking index " + std::to_string(index));
@@ -178,14 +178,6 @@ TEST_F(ConsumerTest, BatchInterval) {
 TEST_F(ConsumerTest, StartStop) {
   Consumer consumer{CreateDefaultConsumerInfo(), kDummyConsumerFunction};
 
-  auto start = [&consumer](const bool use_conditional) {
-    if (use_conditional) {
-      consumer.StartIfStopped();
-    } else {
-      consumer.Start();
-    }
-  };
-
   auto stop = [&consumer](const bool use_conditional) {
     if (use_conditional) {
       consumer.StopIfRunning();
@@ -194,34 +186,28 @@ TEST_F(ConsumerTest, StartStop) {
     }
   };
 
-  auto check_config = [&start, &stop, &consumer](const bool use_conditional_start,
-                                                 const bool use_conditional_stop) mutable {
-    SCOPED_TRACE(
-        fmt::format("Conditional start {} and conditional stop {}", use_conditional_start, use_conditional_stop));
+  auto check_config = [&stop, &consumer](const bool use_conditional_stop) mutable {
+    SCOPED_TRACE(fmt::format("Start and conditionally stop {}", use_conditional_stop));
     EXPECT_FALSE(consumer.IsRunning());
     EXPECT_THROW(consumer.Stop(), ConsumerStoppedException);
     consumer.StopIfRunning();
     EXPECT_FALSE(consumer.IsRunning());
 
-    start(use_conditional_start);
+    consumer.Start();
     EXPECT_TRUE(consumer.IsRunning());
     EXPECT_THROW(consumer.Start(), ConsumerRunningException);
-    consumer.StartIfStopped();
+
     EXPECT_TRUE(consumer.IsRunning());
 
     stop(use_conditional_stop);
     EXPECT_FALSE(consumer.IsRunning());
   };
 
-  static constexpr auto kSimpleStart = false;
   static constexpr auto kSimpleStop = false;
-  static constexpr auto kConditionalStart = true;
   static constexpr auto kConditionalStop = true;
 
-  check_config(kSimpleStart, kSimpleStop);
-  check_config(kSimpleStart, kConditionalStop);
-  check_config(kConditionalStart, kSimpleStop);
-  check_config(kConditionalStart, kConditionalStop);
+  check_config(kSimpleStop);
+  check_config(kConditionalStop);
 }
 
 TEST_F(ConsumerTest, BatchSize) {
@@ -252,7 +238,7 @@ TEST_F(ConsumerTest, BatchSize) {
   }
   std::this_thread::sleep_for(kBatchInterval * 2);
   consumer->Stop();
-  EXPECT_TRUE(expected_messages_received) << "Some unexpected message have been received";
+  EXPECT_TRUE(expected_messages_received) << "Some unexpected message has been received";
 
   auto check_received_timestamp = [&received_timestamps](size_t index, size_t expected_message_count) {
     SCOPED_TRACE("Checking index " + std::to_string(index));
@@ -371,7 +357,7 @@ TEST_F(ConsumerTest, DISABLED_StartsFromPreviousOffset) {
     EXPECT_EQ(expected_total_messages, received_message_count);
     EXPECT_NO_THROW(consumer->Stop());
     ASSERT_FALSE(consumer->IsRunning());
-    EXPECT_TRUE(expected_messages_received) << "Some unexpected message have been received";
+    EXPECT_TRUE(expected_messages_received) << "Some unexpected message has been received";
   };
 
   ASSERT_NO_FATAL_FAILURE(send_and_consume_messages(2));
@@ -383,10 +369,9 @@ TEST_F(ConsumerTest, CheckMethodWorks) {
   auto info = CreateDefaultConsumerInfo();
   info.batch_size = kBatchSize;
   const std::string kMessagePrefix{"Message"};
-  auto consumer_function = [](const std::vector<Message> &messages) mutable {};
 
   // This test depends on CreateConsumer starts and stops the consumer, so the offset is stored
-  auto consumer = CreateConsumer(std::move(info), std::move(consumer_function));
+  auto consumer = CreateConsumer(std::move(info), kDummyConsumerFunction);
 
   static constexpr auto kMessageCount = 4;
   for (auto sent_messages = 0; sent_messages < kMessageCount; ++sent_messages) {
@@ -411,7 +396,7 @@ TEST_F(ConsumerTest, CheckMethodWorks) {
     });
     ASSERT_FALSE(consumer->IsRunning());
 
-    EXPECT_TRUE(expected_messages_received) << "Some unexpected message have been received";
+    EXPECT_TRUE(expected_messages_received) << "Some unexpected message has been received";
     EXPECT_EQ(received_message_count, kMessageCount);
   };
 
@@ -445,8 +430,6 @@ TEST_F(ConsumerTest, CheckWithInvalidTimeout) {
   const auto start = std::chrono::steady_clock::now();
   EXPECT_THROW(consumer.Check(std::chrono::milliseconds{0}, std::nullopt, kDummyConsumerFunction),
                ConsumerCheckFailedException);
-  EXPECT_THROW(consumer.Check(std::chrono::milliseconds{-1}, std::nullopt, kDummyConsumerFunction),
-               ConsumerCheckFailedException);
   const auto end = std::chrono::steady_clock::now();
 
   static constexpr std::chrono::seconds kMaxExpectedTimeout{2};
@@ -459,7 +442,6 @@ TEST_F(ConsumerTest, CheckWithInvalidBatchSize) {
 
   const auto start = std::chrono::steady_clock::now();
   EXPECT_THROW(consumer.Check(std::nullopt, 0, kDummyConsumerFunction), ConsumerCheckFailedException);
-  EXPECT_THROW(consumer.Check(std::nullopt, -1, kDummyConsumerFunction), ConsumerCheckFailedException);
   const auto end = std::chrono::steady_clock::now();
 
   static constexpr std::chrono::seconds kMaxExpectedTimeout{2};
@@ -496,8 +478,85 @@ TEST_F(ConsumerTest, ConsumerStatus) {
   check_info(consumer.Info());
   consumer.Start();
   check_info(consumer.Info());
-  consumer.StartIfStopped();
-  check_info(consumer.Info());
   consumer.StopIfRunning();
   check_info(consumer.Info());
 }
+
+TEST_F(ConsumerTest, LimitBatches_CannotStartIfAlreadyRunning) {
+  static constexpr auto kLimitBatches = 3;
+
+  auto info = CreateDefaultConsumerInfo();
+
+  auto consumer = CreateConsumer(std::move(info), kDummyConsumerFunction);
+
+  consumer->Start();
+  ASSERT_TRUE(consumer->IsRunning());
+
+  EXPECT_THROW(consumer->StartWithLimit(kLimitBatches, std::nullopt /*timeout*/), ConsumerRunningException);
+
+  EXPECT_TRUE(consumer->IsRunning());
+
+  consumer->Stop();
+  EXPECT_FALSE(consumer->IsRunning());
+}
+
+TEST_F(ConsumerTest, LimitBatches_SendingMoreThanLimit) {
+  /*
+  We send more messages than the BatchSize*LimitBatches:
+  -Consumer should receive 2*3=6 messages.
+  -Consumer should not be running afterwards.
+  */
+  static constexpr auto kBatchSize = 2;
+  static constexpr auto kLimitBatches = 3;
+  static constexpr auto kNumberOfMessagesToSend = 20;
+  static constexpr auto kNumberOfMessagesExpected = kBatchSize * kLimitBatches;
+  static constexpr auto kBatchInterval =
+      std::chrono::seconds{2};  // We do not want the batch interval to be the limiting factor here.
+
+  auto info = CreateDefaultConsumerInfo();
+  info.batch_size = kBatchSize;
+  info.batch_interval = kBatchInterval;
+
+  static constexpr std::string_view kMessage = "LimitBatchesTestMessage";
+
+  auto expected_messages_received = true;
+  auto number_of_messages_received = 0;
+  auto consumer_function = [&expected_messages_received,
+                            &number_of_messages_received](const std::vector<Message> &messages) mutable {
+    number_of_messages_received += messages.size();
+    for (const auto &message : messages) {
+      expected_messages_received &= (kMessage == std::string_view(message.Payload().data(), message.Payload().size()));
+    }
+  };
+
+  auto consumer = CreateConsumer(std::move(info), consumer_function);
+
+  for (auto sent_messages = 0; sent_messages <= kNumberOfMessagesToSend; ++sent_messages) {
+    cluster.SeedTopic(kTopicName, kMessage);
+  }
+
+  consumer->StartWithLimit(kLimitBatches, kDontCareTimeout);
+
+  EXPECT_FALSE(consumer->IsRunning());
+  EXPECT_EQ(number_of_messages_received, kNumberOfMessagesExpected);
+  EXPECT_TRUE(expected_messages_received) << "Some unexpected message has been received";
+}
+
+TEST_F(ConsumerTest, LimitBatches_Timeout_Reached) {
+  // We do not send any messages, we expect an exeption to be thrown.
+  static constexpr auto kLimitBatches = 3;
+
+  auto info = CreateDefaultConsumerInfo();
+
+  auto consumer = CreateConsumer(std::move(info), kDummyConsumerFunction);
+
+  std::chrono::milliseconds timeout{3000};
+
+  const auto start = std::chrono::steady_clock::now();
+  EXPECT_THROW(consumer->StartWithLimit(kLimitBatches, timeout), ConsumerStartFailedException);
+  const auto end = std::chrono::steady_clock::now();
+  const auto elapsed = (end - start);
+
+  EXPECT_LE(timeout, elapsed);
+  EXPECT_LE(elapsed, timeout * 1.2);
+}

tests/unit/storage_v2_constraints.cpp

@@ -137,10 +137,8 @@ TEST_F(ConstraintsTest, ExistenceConstraintsViolationOnCommit) {
 
     auto res = acc.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(
-        res.GetError(),
-        (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                     ConstraintViolation{ConstraintViolation::Type::EXISTENCE, label1, std::set<PropertyId>{prop1}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::EXISTENCE, label1, std::set<PropertyId>{prop1}}));
   }
 
   {
@@ -159,10 +157,8 @@ TEST_F(ConstraintsTest, ExistenceConstraintsViolationOnCommit) {
 
     auto res = acc.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(
-        res.GetError(),
-        (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                     ConstraintViolation{ConstraintViolation::Type::EXISTENCE, label1, std::set<PropertyId>{prop1}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::EXISTENCE, label1, std::set<PropertyId>{prop1}}));
   }
 
   {
@@ -462,10 +458,8 @@ TEST_F(ConstraintsTest, UniqueConstraintsViolationOnCommit1) {
     ASSERT_NO_ERROR(vertex2.SetProperty(prop1, PropertyValue(1)));
     auto res = acc.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(
-        res.GetError(),
-        (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                     ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}));
   }
 }
 
@@ -506,10 +500,8 @@ TEST_F(ConstraintsTest, UniqueConstraintsViolationOnCommit2) {
     ASSERT_NO_ERROR(acc2.Commit());
     auto res = acc3.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(
-        res.GetError(),
-        (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                     ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}));
   }
 }
 
@@ -553,16 +545,12 @@ TEST_F(ConstraintsTest, UniqueConstraintsViolationOnCommit3) {
 
     auto res = acc2.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(
-        res.GetError(),
-        (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                     ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}));
     res = acc3.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(
-        res.GetError(),
-        (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                     ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1}}));
   }
 }
 
@@ -632,9 +620,7 @@ TEST_F(ConstraintsTest, UniqueConstraintsLabelAlteration) {
 
     auto res = acc.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(res.GetError(),
-              (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                           ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1}}}));
+    EXPECT_EQ(res.GetError(), (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1}}));
   }
 
   {
@@ -668,9 +654,7 @@ TEST_F(ConstraintsTest, UniqueConstraintsLabelAlteration) {
 
     auto res = acc1.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(res.GetError(),
-              (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                           ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1}}}));
+    EXPECT_EQ(res.GetError(), (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1}}));
   }
 }
 
@@ -765,9 +749,8 @@ TEST_F(ConstraintsTest, UniqueConstraintsMultipleProperties) {
     ASSERT_NO_ERROR(vertex2->SetProperty(prop2, PropertyValue(2)));
     auto res = acc.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(res.GetError(), (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                                           ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1,
-                                                               std::set<PropertyId>{prop1, prop2}}}));
+    EXPECT_EQ(res.GetError(),
+              (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set<PropertyId>{prop1, prop2}}));
   }
 
   // Then change the second property of both vertex to null. Property values of
@@ -878,9 +861,7 @@ TEST_F(ConstraintsTest, UniqueConstraintsInsertRemoveAbortInsert) {
 
     auto res = acc.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(res.GetError(),
-              (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                           ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1, prop2}}}));
+    EXPECT_EQ(res.GetError(), (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1, prop2}}));
   }
 }
 
@@ -919,9 +900,7 @@ TEST_F(ConstraintsTest, UniqueConstraintsDeleteVertexSetProperty) {
 
     auto res = acc1.Commit();
     ASSERT_TRUE(res.HasError());
-    EXPECT_EQ(res.GetError(),
-              (CommitError{CommitError::Type::CONSTRAINT_VIOLATION,
-                           ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1}}}));
+    EXPECT_EQ(res.GetError(), (ConstraintViolation{ConstraintViolation::Type::UNIQUE, label1, std::set{prop1}}));
 
     ASSERT_NO_ERROR(acc2.Commit());
   }

tests/unit/storage_v2_replication.cpp

@@ -32,6 +32,13 @@ class ReplicationTest : public ::testing::Test {
 
   void TearDown() override { Clear(); }
 
+  memgraph::storage::Config configuration{
+      .items = {.properties_on_edges = true},
+      .durability = {
+          .storage_directory = storage_directory,
+          .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
+      }};
+
  private:
   void Clear() {
     if (!std::filesystem::exists(storage_directory)) return;
@@ -40,19 +47,9 @@ class ReplicationTest : public ::testing::Test {
 };
 
 TEST_F(ReplicationTest, BasicSynchronousReplicationTest) {
-  memgraph::storage::Storage main_store(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage main_store(configuration);
 
-  memgraph::storage::Storage replica_store(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage replica_store(configuration);
   replica_store.SetReplicaRole(memgraph::io::network::Endpoint{"127.0.0.1", 10000});
 
   ASSERT_FALSE(main_store
@@ -483,19 +480,9 @@ TEST_F(ReplicationTest, RecoveryProcess) {
 }
 
 TEST_F(ReplicationTest, BasicAsynchronousReplicationTest) {
-  memgraph::storage::Storage main_store(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage main_store(configuration);
 
-  memgraph::storage::Storage replica_store_async(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage replica_store_async(configuration);
 
   replica_store_async.SetReplicaRole(memgraph::io::network::Endpoint{"127.0.0.1", 20000});
 
@@ -533,28 +520,13 @@ TEST_F(ReplicationTest, BasicAsynchronousReplicationTest) {
 }
 
 TEST_F(ReplicationTest, EpochTest) {
-  memgraph::storage::Storage main_store(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage main_store(configuration);
 
-  memgraph::storage::Storage replica_store1(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage replica_store1(configuration);
 
   replica_store1.SetReplicaRole(memgraph::io::network::Endpoint{"127.0.0.1", 10000});
 
-  memgraph::storage::Storage replica_store2(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage replica_store2(configuration);
 
   replica_store2.SetReplicaRole(memgraph::io::network::Endpoint{"127.0.0.1", 10001});
 
@@ -639,30 +611,15 @@ TEST_F(ReplicationTest, EpochTest) {
 }
 
 TEST_F(ReplicationTest, ReplicationInformation) {
-  memgraph::storage::Storage main_store(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage main_store(configuration);
 
-  memgraph::storage::Storage replica_store1(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  memgraph::storage::Storage replica_store1(configuration);
 
-  const memgraph::io::network::Endpoint replica1_endpoint{"127.0.0.1", 10000};
+  const memgraph::io::network::Endpoint replica1_endpoint{"127.0.0.1", 10001};
   replica_store1.SetReplicaRole(replica1_endpoint);
 
-  const memgraph::io::network::Endpoint replica2_endpoint{"127.0.0.1", 10000};
-  memgraph::storage::Storage replica_store2(
-      {.items = {.properties_on_edges = true},
-       .durability = {
-           .storage_directory = storage_directory,
-           .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::PERIODIC_SNAPSHOT_WITH_WAL,
-       }});
+  const memgraph::io::network::Endpoint replica2_endpoint{"127.0.0.1", 10002};
+  memgraph::storage::Storage replica_store2(configuration);
 
   replica_store2.SetReplicaRole(replica2_endpoint);
 
@@ -700,3 +657,55 @@ TEST_F(ReplicationTest, ReplicationInformation) {
   ASSERT_EQ(second_info.endpoint, replica2_endpoint);
   ASSERT_EQ(second_info.state, memgraph::storage::replication::ReplicaState::READY);
 }
+
+TEST_F(ReplicationTest, ReplicationReplicaWithExistingName) {
+  memgraph::storage::Storage main_store(configuration);
+
+  memgraph::storage::Storage replica_store1(configuration);
+
+  const memgraph::io::network::Endpoint replica1_endpoint{"127.0.0.1", 10001};
+  replica_store1.SetReplicaRole(replica1_endpoint);
+
+  const memgraph::io::network::Endpoint replica2_endpoint{"127.0.0.1", 10002};
+  memgraph::storage::Storage replica_store2(configuration);
+
+  replica_store2.SetReplicaRole(replica2_endpoint);
+
+  const std::string replica1_name{"REPLICA1"};
+  ASSERT_FALSE(main_store
+                   .RegisterReplica(replica1_name, replica1_endpoint,
+                                    memgraph::storage::replication::ReplicationMode::SYNC, {.timeout = 2.0})
+                   .HasError());
+
+  const std::string replica2_name{"REPLICA1"};
+  ASSERT_TRUE(
+      main_store
+          .RegisterReplica(replica2_name, replica2_endpoint, memgraph::storage::replication::ReplicationMode::ASYNC)
+          .GetError() == memgraph::storage::Storage::RegisterReplicaError::NAME_EXISTS);
+}
+
+TEST_F(ReplicationTest, ReplicationReplicaWithExistingEndPoint) {
+  memgraph::storage::Storage main_store(configuration);
+
+  memgraph::storage::Storage replica_store1(configuration);
+
+  const memgraph::io::network::Endpoint replica1_endpoint{"127.0.0.1", 10001};
+  replica_store1.SetReplicaRole(replica1_endpoint);
+
+  const memgraph::io::network::Endpoint replica2_endpoint{"127.0.0.1", 10001};
+  memgraph::storage::Storage replica_store2(configuration);
+
+  replica_store2.SetReplicaRole(replica2_endpoint);
+
+  const std::string replica1_name{"REPLICA1"};
+  ASSERT_FALSE(main_store
+                   .RegisterReplica(replica1_name, replica1_endpoint,
+                                    memgraph::storage::replication::ReplicationMode::SYNC, {.timeout = 2.0})
+                   .HasError());
+
+  const std::string replica2_name{"REPLICA2"};
+  ASSERT_TRUE(
+      main_store
+          .RegisterReplica(replica2_name, replica2_endpoint, memgraph::storage::replication::ReplicationMode::ASYNC)
+          .GetError() == memgraph::storage::Storage::RegisterReplicaError::END_POINT_EXISTS);
+}