From 61136e473d827785b35d6accc2c078766b7eb485 Mon Sep 17 00:00:00 2001
From: Jonathan Simon <jsimonweb@gmail.com>
Date: Mon, 25 Mar 2019 15:32:30 -0700
Subject: [PATCH] Add Cloud Spanner Batch DML sample
 [(#2068)](https://github.com/GoogleCloudPlatform/python-docs-samples/issues/2068)

* Add Cloud Spanner Batch DML sample

* Fix test.

* Lint.

* More Lint.
---
 samples/samples/snippets.py      | 48 +++++++++++++++++++++++++++++---
 samples/samples/snippets_test.py | 10 +++++--
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py
index 1b065bb4cf..9bbe0fbd8c 100644
--- a/samples/samples/snippets.py
+++ b/samples/samples/snippets.py
@@ -701,7 +701,7 @@ def query_with_array_of_struct(instance_id, database_id):
             param_types={'names': param_types.Array(name_type)})
 
     for row in results:
-            print(u'SingerId: {}'.format(*row))
+        print(u'SingerId: {}'.format(*row))
     # [END spanner_query_data_with_array_of_struct]
 
 
@@ -725,7 +725,7 @@ def query_struct_field(instance_id, database_id):
             param_types={'name': name_type})
 
     for row in results:
-            print(u'SingerId: {}'.format(*row))
+        print(u'SingerId: {}'.format(*row))
 # [START spanner_field_access_on_struct_parameters]
 
 
@@ -768,7 +768,7 @@ def query_nested_struct_field(instance_id, database_id):
         )
 
     for row in results:
-            print(u'SingerId: {} SongName: {}'.format(*row))
+        print(u'SingerId: {} SongName: {}'.format(*row))
 # [END spanner_field_access_on_nested_struct_parameters]
 
 
@@ -998,7 +998,7 @@ def transfer_budget(transaction):
             )
 
             print("Transferred {} from Album1's budget to Album2's".format(
-                    transfer_amount))
+                transfer_amount))
 
     database.run_in_transaction(transfer_budget)
     # [END spanner_dml_getting_started_update]
@@ -1039,6 +1039,41 @@ def delete_data_with_partitioned_dml(instance_id, database_id):
     # [END spanner_dml_partitioned_delete]
 
 
+def update_with_batch_dml(instance_id, database_id):
+    """Updates sample data in the database using Batch DML. """
+    # [START spanner_dml_batch_update]
+    # instance_id = "your-spanner-instance"
+    # database_id = "your-spanner-db-id"
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    insert_statement = (
+        "INSERT INTO Albums "
+        "(SingerId, AlbumId, AlbumTitle, MarketingBudget) "
+        "VALUES (1, 3, 'Test Album Title', 10000)"
+    )
+
+    update_statement = (
+        "UPDATE Albums "
+        "SET MarketingBudget = MarketingBudget * 2 "
+        "WHERE SingerId = 1 and AlbumId = 3"
+    )
+
+    def update_albums(transaction):
+        row_cts = transaction.batch_update([
+            insert_statement,
+            update_statement,
+        ])
+
+        print("Executed {} SQL statements using Batch DML.".format(
+            len(row_cts)))
+
+    database.run_in_transaction(update_albums)
+    # [END spanner_dml_batch_update]
+
+
 if __name__ == '__main__':  # noqa: C901
     parser = argparse.ArgumentParser(
         description=__doc__,
@@ -1118,6 +1153,9 @@ def delete_data_with_partitioned_dml(instance_id, database_id):
     subparsers.add_parser(
         'delete_data_with_partitioned_dml',
         help=delete_data_with_partitioned_dml.__doc__)
+    subparsers.add_parser(
+        'update_with_batch_dml',
+        help=update_with_batch_dml.__doc__)
 
     args = parser.parse_args()
 
@@ -1195,3 +1233,5 @@ def delete_data_with_partitioned_dml(instance_id, database_id):
         update_data_with_partitioned_dml(args.instance_id, args.database_id)
     elif args.command == 'delete_data_with_partitioned_dml':
         delete_data_with_partitioned_dml(args.instance_id, args.database_id)
+    elif args.command == 'update_with_batch_dml':
+        update_with_batch_dml(args.instance_id, args.database_id)
diff --git a/samples/samples/snippets_test.py b/samples/samples/snippets_test.py
index 2ea6aec642..9bd39d27a5 100644
--- a/samples/samples/snippets_test.py
+++ b/samples/samples/snippets_test.py
@@ -190,13 +190,13 @@ def test_insert_data_with_timestamp(capsys):
 def test_write_struct_data(capsys):
     snippets.write_struct_data(INSTANCE_ID, DATABASE_ID)
     out, _ = capsys.readouterr()
-    assert 'Inserted sample data for STRUCT queries'
+    assert 'Inserted sample data for STRUCT queries' in out
 
 
 def test_query_with_struct(capsys):
     snippets.query_with_struct(INSTANCE_ID, DATABASE_ID)
     out, _ = capsys.readouterr()
-    assert 'SingerId: Elena'
+    assert 'SingerId: 6' in out
 
 
 def test_query_with_array_of_struct(capsys):
@@ -277,3 +277,9 @@ def delete_data_with_partitioned_dml(capsys):
     snippets.delete_data_with_partitioned_dml(INSTANCE_ID, DATABASE_ID)
     out, _ = capsys.readouterr()
     assert "5 record(s) deleted" in out
+
+
+def update_with_batch_dml(capsys):
+    snippets.update_with_batch_dml(INSTANCE_ID, DATABASE_ID)
+    out, _ = capsys.readouterr()
+    assert "Executed 2 SQL statements using Batch DML" in out