feat(spanner): tests and samples for DML RETURNING

google/cloud/spanner/client.h

@@ -341,12 +341,18 @@ class Client {
    * returned/ignored, and the column order is known. This enables more
    * efficient and simpler code.
    *
+   * Can also execute a DML statement with a returning clause in a read/write
+   * transaction.
+   *
    * @par Example with explicitly selected columns.
    * @snippet samples.cc spanner-query-data
    *
-   * @par Example using SELECT *
+   * @par Example using `SELECT *`.
    * @snippet samples.cc spanner-query-data-select-star
    *
+   * @par Example using a DML statement with `THEN RETURN`.
+   * @snippet samples.cc spanner-update-dml-returning
+   *
    * @param statement The SQL statement to execute.
    * @param opts (optional) The `Options` to use for this call. If given,
    *     these will take precedence over the options set at the client and

google/cloud/spanner/integration_tests/data_types_integration_test.cc

@@ -22,6 +22,7 @@
 #include "absl/memory/memory.h"
 #include "absl/time/time.h"
 #include <gmock/gmock.h>
+#include <vector>
 
 namespace google {
 namespace cloud {
@@ -33,6 +34,7 @@ using ::google::cloud::testing_util::IsOk;
 using ::google::cloud::testing_util::StatusIs;
 using ::testing::AllOf;
 using ::testing::AnyOf;
+using ::testing::ElementsAre;
 using ::testing::HasSubstr;
 using ::testing::UnorderedElementsAreArray;
 
@@ -600,6 +602,148 @@ TEST_F(PgDataTypeIntegrationTest, NumericPrimaryKey) {
                                        HasSubstr("part of the primary key"))));
 }
 
