Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 434967e

Browse files
larkeelarkee
authored
feat: add support to log commit stats (googleapis#205)
* feat: add support for logging commit stats * test: add commit stats to CommitResponse * style: fix lint errors * refactor: remove log formatting * test: update info arg assertions * docs: document logger param * refactor: pass CommitStats via extra kwarg * fix: ensure logger is unused if commit fails Co-authored-by: larkee <[email protected]>
1 parent fd14b13 commit 434967e

File tree

10 files changed

+410
-56
lines changed

10 files changed

+410
-56
lines changed

‎google/cloud/spanner_v1/batch.py‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
"""Context manager for Cloud Spanner batched writes."""
1616

17+
from google.cloud.spanner_v1 import CommitRequest
1718
from google.cloud.spanner_v1 import Mutation
1819
from google.cloud.spanner_v1 import TransactionOptions
1920

@@ -123,6 +124,7 @@ class Batch(_BatchBase):
123124
"""
124125

125126
committed = None
127+
commit_stats = None
126128
"""Timestamp at which the batch was successfully committed."""
127129

128130
def _check_state(self):
@@ -136,9 +138,13 @@ def _check_state(self):
136138
if self.committed is not None:
137139
raise ValueError("Batch already committed")
138140

139-
def commit(self):
141+
def commit(self, return_commit_stats=False):
140142
"""Commit mutations to the database.
141143
144+
:type return_commit_stats: bool
145+
:param return_commit_stats:
146+
If true, the response will return commit stats which can be accessed though commit_stats.
147+
142148
:rtype: datetime
143149
:returns: timestamp of the committed changes.
144150
"""
@@ -148,14 +154,16 @@ def commit(self):
148154
metadata = _metadata_with_prefix(database.name)
149155
txn_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
150156
trace_attributes = {"num_mutations": len(self._mutations)}
157+
request = CommitRequest(
158+
session=self._session.name,
159+
mutations=self._mutations,
160+
single_use_transaction=txn_options,
161+
return_commit_stats=return_commit_stats,
162+
)
151163
with trace_call("CloudSpanner.Commit", self._session, trace_attributes):
152-
response = api.commit(
153-
session=self._session.name,
154-
mutations=self._mutations,
155-
single_use_transaction=txn_options,
156-
metadata=metadata,
157-
)
164+
response = api.commit(request=request, metadata=metadata,)
158165
self.committed = response.commit_timestamp
166+
self.commit_stats = response.commit_stats
159167
return self.committed
160168

161169
def __enter__(self):

‎google/cloud/spanner_v1/database.py‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import copy
1818
import functools
1919
import grpc
20+
import logging
2021
import re
2122
import threading
2223

@@ -95,11 +96,19 @@ class Database(object):
9596
:param pool: (Optional) session pool to be used by database. If not
9697
passed, the database will construct an instance of
9798
:class:`~google.cloud.spanner_v1.pool.BurstyPool`.
99+
100+
:type logger: `logging.Logger`
101+
:param logger: (Optional) a custom logger that is used if `log_commit_stats`
102+
is `True` to log commit statistics. If not passed, a logger
103+
will be created when needed that will log the commit statistics
104+
to stdout.
98105
"""
99106

100107
_spanner_api = None
101108

102-
def __init__(self, database_id, instance, ddl_statements=(), pool=None):
109+
def __init__(
110+
self, database_id, instance, ddl_statements=(), pool=None, logger=None
111+
):
103112
self.database_id = database_id
104113
self._instance = instance
105114
self._ddl_statements = _check_ddl_statements(ddl_statements)
@@ -109,6 +118,8 @@ def __init__(self, database_id, instance, ddl_statements=(), pool=None):
109118
self._restore_info = None
110119
self._version_retention_period = None
111120
self._earliest_version_time = None
121+
self.log_commit_stats = False
122+
self._logger = logger
112123

113124
if pool is None:
114125
pool = BurstyPool()
@@ -237,6 +248,25 @@ def ddl_statements(self):
237248
"""
238249
return self._ddl_statements
239250

251+
@property
252+
def logger(self):
253+
"""Logger used by the database.
254+
255+
The default logger will log commit stats at the log level INFO using
256+
`sys.stderr`.
257+
258+
:rtype: :class:`logging.Logger` or `None`
259+
:returns: the logger
260+
"""
261+
if self._logger is None:
262+
self._logger = logging.getLogger(self.name)
263+
self._logger.setLevel(logging.INFO)
264+
265+
ch = logging.StreamHandler()
266+
ch.setLevel(logging.INFO)
267+
self._logger.addHandler(ch)
268+
return self._logger
269+
240270
@property
241271
def spanner_api(self):
242272
"""Helper for session-related API calls."""
@@ -647,8 +677,13 @@ def __exit__(self, exc_type, exc_val, exc_tb):
647677
"""End ``with`` block."""
648678
try:
649679
if exc_type is None:
650-
self._batch.commit()
680+
self._batch.commit(return_commit_stats=self._database.log_commit_stats)
651681
finally:
682+
if self._database.log_commit_stats and self._batch.commit_stats:
683+
self._database.logger.info(
684+
"CommitStats: {}".format(self._batch.commit_stats),
685+
extra={"commit_stats": self._batch.commit_stats},
686+
)
652687
self._database._pool.put(self._session)
653688

654689

‎google/cloud/spanner_v1/instance.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,7 @@ def delete(self):
357357

358358
api.delete_instance(name=self.name, metadata=metadata)
359359

360-
def database(self, database_id, ddl_statements=(), pool=None):
360+
def database(self, database_id, ddl_statements=(), pool=None, logger=None):
361361
"""Factory to create a database within this instance.
362362
363363
:type database_id: str
@@ -371,10 +371,18 @@ def database(self, database_id, ddl_statements=(), pool=None):
371371
:class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`.
372372
:param pool: (Optional) session pool to be used by database.
373373
374+
:type logger: `logging.Logger`
375+
:param logger: (Optional) a custom logger that is used if `log_commit_stats`
376+
is `True` to log commit statistics. If not passed, a logger
377+
will be created when needed that will log the commit statistics
378+
to stdout.
379+
374380
:rtype: :class:`~google.cloud.spanner_v1.database.Database`
375381
:returns: a database owned by this instance.
376382
"""
377-
return Database(database_id, self, ddl_statements=ddl_statements, pool=pool)
383+
return Database(
384+
database_id, self, ddl_statements=ddl_statements, pool=pool, logger=logger
385+
)
378386

379387
def list_databases(self, page_size=None):
380388
"""List databases for the instance.

‎google/cloud/spanner_v1/session.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -349,14 +349,19 @@ def run_in_transaction(self, func, *args, **kw):
349349
raise
350350

351351
try:
352-
txn.commit()
352+
txn.commit(return_commit_stats=self._database.log_commit_stats)
353353
except Aborted as exc:
354354
del self._transaction
355355
_delay_until_retry(exc, deadline, attempts)
356356
except GoogleAPICallError:
357357
del self._transaction
358358
raise
359359
else:
360+
if self._database.log_commit_stats and txn.commit_stats:
361+
self._database.logger.info(
362+
"CommitStats: {}".format(txn.commit_stats),
363+
extra={"commit_stats": txn.commit_stats},
364+
)
360365
return return_value
361366

362367

‎google/cloud/spanner_v1/transaction.py‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
_merge_query_options,
2222
_metadata_with_prefix,
2323
)
24+
from google.cloud.spanner_v1 import CommitRequest
2425
from google.cloud.spanner_v1 import ExecuteBatchDmlRequest
2526
from google.cloud.spanner_v1 import ExecuteSqlRequest
2627
from google.cloud.spanner_v1 import TransactionSelector
@@ -42,6 +43,7 @@ class Transaction(_SnapshotBase, _BatchBase):
4243
committed = None
4344
"""Timestamp at which the transaction was successfully committed."""
4445
rolled_back = False
46+
commit_stats = None
4547
_multi_use = True
4648
_execute_sql_count = 0
4749

@@ -119,9 +121,13 @@ def rollback(self):
119121
self.rolled_back = True
120122
del self._session._transaction
121123

122-
def commit(self):
124+
def commit(self, return_commit_stats=False):
123125
"""Commit mutations to the database.
124126
127+
:type return_commit_stats: bool
128+
:param return_commit_stats:
129+
If true, the response will return commit stats which can be accessed though commit_stats.
130+
125131
:rtype: datetime
126132
:returns: timestamp of the committed changes.
127133
:raises ValueError: if there are no mutations to commit.
@@ -132,14 +138,17 @@ def commit(self):
132138
api = database.spanner_api
133139
metadata = _metadata_with_prefix(database.name)
134140
trace_attributes = {"num_mutations": len(self._mutations)}
141+
request = CommitRequest(
142+
session=self._session.name,
143+
mutations=self._mutations,
144+
transaction_id=self._transaction_id,
145+
return_commit_stats=return_commit_stats,
146+
)
135147
with trace_call("CloudSpanner.Commit", self._session, trace_attributes):
136-
response = api.commit(
137-
session=self._session.name,
138-
mutations=self._mutations,
139-
transaction_id=self._transaction_id,
140-
metadata=metadata,
141-
)
148+
response = api.commit(request=request, metadata=metadata,)
142149
self.committed = response.commit_timestamp
150+
if return_commit_stats:
151+
self.commit_stats = response.commit_stats
143152
del self._session._transaction
144153
return self.committed
145154

‎tests/unit/test_batch.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -339,17 +339,17 @@ def __init__(self, **kwargs):
339339
self.__dict__.update(**kwargs)
340340

341341
def commit(
342-
self,
343-
session,
344-
mutations,
345-
transaction_id="",
346-
single_use_transaction=None,
347-
metadata=None,
342+
self, request=None, metadata=None,
348343
):
349344
from google.api_core.exceptions import Unknown
350345

351-
assert transaction_id == ""
352-
self._committed = (session, mutations, single_use_transaction, metadata)
346+
assert request.transaction_id == b""
347+
self._committed = (
348+
request.session,
349+
request.mutations,
350+
request.single_use_transaction,
351+
metadata,
352+
)
353353
if self._rpc_error:
354354
raise Unknown("error")
355355
return self._commit_response

0 commit comments

Comments
 (0)