From 1343656ad43dbc41c119b652d8fe9360fa2b0e78 Mon Sep 17 00:00:00 2001 From: larkee <31196561+larkee@users.noreply.github.com> Date: Fri, 26 Feb 2021 14:34:11 +1100 Subject: [PATCH] feat: add sample for commit stats (#241) * 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 <8822365+busunkim96@users.noreply.github.com> Co-authored-by: skuruppu * refactor: make last_commit_stats public Co-authored-by: larkee Co-authored-by: Bu Sun Kim <8822365+busunkim96@users.noreply.github.com> Co-authored-by: skuruppu --- samples/samples/snippets.py | 42 +++++++++++++++++++++++++ samples/samples/snippets_test.py | 7 +++++ test.py | 53 -------------------------------- 3 files changed, 49 insertions(+), 53 deletions(-) delete mode 100644 test.py diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index f0379c0210..9a94e85a9b 100644 --- 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 index 237389c8b1..ee8c6ebe23 100644 --- 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 deleted file mode 100644 index 7888bbd090..0000000000 --- a/test.py +++ /dev/null @@ -1,53 +0,0 @@ -import base64 -import time -from google.cloud import spanner -from google.auth.credentials import AnonymousCredentials - -instance_id = 'test-instance' -database_id = 'test-db' - -spanner_client = spanner.Client( - project='test-project', - client_options={"api_endpoint": 'localhost:9010'}, - credentials=AnonymousCredentials() -) - -instance = spanner_client.instance(instance_id) -op = instance.create() -op.result() - -database = instance.database(database_id, ddl_statements=[ - "CREATE TABLE Test (id STRING(36) NOT NULL, megafield BYTES(MAX)) PRIMARY KEY (id)" -]) -op = database.create() -op.result() - -# This must be large enough that the SDK will split the megafield payload across two query chunks -# and try to recombine them, causing the error: -data = base64.standard_b64encode(("a" * 1000000).encode("utf8")) - -try: - with database.batch() as batch: - batch.insert( - table="Test", - columns=("id", "megafield"), - values=[ - (1, data), - ], - ) - - with database.snapshot() as snapshot: - toc = time.time() - results = snapshot.execute_sql( - "SELECT * FROM Test" - ) - tic = time.time() - - print("TIME: ", tic - toc) - - for row in results: - print("Id: ", row[0]) - print("Megafield: ", row[1][:100]) -finally: - database.drop() - instance.delete() \ No newline at end of file