feat: add sample for commit stats (googleapis#241)

larkee · busunkim96 · skuruppu · web-flow · commit 1343656ad43d · 2021-02-26T14:34:11.000+11:00
* feat: add sample for commit stats

* fix: use correct kwarg

* fix: correct super call

* fix: add missing super init call

* fix: update super init call

* fix: use correct key

* refactor: remove testing file

* test: fix typo

* Apply suggestions from code review

Co-authored-by: Bu Sun Kim &lt;8822365+busunkim96@users.noreply.github.com&gt;
Co-authored-by: skuruppu &lt;skuruppu@google.com&gt;

* refactor: make last_commit_stats public

Co-authored-by: larkee &lt;larkee@users.noreply.github.com&gt;
Co-authored-by: Bu Sun Kim &lt;8822365+busunkim96@users.noreply.github.com&gt;
Co-authored-by: skuruppu &lt;skuruppu@google.com&gt;
diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py
@@ -24,6 +24,7 @@
 import base64
 import datetime
 import decimal
+import logging
 
 from google.cloud import spanner
 from google.cloud.spanner_v1 import param_types
@@ -969,6 +970,44 @@ def insert_singers(transaction):
     # [END spanner_dml_standard_insert]
 
 
+# [START spanner_get_commit_stats]
+def log_commit_stats(instance_id, database_id):
+    """Inserts sample data using DML and displays the commit statistics. """
+    # By default, commit statistics are logged via stdout at level Info.
+    # This sample uses a custom logger to access the commit statistics.
+    class CommitStatsSampleLogger(logging.Logger):
+        def __init__(self):
+            self.last_commit_stats = None
+            super().__init__("commit_stats_sample")
+
+        def info(self, msg, *args, **kwargs):
+            if kwargs["extra"] and "commit_stats" in kwargs["extra"]:
+                self.last_commit_stats = kwargs["extra"]["commit_stats"]
+            super().info(msg)
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id, logger=CommitStatsSampleLogger())
+    database.log_commit_stats = True
+
+    def insert_singers(transaction):
+        row_ct = transaction.execute_update(
+            "INSERT Singers (SingerId, FirstName, LastName) "
+            " VALUES (110, 'Virginia', 'Watson')"
+        )
+
+        print("{} record(s) inserted.".format(row_ct))
+
+    database.run_in_transaction(insert_singers)
+    commit_stats = database.logger.last_commit_stats
+    print(
+        "{} mutation(s) in transaction.".format(
+            commit_stats.mutation_count
+        )
+    )
+# [END spanner_get_commit_stats]
+
+
 def update_data_with_dml(instance_id, database_id):
     """Updates sample data from the database using a DML statement. """
     # [START spanner_dml_standard_update]
@@ -1710,6 +1749,7 @@ def create_client_with_query_options(instance_id, database_id):
         "query_nested_struct_field", help=query_nested_struct_field.__doc__
     )
     subparsers.add_parser("insert_data_with_dml", help=insert_data_with_dml.__doc__)
+    subparsers.add_parser("log_commit_stats", help=log_commit_stats.__doc__)
     subparsers.add_parser("update_data_with_dml", help=update_data_with_dml.__doc__)
     subparsers.add_parser("delete_data_with_dml", help=delete_data_with_dml.__doc__)
     subparsers.add_parser(
@@ -1820,6 +1860,8 @@ def create_client_with_query_options(instance_id, database_id):
         query_nested_struct_field(args.instance_id, args.database_id)
     elif args.command == "insert_data_with_dml":
         insert_data_with_dml(args.instance_id, args.database_id)
+    elif args.command == "log_commit_stats":
+        log_commit_stats(args.instance_id, args.database_id)
     elif args.command == "update_data_with_dml":
         update_data_with_dml(args.instance_id, args.database_id)
     elif args.command == "delete_data_with_dml":
diff --git a/samples/samples/snippets_test.py b/samples/samples/snippets_test.py
@@ -236,6 +236,13 @@ def test_insert_data_with_dml(capsys):
     assert "1 record(s) inserted." in out
 
 
+def test_log_commit_stats(capsys):
+    snippets.log_commit_stats(INSTANCE_ID, DATABASE_ID)
+    out, _ = capsys.readouterr()
+    assert "1 record(s) inserted." in out
+    assert "3 mutation(s) in transaction." in out
+
+
 def test_update_data_with_dml(capsys):
     snippets.update_data_with_dml(INSTANCE_ID, DATABASE_ID)
     out, _ = capsys.readouterr()
diff --git a/test.py b/test.py
