Skip to content

feat(google-cloud-ndb): add x-goog-is-transaction routing header #18847

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nikhilchowdary wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(google-cloud-ndb): add x-goog-is-transaction routing header #18847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 46 additions & 8 deletions packages/google-cloud-ndb/google/cloud/ndb/_datastore_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,11 @@ def _datastore_lookup(keys, read_options, retries=None, timeout=None, metadata=(
keys=[key for key in keys],
read_options=read_options,
)
metadata = _add_routing_info(metadata, request)
metadata = _add_routing_info(
metadata,
request,
is_transaction=_is_read_options_transactional(read_options),
)

return make_call(
"lookup", request, retries=retries, timeout=timeout, metadata=metadata
Expand Down Expand Up @@ -889,7 +893,9 @@ def _datastore_commit(mutations, transaction, retries=None, timeout=None, metada
mutations=mutations,
transaction=transaction,
)
metadata = _add_routing_info(metadata, request)
metadata = _add_routing_info(
metadata, request, is_transaction=bool(transaction)
)

return make_call(
"commit", request, retries=retries, timeout=timeout, metadata=metadata
Expand Down Expand Up @@ -1013,7 +1019,7 @@ def _datastore_allocate_ids(keys, retries=None, timeout=None, metadata=()):
request = datastore_pb2.AllocateIdsRequest(
project_id=client.project, database_id=client.database, keys=keys
)
metadata = _add_routing_info(metadata, request)
metadata = _add_routing_info(metadata, request, is_transaction=False)

return make_call(
"allocate_ids", request, retries=retries, timeout=timeout, metadata=metadata
Expand Down Expand Up @@ -1076,7 +1082,7 @@ def _datastore_begin_transaction(read_only, retries=None, timeout=None, metadata
database_id=client.database,
transaction_options=options,
)
metadata = _add_routing_info(metadata, request)
metadata = _add_routing_info(metadata, request, is_transaction=True)

return make_call(
"begin_transaction",
Expand Down Expand Up @@ -1128,21 +1134,47 @@ def _datastore_rollback(transaction, retries=None, timeout=None, metadata=()):
database_id=client.database,
transaction=transaction,
)
metadata = _add_routing_info(metadata, request)
metadata = _add_routing_info(metadata, request, is_transaction=True)

return make_call(
"rollback", request, retries=retries, timeout=timeout, metadata=metadata
)


def _add_routing_info(metadata, request):
def _is_read_options_transactional(read_options):
"""Determines whether a ``ReadOptions`` message specifies a transaction."""
if read_options is None:
return False
return bool(read_options.transaction or "new_transaction" in read_options)


def _is_transaction_request(request):
"""Determines whether a Datastore RPC request is transactional."""
request_type = type(request).__name__
if request_type in ("BeginTransactionRequest", "RollbackRequest"):
return True
if request_type in ("AllocateIdsRequest", "ReserveIdsRequest"):
return False
if request_type == "CommitRequest":
return bool(
request.transaction
or "single_use_transaction" in request
or request.mode == datastore_pb2.CommitRequest.Mode.TRANSACTIONAL
)
return _is_read_options_transactional(getattr(request, "read_options", None))


def _add_routing_info(metadata, request, is_transaction=None):
"""Adds routing header info to the given metadata.

Args:
metadata (Sequence[Tuple[str, str]]): Strings which should be
sent along with the request as metadata. Not modified.
request (Any): An appropriate request object for the call, eg,
`entity_pb2.LookupRequest` for calling ``Lookup``.
is_transaction (Optional[bool]): Explicit indicator of whether the
request is part of a transaction. If ``None``, inferred from
``request``.

Returns:
Sequence[Tuple[str, str]]: Sequence with routing info added,
Expand All @@ -1156,9 +1188,15 @@ def _add_routing_info(metadata, request):
if request.database_id:
header_params["database_id"] = request.database_id

result = tuple(metadata)
if header_params:
return tuple(metadata) + (
result = result + (
gapic_v1.routing_header.to_grpc_metadata(header_params),
)

return tuple(metadata)
if is_transaction is None:
is_transaction = _is_transaction_request(request)

return result + (
("x-goog-is-transaction", "true" if is_transaction else "false"),
)
Comment on lines +1191 to +1202

@gemini-code-assist gemini-code-assist Bot Oct 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In performance-critical paths like Datastore RPC metadata generation, multiple tuple concatenations (e.g., result = result + (...)) introduce unnecessary copying overhead. Using itertools.chain to combine the metadata collections and converting to a tuple once at the end is much more efficient.

Suggested change
result = tuple(metadata)
if header_params:
return tuple(metadata) + (
result = result + (
gapic_v1.routing_header.to_grpc_metadata(header_params),
)
return tuple(metadata)
if is_transaction is None:
is_transaction = _is_transaction_request(request)
return result + (
("x-goog-is-transaction", "true" if is_transaction else "false"),
)
additional = []
if header_params:
additional.append(
gapic_v1.routing_header.to_grpc_metadata(header_params)
)
if is_transaction is None:
is_transaction = _is_transaction_request(request)
additional.append(
("x-goog-is-transaction", "true" if is_transaction else "false")
)
import itertools
return tuple(itertools.chain(metadata, additional))
References
  1. In performance-critical paths, avoid combining collections by converting them to tuples and concatenating them (e.g., tuple(a) + tuple(b)), as this introduces copying overhead. Instead, use lazy iterators like itertools.chain to iterate over multiple collections efficiently.

81 changes: 79 additions & 2 deletions packages/google-cloud-ndb/tests/unit/test__datastore_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -579,6 +579,7 @@ def test__datastore_lookup(datastore_pb2, context):
timeout=_api._DEFAULT_TIMEOUT,
metadata=(
("x-goog-request-params", "project_id=theproject&database_id=testdb"),
("x-goog-is-transaction", "false"),
),
)

Expand Down Expand Up @@ -1539,16 +1540,92 @@ def __init__(self, id=None, name=None):
)
def test__add_routing_info(project_id, database_id, expected):
expected_new_metadata = ("x-goog-request-params", expected)
expected_tx_metadata = ("x-goog-is-transaction", "false")
request = datastore_pb2.LookupRequest(
project_id=project_id, database_id=database_id
)
assert _api._add_routing_info((), request) == (expected_new_metadata,)
assert _api._add_routing_info((), request) == (
expected_new_metadata,
expected_tx_metadata,
)
assert _api._add_routing_info(("already=there",), request) == (
"already=there",
expected_new_metadata,
expected_tx_metadata,
)


def test__add_routing_info_no_request_info():
request = datastore_pb2.LookupRequest()
assert _api._add_routing_info((), request) == ()
assert _api._add_routing_info((), request) == (
("x-goog-is-transaction", "false"),
)


@pytest.mark.parametrize(
"request_pb,expected_is_tx",
[
(datastore_pb2.LookupRequest(), "false"),
(
datastore_pb2.LookupRequest(
read_options=datastore_pb2.ReadOptions(transaction=b"tx1")
),
"true",
),
(
datastore_pb2.LookupRequest(
read_options=datastore_pb2.ReadOptions(
new_transaction=datastore_pb2.TransactionOptions()
)
),
"true",
),
(datastore_pb2.RunQueryRequest(), "false"),
(
datastore_pb2.RunQueryRequest(
read_options=datastore_pb2.ReadOptions(transaction=b"tx1")
),
"true",
),
(datastore_pb2.RunAggregationQueryRequest(), "false"),
(
datastore_pb2.RunAggregationQueryRequest(
read_options=datastore_pb2.ReadOptions(transaction=b"tx1")
),
"true",
),
(datastore_pb2.BeginTransactionRequest(), "true"),
(datastore_pb2.RollbackRequest(), "true"),
(
datastore_pb2.CommitRequest(
mode=datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL
),
"false",
),
(
datastore_pb2.CommitRequest(
mode=datastore_pb2.CommitRequest.Mode.TRANSACTIONAL,
),
"true",
),
(
datastore_pb2.CommitRequest(
mode=datastore_pb2.CommitRequest.Mode.TRANSACTIONAL,
transaction=b"tx1",
),
"true",
),
(
datastore_pb2.CommitRequest(
single_use_transaction=datastore_pb2.TransactionOptions()
),
"true",
),
(datastore_pb2.AllocateIdsRequest(), "false"),
(datastore_pb2.ReserveIdsRequest(), "false"),
],
)
def test__add_routing_info_is_transaction_header(request_pb, expected_is_tx):
assert _api._add_routing_info((), request_pb) == (
("x-goog-is-transaction", expected_is_tx),
)
Loading