Skip to content

Commit f9fefad

Browse files
asthamohtaasthamohta
authored
feat: Leader Aware Routing (#899)
* changes * tests * Update client.py * Update test_client.py * Update connection.py * setting feature false * changes
1 parent 10a1351 commit f9fefad

21 files changed

+543
-93
lines changed

‎google/cloud/spanner_dbapi/connection.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,7 @@ def connect(
508508
pool=None,
509509
user_agent=None,
510510
client=None,
511+
route_to_leader_enabled=False,
511512
):
512513
"""Creates a connection to a Google Cloud Spanner database.
513514
@@ -544,6 +545,14 @@ def connect(
544545
:class:`~google.cloud.spanner_v1.Client`.
545546
:param client: (Optional) Custom user provided Client Object
546547
548+
:type route_to_leader_enabled: boolean
549+
:param route_to_leader_enabled:
550+
(Optional) Default False. Set route_to_leader_enabled as True to
551+
Enable leader aware routing. Enabling leader aware routing
552+
would route all requests in RW/PDML transactions to the
553+
leader region.
554+
555+
547556
:rtype: :class:`google.cloud.spanner_dbapi.connection.Connection`
548557
:returns: Connection object associated with the given Google Cloud Spanner
549558
resource.
@@ -556,11 +565,17 @@ def connect(
556565
)
557566
if isinstance(credentials, str):
558567
client = spanner.Client.from_service_account_json(
559-
credentials, project=project, client_info=client_info
568+
credentials,
569+
project=project,
570+
client_info=client_info,
571+
route_to_leader_enabled=False,
560572
)
561573
else:
562574
client = spanner.Client(
563-
project=project, credentials=credentials, client_info=client_info
575+
project=project,
576+
credentials=credentials,
577+
client_info=client_info,
578+
route_to_leader_enabled=False,
564579
)
565580
else:
566581
if project is not None and client.project != project:

‎google/cloud/spanner_v1/_helpers.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -292,3 +292,15 @@ def _metadata_with_prefix(prefix, **kw):
292292
List[Tuple[str, str]]: RPC metadata with supplied prefix
293293
"""
294294
return [("google-cloud-resource-prefix", prefix)]
295+
296+
297+
def _metadata_with_leader_aware_routing(value, **kw):
298+
"""Create RPC metadata containing a leader aware routing header
299+
300+
Args:
301+
value (bool): header value
302+
303+
Returns:
304+
List[Tuple[str, str]]: RPC metadata with leader aware routing header
305+
"""
306+
return ("x-goog-spanner-route-to-leader", str(value).lower())

‎google/cloud/spanner_v1/batch.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@
2020

2121
from google.cloud.spanner_v1._helpers import _SessionWrapper
2222
from google.cloud.spanner_v1._helpers import _make_list_value_pbs
23-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
23+
from google.cloud.spanner_v1._helpers import (
24+
_metadata_with_prefix,
25+
_metadata_with_leader_aware_routing,
26+
)
2427
from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
2528
from google.cloud.spanner_v1 import RequestOptions
2629

@@ -159,6 +162,10 @@ def commit(self, return_commit_stats=False, request_options=None):
159162
database = self._session._database
160163
api = database.spanner_api
161164
metadata = _metadata_with_prefix(database.name)
165+
if database._route_to_leader_enabled:
166+
metadata.append(
167+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
168+
)
162169
txn_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
163170
trace_attributes = {"num_mutations": len(self._mutations)}
164171

‎google/cloud/spanner_v1/client.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,13 @@ class Client(ClientWithProject):
114114
If a dict is provided, it must be of the same form as the protobuf
115115
message :class:`~google.cloud.spanner_v1.types.QueryOptions`
116116
117+
:type route_to_leader_enabled: boolean
118+
:param route_to_leader_enabled:
119+
(Optional) Default False. Set route_to_leader_enabled as True to
120+
Enable leader aware routing. Enabling leader aware routing
121+
would route all requests in RW/PDML transactions to the
122+
leader region.
123+
117124
:raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
118125
and ``admin`` are :data:`True`
119126
"""
@@ -132,6 +139,7 @@ def __init__(
132139
client_info=_CLIENT_INFO,
133140
client_options=None,
134141
query_options=None,
142+
route_to_leader_enabled=False,
135143
):
136144
self._emulator_host = _get_spanner_emulator_host()
137145

@@ -171,6 +179,8 @@ def __init__(
171179
):
172180
warnings.warn(_EMULATOR_HOST_HTTP_SCHEME)
173181

182+
self._route_to_leader_enabled = route_to_leader_enabled
183+
174184
@property
175185
def credentials(self):
176186
"""Getter for client's credentials.
@@ -242,6 +252,15 @@ def database_admin_api(self):
242252
)
243253
return self._database_admin_api
244254

