From 96e6e15effe958f52b21e52ffa1bb5f67ea8b128 Mon Sep 17 00:00:00 2001
From: Sunny Singh <sunnsing@google.com>
Date: Thu, 2 Nov 2023 11:11:36 +0000
Subject: [PATCH] Update sample

---
 samples/samples/snippets.py      | 23 +++++++++++++++++------
 samples/samples/snippets_test.py |  2 +-
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py
index 8593f6393e..f7c403cfc4 100644
--- a/samples/samples/snippets.py
+++ b/samples/samples/snippets.py
@@ -403,13 +403,15 @@ def insert_data(instance_id, database_id):
 # [END spanner_insert_data]
 
 
-# [START spanner_batch_write]
+# [START spanner_batch_write_at_least_once]
 def batch_write(instance_id, database_id):
     """Inserts sample data into the given database via BatchWrite API.
 
     The database and table must already exist and can be created using
     `create_database`.
     """
+    from google.rpc.code_pb2 import OK
+
     spanner_client = spanner.Client()
     instance = spanner_client.instance(instance_id)
     database = instance.database(database_id)
@@ -429,7 +431,7 @@ def batch_write(instance_id, database_id):
             table="Singers",
             columns=("SingerId", "FirstName", "LastName"),
             values=[
-                (17, "Marc", "Richards"),
+                (17, "Marc", ""),
                 (18, "Catalina", "Smith"),
             ],
         )
@@ -443,12 +445,21 @@ def batch_write(instance_id, database_id):
         )
 
         for response in groups.batch_write():
-            print(response)
-
-    print("Inserted data.")
+            if response.status.code == OK:
+                print(
+                    "Mutation group indexes {} have been applied with commit timestamp {}".format(
+                        response.indexes, response.commit_timestamp
+                    )
+                )
+            else:
+                print(
+                    "Mutation group indexes {} could not be applied with error {}".format(
+                        response.indexes, response.status
+                    )
+                )
 
 
-# [END spanner_batch_write]
+# [END spanner_batch_write_at_least_once]
 
 
 # [START spanner_delete_data]
diff --git a/samples/samples/snippets_test.py b/samples/samples/snippets_test.py
index 5c6fa0eacd..85999363bb 100644
--- a/samples/samples/snippets_test.py
+++ b/samples/samples/snippets_test.py
@@ -294,7 +294,7 @@ def test_insert_data(capsys, instance_id, sample_database):
 def test_batch_write(capsys, instance_id, sample_database):
     snippets.batch_write(instance_id, sample_database.database_id)
     out, _ = capsys.readouterr()
-    assert "Inserted data" in out
+    assert "could not be applied with error" not in out
 
 
 @pytest.mark.dependency(depends=["insert_data"])