couchbase 3.8.2 → 3.9.1

data/ext/couchbase/core/query_stream.cxx

@@ -0,0 +1,447 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *   Copyright 2026. Couchbase, Inc.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ */
+
+#include "query_stream.hxx"
+
+#include <couchbase/error_codes.hxx>
+
+#include "core/io/http_error.hxx"
+#include "core/operations/query_response_parsing.hxx"
+#include "free_form_http_request.hxx"
+#include "logger/logger.hxx"
+#include "row_streamer.hxx"
+#include "utils/json.hxx"
+#include "utils/movable_function.hxx"
+
+#include <asio/io_context.hpp>
+#include <asio/post.hpp>
+#include <tao/json/value.hpp>
+
+#include <atomic>
+#include <cstddef>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <system_error>
+#include <utility>
+#include <vector>
+
+namespace couchbase::core
+{
+namespace
+{
+constexpr auto N1QL_RESULTS_POINTER = "/results/^";
+
+// The bounded streaming lexer surfaces malformed- or oversized-JSON failures as
+// errc::streaming_json_lexer::* codes. The buffered query() path normalizes any body-parse failure
+// to errc::common::parsing_failure (document_query.cxx), and so does row_streamer's own
+// truncated-body terminal; normalize the lexer codes here too so the public streaming error
+// contract matches the buffered one (a consumer matching parsing_failure never silently misses a
+// malformed response).
+auto
+normalize_stream_error(std::error_code ec) -> std::error_code
+{
+  if (ec.category() == couchbase::core::impl::streaming_json_lexer_category()) {
+    return errc::common::parsing_failure;
+  }
+  return ec;
+}
+} // namespace
+
+// Abstract backing for a query_stream. Two implementations: one that lazily pulls rows off a live
+// socket (streaming_query_stream_impl), and one that replays an already-buffered response without
+// re-parsing it (buffered_query_stream_impl), used for the prepared-statement path.
+class query_stream_impl
+{
+public:
+  query_stream_impl() = default;
+  query_stream_impl(const query_stream_impl&) = delete;
+  query_stream_impl(query_stream_impl&&) = delete;
+  auto operator=(const query_stream_impl&) -> query_stream_impl& = delete;
+  auto operator=(query_stream_impl&&) -> query_stream_impl& = delete;
+  virtual ~query_stream_impl() = default;
+
+  virtual void start(utils::movable_function<void(std::error_code)>&& on_ready) = 0;
+  virtual void next_row(
+    utils::movable_function<void(std::optional<std::string>, std::error_code)>&& handler) = 0;
+  [[nodiscard]] virtual auto signature() const -> std::optional<std::string> = 0;
+  [[nodiscard]] virtual auto meta_data() const
+    -> std::optional<operations::query_response::query_meta_data> = 0;
+  [[nodiscard]] virtual auto error_details() const -> stream_error_details = 0;
+  virtual auto set_deadline(std::chrono::time_point<std::chrono::steady_clock> deadline_tp)
+    -> io::deadline_state = 0;
+  virtual void cancel() = 0;
+};
+
+namespace
+{
+// Lift the service-reported diagnostics out of a parsed preamble/trailer. Called for both sections,
+// so whichever one carried the failure is what the error context ends up describing.
+auto
+extract_error_details(const operations::query_response::query_meta_data& meta, std::string raw)
+  -> stream_error_details
+{
+  stream_error_details details{};
+  details.client_context_id = meta.client_context_id;
+  // Retained whether or not the section reports an error: it becomes error_context::http_body,
+  // which the buffered path fills for every response, and even an error-free envelope names the
+  // request the service logged.
+  details.raw = std::move(raw);
+  if (meta.errors && !meta.errors->empty()) {
+    details.first_error_code = meta.errors->front().code;
+    details.first_error_message = meta.errors->front().message;
+  }
+  return details;
+}
+} // namespace
+
+namespace
+{
+class streaming_query_stream_impl
+  : public query_stream_impl
+  , public std::enable_shared_from_this<streaming_query_stream_impl>
+{
+public:
+  streaming_query_stream_impl(asio::io_context& io,
+                              http_response_body body,
+                              row_streamer_options options,
+                              std::uint32_t http_status)
+    : streamer_{ io, std::move(body), N1QL_RESULTS_POINTER, options }
+    , http_status_{ http_status }
+  {
+  }
+  streaming_query_stream_impl(const streaming_query_stream_impl&) = delete;
+  streaming_query_stream_impl(streaming_query_stream_impl&&) = delete;
+  auto operator=(const streaming_query_stream_impl&) -> streaming_query_stream_impl& = delete;
+  auto operator=(streaming_query_stream_impl&&) -> streaming_query_stream_impl& = delete;
+
+  // Runs with no handle left, no pull in flight and the preamble delivered: each pull, and start()
+  // until the preamble, holds this object.
+  //
+  // Cancels unless a pull delivered the terminal of a completed JSON document. Without the cancel,
+  // the read-ahead keeps the body and its session alive:
+  //  - through a read in flight, until the idle timer fires or the cluster closes;
+  //  - through a row send queued on a full channel, indefinitely.
+  // The cancel stops the connection only while the body is still being read. A body read to its
+  // end has already checked the connection in.
+  //
+  // After a completed document the read-ahead reads the trailing bytes and the connection is
+  // checked in, or the idle timer ends the body.
+  ~streaming_query_stream_impl() override
+  {
+    if (!document_completed_) {
+      streamer_.cancel();
+    }
+  }
+
+  void start(utils::movable_function<void(std::error_code)>&& on_ready) override
+  {
+    streamer_.start([self = shared_from_this(), on_ready = std::move(on_ready)](
+                      std::string preamble, std::error_code ec) mutable {
+      if (ec) {
+        return on_ready(normalize_stream_error(
+          io::invalid_argument_for_undecodable_400(ec, self->http_status_, {})));
+      }
+      operations::query_response::query_meta_data meta{};
+      try {
+        meta = operations::parse_query_meta(utils::json::parse(preamble));
+      } catch (const std::exception&) {
+        {
+          // Keep the unparseable preamble: it is the only diagnostic there is for a malformed
+          // response, and the buffered path reports the raw body in the same situation.
+          const std::scoped_lock lock{ self->mutex_ };
+          self->error_details_.raw = std::move(preamble);
+        }
+        return on_ready(io::invalid_argument_for_undecodable_400(
+          errc::common::parsing_failure, self->http_status_, {}));
+      }
+      {
+        const std::scoped_lock lock{ self->mutex_ };
+        self->signature_ = meta.signature;
+        self->error_details_ = extract_error_details(meta, std::move(preamble));
+      }
+      // An upfront error is rare for the streaming path (status/errors usually arrive in the
+      // trailer). Mirror the analytics preamble check and gate on error entries rather than status:
+      // a benign non-success preamble status (e.g. "running") would otherwise be handed to
+      // map_query_error, whose no-errors fallback reports an error — a spurious
+      // error. Checking errors also surfaces a genuine preamble error now rather than one next()
+      // call later.
+      std::error_code early_error{};
+      if (meta.errors && !meta.errors->empty()) {
+        early_error = operations::map_query_error(meta, self->http_status_);
+      }
+      on_ready(early_error);
+    });
+  }
+
+  void next_row(
+    utils::movable_function<void(std::optional<std::string>, std::error_code)>&& handler) override
+  {
+    // Terminal-sticky: once the stream has ended, re-deliver the stored terminal on every
+    // subsequent pull instead of issuing another receive that would park forever on the
+    // drained-but-open row channel. The public handle guards this too, but a core::-direct consumer
+    // (e.g. a wrapper) relies on this being safe.
+    std::optional<std::error_code> sticky_terminal;
+    {
+      const std::scoped_lock lock{ mutex_ };
+      if (terminal_reached_) {
+        sticky_terminal = terminal_ec_;
+      }
+    }
+    if (sticky_terminal) {
+      // Deliver outside the lock: the handler typically issues the next pull, which re-enters
+      // next_row and would re-lock the non-recursive mutex_ and deadlock if invoked while held.
+      return handler(std::nullopt, *sticky_terminal);
+    }
+    streamer_.next_row([self = shared_from_this(),
+                        handler = std::move(handler)](std::string row, std::error_code ec) mutable {
+      if (!row.empty()) {
+        return handler(std::move(row), {});
+      }
+      // Empty row marks the end of the stream. Classify the terminal, remember it (so later pulls
+      // are served the sticky terminal above), then deliver.
+      std::error_code terminal_ec{};
+      if (ec) {
+        // Normalize a mid-document lexer failure to parsing_failure, matching the buffered path.
+        terminal_ec = normalize_stream_error(
+          io::invalid_argument_for_undecodable_400(ec, self->http_status_, {}));
+      } else if (auto raw_meta = self->streamer_.metadata(); !raw_meta.has_value()) {
+        // A clean end without reconstructed metadata violates the row_streamer invariant; surface
+        // it rather than silently reporting success with no metadata.
+        CB_LOG_WARNING("query_stream: end reached but row_streamer returned no metadata");
+        terminal_ec = errc::common::internal_server_failure;
+      } else {
+        try {
+          auto meta = operations::parse_query_meta(utils::json::parse(raw_meta.value()));
+          terminal_ec = operations::map_query_error(meta, self->http_status_);
+          const std::scoped_lock lock{ self->mutex_ };
+          // The trailer is where a request that failed part-way reports itself, so its diagnostics
+          // supersede the preamble's for the terminal error.
+          self->error_details_ = extract_error_details(meta, std::move(raw_meta.value()));
+          self->meta_data_ = std::move(meta);
+        } catch (const std::exception&) {
+          terminal_ec = io::invalid_argument_for_undecodable_400(
+            errc::common::parsing_failure, self->http_status_, {});
+          const std::scoped_lock lock{ self->mutex_ };
+          self->error_details_.raw = std::move(raw_meta.value());
+        }
+      }
+      {
+        const std::scoped_lock lock{ self->mutex_ };
+        self->terminal_reached_ = true;
+        self->terminal_ec_ = terminal_ec;
+        self->document_completed_ = !ec;
+      }
+      handler(std::nullopt, terminal_ec);
+    });
+  }
+
+  [[nodiscard]] auto signature() const -> std::optional<std::string> override
+  {
+    const std::scoped_lock lock{ mutex_ };
+    return signature_;
+  }
+
+  [[nodiscard]] auto meta_data() const
+    -> std::optional<operations::query_response::query_meta_data> override
+  {
+    const std::scoped_lock lock{ mutex_ };
+    return meta_data_;
+  }
+
+  [[nodiscard]] auto error_details() const -> stream_error_details override
+  {
+    const std::scoped_lock lock{ mutex_ };
+    return error_details_;
+  }
+
+  auto set_deadline(std::chrono::time_point<std::chrono::steady_clock> deadline_tp)
+    -> io::deadline_state override
+  {
+    return streamer_.set_deadline(deadline_tp);
+  }
+
+  void cancel() override
+  {
+    streamer_.cancel();
+  }
+
+private:
+  row_streamer streamer_;
+  std::uint32_t http_status_{};
+  mutable std::mutex mutex_{};
+  std::optional<std::string> signature_{};
+  std::optional<operations::query_response::query_meta_data> meta_data_{};
+  stream_error_details error_details_{};
+  // Terminal-sticky state (guarded by mutex_): once next_row reports the terminal, it is remembered
+  // so every later pull re-delivers it instead of parking on the drained channel.
+  bool terminal_reached_{ false };
+  std::error_code terminal_ec_{};
+  // Set when a pull delivers a terminal with no streamer error: the lexer completed the JSON
+  // document. Written under mutex_. The destructor reads it unlocked: the release of the last
+  // shared_ptr orders that read after every write.
+  bool document_completed_{ false };
+};
+
+// Replays an already-buffered query_response as a stream without any JSON re-parsing: the rows are
+// already row-JSON strings and the metadata is already parsed. Used for the prepared-statement
+// (adhoc == false) path, which must run buffered but presents the uniform streaming handle.
+class buffered_query_stream_impl
+  : public query_stream_impl
+  , public std::enable_shared_from_this<buffered_query_stream_impl>
+{
+public:
+  buffered_query_stream_impl(asio::io_context& io,
+                             std::vector<std::string> rows,
+                             operations::query_response::query_meta_data meta)
+    : io_{ io }
+    , rows_{ std::move(rows) }
+    , meta_{ std::move(meta) }
+  {
+  }
+
+  void start(utils::movable_function<void(std::error_code)>&& on_ready) override
+  {
+    // The whole response is already in hand; there is no preamble to parse. Any query error is
+    // delivered at the terminal (mirroring the streaming path's trailing-error semantics).
+    asio::post(io_, [self = shared_from_this(), on_ready = std::move(on_ready)]() mutable {
+      on_ready({});
+    });
+  }
+
+  void next_row(
+    utils::movable_function<void(std::optional<std::string>, std::error_code)>&& handler) override
+  {
+    // Post the delivery so a consumer that pulls the next row from within its handler does not
+    // recurse synchronously through a large buffered result and overflow the stack.
+    asio::post(io_, [self = shared_from_this(), handler = std::move(handler)]() mutable {
+      if (self->cancelled_) {
+        return handler(std::nullopt, errc::common::request_canceled);
+      }
+      if (self->index_ < self->rows_.size()) {
+        // Each buffered row is consumed at most once, so hand it to the handler by move rather than
+        // copying what can be a large prepared-statement result out of the vector.
+        return handler(std::move(self->rows_[self->index_++]), {});
+      }
+      return handler(std::nullopt, operations::map_query_error(self->meta_));
+    });
+  }
+
+  [[nodiscard]] auto signature() const -> std::optional<std::string> override
+  {
+    return meta_.signature;
+  }
+
+  [[nodiscard]] auto meta_data() const
+    -> std::optional<operations::query_response::query_meta_data> override
+  {
+    return meta_;
+  }
+
+  [[nodiscard]] auto error_details() const -> stream_error_details override
+  {
+    // The response was parsed by the buffered path, so there is no raw JSON left to report; the
+    // service's error itself still is.
+    return extract_error_details(meta_, {});
+  }
+
+  auto set_deadline(std::chrono::time_point<std::chrono::steady_clock> /* deadline_tp */)
+    -> io::deadline_state override
+  {
+    // The prepared-statement path replays a buffered response. There is no socket to reclaim and
+    // next_row() always completes, so a deadline has nothing to bound.
+    return io::deadline_state::body_already_ended;
+  }
+
+  void cancel() override
+  {
+    cancelled_ = true;
+  }
+
+private:
+  asio::io_context& io_;
+  std::vector<std::string> rows_;
+  operations::query_response::query_meta_data meta_;
+  std::size_t index_{ 0 };
+  // cancel() may be called from the consumer thread while the posted delivery reads it on the io
+  // thread, so it must be atomic. index_/rows_ are only touched on the io thread (posted handlers).
+  std::atomic_bool cancelled_{ false };
+};
+} // namespace
+
+query_stream::query_stream(asio::io_context& io,
+                           http_response_body body,
+                           row_streamer_options options,
+                           std::uint32_t http_status)
+  : impl_{
+    std::make_shared<streaming_query_stream_impl>(io, std::move(body), options, http_status)
+  }
+{
+}
+
+query_stream::query_stream(asio::io_context& io,
+                           std::vector<std::string> rows,
+                           operations::query_response::query_meta_data meta)
+  : impl_{ std::make_shared<buffered_query_stream_impl>(io, std::move(rows), std::move(meta)) }
+{
+}
+
+void
+query_stream::start(utils::movable_function<void(std::error_code)>&& on_ready)
+{
+  impl_->start(std::move(on_ready));
+}
+
+auto
+query_stream::error_details() const -> stream_error_details
+{
+  return impl_->error_details();
+}
+
+void
+query_stream::next_row(
+  utils::movable_function<void(std::optional<std::string>, std::error_code)>&& handler)
+{
+  impl_->next_row(std::move(handler));
+}
+
+auto
+query_stream::signature() const -> std::optional<std::string>
+{
+  return impl_->signature();
+}
+
+auto
+query_stream::meta_data() const -> std::optional<operations::query_response::query_meta_data>
+{
+  return impl_->meta_data();
+}
+
+auto
+query_stream::set_deadline(std::chrono::time_point<std::chrono::steady_clock> deadline_tp)
+  -> io::deadline_state
+{
+  return impl_->set_deadline(deadline_tp);
+}
+
+void
+query_stream::cancel()
+{
+  impl_->cancel();
+}
+} // namespace couchbase::core