255+
@property
256+
def route_to_leader_enabled(self):
257+
"""Getter for if read-write or pdml requests will be routed to leader.
258+
259+
:rtype: boolean
260+
:returns: If read-write requests will be routed to leader.
261+
"""
262+
return self._route_to_leader_enabled
263+
245264
def copy(self):
246265
"""Make a copy of this client.
247266

‎google/cloud/spanner_v1/database.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,10 @@
4444
from google.cloud.spanner_v1 import RequestOptions
4545
from google.cloud.spanner_v1 import SpannerClient
4646
from google.cloud.spanner_v1._helpers import _merge_query_options
47-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
47+
from google.cloud.spanner_v1._helpers import (
48+
_metadata_with_prefix,
49+
_metadata_with_leader_aware_routing,
50+
)
4851
from google.cloud.spanner_v1.batch import Batch
4952
from google.cloud.spanner_v1.keyset import KeySet
5053
from google.cloud.spanner_v1.pool import BurstyPool
@@ -155,6 +158,7 @@ def __init__(
155158
self._encryption_config = encryption_config
156159
self._database_dialect = database_dialect
157160
self._database_role = database_role
161+
self._route_to_leader_enabled = self._instance._client.route_to_leader_enabled
158162

159163
if pool is None:
160164
pool = BurstyPool(database_role=database_role)
@@ -565,6 +569,10 @@ def execute_partitioned_dml(
565569
)
566570

567571
metadata = _metadata_with_prefix(self.name)
572+
if self._route_to_leader_enabled:
573+
metadata.append(
574+
_metadata_with_leader_aware_routing(self._route_to_leader_enabled)
575+
)
568576

569577
def execute_pdml():
570578
with SessionCheckout(self._pool) as session:

‎google/cloud/spanner_v1/pool.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@
2020
from google.cloud.exceptions import NotFound
2121
from google.cloud.spanner_v1 import BatchCreateSessionsRequest
2222
from google.cloud.spanner_v1 import Session
23-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
23+
from google.cloud.spanner_v1._helpers import (
24+
_metadata_with_prefix,
25+
_metadata_with_leader_aware_routing,
26+
)
2427
from warnings import warn
2528

2629
_NOW = datetime.datetime.utcnow # unit tests may replace
@@ -191,6 +194,10 @@ def bind(self, database):
191194
self._database = database
192195
api = database.spanner_api
193196
metadata = _metadata_with_prefix(database.name)
197+
if database._route_to_leader_enabled:
198+
metadata.append(
199+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
200+
)
194201
self._database_role = self._database_role or self._database.database_role
195202
request = BatchCreateSessionsRequest(
196203
database=database.name,
@@ -402,6 +409,10 @@ def bind(self, database):
402409
self._database = database
403410
api = database.spanner_api
404411
metadata = _metadata_with_prefix(database.name)
412+
if database._route_to_leader_enabled:
413+
metadata.append(
414+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
415+
)
405416
created_session_count = 0
406417
self._database_role = self._database_role or self._database.database_role
407418

‎google/cloud/spanner_v1/session.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,10 @@
2626

2727
from google.cloud.spanner_v1 import ExecuteSqlRequest
2828
from google.cloud.spanner_v1 import CreateSessionRequest
29-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
29+
from google.cloud.spanner_v1._helpers import (
30+
_metadata_with_prefix,
31+
_metadata_with_leader_aware_routing,
32+
)
3033
from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
3134
from google.cloud.spanner_v1.batch import Batch
3235
from google.cloud.spanner_v1.snapshot import Snapshot
@@ -125,6 +128,12 @@ def create(self):
125128
raise ValueError("Session ID already set by back-end")
126129
api = self._database.spanner_api
127130
metadata = _metadata_with_prefix(self._database.name)
131+
if self._database._route_to_leader_enabled:
132+
metadata.append(
133+
_metadata_with_leader_aware_routing(
134+
self._database._route_to_leader_enabled
135+
)
136+
)
128137

129138
request = CreateSessionRequest(database=self._database.name)
130139
if self._database.database_role is not None:
@@ -153,6 +162,12 @@ def exists(self):
153162
return False
154163
api = self._database.spanner_api
155164
metadata = _metadata_with_prefix(self._database.name)
165+
if self._database._route_to_leader_enabled:
166+
metadata.append(
167+
_metadata_with_leader_aware_routing(
168+
self._database._route_to_leader_enabled
169+
)
170+
)
156171

157172
with trace_call("CloudSpanner.GetSession", self) as span:
158173
try:

‎google/cloud/spanner_v1/snapshot.py

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,10 @@
3131
from google.api_core import gapic_v1
3232
from google.cloud.spanner_v1._helpers import _make_value_pb
3333
from google.cloud.spanner_v1._helpers import _merge_query_options
34-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
34+
from google.cloud.spanner_v1._helpers import (
35+
_metadata_with_prefix,
36+
_metadata_with_leader_aware_routing,
37+
)
3538
from google.cloud.spanner_v1._helpers import _SessionWrapper
3639
from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
3740
from google.cloud.spanner_v1.streamed import StreamedResultSet
@@ -235,6 +238,10 @@ def read(
235238
database = self._session._database
236239
api = database.spanner_api
237240
metadata = _metadata_with_prefix(database.name)
241+
if not self._read_only and database._route_to_leader_enabled:
242+
metadata.append(
243+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
244+
)
238245

239246
if request_options is None:
240247
request_options = RequestOptions()
@@ -244,7 +251,7 @@ def read(
244251
if self._read_only:
245252
# Transaction tags are not supported for read only transactions.
246253
request_options.transaction_tag = None
247-
else:
254+
elif self.transaction_tag is not None:
248255
request_options.transaction_tag = self.transaction_tag
249256

250257
request = ReadRequest(
@@ -391,6 +398,10 @@ def execute_sql(
391398

392399
database = self._session._database
393400
metadata = _metadata_with_prefix(database.name)
401+
if not self._read_only and database._route_to_leader_enabled:
402+
metadata.append(
403+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
404+
)
394405

395406
api = database.spanner_api
396407

@@ -406,7 +417,7 @@ def execute_sql(
406417
if self._read_only:
407418
# Transaction tags are not supported for read only transactions.
408419
request_options.transaction_tag = None
409-
else:
420+
elif self.transaction_tag is not None:
410421
request_options.transaction_tag = self.transaction_tag
411422

412423
request = ExecuteSqlRequest(
@@ -527,6 +538,10 @@ def partition_read(
527538
database = self._session._database
528539
api = database.spanner_api
529540
metadata = _metadata_with_prefix(database.name)
541+
if database._route_to_leader_enabled:
542+
metadata.append(
543+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
544+
)
530545
transaction = self._make_txn_selector()
531546
partition_options = PartitionOptions(
532547
partition_size_bytes=partition_size_bytes, max_partitions=max_partitions
@@ -621,6 +636,10 @@ def partition_query(
621636
database = self._session._database
622637
api = database.spanner_api
623638
metadata = _metadata_with_prefix(database.name)
639+
if database._route_to_leader_enabled:
640+
metadata.append(
641+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
642+
)
624643
transaction = self._make_txn_selector()
625644
partition_options = PartitionOptions(
626645
partition_size_bytes=partition_size_bytes, max_partitions=max_partitions
@@ -766,6 +785,10 @@ def begin(self):
766785
database = self._session._database
767786
api = database.spanner_api
768787
metadata = _metadata_with_prefix(database.name)
788+
if not self._read_only and database._route_to_leader_enabled:
789+
metadata.append(
790+
(_metadata_with_leader_aware_routing(database._route_to_leader_enabled))
791+
)
769792
txn_selector = self._make_txn_selector()
770793
with trace_call("CloudSpanner.BeginTransaction", self._session):
771794
response = api.begin_transaction(

‎google/cloud/spanner_v1/transaction.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
_make_value_pb,
2222
_merge_query_options,
2323
_metadata_with_prefix,
24+
_metadata_with_leader_aware_routing,
2425
)
2526
from google.cloud.spanner_v1 import CommitRequest
2627
from google.cloud.spanner_v1 import ExecuteBatchDmlRequest
@@ -50,6 +51,7 @@ class Transaction(_SnapshotBase, _BatchBase):
5051
_multi_use = True
5152
_execute_sql_count = 0
5253
_lock = threading.Lock()
54+
_read_only = False
5355

5456
def __init__(self, session):
5557
if session._transaction is not None:
@@ -124,6 +126,10 @@ def begin(self):
124126
database = self._session._database
125127
api = database.spanner_api
126128
metadata = _metadata_with_prefix(database.name)
129+
if database._route_to_leader_enabled:
130+
metadata.append(
131+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
132+
)
127133
txn_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
128134
with trace_call("CloudSpanner.BeginTransaction", self._session):
129135
response = api.begin_transaction(
@@ -140,6 +146,12 @@ def rollback(self):
140146
database = self._session._database
141147
api = database.spanner_api
142148
metadata = _metadata_with_prefix(database.name)
149+
if database._route_to_leader_enabled:
150+
metadata.append(
151+
_metadata_with_leader_aware_routing(
152+
database._route_to_leader_enabled
153+
)
154+
)
143155
with trace_call("CloudSpanner.Rollback", self._session):
144156
api.rollback(
145157
session=self._session.name,
@@ -176,6 +188,10 @@ def commit(self, return_commit_stats=False, request_options=None):
176188
database = self._session._database
177189
api = database.spanner_api
178190
metadata = _metadata_with_prefix(database.name)
191+
if database._route_to_leader_enabled:
192+
metadata.append(
193+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
194+
)
179195
trace_attributes = {"num_mutations": len(self._mutations)}
180196

181197
if request_options is None:
@@ -294,6 +310,10 @@ def execute_update(
294310
params_pb = self._make_params_pb(params, param_types)
295311
database = self._session._database
296312
metadata = _metadata_with_prefix(database.name)
313+
if database._route_to_leader_enabled:
314+
metadata.append(
315+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
316+
)
297317
api = database.spanner_api
298318

299319
seqno, self._execute_sql_count = (
@@ -406,6 +426,10 @@ def batch_update(self, statements, request_options=None):
406426

407427
database = self._session._database
408428
metadata = _metadata_with_prefix(database.name)
429+
if database._route_to_leader_enabled:
430+
metadata.append(
431+
_metadata_with_leader_aware_routing(database._route_to_leader_enabled)
432+
)
409433
api = database.spanner_api
410434

411435
seqno, self._execute_sql_count = (

0 commit comments

Comments
 (0)