Merge branch 'upstream_emulator' into sample_row_keys

Author: marete

google/cloud/bigtable/emulator/column_family.cc

@@ -14,10 +14,15 @@
 
 #include "google/cloud/bigtable/emulator/column_family.h"
 #include "google/cloud/internal/big_endian.h"
+#include <google/bigtable/v2/types.pb.h>
 #include <absl/types/optional.h>
+#include <array>
 #include <chrono>
 #include <cstdint>
+#include <iterator>
 #include <map>
+#include <string>
+#include <utility>
 
 namespace google {
 namespace cloud {
@@ -118,9 +123,32 @@ ReadModifyWriteCellResult ColumnRow::ReadModifyWrite(
 absl::optional<std::string> ColumnRow::SetCell(
     std::chrono::milliseconds timestamp, std::string const& value) {
   absl::optional<std::string> ret = absl::nullopt;
+
+  auto cell_it = cells_.find(timestamp);
+  if (!(cell_it == cells_.end())) {
+    ret = std::move(cell_it->second);
+  }
+
+  cells_[timestamp] = value;
+
+  return ret;
+}
+
+StatusOr<absl::optional<std::string>> ColumnRow::UpdateCell(
+    std::chrono::milliseconds timestamp, std::string& value,
+    std::function<StatusOr<std::string>(std::string const&,
+                                        std::string&&)> const& update_fn) {
+  absl::optional<std::string> ret = absl::nullopt;
+
   auto cell_it = cells_.find(timestamp);
   if (!(cell_it == cells_.end())) {
+    auto maybe_update_value = update_fn(cell_it->second, std::move(value));
+    if (!maybe_update_value) {
+      return maybe_update_value.status();
+    }
     ret = std::move(cell_it->second);
+    maybe_update_value.value().swap(cell_it->second);
+    return ret;
   }
 
   cells_[timestamp] = value;
@@ -166,6 +194,15 @@ absl::optional<std::string> ColumnFamilyRow::SetCell(
   return columns_[column_qualifier].SetCell(timestamp, value);
 }
 
+StatusOr<absl::optional<std::string>> ColumnFamilyRow::UpdateCell(
+    std::string const& column_qualifier, std::chrono::milliseconds timestamp,
+    std::string& value,
+    std::function<StatusOr<std::string>(std::string const&,
+                                        std::string&&)> const& update_fn) {
+  return columns_[column_qualifier].UpdateCell(timestamp, value,
+                                               std::move(update_fn));
+}
+
 std::vector<Cell> ColumnFamilyRow::DeleteColumn(
     std::string const& column_qualifier,
     ::google::bigtable::v2::TimestampRange const& time_range) {
@@ -201,6 +238,13 @@ absl::optional<std::string> ColumnFamily::SetCell(
   return rows_[row_key].SetCell(column_qualifier, timestamp, value);
 }
 
+StatusOr<absl::optional<std::string>> ColumnFamily::UpdateCell(
+    std::string const& row_key, std::string const& column_qualifier,
+    std::chrono::milliseconds timestamp, std::string& value) {
+  return rows_[row_key].UpdateCell(column_qualifier, timestamp, value,
+                                   update_cell_);
+}
+
 std::map<std::string, std::vector<Cell>> ColumnFamily::DeleteRow(
     std::string const& row_key) {
   std::map<std::string, std::vector<Cell>> res;
@@ -215,11 +259,11 @@ std::map<std::string, std::vector<Cell>> ColumnFamily::DeleteRow(
     ::google::bigtable::v2::TimestampRange time_range;
     auto deleted_cells = column.second.DeleteTimeRange(time_range);
     if (!deleted_cells.empty()) {
-      res[std::move(column.first)] = std::move(deleted_cells);
+      res[column.first] = std::move(deleted_cells);
     }
   }
 
-  rows_.erase(row_key);
+  rows_.erase(row_it);
 
   return res;
 }
@@ -385,6 +429,41 @@ bool FilteredColumnFamilyStream::PointToFirstCellAfterRowChange() const {
   return false;
 }
 
+StatusOr<std::shared_ptr<ColumnFamily>>
+ColumnFamily::ConstructAggregateColumnFamily(
+    google::bigtable::admin::v2::Type value_type) {
+  auto cf = std::make_shared<ColumnFamily>();
+
+  if (value_type.has_aggregate_type()) {
+    auto const& aggregate_type = value_type.aggregate_type();
+    switch (aggregate_type.aggregator_case()) {
+      case google::bigtable::admin::v2::Type::Aggregate::kSum:
+        cf->update_cell_ = cf->SumUpdateCellBEInt64;
+        break;
+      case google::bigtable::admin::v2::Type::Aggregate::kMin:
+        cf->update_cell_ = cf->MinUpdateCellBEInt64;
+        break;
+      case google::bigtable::admin::v2::Type::Aggregate::kMax:
+        cf->update_cell_ = cf->MaxUpdateCellBEInt64;
+        break;
+      default:
+        return InvalidArgumentError(
+            "unsupported aggregation type",
+            GCP_ERROR_INFO().WithMetadata(
+                "aggregation case",
+                absl::StrFormat("%d", aggregate_type.aggregator_case())));
+    }
+
+    cf->value_type_ = std::move(value_type);
+
+    return cf;
+  }
+
+  return InvalidArgumentError(
+      "no aggregate type set in the supplied value_type",
+      GCP_ERROR_INFO().WithMetadata("supplied value type",
+                                    value_type.DebugString()));
+}
 }  // namespace emulator
 }  // namespace bigtable
 }  // namespace cloud

google/cloud/bigtable/emulator/column_family.h

@@ -15,18 +15,23 @@
 #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_EMULATOR_COLUMN_FAMILY_H
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_EMULATOR_COLUMN_FAMILY_H
 
+#include "google/cloud/bigtable/cell.h"
 #include "google/cloud/bigtable/emulator/cell_view.h"
 #include "google/cloud/bigtable/emulator/filter.h"
 #include "google/cloud/bigtable/emulator/filtered_map.h"
 #include "google/cloud/bigtable/emulator/range_set.h"
 #include "google/cloud/bigtable/read_modify_write_rule.h"
 #include "google/cloud/internal/big_endian.h"
 #include "google/cloud/internal/make_status.h"
+#include "google/cloud/status_or.h"
 #include "absl/types/optional.h"
 #include <google/bigtable/admin/v2/table.pb.h>
 #include <google/bigtable/v2/data.pb.h>
+#include <google/bigtable/v2/types.pb.h>
+#include <absl/strings/str_format.h>
 #include <chrono>
 #include <cstddef>
+#include <cstdint>
 #include <map>
 #include <memory>
 #include <optional>
@@ -89,6 +94,12 @@ class ColumnRow {
    */
   absl::optional<std::string> SetCell(std::chrono::milliseconds timestamp,
                                       std::string const& value);
+
+  StatusOr<absl::optional<std::string>> UpdateCell(
+      std::chrono::milliseconds timestamp, std::string& value,
+      std::function<StatusOr<std::string>(std::string const&,
+                                          std::string&&)> const& update_fn);
+
   /**
    * Delete cells falling into a given timestamp range.
    *
@@ -174,6 +185,13 @@ class ColumnFamilyRow {
   absl::optional<std::string> SetCell(std::string const& column_qualifier,
                                       std::chrono::milliseconds timestamp,
                                       std::string const& value);
+
+  StatusOr<absl::optional<std::string>> UpdateCell(
+      std::string const& column_qualifier, std::chrono::milliseconds timestamp,
+      std::string& value,
+      std::function<StatusOr<std::string>(std::string const&,
+                                          std::string&&)> const& update_fn);
+
   /**
    * Delete cells falling into a given timestamp range in one column.
    *
@@ -238,6 +256,13 @@ class ColumnFamilyRow {
 class ColumnFamily {
  public:
   ColumnFamily() = default;
+  // ConstructAggregateColumnFamily can be used to return an aggregate
+  // ColumnFamily that can support AddToCell or MergeToCell and
+  // similar aggregate complex types. To construct an ordinary
+  // ColumnFamily, use the default constructor ColumnFamily().
+  static StatusOr<std::shared_ptr<ColumnFamily>> ConstructAggregateColumnFamily(
+      google::bigtable::admin::v2::Type value_type);
+
   // Disable copying.
   ColumnFamily(ColumnFamily const&) = delete;
   ColumnFamily& operator=(ColumnFamily const&) = delete;
@@ -275,6 +300,21 @@ class ColumnFamily {
                                       std::string const& column_qualifier,
                                       std::chrono::milliseconds timestamp,
                                       std::string const& value);
+
+  /**
+   * UpdateCell is like SetCell except that, when a cell exists with
+   * the same timestamp, an update function (that depends on the column
+   * family type) is called to derive a new value from the new and
+   * existing value, and that is the value that is written.
+   *
+   * Simple (non-aggregate) column families have a default update
+   * function that just returns the new value.
+   *
+   */
+  StatusOr<absl::optional<std::string>> UpdateCell(
+      std::string const& row_key, std::string const& column_qualifier,
+      std::chrono::milliseconds timestamp, std::string& value);
+
   /**
    * Delete the whole row from this column family.
    *
@@ -352,9 +392,81 @@ class ColumnFamily {
   }
 
   void clear() { rows_.clear(); }
+  absl::optional<google::bigtable::admin::v2::Type> GetValueType() {
+    return value_type_;
+  };
 
  private:
   std::map<std::string, ColumnFamilyRow> rows_;
+
+  // Support for aggregate and other complex types.
+  absl::optional<google::bigtable::admin::v2::Type> value_type_ = absl::nullopt;
+
+  static StatusOr<std::string> DefaultUpdateCell(
+      std::string const& /*existing_value*/, std::string&& new_value) {
+    return new_value;
+  };
+
+  static StatusOr<std::string> SumUpdateCellBEInt64(
+      std::string const& existing_value, std::string&& new_value) {
+    auto existing_value_int =
+        google::cloud::internal::DecodeBigEndian<std::int64_t>(existing_value);
+    if (!existing_value_int) {
+      return existing_value_int.status();
+    }
+
+    auto new_value_int =
+        google::cloud::internal::DecodeBigEndian<std::int64_t>(new_value);
+    if (!new_value_int) {
+      return new_value_int.status();
+    }
+
+    return google::cloud::internal::EncodeBigEndian(existing_value_int.value() +
+                                                    new_value_int.value());
+  };
+
+  static StatusOr<std::string> MaxUpdateCellBEInt64(
+      std::string const& existing_value, std::string&& new_value) {
+    auto existing_int =
+        google::cloud::internal::DecodeBigEndian<std::int64_t>(existing_value);
+    if (!existing_int) {
+      return existing_int.status();
+    }
+    auto new_int = google::cloud::internal::DecodeBigEndian<std::int64_t>(
+        std::move(new_value));
+    if (!new_int) {
+      return new_int.status();
+    }
+
+    if (existing_int.value() > new_int.value()) {
+      return existing_value;
+    }
+
+    return new_value;
+  };
+
+  static StatusOr<std::string> MinUpdateCellBEInt64(
+      std::string const& existing_value, std::string&& new_value) {
+    auto existing_int =
+        google::cloud::internal::DecodeBigEndian<std::int64_t>(existing_value);
+    if (!existing_int) {
+      return existing_int.status();
+    }
+    auto new_int = google::cloud::internal::DecodeBigEndian<std::int64_t>(
+        std::move(new_value));
+    if (!new_int) {
+      return new_int.status();
+    }
+
+    if (existing_int.value() < new_int.value()) {
+      return existing_value;
+    }
+
+    return new_value;
+  };
+
+  std::function<StatusOr<std::string>(std::string const&, std::string&&)>
+      update_cell_ = DefaultUpdateCell;
 };
 
 /**

google/cloud/bigtable/emulator/column_family_test.cc

@@ -18,6 +18,8 @@
 #include "google/cloud/testing_util/is_proto_equal.h"
 #include <google/protobuf/text_format.h>
 #include <gmock/gmock.h>
+#include <array>
+#include <cstdint>
 
 namespace google {
 namespace cloud {