feat(spanner): support for the PG.JSONB data type (#10098)

Introduce `google::cloud::spanner::JsonB` for PostgreSQL databases.
The interface is the same as `Json` (just a strongly-typed string),
but `JsonB` values are stored in the database in a decomposed, binary
format, which affects performance: `JsonB` is slower to input, but
faster to process as it avoids reparsing.

Note: Support for arrays of `JSONB` will come later.

google/cloud/spanner/integration_tests/data_types_integration_test.cc

@@ -50,7 +50,7 @@ StatusOr<T> WriteReadData(Client& client, T const& data,
   int id = 0;
   for (auto&& x : data) {
     mutations.push_back(MakeInsertMutation("DataTypes", {"Id", column},
-                                           "Id-" + std::to_string(id++), x));
+                                           "Id-" + std::to_string(++id), x));
   }
   auto commit_result = client.Commit(std::move(mutations));
   if (!commit_result) return commit_result.status();
@@ -241,6 +241,20 @@ TEST_F(DataTypeIntegrationTest, WriteReadJson) {
   EXPECT_THAT(*result, UnorderedElementsAreArray(data));
 }
 
+TEST_F(PgDataTypeIntegrationTest, WriteReadJson) {
+  if (UsingEmulator()) GTEST_SKIP() << "emulator does not support PostgreSQL";
+
+  std::vector<JsonB> const data = {
+      JsonB(),                     //
+      JsonB(R"("Hello world!")"),  //
+      JsonB("42"),                 //
+      JsonB("true"),               //
+  };
+  auto result = WriteReadData(*client_, data, "JsonValue");
+  ASSERT_STATUS_OK(result);
+  EXPECT_THAT(*result, UnorderedElementsAreArray(data));
+}
+
 TEST_F(DataTypeIntegrationTest, WriteReadNumeric) {
   auto min = MakeNumeric("-99999999999999999999999999999.999999999");
   ASSERT_STATUS_OK(min);
@@ -387,6 +401,33 @@ TEST_F(DataTypeIntegrationTest, WriteReadArrayJson) {
   EXPECT_THAT(*result, UnorderedElementsAreArray(data));
 }
 
+TEST_F(PgDataTypeIntegrationTest, WriteReadArrayJson) {
+  if (UsingEmulator()) GTEST_SKIP() << "emulator does not support PostgreSQL";
+
+  std::vector<std::vector<JsonB>> const data = {
+      std::vector<JsonB>{},
+      std::vector<JsonB>{JsonB()},
+      std::vector<JsonB>{
+          JsonB(),
+          JsonB(R"("Hello world!")"),
+          JsonB("42"),
+          JsonB("true"),
+      },
+  };
+  auto result = WriteReadData(*client_, data, "ArrayJsonValue");
+  {
+    // TODO(#10095): Remove this when JSONB[] is supported.
+    auto matcher = testing_util::StatusIs(
+        StatusCode::kNotFound, testing::HasSubstr("Column not found in table"));
+    testing::StringMatchResultListener listener;
+    if (matcher.impl().MatchAndExplain(result, &listener)) {
+      GTEST_SKIP();
+    }
+  }
+  ASSERT_STATUS_OK(result);
+  EXPECT_THAT(*result, UnorderedElementsAreArray(data));
+}
+
 TEST_F(DataTypeIntegrationTest, WriteReadArrayNumeric) {
   std::vector<std::vector<Numeric>> const data = {
       std::vector<Numeric>{},
@@ -432,8 +473,12 @@ TEST_F(DataTypeIntegrationTest, JsonIndexAndPrimaryKey) {
   auto metadata =
       admin_client.UpdateDatabaseDdl(GetDatabase().FullName(), statements)
           .get();
-  EXPECT_THAT(metadata, StatusIs(StatusCode::kFailedPrecondition,
-                                 HasSubstr("DataTypesByJsonValue")));
+  EXPECT_THAT(metadata,
+              StatusIs(StatusCode::kFailedPrecondition,
+                       AnyOf(AllOf(HasSubstr("Index DataTypesByJsonValue"),
+                                   HasSubstr("unsupported type JSON")),
+                             AllOf(HasSubstr("index DataTypesByJsonValue"),
+                                   HasSubstr("Cannot reference JSON")))));
 
   // Verify that a JSON column cannot be used as a primary key.
   statements.clear();
@@ -446,10 +491,75 @@ TEST_F(DataTypeIntegrationTest, JsonIndexAndPrimaryKey) {
       admin_client.UpdateDatabaseDdl(GetDatabase().FullName(), statements)
           .get();
   EXPECT_THAT(metadata, StatusIs(StatusCode::kInvalidArgument,
-                                 AllOf(HasSubstr("Key has type JSON"),
+                                 AllOf(HasSubstr("has type JSON"),
+                                       HasSubstr("part of the primary key"))));
+}
+
+TEST_F(PgDataTypeIntegrationTest, JsonIndexAndPrimaryKey) {
+  if (UsingEmulator()) GTEST_SKIP() << "emulator does not support PostgreSQL";
+
+  spanner_admin::DatabaseAdminClient admin_client(
+      spanner_admin::MakeDatabaseAdminConnection());
+
+  // Verify that a JSONB column cannot be used as an index.
+  std::vector<std::string> statements;
+  statements.emplace_back(R"""(
+        CREATE INDEX DataTypesByJsonValue
+          ON DataTypes(JsonValue)
+      )""");
+  auto metadata =
+      admin_client.UpdateDatabaseDdl(GetDatabase().FullName(), statements)
+          .get();
+  EXPECT_THAT(metadata,
+              StatusIs(StatusCode::kFailedPrecondition,
+                       AllOf(HasSubstr("Index datatypesbyjsonvalue"),
+                             HasSubstr("unsupported type PG.JSONB"))));
+
+  // Verify that a JSONB column cannot be used as a primary key.
+  statements.clear();
+  statements.emplace_back(R"""(
+        CREATE TABLE JsonKey (
+          Key JSONB NOT NULL,
+          PRIMARY KEY (Key)
+        )
+      )""");
+  metadata =
+      admin_client.UpdateDatabaseDdl(GetDatabase().FullName(), statements)
+          .get();
+  EXPECT_THAT(metadata, StatusIs(StatusCode::kInvalidArgument,
+                                 AllOf(HasSubstr("has type PG.JSONB"),
                                        HasSubstr("part of the primary key"))));
 }
 
+TEST_F(PgDataTypeIntegrationTest, InsertAndQueryWithJson) {
+  if (UsingEmulator()) GTEST_SKIP() << "emulator does not support PostgreSQL";
+
+  auto& client = *client_;
+  auto commit_result =
+      client.Commit([&client](Transaction const& txn) -> StatusOr<Mutations> {
+        auto dml_result = client.ExecuteDml(
+            txn,
+            SqlStatement("INSERT INTO DataTypes (Id, JsonValue)"
+                         "  VALUES($1, $2)",
+                         {{"p1", Value("Id-1")}, {"p2", Value(JsonB("42"))}}));
+        if (!dml_result) return dml_result.status();
+        return Mutations{};
+      });
+  ASSERT_STATUS_OK(commit_result);
+
+  auto rows =
+      client_->ExecuteQuery(SqlStatement("SELECT Id, JsonValue FROM DataTypes"
+                                         "  WHERE Id = $1",
+                                         {{"p1", Value("Id-1")}}));
+  using RowType = std::tuple<std::string, absl::optional<JsonB>>;
+  auto row = GetSingularRow(StreamOf<RowType>(rows));
+  ASSERT_STATUS_OK(row);
+
+  auto const& v = std::get<1>(*row);
+  ASSERT_TRUE(v.has_value());
+  EXPECT_EQ(*v, JsonB("42"));
+}
+
 TEST_F(DataTypeIntegrationTest, InsertAndQueryWithNumericKey) {
   auto& client = *client_;
   auto const key = MakeNumeric(42).value();
@@ -513,8 +623,8 @@ TEST_F(DataTypeIntegrationTest, InsertAndQueryWithStruct) {
             txn,
             SqlStatement(
                 "INSERT INTO DataTypes (Id, StringValue, ArrayInt64Value)"
-                "VALUES(@id, @struct.StringValue, @struct.ArrayInt64Value)",
-                {{"id", Value("id-1")}, {"struct", Value(data)}}));
+                "  VALUES(@id, @struct.StringValue, @struct.ArrayInt64Value)",
+                {{"id", Value("Id-1")}, {"struct", Value(data)}}));
         if (!dml_result) return dml_result.status();
         return Mutations{};
       });

google/cloud/spanner/internal/connection_impl_test.cc

@@ -853,6 +853,53 @@ TEST(ConnectionImplTest, ExecuteQueryPgNumericResult) {
   EXPECT_EQ(row_number, expected.size());
 }
 
+TEST(ConnectionImplTest, ExecuteQueryJsonBResult) {
+  auto mock = std::make_shared<spanner_testing::MockSpannerStub>();
+  auto db = spanner::Database("placeholder_project", "placeholder_instance",
+                              "placeholder_database_id");
+  EXPECT_CALL(*mock, BatchCreateSessions(_, HasDatabase(db)))
+      .WillOnce(Return(MakeSessionsResponse({"test-session-name"})));
+  auto constexpr kText = R"pb(
+    metadata: {
+      row_type: {
+        fields: {
+          name: "ColumnA",
+          type: { code: JSON type_annotation: PG_JSONB }
+        }
+        fields: {
+          name: "ColumnB",
+          type: { code: JSON type_annotation: PG_JSONB }
+        }
+      }
+    }
+    values: { string_value: "42" }
+    values: { null_value: NULL_VALUE }
+    values: { string_value: "[null, null]" }
+    values: { string_value: "{\"a\": 1, \"b\": 2}" }
+  )pb";
+  EXPECT_CALL(*mock, ExecuteStreamingSql)
+      .WillOnce(Return(ByMove(MakeReader({kText}))));
+
+  auto conn = MakeConnectionImpl(db, mock);
+  internal::OptionsSpan span(MakeLimitedTimeOptions());
+  auto rows = conn->ExecuteQuery(
+      {MakeSingleUseTransaction(spanner::Transaction::ReadOnlyOptions()),
+       spanner::SqlStatement("SELECT * FROM Table")});
+  using RowType = std::tuple<spanner::JsonB, absl::optional<spanner::JsonB>>;
+  auto expected = std::vector<RowType>{
+      RowType(spanner::JsonB("42"), absl::nullopt),
+      RowType(spanner::JsonB("[null, null]"),
+              spanner::JsonB(R"({"a": 1, "b": 2})")),
+  };
+  int row_number = 0;
+  for (auto& row : spanner::StreamOf<RowType>(rows)) {
+    EXPECT_STATUS_OK(row);
+    EXPECT_EQ(*row, expected[row_number]);
+    ++row_number;
+  }
+  EXPECT_EQ(row_number, expected.size());
+}
+
 TEST(ConnectionImplTest, ExecuteQueryNumericParameter) {
   auto mock = std::make_shared<spanner_testing::MockSpannerStub>();
   auto db = spanner::Database("placeholder_project", "placeholder_instance",

google/cloud/spanner/json.h

@@ -81,6 +81,68 @@ inline std::ostream& operator<<(std::ostream& os, Json const& json) {
   return os << std::string(json);
 }
 
+/**
+ * `JsonB` is a variant of `Json` (see above). While both classes share the
+ * same, thin client-side API, `JsonB` stores the data in a decomposed,
+ * binary format, whereas `Json` stores an exact copy of the RFC 7159 text.
+ *
+ * This means that `JsonB` is slower to input, but faster to process as it
+ * avoids reparsing. Therefore, applications that utilize the structured
+ * state of a JSON value should prefer `JsonB`.
+ *
+ * It also means that the `JsonB` stored representation does NOT preserve:
+ *   - white space,
+ *   - the order of object keys, or
+ *   - duplicate object keys.
+ *
+ * Note: `JsonB` is only applicable to PostgreSQL databases (i.e., those
+ * created using `DatabaseDialect::POSTGRESQL`).
+ */
+class JsonB {
+ public:
+  /// A null value.
+  JsonB() : rep_("null") {}
+
+  /// Regular value type, supporting copy, assign, move.
+  ///@{
+  JsonB(JsonB const&) = default;
+  JsonB& operator=(JsonB const&) = default;
+  JsonB(JsonB&&) = default;
+  JsonB& operator=(JsonB&&) = default;
+  ///@}
+
+  /**
+   * Construction from a JSON-formatted string. Note that there is no check
+   * here that the argument string is indeed well-formatted. Error detection
+   * will be delayed until the value is passed to Spanner.
+   */
+  explicit JsonB(std::string s) : rep_(std::move(s)) {}
+
+  /// Conversion to a JSON-formatted string.
+  ///@{
+  explicit operator std::string() const& { return rep_; }
+  explicit operator std::string() && { return std::move(rep_); }
+  ///@}
+
+ private:
+  std::string rep_;  // a (presumably) JSON-formatted string
+};
+
+/// @name Relational operators
+///@{
+inline bool operator==(JsonB const& lhs, JsonB const& rhs) {
+  return std::string(lhs) == std::string(rhs);
+}
+inline bool operator!=(JsonB const& lhs, JsonB const& rhs) {
+  return !(lhs == rhs);
+}
+///@}
+
+/// Outputs a JSON-formatted string to the provided stream.
+inline std::ostream& operator<<(std::ostream& os, JsonB const& json) {
+  return os << std::string(json);
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace spanner
 }  // namespace cloud