+TEST_F(DataTypeIntegrationTest, DmlReturning) {
+  if (UsingEmulator()) GTEST_SKIP() << "emulator does not support THEN RETURN";
+
+  auto& client = *client_;
+  using RowType = std::tuple<std::string, std::int64_t>;
+
+  std::vector<RowType> insert_actual;
+  auto insert_result = client.Commit(
+      [&client, &insert_actual](Transaction const& txn) -> StatusOr<Mutations> {
+        auto sql = SqlStatement(R"""(
+            INSERT INTO DataTypes (Id, Int64Value)
+              VALUES ('Id-Ret-1', 1),
+                     ('Id-Ret-2', 2),
+                     ('Id-Ret-3', 3),
+                     ('Id-Ret-4', 4)
+              THEN RETURN Id, Int64Value
+        )""");
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        EXPECT_EQ(rows.RowsModified(), 4);
+        for (auto& row : StreamOf<RowType>(rows)) {
+          if (row) insert_actual.push_back(*std::move(row));
+        }
+        return Mutations{};
+      });
+  ASSERT_THAT(insert_result, IsOk());
+  EXPECT_THAT(insert_actual,
+              ElementsAre(RowType{"Id-Ret-1", 1}, RowType{"Id-Ret-2", 2},
+                          RowType{"Id-Ret-3", 3}, RowType{"Id-Ret-4", 4}));
+
+  std::vector<RowType> update_actual;
+  auto update_result = client.Commit(
+      [&client, &update_actual](Transaction const& txn) -> StatusOr<Mutations> {
+        auto sql = SqlStatement(R"""(
+            UPDATE DataTypes SET Int64Value = 100
+              WHERE Id LIKE 'Id-Ret-%%'
+              THEN RETURN Id, Int64Value
+        )""");
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        EXPECT_EQ(rows.RowsModified(), 4);
+        for (auto& row : StreamOf<RowType>(rows)) {
+          if (row) update_actual.push_back(*std::move(row));
+        }
+        return Mutations{};
+      });
+  ASSERT_THAT(update_result, IsOk());
+  EXPECT_THAT(update_actual,
+              ElementsAre(RowType{"Id-Ret-1", 100}, RowType{"Id-Ret-2", 100},
+                          RowType{"Id-Ret-3", 100}, RowType{"Id-Ret-4", 100}));
+
+  std::vector<RowType> delete_actual;
+  auto delete_result = client.Commit(
+      [&client, &delete_actual](Transaction const& txn) -> StatusOr<Mutations> {
+        auto sql = SqlStatement(R"""(
+            DELETE FROM DataTypes
+              WHERE Id LIKE 'Id-Ret-%%'
+              THEN RETURN Id, Int64Value
+        )""");
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        EXPECT_EQ(rows.RowsModified(), 4);
+        for (auto& row : StreamOf<RowType>(rows)) {
+          if (row) delete_actual.push_back(*std::move(row));
+        }
+        return Mutations{};
+      });
+  ASSERT_THAT(delete_result, IsOk());
+  EXPECT_THAT(delete_actual,
+              ElementsAre(RowType{"Id-Ret-1", 100}, RowType{"Id-Ret-2", 100},
+                          RowType{"Id-Ret-3", 100}, RowType{"Id-Ret-4", 100}));
+}
+
+TEST_F(PgDataTypeIntegrationTest, DmlReturning) {
+  if (UsingEmulator()) {
+    GTEST_SKIP() << "emulator does not support PostgreSQL or RETURNING";
+  }
+
+  auto& client = *client_;
+  using RowType = std::tuple<std::string, std::int64_t>;
+
+  std::vector<RowType> insert_actual;
+  auto insert_result = client.Commit(
+      [&client, &insert_actual](Transaction const& txn) -> StatusOr<Mutations> {
+        auto sql = SqlStatement(R"""(
+            INSERT INTO DataTypes (Id, Int64Value)
+              VALUES ('Id-Ret-1', 1),
+                     ('Id-Ret-2', 2),
+                     ('Id-Ret-3', 3),
+                     ('Id-Ret-4', 4)
+              RETURNING Id, Int64Value
+        )""");
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        EXPECT_EQ(rows.RowsModified(), 4);
+        for (auto& row : StreamOf<RowType>(rows)) {
+          if (row) insert_actual.push_back(*std::move(row));
+        }
+        return Mutations{};
+      });
+  ASSERT_THAT(insert_result, IsOk());
+  EXPECT_THAT(insert_actual,
+              ElementsAre(RowType{"Id-Ret-1", 1}, RowType{"Id-Ret-2", 2},
+                          RowType{"Id-Ret-3", 3}, RowType{"Id-Ret-4", 4}));
+
+  std::vector<RowType> update_actual;
+  auto update_result = client.Commit(
+      [&client, &update_actual](Transaction const& txn) -> StatusOr<Mutations> {
+        auto sql = SqlStatement(R"""(
+            UPDATE DataTypes SET Int64Value = 100
+              WHERE Id LIKE 'Id-Ret-%%'
+              RETURNING Id, Int64Value
+        )""");
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        EXPECT_EQ(rows.RowsModified(), 4);
+        for (auto& row : StreamOf<RowType>(rows)) {
+          if (row) update_actual.push_back(*std::move(row));
+        }
+        return Mutations{};
+      });
+  ASSERT_THAT(update_result, IsOk());
+  EXPECT_THAT(update_actual,
+              ElementsAre(RowType{"Id-Ret-1", 100}, RowType{"Id-Ret-2", 100},
+                          RowType{"Id-Ret-3", 100}, RowType{"Id-Ret-4", 100}));
+
+  std::vector<RowType> delete_actual;
+  auto delete_result = client.Commit(
+      [&client, &delete_actual](Transaction const& txn) -> StatusOr<Mutations> {
+        auto sql = SqlStatement(R"""(
+            DELETE FROM DataTypes
+              WHERE Id LIKE 'Id-Ret-%%'
+              RETURNING Id, Int64Value
+        )""");
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        EXPECT_EQ(rows.RowsModified(), 4);
+        for (auto& row : StreamOf<RowType>(rows)) {
+          if (row) delete_actual.push_back(*std::move(row));
+        }
+        return Mutations{};
+      });
+  ASSERT_THAT(delete_result, IsOk());
+  EXPECT_THAT(delete_actual,
+              ElementsAre(RowType{"Id-Ret-1", 100}, RowType{"Id-Ret-2", 100},
+                          RowType{"Id-Ret-3", 100}, RowType{"Id-Ret-4", 100}));
+}
+
 // This test differs a lot from the other tests since Spanner STRUCT types may
 // not be used as column types, and they may not be returned as top-level
 // objects in a select statement. See

google/cloud/spanner/samples/postgresql_samples.cc

@@ -219,6 +219,93 @@ void BatchDml(google::cloud::spanner::Client client) {
 }
 // [END spanner_postgresql_batch_dml]
 