data/ext/couchbase/core/query_stream.hxx

@@ -0,0 +1,141 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *   Copyright 2026. Couchbase, Inc.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ */
+
+#pragma once
+
+#include "core/io/stream_deadline.hxx"
+#include "core/operations/document_query.hxx"
+#include "row_streamer.hxx"
+#include "stream_error_details.hxx"
+#include "utils/movable_function.hxx"
+
+#include <chrono>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <system_error>
+#include <vector>
+
+namespace asio
+{
+class io_context;
+} // namespace asio
+
+namespace couchbase::core
+{
+class query_stream_impl;
+class http_response_body;
+
+/**
+ * L3 streaming state machine for N1QL query responses.
+ *
+ * Wraps a row_streamer (pointed at "/results/^") and adds query-specific semantics:
+ * preamble parsing (signature + upfront error detection), late trailer metadata, and
+ * terminal error classification via map_query_error.
+ */
+class query_stream
+{
+public:
+  /**
+   * Constructs an empty handle (no underlying stream). Used to carry the "no stream" outcome on
+   * the error path of the dispatch; none of the streaming methods may be called on it.
+   */
+  query_stream() = default;
+
+  /**
+   * http_status is the status of the response the body belongs to; an HTTP 400 that matches no
+   * specific error is reported as invalid_argument.
+   */
+  query_stream(asio::io_context& io,
+               http_response_body body,
+               row_streamer_options options = {},
+               std::uint32_t http_status = 0);
+
+  /**
+   * Constructs a stream that replays an already-buffered response (rows are row-JSON strings,
+   * metadata already parsed) without any JSON re-parsing. Used for the prepared-statement path,
+   * which must run buffered but presents the uniform streaming handle.
+   */
+  query_stream(asio::io_context& io,
+               std::vector<std::string> rows,
+               operations::query_response::query_meta_data meta);
+
+  /**
+   * Starts the stream. Resolves once the preamble has been parsed. The early_error is set when
+   * the response carried an upfront error (or the preamble failed to parse).
+   */
+  void start(utils::movable_function<void(std::error_code early_error)>&& on_ready);
+
+  /**
+   * Service-reported diagnostics from the most recently parsed JSON section — the preamble after
+   * start(), the trailer once the terminal has been reached. The dispatch layer stamps these into
+   * the error context so a streaming failure reports the same detail as a buffered one.
+   */
+  [[nodiscard]] auto error_details() const -> stream_error_details;
+
+  /**
+   * Retrieves the next row. A populated row is delivered with a falsy error_code. An empty row
+   * (std::nullopt) signals the end: a falsy error_code means a clean success end, a truthy
+   * error_code means a trailing query/transport error.
+   */
+  void next_row(
+    utils::movable_function<void(std::optional<std::string> row, std::error_code)>&& handler);
+
+  /**
+   * Signature captured from the preamble (available after start resolves).
+   */
+  [[nodiscard]] auto signature() const -> std::optional<std::string>;
+
+  /**
+   * Trailer metadata. Valid only after the stream has reached its end.
+   */
+  [[nodiscard]] auto meta_data() const
+    -> std::optional<operations::query_response::query_meta_data>;
+
+  /**
+   * Closes the HTTP body at `deadline_tp`, whether or not a row is being pulled, releasing the
+   * connection where one is still held. A response received in full holds none: once its bytes
+   * have also been parsed the call arms nothing and reports body_already_ended, and while bytes
+   * remain unparsed it ends the body, dropping those and reporting the timeout. Rows already
+   * parsed stay consumable either way. The inter-read idle timer is armed only while a socket
+   * read is in flight, and a consumer stopped above the high-water mark has none.
+   *
+   * The terminal reaches the consumer through a pull. A pull outstanding at expiry completes
+   * with the timeout. A read that completes the response is the exception: the whole response
+   * arrived, so those bytes are delivered as a clean end. Otherwise the next pull reports the
+   * terminal, after any rows already buffered.
+   *
+   * Optional. The terminal is unambiguous_timeout for a read-only request, ambiguous_timeout
+   * otherwise. Re-arming before expiry replaces the deadline. Once the stream has ended a
+   * further call arms nothing and reports body_already_ended. Callable from any thread.
+   *
+   * A stream replaying an already-buffered response holds no socket, so a deadline on it is a
+   * no-op.
+   */
+  [[nodiscard]] auto set_deadline(std::chrono::time_point<std::chrono::steady_clock> deadline_tp)
+    -> io::deadline_state;
+
+  /**
+   * Cancels the stream & closes the HTTP connection.
+   */
+  void cancel();
+
+private:
+  std::shared_ptr<query_stream_impl> impl_;
+};
+} // namespace couchbase::core