+// [START spanner_postgresql_update_dml_returning]
+void UpdateUsingDmlReturning(google::cloud::spanner::Client client) {
+  // Update MarketingBudget column for records satisfying a particular
+  // condition and return the modified MarketingBudget column of the
+  // updated records using `RETURNING MarketingBudget`.
+  auto commit = client.Commit(
+      [&client](google::cloud::spanner::Transaction txn)
+          -> google::cloud::StatusOr<google::cloud::spanner::Mutations> {
+        auto sql = google::cloud::spanner::SqlStatement(R"""(
+            UPDATE Albums SET MarketingBudget = MarketingBudget * 2
+              WHERE SingerId = 1 AND AlbumId = 1
+              RETURNING MarketingBudget
+        )""");
+        using RowType = std::tuple<absl::optional<std::int64_t>>;
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        for (auto& row : google::cloud::spanner::StreamOf<RowType>(rows)) {
+          if (!row) return std::move(row).status();
+          std::cout << "MarketingBudget: ";
+          if (std::get<0>(*row).has_value()) {
+            std::cout << *std::get<0>(*row);
+          } else {
+            std::cout << "NULL";
+          }
+          std::cout << "\n";
+        }
+        std::cout << "Updated row(s) count: " << rows.RowsModified() << "\n";
+        return google::cloud::spanner::Mutations{};
+      });
+  if (!commit) throw std::move(commit).status();
+}
+// [END spanner_postgresql_update_dml_returning]
+
+// [START spanner_postgresql_insert_dml_returning]
+void InsertUsingDmlReturning(google::cloud::spanner::Client client) {
+  // Insert records into SINGERS table and return the generated column
+  // FullName of the inserted records using `RETURNING FullName`.
+  auto commit = client.Commit(
+      [&client](google::cloud::spanner::Transaction txn)
+          -> google::cloud::StatusOr<google::cloud::spanner::Mutations> {
+        auto sql = google::cloud::spanner::SqlStatement(R"""(
+            INSERT INTO Singers (SingerId, FirstName, LastName)
+              VALUES (12, 'Melissa', 'Garcia'),
+                     (13, 'Russell', 'Morales'),
+                     (14, 'Jacqueline', 'Long'),
+                     (15, 'Dylan', 'Shaw')
+              RETURNING FullName
+        )""");
+        using RowType = std::tuple<std::string>;
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        for (auto& row : google::cloud::spanner::StreamOf<RowType>(rows)) {
+          if (!row) return std::move(row).status();
+          std::cout << "FullName: " << std::get<0>(*row) << "\n";
+        }
+        std::cout << "Inserted row(s) count: " << rows.RowsModified() << "\n";
+        return google::cloud::spanner::Mutations{};
+      });
+  if (!commit) throw std::move(commit).status();
+}
+// [END spanner_postgresql_insert_dml_returning]
+
+// [START spanner_postgresql_delete_dml_returning]
+void DeleteUsingDmlReturning(google::cloud::spanner::Client client) {
+  // Delete records from SINGERS table satisfying a particular condition
+  // and return the SingerId and FullName column of the deleted records
+  // using `RETURNING SingerId, FullName'.
+  auto commit = client.Commit(
+      [&client](google::cloud::spanner::Transaction txn)
+          -> google::cloud::StatusOr<google::cloud::spanner::Mutations> {
+        auto sql = google::cloud::spanner::SqlStatement(R"""(
+            DELETE FROM Singers
+              WHERE FirstName = 'Alice'
+              RETURNING SingerId, FullName
+        )""");
+        using RowType = std::tuple<std::int64_t, std::string>;
+        auto rows = client.ExecuteQuery(std::move(txn), std::move(sql));
+        for (auto& row : google::cloud::spanner::StreamOf<RowType>(rows)) {
+          if (!row) return std::move(row).status();
+          std::cout << "SingerId: " << std::get<0>(*row) << " ";
+          std::cout << "FullName: " << std::get<1>(*row) << "\n";
+        }
+        std::cout << "Deleted row(s) count: " << rows.RowsModified() << "\n";
+        return google::cloud::spanner::Mutations{};
+      });
+  if (!commit) throw std::move(commit).status();
+}
+// [END spanner_postgresql_delete_dml_returning]
+
 // [START spanner_postgresql_case_sensitivity]
 void CaseSensitivity(
     google::cloud::spanner_admin::DatabaseAdminClient admin_client,
@@ -782,6 +869,8 @@ void CreateTables(google::cloud::spanner_admin::DatabaseAdminClient client,
             SingerId   BIGINT NOT NULL,
             FirstName  CHARACTER VARYING(1024),
             LastName   CHARACTER VARYING(1024),
+            FullName   CHARACTER VARYING(2049)
+                GENERATED ALWAYS AS (FirstName || ' ' || LastName) STORED,
             SingerInfo BYTEA,
             PRIMARY KEY(singerid)
         )
@@ -904,6 +993,9 @@ int RunOneCommand(std::vector<std::string> argv,
       {"query-with-parameter", Command(samples::QueryWithParameter)},
       {"dml-getting-started-update", Command(samples::DmlGettingStartedUpdate)},
       {"batch-dml", Command(samples::BatchDml)},
+      {"update-dml-returning", Command(samples::UpdateUsingDmlReturning)},
+      {"insert-dml-returning", Command(samples::InsertUsingDmlReturning)},
+      {"delete-dml-returning", Command(samples::DeleteUsingDmlReturning)},
       {"drop-tables", Command(helpers::DropTables)},
       {"case-sensitivity", Command(samples::CaseSensitivity)},
       {"cast-data-type", Command(samples::CastDataType)},
@@ -992,6 +1084,15 @@ int RunAll() {
     SampleBanner("spanner_postgresql_batch_dml");
     samples::BatchDml(client);
 
+    SampleBanner("spanner_postgresql_update_dml_returning");
+    samples::UpdateUsingDmlReturning(client);
+
+    SampleBanner("spanner_postgresql_insert_dml_returning");
+    samples::InsertUsingDmlReturning(client);
+
+    SampleBanner("spanner_postgresql_delete_dml_returning");
+    samples::DeleteUsingDmlReturning(client);
+
     helpers::DropTables(database_admin_client, database);
 
     SampleBanner("spanner_postgresql_case_sensitivity");