From 2f70017fed4ab59192e80e1a224be5395e4b9912 Mon Sep 17 00:00:00 2001 From: Astha Mohta Date: Fri, 24 Feb 2023 02:04:31 +0530 Subject: [PATCH 01/19] changes --- samples/samples/snippets.py | 78 ++++++++++++++++++++++++++++++++ samples/samples/snippets_test.py | 7 +++ 2 files changed, 85 insertions(+) diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 82fb95a0dd..47f48aa8ce 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -31,6 +31,7 @@ from google.cloud import spanner from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin from google.cloud.spanner_v1 import param_types +from google.cloud.spanner_v1 import TransactionTypes from google.type import expr_pb2 from google.iam.v1 import policy_pb2 from google.cloud.spanner_v1.data_types import JsonObject @@ -2664,6 +2665,78 @@ def drop_sequence(instance_id, database_id): # [END spanner_drop_sequence] + +def directed_read_options( + instance_id, + database_id, +): + """ + Shows how to run an execute sql request with directed read options. + Only one of exclude_replicas or include_replicas can be set + Each accepts a list of replicaSelections which contains location and type + * `location` - The location must be one of the regions within the + multi-region configuration of your database. + * `type_` - The type of the replica + Some examples of using replica_selectors are: + * `location:us-east1` --> The "us-east1" replica(s) of any available type + will be used to process the request. + * `type:READ_ONLY` --> The "READ_ONLY" type replica(s) in nearest + available location will be used to process the + request. + * `location:us-east1 type:READ_ONLY` --> The "READ_ONLY" type replica(s) + in location "us-east1" will be used to process + the request. + include_replicas also contains an option for auto_failover which when set + Spanner will not route requests to a replica outside the + include_replicas list when all the specified replicas are unavailable + or unhealthy. The default value is `false` + """ + # [START spanner_directed_read] + # instance_id = "your-spanner-instance" + # database_id = "your-spanner-db-id" + + directed_read_options_for_client = { + "exclude_replicas": { + "replica_selections": [ + { + "location": "us-east4", + }, + ], + }, + } + + # directed_read_options can be set at client level and will be used in all + # read-only transaction requests + spanner_client = spanner.Client( + directed_read_options=directed_read_options_for_client + ) + instance = spanner_client.instance(instance_id) + database = instance.database(database_id) + + directed_read_options_for_request = { + "include_replicas": { + "replica_selections": [ + { + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover": True, + }, + } + + with database.snapshot() as snapshot: + # Read rows while passing directed_read_options directly to the query. + # These will override the options passed at Client level. + results = snapshot.execute_sql( + "SELECT SingerId, AlbumId, AlbumTitle FROM Albums", + directed_read_options=directed_read_options_for_request, + ) + + for row in results: + print("SingerId: {}, AlbumId: {}, AlbumTitle: {}".format(*row)) + # [END spanner_directed_read] + + if __name__ == "__main__": # noqa: C901 parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter @@ -2802,6 +2875,9 @@ def drop_sequence(instance_id, database_id): "--database_role", default="new_parent" ) enable_fine_grained_access_parser.add_argument("--title", default="condition title") + subparsers.add_parser( + "directed_read_options", help=directed_read_options.__doc__ + ) args = parser.parse_args() @@ -2931,3 +3007,5 @@ def drop_sequence(instance_id, database_id): args.database_role, args.title, ) + elif args.command == "directed_read_options": + directed_read_options(args.instance_id, args.database_id) diff --git a/samples/samples/snippets_test.py b/samples/samples/snippets_test.py index 22b5b6f944..99eed700bf 100644 --- a/samples/samples/snippets_test.py +++ b/samples/samples/snippets_test.py @@ -845,3 +845,10 @@ def test_drop_sequence(capsys, instance_id, bit_reverse_sequence_database): "Altered Customers table to drop DEFAULT from CustomerId column and dropped the Seq sequence on database" in out ) + + +@pytest.mark.dependency(depends=["insert_data"]) +def test_directed_read_options(capsys, instance_id, sample_database): + snippets.directed_read_options(instance_id, sample_database.database_id) + out, _ = capsys.readouterr() + assert "SingerId: 1, AlbumId: 1, AlbumTitle: Total Junk" in out From 88c70ceaa81f90e917fe8bcbf2eb56ce3af4d99b Mon Sep 17 00:00:00 2001 From: Astha Mohta Date: Fri, 24 Feb 2023 02:04:43 +0530 Subject: [PATCH 02/19] changes --- google/cloud/spanner_v1/__init__.py | 5 +- google/cloud/spanner_v1/_helpers.py | 43 ++++++ google/cloud/spanner_v1/client.py | 17 +++ google/cloud/spanner_v1/database.py | 1 + google/cloud/spanner_v1/snapshot.py | 32 ++++ tests/system/test_session_api.py | 48 ++++++ tests/unit/spanner_dbapi/test_connection.py | 3 +- tests/unit/test__helpers.py | 153 ++++++++++++++++++++ tests/unit/test_client.py | 25 ++++ tests/unit/test_database.py | 27 +++- tests/unit/test_instance.py | 1 + tests/unit/test_snapshot.py | 88 ++++++++++- tests/unit/test_spanner.py | 42 +++++- tests/unit/test_transaction.py | 2 + 14 files changed, 478 insertions(+), 9 deletions(-) diff --git a/google/cloud/spanner_v1/__init__.py b/google/cloud/spanner_v1/__init__.py index 039919563f..1552213f3b 100644 --- a/google/cloud/spanner_v1/__init__.py +++ b/google/cloud/spanner_v1/__init__.py @@ -38,6 +38,7 @@ from .types.spanner import CommitRequest from .types.spanner import CreateSessionRequest from .types.spanner import DeleteSessionRequest +from .types.spanner import DirectedReadOptions from .types.spanner import ExecuteBatchDmlRequest from .types.spanner import ExecuteBatchDmlResponse from .types.spanner import ExecuteSqlRequest @@ -77,7 +78,7 @@ This value can only be used for timestamp columns that have set the option ``(allow_commit_timestamp=true)`` in the schema. """ - +TransactionTypes = DirectedReadOptions.ReplicaSelection.Type __all__ = ( # google.cloud.spanner_v1 @@ -96,6 +97,7 @@ "TransactionPingingPool", # local "COMMIT_TIMESTAMP", + "TransactionTypes", # google.cloud.spanner_v1.types "BatchCreateSessionsRequest", "BatchCreateSessionsResponse", @@ -104,6 +106,7 @@ "CommitResponse", "CreateSessionRequest", "DeleteSessionRequest", + "DirectedReadOptions", "ExecuteBatchDmlRequest", "ExecuteBatchDmlResponse", "ExecuteSqlRequest", diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py index e0e2bfdbd0..3b62515da0 100644 --- a/google/cloud/spanner_v1/_helpers.py +++ b/google/cloud/spanner_v1/_helpers.py @@ -28,6 +28,8 @@ from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import JsonObject +from google.cloud.spanner_v1 import DirectedReadOptions +from google.api_core.exceptions import InvalidArgument # Validation error messages NUMERIC_MAX_SCALE_ERR_MSG = ( @@ -358,3 +360,44 @@ def _metadata_with_leader_aware_routing(value, **kw): List[Tuple[str, str]]: RPC metadata with leader aware routing header """ return ("x-goog-spanner-route-to-leader", str(value).lower()) + + +def verify_directed_read_options(directed_read_options): + if type(directed_read_options) == dict: + if ( + "include_replicas" in directed_read_options.keys() + and "exclude_replicas" in directed_read_options.keys() + ): + raise InvalidArgument( + "Only one of include_replicas or exclude_replicas can be set" + ) + if ( + "include_replicas" in directed_read_options.keys() + and "replica_selections" in directed_read_options["include_replicas"].keys() + and len(directed_read_options["include_replicas"]["replica_selections"]) + > 10 + ) or ( + "exclude_replicas" in directed_read_options.keys() + and "replica_selections" in directed_read_options["exclude_replicas"].keys() + and len(directed_read_options["exclude_replicas"]["replica_selections"]) + > 10 + ): + raise InvalidArgument("Maximum length of replica selection allowed is 10") + elif isinstance(directed_read_options, DirectedReadOptions): + if ( + directed_read_options.include_replicas is not None + and directed_read_options.exclude_replicas is not None + ): + raise InvalidArgument( + "Only one of include_replicas or exclude_replicas can be set" + ) + if ( + directed_read_options.include_replicas is not None + and directed_read_options.include_replicas.replica_selections is not None + and len(directed_read_options.include_replicas.replica_selections) > 10 + ) or ( + directed_read_options.include_replicas is not None + and directed_read_options.exclude_replicas.replica_selections is not None + and len(directed_read_options.exclude_replicas.replica_selections) > 10 + ): + raise InvalidArgument("Maximum length of replica selection allowed is 10") diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index a0e848228b..1b5d971e4a 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -47,6 +47,7 @@ from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1._helpers import _merge_query_options from google.cloud.spanner_v1._helpers import _metadata_with_prefix +from google.cloud.spanner_v1._helpers import verify_directed_read_options from google.cloud.spanner_v1.instance import Instance _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__) @@ -139,6 +140,7 @@ def __init__( client_options=None, query_options=None, route_to_leader_enabled=True, + directed_read_options=None, ): self._emulator_host = _get_spanner_emulator_host() @@ -179,6 +181,8 @@ def __init__( warnings.warn(_EMULATOR_HOST_HTTP_SCHEME) self._route_to_leader_enabled = route_to_leader_enabled + verify_directed_read_options(directed_read_options) + self._directed_read_options = directed_read_options @property def credentials(self): @@ -260,6 +264,16 @@ def route_to_leader_enabled(self): """ return self._route_to_leader_enabled + + def directed_read_options(self): + """Getter for client's credentials. + + :rtype: + :class:`Credentials ` + :returns: The credentials stored on the client. + """ + return self._directed_read_options + def copy(self): """Make a copy of this client. @@ -383,3 +397,6 @@ def list_instances(self, filter_="", page_size=None): request=request, metadata=metadata ) return page_iter + + def set_directed_read_options(self, directed_read_options): + self._directed_read_options = directed_read_options diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 1d211f7d6d..81c144833e 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -166,6 +166,7 @@ def __init__( self._route_to_leader_enabled = self._instance._client.route_to_leader_enabled self._enable_drop_protection = enable_drop_protection self._reconciling = False + self._directed_read_options = self._instance._client.directed_read_options if pool is None: pool = BurstyPool(database_role=database_role) diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index 573042aa11..b2223b1831 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -28,6 +28,7 @@ from google.api_core.exceptions import InternalServerError from google.api_core.exceptions import ServiceUnavailable from google.api_core.exceptions import InvalidArgument +from google.api_core.exceptions import BadRequest from google.api_core import gapic_v1 from google.cloud.spanner_v1._helpers import ( _make_value_pb, @@ -41,6 +42,7 @@ from google.cloud.spanner_v1._opentelemetry_tracing import trace_call from google.cloud.spanner_v1.streamed import StreamedResultSet from google.cloud.spanner_v1 import RequestOptions +from google.cloud.spanner_v1._helpers import verify_directed_read_options _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES = ( "RST_STREAM", @@ -176,6 +178,7 @@ def read( *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, + directed_read_options=None, ): """Perform a ``StreamingRead`` API request for rows in a table. @@ -255,6 +258,19 @@ def read( request_options.transaction_tag = None elif self.transaction_tag is not None: request_options.transaction_tag = self.transaction_tag + if ( + directed_read_options is None + and database._directed_read_options is not None + ): + directed_read_options = database._directed_read_options + else: + if self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag + if directed_read_options is not None: + raise BadRequest( + "directed_read_options can't be set for readWrite transactions or partitioned dml requests" + ) + verify_directed_read_options(directed_read_options) request = ReadRequest( session=self._session.name, @@ -266,6 +282,7 @@ def read( partition_token=partition, request_options=request_options, data_boost_enabled=data_boost_enabled, + directed_read_options=directed_read_options, ) restart = functools.partial( api.streaming_read, @@ -322,6 +339,7 @@ def execute_sql( retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, data_boost_enabled=False, + directed_read_options=None, ): """Perform an ``ExecuteStreamingSql`` API request. @@ -421,6 +439,19 @@ def execute_sql( request_options.transaction_tag = None elif self.transaction_tag is not None: request_options.transaction_tag = self.transaction_tag + if ( + directed_read_options is None + and database._directed_read_options is not None + ): + directed_read_options = database._directed_read_options + else: + if self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag + if directed_read_options is not None: + raise BadRequest( + "directed_read_options can't be set for readWrite transactions or partitioned dml requests" + ) + verify_directed_read_options(directed_read_options) request = ExecuteSqlRequest( session=self._session.name, @@ -433,6 +464,7 @@ def execute_sql( query_options=query_options, request_options=request_options, data_boost_enabled=data_boost_enabled, + directed_read_options=directed_read_options, ) restart = functools.partial( api.execute_streaming_sql, diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index 7d58324b04..b6c474af78 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -28,6 +28,7 @@ from google.cloud import spanner_v1 from google.cloud.spanner_admin_database_v1 import DatabaseDialect from google.cloud._helpers import UTC +from google.cloud.spanner_v1 import TransactionTypes from google.cloud.spanner_v1.data_types import JsonObject from tests import _helpers as ot_helpers from . import _helpers @@ -57,6 +58,17 @@ JSON_2 = JsonObject( {"sample_object": {"name": "Anamika", "id": 2635}}, ) +DIRECTED_READ_OPTIONS = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover": True, + }, +} COUNTERS_TABLE = "counters" COUNTERS_COLUMNS = ("name", "value") @@ -502,6 +514,29 @@ def test_batch_insert_w_commit_timestamp(sessions_database, not_postgres): assert not deleted +def test_snapshot_read_w_directed_read_options(sessions_database, not_postgres): + table = "users_history" + columns = ["id", "commit_ts", "name", "email", "deleted"] + user_id = 1234 + name = "phred" + email = "phred@example.com" + row_data = [[user_id, spanner_v1.COMMIT_TIMESTAMP, name, email, False]] + sd = _sample_data + + with sessions_database.batch() as batch: + batch.delete(table, sd.ALL) + batch.insert(table, columns, row_data) + + with sessions_database.snapshot() as snapshot: + rows = list( + snapshot.read( + table, columns, sd.ALL, directed_read_options=DIRECTED_READ_OPTIONS + ) + ) + + assert len(rows) == 1 + + @_helpers.retry_mabye_aborted_txn def test_transaction_read_and_insert_then_rollback( sessions_database, @@ -1918,6 +1953,19 @@ def test_execute_sql_w_manual_consume(sessions_database): assert streamed._pending_chunk is None +def test_execute_sql_w_directed_read_options(sessions_database, not_postgres): + sd = _sample_data + row_count = 3000 + committed = _set_up_table(sessions_database, row_count) + + with sessions_database.snapshot() as snapshot: + streamed = snapshot.execute_sql( + sd.SQL, directed_read_options=DIRECTED_READ_OPTIONS + ) + + assert len(list(streamed)) == 3000 + + def _check_sql_results( database, sql, diff --git a/tests/unit/spanner_dbapi/test_connection.py b/tests/unit/spanner_dbapi/test_connection.py index 1628f84062..a6fbd323fc 100644 --- a/tests/unit/spanner_dbapi/test_connection.py +++ b/tests/unit/spanner_dbapi/test_connection.py @@ -47,7 +47,8 @@ def _make_connection(self, **kwargs): from google.cloud.spanner_v1.client import Client # We don't need a real Client object to test the constructor - instance = Instance(INSTANCE, client=Client) + client = Client() + instance = Instance(INSTANCE, client=client) database = instance.database(DATABASE) return Connection(instance, database, **kwargs) diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 0e0ec903a2..54499738e7 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -761,3 +761,156 @@ def test(self): self.assertEqual( metadata, ("x-goog-spanner-route-to-leader", str(value).lower()) ) + + +class Test_verify_directed_read_options(unittest.TestCase): + def _call_fut(self, directed_read_options): + from google.cloud.spanner_v1._helpers import verify_directed_read_options + + verify_directed_read_options(directed_read_options) + + def test_dict_include_exclude_replica_set_error(self): + from google.api_core.exceptions import InvalidArgument + + directed_read_options = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + }, + ], + }, + "exclude_replicas": { + "replica_selections": [ + { + "location": "us-east1", + }, + ], + }, + } + with self.assertRaises(InvalidArgument): + self._call_fut(directed_read_options) + + def test_dict_greater_than_ten_replica_selections_error(self): + from google.api_core.exceptions import InvalidArgument + + directed_read_options = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + ], + }, + } + with self.assertRaises(InvalidArgument): + self._call_fut(directed_read_options) + + def test_object_include_exclude_replica_set_error(self): + from google.api_core.exceptions import InvalidArgument + from google.cloud.spanner_v1 import DirectedReadOptions + + directed_read_options = DirectedReadOptions( + { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + }, + ], + }, + "exclude_replicas": { + "replica_selections": [ + { + "location": "us-east1", + }, + ], + }, + } + ) + with self.assertRaises(InvalidArgument): + self._call_fut(directed_read_options) + + def test_dict_greater_than_ten_replica_selections_error(self): + from google.api_core.exceptions import InvalidArgument + from google.cloud.spanner_v1 import DirectedReadOptions + + directed_read_options = DirectedReadOptions( + { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + ], + }, + } + ) + with self.assertRaises(InvalidArgument): + self._call_fut(directed_read_options) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index ed79271a96..4c8f1449c1 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -15,6 +15,7 @@ import unittest import mock +from google.cloud.spanner_v1 import TransactionTypes def _make_credentials(): @@ -41,6 +42,17 @@ class TestClient(unittest.TestCase): LABELS = {"test": "true"} TIMEOUT_SECONDS = 80 LEADER_OPTIONS = ["leader1", "leader2"] + DIRECTED_READ_OPTIONS = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover": True, + }, + } def _get_target_class(self): from google.cloud import spanner @@ -60,6 +72,7 @@ def _constructor_test_helper( query_options=None, expected_query_options=None, route_to_leader_enabled=True, + directed_read_options=None, ): import google.api_core.client_options from google.cloud.spanner_v1 import client as MUT @@ -85,6 +98,7 @@ def _constructor_test_helper( project=self.PROJECT, credentials=creds, query_options=query_options, + directed_read_options=directed_read_options, **kwargs ) @@ -113,6 +127,8 @@ def _constructor_test_helper( self.assertEqual(client.route_to_leader_enabled, route_to_leader_enabled) else: self.assertFalse(client.route_to_leader_enabled) + if directed_read_options is not None: + self.assertEqual(client.directed_read_options, directed_read_options) @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host") @mock.patch("warnings.warn") @@ -226,6 +242,15 @@ def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats) expected_query_options=expected_query_options, ) + def test_constructor_w_directed_read_options(self): + from google.cloud.spanner_v1 import client as MUT + + expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,) + creds = _make_credentials() + self._constructor_test_helper( + expected_scopes, creds, directed_read_options=self.DIRECTED_READ_OPTIONS + ) + def test_constructor_route_to_leader_disbled(self): from google.cloud.spanner_v1 import client as MUT diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 5a6abf8084..7611904f25 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -22,7 +22,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask -from google.cloud.spanner_v1 import RequestOptions +from google.cloud.spanner_v1 import RequestOptions, TransactionTypes DML_WO_PARAM = """ DELETE FROM citizens @@ -35,6 +35,17 @@ PARAMS = {"age": 30} PARAM_TYPES = {"age": INT64} MODE = 2 # PROFILE +DIRECTED_READ_OPTIONS = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover": True, + }, +} def _make_credentials(): # pragma: NO COVER @@ -199,6 +210,16 @@ def test_ctor_w_encryption_config(self): self.assertIs(database._instance, instance) self.assertEqual(database._encryption_config, encryption_config) + def test_ctor_w_directed_read_options(self): + client = _Client(directed_read_options=DIRECTED_READ_OPTIONS) + instance = _Instance(self.INSTANCE_NAME, client=client) + database = self._make_one( + self.DATABASE_ID, instance, database_role=self.DATABASE_ROLE + ) + self.assertEqual(database.database_id, self.DATABASE_ID) + self.assertIs(database._instance, instance) + self.assertEqual(database._directed_read_options, DIRECTED_READ_OPTIONS) + def test_from_pb_bad_database_name(self): from google.cloud.spanner_admin_database_v1 import Database @@ -2690,7 +2711,7 @@ def _make_instance_api(): class _Client(object): - def __init__(self, project=TestDatabase.PROJECT_ID, route_to_leader_enabled=True): + def __init__(self, project=TestDatabase.PROJECT_ID, route_to_leader_enabled=True, directed_read_options=None): from google.cloud.spanner_v1 import ExecuteSqlRequest self.project = project @@ -2701,6 +2722,7 @@ def __init__(self, project=TestDatabase.PROJECT_ID, route_to_leader_enabled=True self._client_options = mock.Mock() self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1") self.route_to_leader_enabled = route_to_leader_enabled + self.directed_read_options = directed_read_options class _Instance(object): @@ -2727,6 +2749,7 @@ def __init__(self, name, instance=None): from logging import Logger self.logger = mock.create_autospec(Logger, instance=True) + self._directed_read_options = None class _Pool(object): diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 0a7dbccb81..503e8f83bd 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -1016,6 +1016,7 @@ def __init__(self, project, timeout_seconds=None): self.project_name = "projects/" + self.project self.timeout_seconds = timeout_seconds self.route_to_leader_enabled = True + self.directed_read_options = None def copy(self): from copy import deepcopy diff --git a/tests/unit/test_snapshot.py b/tests/unit/test_snapshot.py index 5d2afb4fe6..d80edee3f7 100644 --- a/tests/unit/test_snapshot.py +++ b/tests/unit/test_snapshot.py @@ -16,7 +16,7 @@ from google.api_core import gapic_v1 import mock -from google.cloud.spanner_v1 import RequestOptions +from google.cloud.spanner_v1 import RequestOptions, TransactionTypes from tests._helpers import ( OpenTelemetryBase, StatusCode, @@ -46,6 +46,26 @@ "db.instance": "testing", "net.host.name": "spanner.googleapis.com", } +DIRECTED_READ_OPTIONS = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover": True, + }, +} +DIRECTED_READ_OPTIONS_FOR_CLIENT = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-east1", + }, + ], + }, +} class Test_restart_on_unavailable(OpenTelemetryBase): @@ -603,6 +623,8 @@ def _read_helper( timeout=gapic_v1.method.DEFAULT, retry=gapic_v1.method.DEFAULT, request_options=None, + directed_read_options=None, + directed_read_options_at_client_level=None, ): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ( @@ -642,7 +664,9 @@ def _read_helper( keyset = KeySet(keys=KEYS) INDEX = "email-address-index" LIMIT = 20 - database = _Database() + database = _Database( + directed_read_options=directed_read_options_at_client_level + ) api = database.spanner_api = self._make_spanner_api() api.streaming_read.return_value = _MockIterator(*result_sets) session = _Session(database) @@ -667,6 +691,7 @@ def _read_helper( retry=retry, timeout=timeout, request_options=request_options, + directed_read_options=directed_read_options, ) else: result_set = derived.read( @@ -678,6 +703,7 @@ def _read_helper( retry=retry, timeout=timeout, request_options=request_options, + directed_read_options=directed_read_options, ) self.assertEqual(derived._read_request_count, count + 1) @@ -712,6 +738,12 @@ def _read_helper( expected_request_options = request_options expected_request_options.transaction_tag = None + expected_directed_read_options = ( + directed_read_options + if directed_read_options is not None + else directed_read_options_at_client_level + ) + expected_request = ReadRequest( session=self.SESSION_NAME, table=TABLE_NAME, @@ -722,6 +754,7 @@ def _read_helper( limit=expected_limit, partition_token=partition, request_options=expected_request_options, + directed_read_options=expected_directed_read_options, ) api.streaming_read.assert_called_once_with( request=expected_request, @@ -797,6 +830,22 @@ def test_read_w_timeout_and_retry_params(self): multi_use=True, first=False, retry=Retry(deadline=60), timeout=2.0 ) + def test_read_w_directed_read_options(self): + self._read_helper(multi_use=False, directed_read_options=DIRECTED_READ_OPTIONS) + + def test_read_w_directed_read_options_at_client_level(self): + self._read_helper( + multi_use=False, + directed_read_options_at_client_level=DIRECTED_READ_OPTIONS_FOR_CLIENT, + ) + + def test_read_w_directed_read_options_override(self): + self._read_helper( + multi_use=False, + directed_read_options=DIRECTED_READ_OPTIONS, + directed_read_options_at_client_level=DIRECTED_READ_OPTIONS_FOR_CLIENT, + ) + def test_execute_sql_other_error(self): database = _Database() database.spanner_api = self._make_spanner_api() @@ -836,6 +885,8 @@ def _execute_sql_helper( request_options=None, timeout=gapic_v1.method.DEFAULT, retry=gapic_v1.method.DEFAULT, + directed_read_options=None, + directed_read_options_at_client_level=None, ): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ( @@ -876,7 +927,9 @@ def _execute_sql_helper( for i in range(len(result_sets)): result_sets[i].values.extend(VALUE_PBS[i]) iterator = _MockIterator(*result_sets) - database = _Database() + database = _Database( + directed_read_options=directed_read_options_at_client_level + ) api = database.spanner_api = self._make_spanner_api() api.execute_streaming_sql.return_value = iterator session = _Session(database) @@ -902,6 +955,7 @@ def _execute_sql_helper( partition=partition, retry=retry, timeout=timeout, + directed_read_options=directed_read_options, ) self.assertEqual(derived._read_request_count, count + 1) @@ -942,6 +996,12 @@ def _execute_sql_helper( expected_request_options = request_options expected_request_options.transaction_tag = None + expected_directed_read_options = ( + directed_read_options + if directed_read_options is not None + else directed_read_options_at_client_level + ) + expected_request = ExecuteSqlRequest( session=self.SESSION_NAME, sql=SQL_QUERY_WITH_PARAM, @@ -953,6 +1013,7 @@ def _execute_sql_helper( request_options=expected_request_options, partition_token=partition, seqno=sql_count, + directed_read_options=expected_directed_read_options, ) api.execute_streaming_sql.assert_called_once_with( request=expected_request, @@ -1039,6 +1100,24 @@ def test_execute_sql_w_incorrect_tag_dictionary_error(self): with self.assertRaises(ValueError): self._execute_sql_helper(multi_use=False, request_options=request_options) + def test_execute_sql_w_directed_read_options(self): + self._execute_sql_helper( + multi_use=False, directed_read_options=DIRECTED_READ_OPTIONS + ) + + def test_execute_sql_w_directed_read_options(self): + self._execute_sql_helper( + multi_use=False, + directed_read_options_at_client_level=DIRECTED_READ_OPTIONS_FOR_CLIENT, + ) + + def test_execute_sql_w_directed_read_options(self): + self._execute_sql_helper( + multi_use=False, + directed_read_options=DIRECTED_READ_OPTIONS, + directed_read_options_at_client_level=DIRECTED_READ_OPTIONS_FOR_CLIENT, + ) + def _partition_read_helper( self, multi_use, @@ -1747,10 +1826,11 @@ def __init__(self): class _Database(object): - def __init__(self): + def __init__(self, directed_read_options=None): self.name = "testing" self._instance = _Instance() self._route_to_leader_enabled = True + self._directed_read_options = directed_read_options class _Session(object): diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index e4cd1e84cd..a9c05c563f 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -28,6 +28,7 @@ StructType, TransactionOptions, TransactionSelector, + TransactionTypes, ExecuteBatchDmlRequest, ExecuteBatchDmlResponse, param_types, @@ -74,6 +75,17 @@ RETRY = gapic_v1.method.DEFAULT TIMEOUT = gapic_v1.method.DEFAULT REQUEST_OPTIONS = RequestOptions() +DIRECTED_READ_OPTIONS = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover": True, + }, +} insert_dml = "INSERT INTO table(pkey, desc) VALUES (%pkey, %desc)" insert_params = {"pkey": 12345, "desc": "DESCRIPTION"} insert_param_types = {"pkey": param_types.INT64, "desc": param_types.STRING} @@ -193,6 +205,7 @@ def _execute_sql_helper( partition=None, sql_count=0, query_options=None, + directed_read_options=None, ): VALUES = [["bharney", "rhubbyl", 31], ["phred", "phlyntstone", 32]] VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES] @@ -231,6 +244,7 @@ def _execute_sql_helper( partition=partition, retry=RETRY, timeout=TIMEOUT, + directed_read_options=directed_read_options, ) self.assertEqual(transaction._read_request_count, count + 1) @@ -284,6 +298,7 @@ def _read_helper( api, count=0, partition=None, + directed_read_options=None, ): VALUES = [["bharney", 31], ["phred", 32]] VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES] @@ -322,6 +337,7 @@ def _read_helper( retry=RETRY, timeout=TIMEOUT, request_options=REQUEST_OPTIONS, + directed_read_options=directed_read_options, ) else: result_set = transaction.read( @@ -333,6 +349,7 @@ def _read_helper( retry=RETRY, timeout=TIMEOUT, request_options=REQUEST_OPTIONS, + directed_read_options=directed_read_options, ) self.assertEqual(transaction._read_request_count, count + 1) @@ -357,7 +374,6 @@ def _read_helper_expected_request(self, partition=None, begin=True, count=0): else: expected_limit = LIMIT - # Transaction tag is ignored for read request. expected_request_options = REQUEST_OPTIONS expected_request_options.transaction_tag = self.TRANSACTION_TAG @@ -606,6 +622,28 @@ def test_transaction_should_use_transaction_id_returned_by_first_update(self): ], ) + def test_transaction_should_throw_error_w_directed_read_options(self): + from google.api_core.exceptions import BadRequest + + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + + with self.assertRaises(BadRequest): + self._execute_sql_helper( + transaction=transaction, + api=api, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + + with self.assertRaises(BadRequest): + self._read_helper( + transaction=transaction, + api=api, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + def test_transaction_should_use_transaction_id_returned_by_first_read(self): database = _Database() session = _Session(database) @@ -924,6 +962,7 @@ def __init__(self): from google.cloud.spanner_v1 import ExecuteSqlRequest self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1") + self.directed_read_options = None class _Instance(object): @@ -936,6 +975,7 @@ def __init__(self): self.name = "testing" self._instance = _Instance() self._route_to_leader_enabled = True + self._directed_read_options = None class _Session(object): diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py index 85359dac19..7d6a4b1ec4 100644 --- a/tests/unit/test_transaction.py +++ b/tests/unit/test_transaction.py @@ -895,6 +895,7 @@ def __init__(self): from google.cloud.spanner_v1 import ExecuteSqlRequest self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1") + self.directed_read_options = None class _Instance(object): @@ -907,6 +908,7 @@ def __init__(self): self.name = "testing" self._instance = _Instance() self._route_to_leader_enabled = True + self._directed_read_options = None class _Session(object): From 69c3f6e432166c901ad7421736c5d58ab51e2df5 Mon Sep 17 00:00:00 2001 From: Astha Mohta Date: Fri, 24 Feb 2023 02:14:36 +0530 Subject: [PATCH 03/19] docs --- google/cloud/spanner_v1/__init__.py | 1 + google/cloud/spanner_v1/client.py | 6 ++++++ google/cloud/spanner_v1/snapshot.py | 12 ++++++++++++ 3 files changed, 19 insertions(+) diff --git a/google/cloud/spanner_v1/__init__.py b/google/cloud/spanner_v1/__init__.py index 1552213f3b..a1b3bb143a 100644 --- a/google/cloud/spanner_v1/__init__.py +++ b/google/cloud/spanner_v1/__init__.py @@ -80,6 +80,7 @@ """ TransactionTypes = DirectedReadOptions.ReplicaSelection.Type + __all__ = ( # google.cloud.spanner_v1 "__version__", diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 1b5d971e4a..9f2915b886 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -121,6 +121,12 @@ class Client(ClientWithProject): disable leader aware routing. Disabling leader aware routing would route all requests in RW/PDML transactions to the closest region. + :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: (Optional) Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + or regions should be used for non-transactional reads or queries. + :raises: :class:`ValueError ` if both ``read_only`` and ``admin`` are :data:`True` """ diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index b2223b1831..fab0bc6cbd 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -227,6 +227,12 @@ def read( ``partition_token``, the API will return an ``INVALID_ARGUMENT`` error. + :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: (Optional) Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + or regions should be used for non-transactional reads or queries. + :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet` :returns: a result set instance which can be used to consume rows. @@ -397,6 +403,12 @@ def execute_sql( ``partition_token``, the API will return an ``INVALID_ARGUMENT`` error. + :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: (Optional) Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + or regions should be used for non-transactional reads or queries. + :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. From aa9f18cc0f87eae8d0d0a8e270ee81707cad0aa1 Mon Sep 17 00:00:00 2001 From: Astha Mohta Date: Fri, 24 Feb 2023 02:22:37 +0530 Subject: [PATCH 04/19] docs --- google/cloud/spanner_v1/_helpers.py | 7 +++++++ google/cloud/spanner_v1/client.py | 14 +++++++++++--- google/cloud/spanner_v1/snapshot.py | 4 ++++ 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py index 3b62515da0..b7b6f0bb9e 100644 --- a/google/cloud/spanner_v1/_helpers.py +++ b/google/cloud/spanner_v1/_helpers.py @@ -363,6 +363,13 @@ def _metadata_with_leader_aware_routing(value, **kw): def verify_directed_read_options(directed_read_options): + """Verifies if value of directed_read_options is correct. + :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: directed_read_options for ReadRequests and ExecuteSqlRequests. + + :raises InvalidArguement: if directed_read_options is incorrect. + """ if type(directed_read_options) == dict: if ( "include_replicas" in directed_read_options.keys() diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 9f2915b886..0000d1cd19 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -272,11 +272,12 @@ def route_to_leader_enabled(self): def directed_read_options(self): - """Getter for client's credentials. + """Getter for directed_read_options. :rtype: - :class:`Credentials ` - :returns: The credentials stored on the client. + :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :returns: The directed_read_options for the client. """ return self._directed_read_options @@ -405,4 +406,11 @@ def list_instances(self, filter_="", page_size=None): return page_iter def set_directed_read_options(self, directed_read_options): + """Sets directed_read_options for the client + :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + or regions should be used for non-transactional reads or queries. + """ self._directed_read_options = directed_read_options diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index fab0bc6cbd..a0f5c47805 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -239,6 +239,8 @@ def read( :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. + + :raises InvalidArguement: if directed_read_options is incorrect. """ if self._read_request_count > 0: if not self._multi_use: @@ -412,6 +414,8 @@ def execute_sql( :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. + + :raises InvalidArguement: if directed_read_options is incorrect. """ if self._read_request_count > 0: if not self._multi_use: From 5a0eadc09f0a4e1507500f261d9c56955900aba1 Mon Sep 17 00:00:00 2001 From: Astha Mohta Date: Fri, 24 Feb 2023 02:29:26 +0530 Subject: [PATCH 05/19] linting --- google/cloud/spanner_v1/client.py | 8 +- google/cloud/spanner_v1/snapshot.py | 14 ++- tests/unit/test__helpers.py | 183 ++++++++++------------------ 3 files changed, 76 insertions(+), 129 deletions(-) diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 0000d1cd19..3e6c862911 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -123,8 +123,8 @@ class Client(ClientWithProject): :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` - :param directed_read_options: (Optional) Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + :param directed_read_options: (Optional) Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas or regions should be used for non-transactional reads or queries. :raises: :class:`ValueError ` if both ``read_only`` @@ -409,8 +409,8 @@ def set_directed_read_options(self, directed_read_options): """Sets directed_read_options for the client :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` - :param directed_read_options: Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + :param directed_read_options: Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas or regions should be used for non-transactional reads or queries. """ self._directed_read_options = directed_read_options diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index a0f5c47805..f2c9ee46b4 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -229,8 +229,8 @@ def read( :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` - :param directed_read_options: (Optional) Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + :param directed_read_options: (Optional) Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas or regions should be used for non-transactional reads or queries. :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet` @@ -239,8 +239,10 @@ def read( :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. - + :raises InvalidArguement: if directed_read_options is incorrect. + + :raises BadRequest: if directed_read_options are set for READ-WRITE Transaction or PDML. """ if self._read_request_count > 0: if not self._multi_use: @@ -407,8 +409,8 @@ def execute_sql( :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` - :param directed_read_options: (Optional) Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + :param directed_read_options: (Optional) Client options used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas or regions should be used for non-transactional reads or queries. :raises ValueError: @@ -416,6 +418,8 @@ def execute_sql( already pending for multiple-use snapshots. :raises InvalidArguement: if directed_read_options is incorrect. + + :raises BadRequest: if directed_read_options are set for READ-WRITE Transaction or PDML. """ if self._read_request_count > 0: if not self._multi_use: diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 54499738e7..46c54bd8a3 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -16,6 +16,65 @@ import unittest import mock +DIRECTED_READ_INCORRECT_OPTIONS1 = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + }, + ], + }, + "exclude_replicas": { + "replica_selections": [ + { + "location": "us-east1", + }, + ], + }, +} +DIRECTED_READ_INCORRECT_OPTIONS2 = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + { + "location": "us-west1", + }, + ], + }, +} + class Test_merge_query_options(unittest.TestCase): def _callFUT(self, *args, **kw): @@ -772,70 +831,14 @@ def _call_fut(self, directed_read_options): def test_dict_include_exclude_replica_set_error(self): from google.api_core.exceptions import InvalidArgument - directed_read_options = { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - }, - ], - }, - "exclude_replicas": { - "replica_selections": [ - { - "location": "us-east1", - }, - ], - }, - } + directed_read_options = DIRECTED_READ_INCORRECT_OPTIONS1 with self.assertRaises(InvalidArgument): self._call_fut(directed_read_options) def test_dict_greater_than_ten_replica_selections_error(self): from google.api_core.exceptions import InvalidArgument - directed_read_options = { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - ], - }, - } + directed_read_options = DIRECTED_READ_INCORRECT_OPTIONS2 with self.assertRaises(InvalidArgument): self._call_fut(directed_read_options) @@ -843,24 +846,7 @@ def test_object_include_exclude_replica_set_error(self): from google.api_core.exceptions import InvalidArgument from google.cloud.spanner_v1 import DirectedReadOptions - directed_read_options = DirectedReadOptions( - { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - }, - ], - }, - "exclude_replicas": { - "replica_selections": [ - { - "location": "us-east1", - }, - ], - }, - } - ) + directed_read_options = DirectedReadOptions(DIRECTED_READ_INCORRECT_OPTIONS1) with self.assertRaises(InvalidArgument): self._call_fut(directed_read_options) @@ -868,49 +854,6 @@ def test_dict_greater_than_ten_replica_selections_error(self): from google.api_core.exceptions import InvalidArgument from google.cloud.spanner_v1 import DirectedReadOptions - directed_read_options = DirectedReadOptions( - { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - ], - }, - } - ) + directed_read_options = DirectedReadOptions(DIRECTED_READ_INCORRECT_OPTIONS2) with self.assertRaises(InvalidArgument): self._call_fut(directed_read_options) From 3036bb6ff94b702ab3faa8ceac02cfc426949613 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 20 Nov 2023 04:57:21 +0000 Subject: [PATCH 06/19] feat(spanner): remove client side validations for directed read options --- google/cloud/spanner_v1/_helpers.py | 48 --------------- google/cloud/spanner_v1/client.py | 7 +-- google/cloud/spanner_v1/snapshot.py | 25 ++------ samples/samples/snippets.py | 4 +- tests/system/test_session_api.py | 2 +- tests/unit/test__helpers.py | 96 ----------------------------- tests/unit/test_database.py | 7 ++- 7 files changed, 15 insertions(+), 174 deletions(-) diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py index b7b6f0bb9e..dc98235f6f 100644 --- a/google/cloud/spanner_v1/_helpers.py +++ b/google/cloud/spanner_v1/_helpers.py @@ -360,51 +360,3 @@ def _metadata_with_leader_aware_routing(value, **kw): List[Tuple[str, str]]: RPC metadata with leader aware routing header """ return ("x-goog-spanner-route-to-leader", str(value).lower()) - - -def verify_directed_read_options(directed_read_options): - """Verifies if value of directed_read_options is correct. - :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` - or :class:`dict` - :param directed_read_options: directed_read_options for ReadRequests and ExecuteSqlRequests. - - :raises InvalidArguement: if directed_read_options is incorrect. - """ - if type(directed_read_options) == dict: - if ( - "include_replicas" in directed_read_options.keys() - and "exclude_replicas" in directed_read_options.keys() - ): - raise InvalidArgument( - "Only one of include_replicas or exclude_replicas can be set" - ) - if ( - "include_replicas" in directed_read_options.keys() - and "replica_selections" in directed_read_options["include_replicas"].keys() - and len(directed_read_options["include_replicas"]["replica_selections"]) - > 10 - ) or ( - "exclude_replicas" in directed_read_options.keys() - and "replica_selections" in directed_read_options["exclude_replicas"].keys() - and len(directed_read_options["exclude_replicas"]["replica_selections"]) - > 10 - ): - raise InvalidArgument("Maximum length of replica selection allowed is 10") - elif isinstance(directed_read_options, DirectedReadOptions): - if ( - directed_read_options.include_replicas is not None - and directed_read_options.exclude_replicas is not None - ): - raise InvalidArgument( - "Only one of include_replicas or exclude_replicas can be set" - ) - if ( - directed_read_options.include_replicas is not None - and directed_read_options.include_replicas.replica_selections is not None - and len(directed_read_options.include_replicas.replica_selections) > 10 - ) or ( - directed_read_options.include_replicas is not None - and directed_read_options.exclude_replicas.replica_selections is not None - and len(directed_read_options.exclude_replicas.replica_selections) > 10 - ): - raise InvalidArgument("Maximum length of replica selection allowed is 10") diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 3e6c862911..29a91c55ea 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -47,7 +47,6 @@ from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1._helpers import _merge_query_options from google.cloud.spanner_v1._helpers import _metadata_with_prefix -from google.cloud.spanner_v1._helpers import verify_directed_read_options from google.cloud.spanner_v1.instance import Instance _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__) @@ -187,7 +186,6 @@ def __init__( warnings.warn(_EMULATOR_HOST_HTTP_SCHEME) self._route_to_leader_enabled = route_to_leader_enabled - verify_directed_read_options(directed_read_options) self._directed_read_options = directed_read_options @property @@ -270,7 +268,7 @@ def route_to_leader_enabled(self): """ return self._route_to_leader_enabled - + @property def directed_read_options(self): """Getter for directed_read_options. @@ -405,7 +403,8 @@ def list_instances(self, filter_="", page_size=None): ) return page_iter - def set_directed_read_options(self, directed_read_options): + @directed_read_options.setter + def directed_read_options(self, directed_read_options): """Sets directed_read_options for the client :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index f2c9ee46b4..bd3024118d 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -42,7 +42,6 @@ from google.cloud.spanner_v1._opentelemetry_tracing import trace_call from google.cloud.spanner_v1.streamed import StreamedResultSet from google.cloud.spanner_v1 import RequestOptions -from google.cloud.spanner_v1._helpers import verify_directed_read_options _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES = ( "RST_STREAM", @@ -266,21 +265,13 @@ def read( if self._read_only: # Transaction tags are not supported for read only transactions. request_options.transaction_tag = None - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag if ( directed_read_options is None and database._directed_read_options is not None ): directed_read_options = database._directed_read_options - else: - if self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag - if directed_read_options is not None: - raise BadRequest( - "directed_read_options can't be set for readWrite transactions or partitioned dml requests" - ) - verify_directed_read_options(directed_read_options) + elif self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag request = ReadRequest( session=self._session.name, @@ -457,21 +448,13 @@ def execute_sql( if self._read_only: # Transaction tags are not supported for read only transactions. request_options.transaction_tag = None - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag if ( directed_read_options is None and database._directed_read_options is not None ): directed_read_options = database._directed_read_options - else: - if self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag - if directed_read_options is not None: - raise BadRequest( - "directed_read_options can't be set for readWrite transactions or partitioned dml requests" - ) - verify_directed_read_options(directed_read_options) + elif self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag request = ExecuteSqlRequest( session=self._session.name, diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 47f48aa8ce..007bdebbee 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -2875,9 +2875,7 @@ def directed_read_options( "--database_role", default="new_parent" ) enable_fine_grained_access_parser.add_argument("--title", default="condition title") - subparsers.add_parser( - "directed_read_options", help=directed_read_options.__doc__ - ) + subparsers.add_parser("directed_read_options", help=directed_read_options.__doc__) args = parser.parse_args() diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index fe6f465cd1..f30cfe1336 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -1960,7 +1960,7 @@ def test_execute_sql_w_directed_read_options(sessions_database, not_postgres): assert len(list(streamed)) == 3000 - + def test_execute_sql_w_to_dict_list(sessions_database): sd = _sample_data row_count = 40 diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 46c54bd8a3..0e0ec903a2 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -16,65 +16,6 @@ import unittest import mock -DIRECTED_READ_INCORRECT_OPTIONS1 = { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - }, - ], - }, - "exclude_replicas": { - "replica_selections": [ - { - "location": "us-east1", - }, - ], - }, -} -DIRECTED_READ_INCORRECT_OPTIONS2 = { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - { - "location": "us-west1", - }, - ], - }, -} - class Test_merge_query_options(unittest.TestCase): def _callFUT(self, *args, **kw): @@ -820,40 +761,3 @@ def test(self): self.assertEqual( metadata, ("x-goog-spanner-route-to-leader", str(value).lower()) ) - - -class Test_verify_directed_read_options(unittest.TestCase): - def _call_fut(self, directed_read_options): - from google.cloud.spanner_v1._helpers import verify_directed_read_options - - verify_directed_read_options(directed_read_options) - - def test_dict_include_exclude_replica_set_error(self): - from google.api_core.exceptions import InvalidArgument - - directed_read_options = DIRECTED_READ_INCORRECT_OPTIONS1 - with self.assertRaises(InvalidArgument): - self._call_fut(directed_read_options) - - def test_dict_greater_than_ten_replica_selections_error(self): - from google.api_core.exceptions import InvalidArgument - - directed_read_options = DIRECTED_READ_INCORRECT_OPTIONS2 - with self.assertRaises(InvalidArgument): - self._call_fut(directed_read_options) - - def test_object_include_exclude_replica_set_error(self): - from google.api_core.exceptions import InvalidArgument - from google.cloud.spanner_v1 import DirectedReadOptions - - directed_read_options = DirectedReadOptions(DIRECTED_READ_INCORRECT_OPTIONS1) - with self.assertRaises(InvalidArgument): - self._call_fut(directed_read_options) - - def test_dict_greater_than_ten_replica_selections_error(self): - from google.api_core.exceptions import InvalidArgument - from google.cloud.spanner_v1 import DirectedReadOptions - - directed_read_options = DirectedReadOptions(DIRECTED_READ_INCORRECT_OPTIONS2) - with self.assertRaises(InvalidArgument): - self._call_fut(directed_read_options) diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 4b138eb4e1..f2da464d7a 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -2707,7 +2707,12 @@ def _make_instance_api(): class _Client(object): - def __init__(self, project=TestDatabase.PROJECT_ID, route_to_leader_enabled=True, directed_read_options=None): + def __init__( + self, + project=TestDatabase.PROJECT_ID, + route_to_leader_enabled=True, + directed_read_options=None, + ): from google.cloud.spanner_v1 import ExecuteSqlRequest self.project = project From 0e6822988ceef0c25dbb3df523bae9bc8725c86f Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 20 Nov 2023 05:27:19 +0000 Subject: [PATCH 07/19] feat(spanner): update the auto_failover_disabled field --- samples/samples/snippets.py | 4 ++-- tests/system/test_session_api.py | 2 +- tests/unit/test_client.py | 2 +- tests/unit/test_database.py | 2 +- tests/unit/test_snapshot.py | 2 +- tests/unit/test_spanner.py | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 007bdebbee..149df64713 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -2686,7 +2686,7 @@ def directed_read_options( * `location:us-east1 type:READ_ONLY` --> The "READ_ONLY" type replica(s) in location "us-east1" will be used to process the request. - include_replicas also contains an option for auto_failover which when set + include_replicas also contains an option for auto_failover_disabled which when set Spanner will not route requests to a replica outside the include_replicas list when all the specified replicas are unavailable or unhealthy. The default value is `false` @@ -2720,7 +2720,7 @@ def directed_read_options( "type_": TransactionTypes.READ_ONLY, }, ], - "auto_failover": True, + "auto_failover_disabled": True, }, } diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index f30cfe1336..3917d67c6b 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -66,7 +66,7 @@ "type_": TransactionTypes.READ_ONLY, }, ], - "auto_failover": True, + "auto_failover_disabled": True, }, } diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 7b3d5d0723..2678802d74 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -49,7 +49,7 @@ class TestClient(unittest.TestCase): "type_": TransactionTypes.READ_ONLY, }, ], - "auto_failover": True, + "auto_failover_disabled": True, }, } diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index f2da464d7a..c6329b3810 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -43,7 +43,7 @@ "type_": TransactionTypes.READ_ONLY, }, ], - "auto_failover": True, + "auto_failover_disabled": True, }, } diff --git a/tests/unit/test_snapshot.py b/tests/unit/test_snapshot.py index 0986d017c3..1c99c5e08a 100644 --- a/tests/unit/test_snapshot.py +++ b/tests/unit/test_snapshot.py @@ -54,7 +54,7 @@ "type_": TransactionTypes.READ_ONLY, }, ], - "auto_failover": True, + "auto_failover_disabled": True, }, } DIRECTED_READ_OPTIONS_FOR_CLIENT = { diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index 713ed125d1..d40e3f6af9 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -83,7 +83,7 @@ "type_": TransactionTypes.READ_ONLY, }, ], - "auto_failover": True, + "auto_failover_disabled": True, }, } insert_dml = "INSERT INTO table(pkey, desc) VALUES (%pkey, %desc)" From dad4a2dfaefb51ef0a0462457c4b43e692db8a4c Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 20 Nov 2023 08:07:35 +0000 Subject: [PATCH 08/19] feat(spanner): update unit tests --- tests/unit/test_snapshot.py | 4 +-- tests/unit/test_spanner.py | 58 +++++++++++++++++++++++++++---------- 2 files changed, 44 insertions(+), 18 deletions(-) diff --git a/tests/unit/test_snapshot.py b/tests/unit/test_snapshot.py index 1c99c5e08a..dfabf71e3c 100644 --- a/tests/unit/test_snapshot.py +++ b/tests/unit/test_snapshot.py @@ -1102,13 +1102,13 @@ def test_execute_sql_w_directed_read_options(self): multi_use=False, directed_read_options=DIRECTED_READ_OPTIONS ) - def test_execute_sql_w_directed_read_options(self): + def test_execute_sql_w_directed_read_options_at_client_level(self): self._execute_sql_helper( multi_use=False, directed_read_options_at_client_level=DIRECTED_READ_OPTIONS_FOR_CLIENT, ) - def test_execute_sql_w_directed_read_options(self): + def test_execute_sql_w_directed_read_options_override(self): self._execute_sql_helper( multi_use=False, directed_read_options=DIRECTED_READ_OPTIONS, diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index d40e3f6af9..a567e58533 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -254,7 +254,13 @@ def _execute_sql_helper( self.assertEqual(transaction._execute_sql_count, sql_count + 1) def _execute_sql_expected_request( - self, database, partition=None, query_options=None, begin=True, sql_count=0 + self, + database, + partition=None, + query_options=None, + begin=True, + sql_count=0, + directed_read_options=None, ): if begin is True: expected_transaction = TransactionSelector( @@ -287,6 +293,7 @@ def _execute_sql_expected_request( request_options=expected_request_options, partition_token=partition, seqno=sql_count, + directed_read_options=directed_read_options, ) return expected_request @@ -359,7 +366,9 @@ def _read_helper( self.assertEqual(result_set.metadata, metadata_pb) self.assertEqual(result_set.stats, stats_pb) - def _read_helper_expected_request(self, partition=None, begin=True, count=0): + def _read_helper_expected_request( + self, partition=None, begin=True, count=0, directed_read_options=None + ): if begin is True: expected_transaction = TransactionSelector( begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()) @@ -385,6 +394,7 @@ def _read_helper_expected_request(self, partition=None, begin=True, count=0): limit=expected_limit, partition_token=partition, request_options=expected_request_options, + directed_read_options=directed_read_options, ) return expected_request @@ -621,26 +631,42 @@ def test_transaction_should_use_transaction_id_returned_by_first_update(self): ) def test_transaction_should_throw_error_w_directed_read_options(self): - from google.api_core.exceptions import BadRequest - database = _Database() session = _Session(database) api = database.spanner_api = self._make_spanner_api() transaction = self._make_one(session) - with self.assertRaises(BadRequest): - self._execute_sql_helper( - transaction=transaction, - api=api, - directed_read_options=DIRECTED_READ_OPTIONS, - ) + self._execute_sql_helper( + transaction=transaction, + api=api, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + api.execute_streaming_sql.assert_called_once_with( + request=self._execute_sql_expected_request( + database=database, directed_read_options=DIRECTED_READ_OPTIONS + ), + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + metadata=[ + ("google-cloud-resource-prefix", database.name), + ], + ) - with self.assertRaises(BadRequest): - self._read_helper( - transaction=transaction, - api=api, - directed_read_options=DIRECTED_READ_OPTIONS, - ) + self._read_helper( + transaction=transaction, + api=api, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + api.streaming_read.assert_called_once_with( + request=self._read_helper_expected_request( + directed_read_options=DIRECTED_READ_OPTIONS + ), + metadata=[ + ("google-cloud-resource-prefix", database.name), + ], + retry=RETRY, + timeout=TIMEOUT, + ) def test_transaction_should_use_transaction_id_returned_by_first_read(self): database = _Database() From c18e0a3ec7b4b5b4ae4b306edb2e2612f673f58e Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Wed, 22 Nov 2023 16:03:04 +0000 Subject: [PATCH 09/19] feat(spanner): update test --- tests/unit/test_spanner.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index a567e58533..703ea85f02 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -630,7 +630,7 @@ def test_transaction_should_use_transaction_id_returned_by_first_update(self): ], ) - def test_transaction_should_throw_error_w_directed_read_options(self): + def test_transaction_w_directed_read_options(self): database = _Database() session = _Session(database) api = database.spanner_api = self._make_spanner_api() @@ -645,11 +645,12 @@ def test_transaction_should_throw_error_w_directed_read_options(self): request=self._execute_sql_expected_request( database=database, directed_read_options=DIRECTED_READ_OPTIONS ), - retry=gapic_v1.method.DEFAULT, - timeout=gapic_v1.method.DEFAULT, metadata=[ ("google-cloud-resource-prefix", database.name), + ("x-goog-spanner-route-to-leader", "true"), ], + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, ) self._read_helper( @@ -663,6 +664,7 @@ def test_transaction_should_throw_error_w_directed_read_options(self): ), metadata=[ ("google-cloud-resource-prefix", database.name), + ("x-goog-spanner-route-to-leader", "true"), ], retry=RETRY, timeout=TIMEOUT, From 76cbdace6604b60d23bc08542166684c24f7f245 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Fri, 8 Dec 2023 09:45:25 +0000 Subject: [PATCH 10/19] feat(spanner): update documentation --- google/cloud/spanner_v1/_helpers.py | 2 -- google/cloud/spanner_v1/client.py | 12 ++++++------ google/cloud/spanner_v1/snapshot.py | 20 ++++++++------------ 3 files changed, 14 insertions(+), 20 deletions(-) diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py index dc98235f6f..e0e2bfdbd0 100644 --- a/google/cloud/spanner_v1/_helpers.py +++ b/google/cloud/spanner_v1/_helpers.py @@ -28,8 +28,6 @@ from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import JsonObject -from google.cloud.spanner_v1 import DirectedReadOptions -from google.api_core.exceptions import InvalidArgument # Validation error messages NUMERIC_MAX_SCALE_ERR_MSG = ( diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 29a91c55ea..c8eb2ac988 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -120,11 +120,11 @@ class Client(ClientWithProject): disable leader aware routing. Disabling leader aware routing would route all requests in RW/PDML transactions to the closest region. - :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` :param directed_read_options: (Optional) Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas - or regions should be used for non-transactional reads or queries. + for all ReadRequests and ExecuteSqlRequests that indicates which replicas + or regions should be used for non-transactional reads or queries. :raises: :class:`ValueError ` if both ``read_only`` and ``admin`` are :data:`True` @@ -273,7 +273,7 @@ def directed_read_options(self): """Getter for directed_read_options. :rtype: - :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` :returns: The directed_read_options for the client. """ @@ -406,10 +406,10 @@ def list_instances(self, filter_="", page_size=None): @directed_read_options.setter def directed_read_options(self, directed_read_options): """Sets directed_read_options for the client - :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` :param directed_read_options: Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + for all ReadRequests and ExecuteSqlRequests that indicates which replicas or regions should be used for non-transactional reads or queries. """ self._directed_read_options = directed_read_options diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index bd3024118d..f0de5a08aa 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -226,10 +226,10 @@ def read( ``partition_token``, the API will return an ``INVALID_ARGUMENT`` error. - :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` - :param directed_read_options: (Optional) Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + :param directed_read_options: (Optional) Request level option used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests that indicates which replicas or regions should be used for non-transactional reads or queries. :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet` @@ -239,9 +239,7 @@ def read( for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. - :raises InvalidArguement: if directed_read_options is incorrect. - - :raises BadRequest: if directed_read_options are set for READ-WRITE Transaction or PDML. + :raises InvalidArgument: if directed_read_options is set for READ-WRITE Transaction or PDML. """ if self._read_request_count > 0: if not self._multi_use: @@ -398,19 +396,17 @@ def execute_sql( ``partition_token``, the API will return an ``INVALID_ARGUMENT`` error. - :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` or :class:`dict` - :param directed_read_options: (Optional) Client options used to set the directed_read_options - for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas + :param directed_read_options: (Optional) Request level option used to set the directed_read_options + for all ReadRequests and ExecuteSqlRequests that indicates which replicas or regions should be used for non-transactional reads or queries. :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. - :raises InvalidArguement: if directed_read_options is incorrect. - - :raises BadRequest: if directed_read_options are set for READ-WRITE Transaction or PDML. + :raises InvalidArgument: if directed_read_options is set for READ-WRITE Transaction or PDML. """ if self._read_request_count > 0: if not self._multi_use: From 71d728d7891de057c9328f1804b8668b4c15d113 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Fri, 8 Dec 2023 11:21:34 +0000 Subject: [PATCH 11/19] feat(spanner): add system test to validate exception in case of RW transaction --- google/cloud/spanner_v1/snapshot.py | 1 - tests/system/test_session_api.py | 19 +++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index f0de5a08aa..8a2f0437ca 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -28,7 +28,6 @@ from google.api_core.exceptions import InternalServerError from google.api_core.exceptions import ServiceUnavailable from google.api_core.exceptions import InvalidArgument -from google.api_core.exceptions import BadRequest from google.api_core import gapic_v1 from google.cloud.spanner_v1._helpers import ( _make_value_pb, diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index 580454e8e8..ac4113d6b4 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -2604,6 +2604,25 @@ def test_mutation_groups_insert_or_update_then_query(not_emulator, sessions_data sd._check_rows_data(rows, sd.BATCH_WRITE_ROW_DATA) +def test_readwrite_transaction_w_directed_read_options_w_error( + sessions_database, not_emulator, not_postgres +): + sd = _sample_data + + def _transaction_read(transaction): + list( + transaction.read( + sd.TABLE, + sd.COLUMNS, + sd.ALL, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + ) + + with pytest.raises(exceptions.InvalidArgument): + sessions_database.run_in_transaction(_transaction_read) + + class FauxCall: def __init__(self, code, details="FauxCall"): self._code = code From d349418312dbe69b4ee6e848cab6df4b0ce2aa15 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Fri, 8 Dec 2023 17:17:04 +0000 Subject: [PATCH 12/19] feat(spanner): update unit test --- tests/unit/test_spanner.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index 703ea85f02..ca14bd665f 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -630,7 +630,7 @@ def test_transaction_should_use_transaction_id_returned_by_first_update(self): ], ) - def test_transaction_w_directed_read_options(self): + def test_transaction_execute_sql_w_directed_read_options(self): database = _Database() session = _Session(database) api = database.spanner_api = self._make_spanner_api() @@ -653,6 +653,12 @@ def test_transaction_w_directed_read_options(self): timeout=gapic_v1.method.DEFAULT, ) + def test_transaction_streaming_read_w_directed_read_options(self): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._read_helper( transaction=transaction, api=api, From 927da677348c726a9cffb34293cadbf2c629c01f Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Sat, 9 Dec 2023 09:31:00 +0000 Subject: [PATCH 13/19] feat(spanner): add dro for batchsnapshot and update system tests --- google/cloud/spanner_v1/database.py | 16 ++++++ google/cloud/spanner_v1/snapshot.py | 2 +- tests/system/test_database_api.py | 79 +++++++++++++++++++++++++++++ tests/system/test_session_api.py | 67 ------------------------ 4 files changed, 96 insertions(+), 68 deletions(-) diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index b196c0771e..83db1409f7 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -1227,6 +1227,7 @@ def generate_read_batches( partition_size_bytes=None, max_partitions=None, data_boost_enabled=False, + directed_read_options=None, *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, @@ -1266,6 +1267,12 @@ def generate_read_batches( (Optional) If this is for a partitioned read and this field is set ``true``, the request will be executed via offline access. + :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: (Optional) Request level option used to set the directed_read_options + for ReadRequests that indicates which replicas + or regions should be used for non-transactional reads. + :type retry: :class:`~google.api_core.retry.Retry` :param retry: (Optional) The retry settings for this request. @@ -1294,6 +1301,7 @@ def generate_read_batches( "keyset": keyset._to_dict(), "index": index, "data_boost_enabled": data_boost_enabled, + "directed_read_options": directed_read_options, } for partition in partitions: yield {"partition": partition, "read": read_info.copy()} @@ -1338,6 +1346,7 @@ def generate_query_batches( max_partitions=None, query_options=None, data_boost_enabled=False, + directed_read_options=None, *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, @@ -1389,6 +1398,12 @@ def generate_query_batches( (Optional) If this is for a partitioned query and this field is set ``true``, the request will be executed via offline access. + :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + or :class:`dict` + :param directed_read_options: (Optional) Request level option used to set the directed_read_options + for ExecuteSqlRequests that indicates which replicas + or regions should be used for non-transactional queries. + :type retry: :class:`~google.api_core.retry.Retry` :param retry: (Optional) The retry settings for this request. @@ -1413,6 +1428,7 @@ def generate_query_batches( query_info = { "sql": sql, "data_boost_enabled": data_boost_enabled, + "directed_read_options": directed_read_options, } if params: query_info["params"] = params diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index 8a2f0437ca..15fcd6a59e 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -173,10 +173,10 @@ def read( partition=None, request_options=None, data_boost_enabled=False, + directed_read_options=None, *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, - directed_read_options=None, ): """Perform a ``StreamingRead`` API request for rows in a table. diff --git a/tests/system/test_database_api.py b/tests/system/test_database_api.py index 153567810a..5bb604994f 100644 --- a/tests/system/test_database_api.py +++ b/tests/system/test_database_api.py @@ -22,6 +22,7 @@ from google.cloud import spanner_v1 from google.cloud.spanner_v1.pool import FixedSizePool, PingingPool from google.cloud.spanner_admin_database_v1 import DatabaseDialect +from google.cloud.spanner_v1 import TransactionTypes from google.type import expr_pb2 from . import _helpers from . import _sample_data @@ -31,6 +32,17 @@ FKADC_CUSTOMERS_COLUMNS = ("CustomerId", "CustomerName") FKADC_SHOPPING_CARTS_COLUMNS = ("CartId", "CustomerId", "CustomerName") ALL_KEYSET = spanner_v1.KeySet(all_=True) +DIRECTED_READ_OPTIONS = { + "include_replicas": { + "replica_selections": [ + { + "location": "us-west1", + "type_": TransactionTypes.READ_ONLY, + }, + ], + "auto_failover_disabled": True, + }, +} @pytest.fixture(scope="module") @@ -740,3 +752,70 @@ def test_update_database_invalid(not_emulator, shared_database): # Empty `fields` is not supported. with pytest.raises(exceptions.InvalidArgument): shared_database.update([]) + + +def test_snapshot_read_w_directed_read_options( + shared_database, not_postgres, not_emulator +): + _helpers.retry_has_all_dll(shared_database.reload)() + table = "users_history" + columns = ["id", "commit_ts", "name", "email", "deleted"] + user_id = 1234 + name = "phred" + email = "phred@example.com" + row_data = [[user_id, spanner_v1.COMMIT_TIMESTAMP, name, email, False]] + sd = _sample_data + + with shared_database.batch() as batch: + batch.delete(table, sd.ALL) + batch.insert(table, columns, row_data) + + with shared_database.snapshot() as snapshot: + rows = list( + snapshot.read( + table, columns, sd.ALL, directed_read_options=DIRECTED_READ_OPTIONS + ) + ) + + assert len(rows) == 1 + + +def test_execute_sql_w_directed_read_options( + shared_database, not_postgres, not_emulator +): + _helpers.retry_has_all_dll(shared_database.reload)() + sd = _sample_data + + with shared_database.batch() as batch: + batch.delete(sd.TABLE, sd.ALL) + + def _unit_of_work(transaction, test): + transaction.insert_or_update(test.TABLE, test.COLUMNS, test.ROW_DATA) + + shared_database.run_in_transaction(_unit_of_work, test=sd) + + with shared_database.snapshot() as snapshot: + rows = list( + snapshot.execute_sql(sd.SQL, directed_read_options=DIRECTED_READ_OPTIONS) + ) + sd._check_rows_data(rows) + + +def test_readwrite_transaction_w_directed_read_options_w_error( + shared_database, not_emulator, not_postgres +): + _helpers.retry_has_all_dll(shared_database.reload)() + sd = _sample_data + + def _transaction_read(transaction): + list( + transaction.read( + sd.TABLE, + sd.COLUMNS, + sd.ALL, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + ) + + with pytest.raises(exceptions.InvalidArgument): + shared_database.run_in_transaction(_transaction_read) diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index ac4113d6b4..30981322cc 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -28,7 +28,6 @@ from google.cloud import spanner_v1 from google.cloud.spanner_admin_database_v1 import DatabaseDialect from google.cloud._helpers import UTC -from google.cloud.spanner_v1 import TransactionTypes from google.cloud.spanner_v1.data_types import JsonObject from tests import _helpers as ot_helpers from . import _helpers @@ -58,17 +57,6 @@ JSON_2 = JsonObject( {"sample_object": {"name": "Anamika", "id": 2635}}, ) -DIRECTED_READ_OPTIONS = { - "include_replicas": { - "replica_selections": [ - { - "location": "us-west1", - "type_": TransactionTypes.READ_ONLY, - }, - ], - "auto_failover_disabled": True, - }, -} COUNTERS_TABLE = "counters" COUNTERS_COLUMNS = ("name", "value") @@ -513,29 +501,6 @@ def test_batch_insert_w_commit_timestamp(sessions_database, not_postgres): assert not deleted -def test_snapshot_read_w_directed_read_options(sessions_database, not_postgres): - table = "users_history" - columns = ["id", "commit_ts", "name", "email", "deleted"] - user_id = 1234 - name = "phred" - email = "phred@example.com" - row_data = [[user_id, spanner_v1.COMMIT_TIMESTAMP, name, email, False]] - sd = _sample_data - - with sessions_database.batch() as batch: - batch.delete(table, sd.ALL) - batch.insert(table, columns, row_data) - - with sessions_database.snapshot() as snapshot: - rows = list( - snapshot.read( - table, columns, sd.ALL, directed_read_options=DIRECTED_READ_OPTIONS - ) - ) - - assert len(rows) == 1 - - @_helpers.retry_mabye_aborted_txn def test_transaction_read_and_insert_then_rollback( sessions_database, @@ -1948,19 +1913,6 @@ def test_execute_sql_w_manual_consume(sessions_database): assert streamed._pending_chunk is None -def test_execute_sql_w_directed_read_options(sessions_database, not_postgres): - sd = _sample_data - row_count = 3000 - committed = _set_up_table(sessions_database, row_count) - - with sessions_database.snapshot() as snapshot: - streamed = snapshot.execute_sql( - sd.SQL, directed_read_options=DIRECTED_READ_OPTIONS - ) - - assert len(list(streamed)) == 3000 - - def test_execute_sql_w_to_dict_list(sessions_database): sd = _sample_data row_count = 40 @@ -2604,25 +2556,6 @@ def test_mutation_groups_insert_or_update_then_query(not_emulator, sessions_data sd._check_rows_data(rows, sd.BATCH_WRITE_ROW_DATA) -def test_readwrite_transaction_w_directed_read_options_w_error( - sessions_database, not_emulator, not_postgres -): - sd = _sample_data - - def _transaction_read(transaction): - list( - transaction.read( - sd.TABLE, - sd.COLUMNS, - sd.ALL, - directed_read_options=DIRECTED_READ_OPTIONS, - ) - ) - - with pytest.raises(exceptions.InvalidArgument): - sessions_database.run_in_transaction(_transaction_read) - - class FauxCall: def __init__(self, code, details="FauxCall"): self._code = code From 2dfc4060ff6b537cd0a1b159f3146a358a2cdd7d Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Sat, 9 Dec 2023 10:43:33 +0000 Subject: [PATCH 14/19] feat(spanner): fix unit tests for batchsnapshot --- tests/unit/test_database.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index af9ae2db5c..18fa21a29f 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -2214,6 +2214,7 @@ def test_generate_read_batches_w_max_partitions(self): "keyset": {"all": True}, "index": "", "data_boost_enabled": False, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2256,6 +2257,7 @@ def test_generate_read_batches_w_retry_and_timeout_params(self): "keyset": {"all": True}, "index": "", "data_boost_enabled": False, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2297,6 +2299,7 @@ def test_generate_read_batches_w_index_w_partition_size_bytes(self): "keyset": {"all": True}, "index": self.INDEX, "data_boost_enabled": False, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2338,6 +2341,7 @@ def test_generate_read_batches_w_data_boost_enabled(self): "keyset": {"all": True}, "index": self.INDEX, "data_boost_enabled": True, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2435,6 +2439,7 @@ def test_generate_query_batches_w_max_partitions(self): "sql": sql, "data_boost_enabled": False, "query_options": client._query_options, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2477,6 +2482,7 @@ def test_generate_query_batches_w_params_w_partition_size_bytes(self): "params": params, "param_types": param_types, "query_options": client._query_options, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2524,6 +2530,7 @@ def test_generate_query_batches_w_retry_and_timeout_params(self): "params": params, "param_types": param_types, "query_options": client._query_options, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): @@ -2555,6 +2562,7 @@ def test_generate_query_batches_w_data_boost_enabled(self): "sql": sql, "data_boost_enabled": True, "query_options": client._query_options, + "directed_read_options": None, } self.assertEqual(len(batches), len(self.TOKENS)) for batch, token in zip(batches, self.TOKENS): From f9af1d833b0d0ddbff98825140198b69c2f811f3 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Sat, 9 Dec 2023 14:09:09 +0000 Subject: [PATCH 15/19] feat(spanner): add unit tests for partition read and query --- tests/unit/test_database.py | 101 ++++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 18fa21a29f..032bdbe9e7 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -2359,6 +2359,47 @@ def test_generate_read_batches_w_data_boost_enabled(self): timeout=gapic_v1.method.DEFAULT, ) + def test_generate_read_batches_w_directed_read_options(self): + keyset = self._make_keyset() + database = self._make_database() + batch_txn = self._make_one(database) + snapshot = batch_txn._snapshot = self._make_snapshot() + snapshot.partition_read.return_value = self.TOKENS + + batches = list( + batch_txn.generate_read_batches( + self.TABLE, + self.COLUMNS, + keyset, + index=self.INDEX, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + ) + + expected_read = { + "table": self.TABLE, + "columns": self.COLUMNS, + "keyset": {"all": True}, + "index": self.INDEX, + "data_boost_enabled": False, + "directed_read_options": DIRECTED_READ_OPTIONS, + } + self.assertEqual(len(batches), len(self.TOKENS)) + for batch, token in zip(batches, self.TOKENS): + self.assertEqual(batch["partition"], token) + self.assertEqual(batch["read"], expected_read) + + snapshot.partition_read.assert_called_once_with( + table=self.TABLE, + columns=self.COLUMNS, + keyset=keyset, + index=self.INDEX, + partition_size_bytes=None, + max_partitions=None, + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + ) + def test_process_read_batch(self): keyset = self._make_keyset() token = b"TOKEN" @@ -2579,6 +2620,42 @@ def test_generate_query_batches_w_data_boost_enabled(self): timeout=gapic_v1.method.DEFAULT, ) + def test_generate_query_batches_w_directed_read_options(self): + sql = "SELECT COUNT(*) FROM table_name" + client = _Client(self.PROJECT_ID) + instance = _Instance(self.INSTANCE_NAME, client=client) + database = _Database(self.DATABASE_NAME, instance=instance) + batch_txn = self._make_one(database) + snapshot = batch_txn._snapshot = self._make_snapshot() + snapshot.partition_query.return_value = self.TOKENS + + batches = list( + batch_txn.generate_query_batches( + sql, directed_read_options=DIRECTED_READ_OPTIONS + ) + ) + + expected_query = { + "sql": sql, + "data_boost_enabled": False, + "query_options": client._query_options, + "directed_read_options": DIRECTED_READ_OPTIONS, + } + self.assertEqual(len(batches), len(self.TOKENS)) + for batch, token in zip(batches, self.TOKENS): + self.assertEqual(batch["partition"], token) + self.assertEqual(batch["query"], expected_query) + + snapshot.partition_query.assert_called_once_with( + sql=sql, + params=None, + param_types=None, + partition_size_bytes=None, + max_partitions=None, + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + ) + def test_process_query_batch(self): sql = ( "SELECT first_name, last_name, email FROM citizens " "WHERE age <= @max_age" @@ -2637,6 +2714,30 @@ def test_process_query_batch_w_retry_timeout(self): timeout=2.0, ) + def test_process_query_batch_w_directed_read_options(self): + sql = "SELECT first_name, last_name, email FROM citizens" + token = b"TOKEN" + batch = { + "partition": token, + "query": {"sql": sql, "directed_read_options": DIRECTED_READ_OPTIONS}, + } + database = self._make_database() + batch_txn = self._make_one(database) + snapshot = batch_txn._snapshot = self._make_snapshot() + expected = snapshot.execute_sql.return_value = object() + + found = batch_txn.process_query_batch(batch) + + self.assertIs(found, expected) + + snapshot.execute_sql.assert_called_once_with( + sql=sql, + partition=token, + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + directed_read_options=DIRECTED_READ_OPTIONS, + ) + def test_close_wo_session(self): database = self._make_database() batch_txn = self._make_one(database) From 12157b4ee90444d95261c68a76af95e6f57eec95 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 8 Jan 2024 04:22:22 +0000 Subject: [PATCH 16/19] feat(spanner): lint fixes --- tests/unit/test_spanner.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index ac3b69a338..741cc93862 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -371,7 +371,12 @@ def _read_helper( self.assertEqual(result_set.stats, stats_pb) def _read_helper_expected_request( - self, partition=None, begin=True, count=0, transaction_tag=False, directed_read_options=None + self, + partition=None, + begin=True, + count=0, + transaction_tag=False, + directed_read_options=None, ): if begin is True: expected_transaction = TransactionSelector( @@ -392,7 +397,7 @@ def _read_helper_expected_request( expected_request_options.transaction_tag = self.TRANSACTION_TAG else: expected_request_options.transaction_tag = None - + expected_request = ReadRequest( session=self.SESSION_NAME, table=TABLE_NAME, From ce21ed8ec5e8ca810be707be8ef62995fd209139 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 8 Jan 2024 05:03:22 +0000 Subject: [PATCH 17/19] feat(spanner): code refactor remove TransactionType --- google/cloud/spanner_v1/__init__.py | 2 -- google/cloud/spanner_v1/snapshot.py | 4 ---- samples/samples/snippets.py | 4 ++-- tests/system/test_database_api.py | 4 ++-- tests/unit/test_client.py | 4 ++-- tests/unit/test_database.py | 4 ++-- tests/unit/test_snapshot.py | 4 ++-- tests/unit/test_spanner.py | 4 ++-- 8 files changed, 12 insertions(+), 18 deletions(-) diff --git a/google/cloud/spanner_v1/__init__.py b/google/cloud/spanner_v1/__init__.py index 2dc05d68a0..47805d4ebc 100644 --- a/google/cloud/spanner_v1/__init__.py +++ b/google/cloud/spanner_v1/__init__.py @@ -80,7 +80,6 @@ This value can only be used for timestamp columns that have set the option ``(allow_commit_timestamp=true)`` in the schema. """ -TransactionTypes = DirectedReadOptions.ReplicaSelection.Type __all__ = ( @@ -100,7 +99,6 @@ "TransactionPingingPool", # local "COMMIT_TIMESTAMP", - "TransactionTypes", # google.cloud.spanner_v1.types "BatchCreateSessionsRequest", "BatchCreateSessionsResponse", diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index c8c560ecf9..37d983e774 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -237,8 +237,6 @@ def read( :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. - - :raises InvalidArgument: if directed_read_options is set for READ-WRITE Transaction or PDML. """ if self._read_request_count > 0: if not self._multi_use: @@ -404,8 +402,6 @@ def execute_sql( :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. - - :raises InvalidArgument: if directed_read_options is set for READ-WRITE Transaction or PDML. """ if self._read_request_count > 0: if not self._multi_use: diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 79bebce434..3ffd579f4a 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -31,7 +31,7 @@ from google.cloud import spanner from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin from google.cloud.spanner_v1 import param_types -from google.cloud.spanner_v1 import TransactionTypes +from google.cloud.spanner_v1 import DirectedReadOptions from google.type import expr_pb2 from google.iam.v1 import policy_pb2 from google.cloud.spanner_v1.data_types import JsonObject @@ -2776,7 +2776,7 @@ def directed_read_options( "include_replicas": { "replica_selections": [ { - "type_": TransactionTypes.READ_ONLY, + "type_": DirectedReadOptions.ReplicaSelection.Type.READ_ONLY, }, ], "auto_failover_disabled": True, diff --git a/tests/system/test_database_api.py b/tests/system/test_database_api.py index 5bb604994f..052e628188 100644 --- a/tests/system/test_database_api.py +++ b/tests/system/test_database_api.py @@ -22,7 +22,7 @@ from google.cloud import spanner_v1 from google.cloud.spanner_v1.pool import FixedSizePool, PingingPool from google.cloud.spanner_admin_database_v1 import DatabaseDialect -from google.cloud.spanner_v1 import TransactionTypes +from google.cloud.spanner_v1 import DirectedReadOptions from google.type import expr_pb2 from . import _helpers from . import _sample_data @@ -37,7 +37,7 @@ "replica_selections": [ { "location": "us-west1", - "type_": TransactionTypes.READ_ONLY, + "type_": DirectedReadOptions.ReplicaSelection.Type.READ_ONLY, }, ], "auto_failover_disabled": True, diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 2678802d74..8fb5b13a9a 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -15,7 +15,7 @@ import unittest import mock -from google.cloud.spanner_v1 import TransactionTypes +from google.cloud.spanner_v1 import DirectedReadOptions def _make_credentials(): @@ -46,7 +46,7 @@ class TestClient(unittest.TestCase): "replica_selections": [ { "location": "us-west1", - "type_": TransactionTypes.READ_ONLY, + "type_": DirectedReadOptions.ReplicaSelection.Type.READ_ONLY, }, ], "auto_failover_disabled": True, diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 032bdbe9e7..5f563773bc 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -22,7 +22,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask -from google.cloud.spanner_v1 import RequestOptions, TransactionTypes +from google.cloud.spanner_v1 import RequestOptions, DirectedReadOptions DML_WO_PARAM = """ DELETE FROM citizens @@ -40,7 +40,7 @@ "replica_selections": [ { "location": "us-west1", - "type_": TransactionTypes.READ_ONLY, + "type_": DirectedReadOptions.ReplicaSelection.Type.READ_ONLY, }, ], "auto_failover_disabled": True, diff --git a/tests/unit/test_snapshot.py b/tests/unit/test_snapshot.py index 5b3e3a1baf..aec20c2f54 100644 --- a/tests/unit/test_snapshot.py +++ b/tests/unit/test_snapshot.py @@ -16,7 +16,7 @@ from google.api_core import gapic_v1 import mock -from google.cloud.spanner_v1 import RequestOptions, TransactionTypes +from google.cloud.spanner_v1 import RequestOptions, DirectedReadOptions from tests._helpers import ( OpenTelemetryBase, StatusCode, @@ -51,7 +51,7 @@ "replica_selections": [ { "location": "us-west1", - "type_": TransactionTypes.READ_ONLY, + "type_": DirectedReadOptions.ReplicaSelection.Type.READ_ONLY, }, ], "auto_failover_disabled": True, diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py index 741cc93862..3663d8bdc9 100644 --- a/tests/unit/test_spanner.py +++ b/tests/unit/test_spanner.py @@ -28,7 +28,7 @@ StructType, TransactionOptions, TransactionSelector, - TransactionTypes, + DirectedReadOptions, ExecuteBatchDmlRequest, ExecuteBatchDmlResponse, param_types, @@ -79,7 +79,7 @@ "replica_selections": [ { "location": "us-west1", - "type_": TransactionTypes.READ_ONLY, + "type_": DirectedReadOptions.ReplicaSelection.Type.READ_ONLY, }, ], "auto_failover_disabled": True, From 37abf815f9df1365c78d4c09d743f8b73f4bf612 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 8 Jan 2024 10:56:54 +0000 Subject: [PATCH 18/19] feat(spanner): comment refactor --- google/cloud/spanner_v1/client.py | 6 +++--- google/cloud/spanner_v1/database.py | 4 ++-- google/cloud/spanner_v1/snapshot.py | 4 ++-- google/cloud/spanner_v1/types/spanner.py | 14 +++++++------- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index c8eb2ac988..f8f3fdb72c 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -120,7 +120,7 @@ class Client(ClientWithProject): disable leader aware routing. Disabling leader aware routing would route all requests in RW/PDML transactions to the closest region. - :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :param directed_read_options: (Optional) Client options used to set the directed_read_options for all ReadRequests and ExecuteSqlRequests that indicates which replicas @@ -273,7 +273,7 @@ def directed_read_options(self): """Getter for directed_read_options. :rtype: - :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :returns: The directed_read_options for the client. """ @@ -406,7 +406,7 @@ def list_instances(self, filter_="", page_size=None): @directed_read_options.setter def directed_read_options(self, directed_read_options): """Sets directed_read_options for the client - :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :param directed_read_options: Client options used to set the directed_read_options for all ReadRequests and ExecuteSqlRequests that indicates which replicas diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 83db1409f7..e5f00c8ebd 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -1267,7 +1267,7 @@ def generate_read_batches( (Optional) If this is for a partitioned read and this field is set ``true``, the request will be executed via offline access. - :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :param directed_read_options: (Optional) Request level option used to set the directed_read_options for ReadRequests that indicates which replicas @@ -1398,7 +1398,7 @@ def generate_query_batches( (Optional) If this is for a partitioned query and this field is set ``true``, the request will be executed via offline access. - :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :param directed_read_options: (Optional) Request level option used to set the directed_read_options for ExecuteSqlRequests that indicates which replicas diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index 37d983e774..37bed11d7e 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -225,7 +225,7 @@ def read( ``partition_token``, the API will return an ``INVALID_ARGUMENT`` error. - :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :param directed_read_options: (Optional) Request level option used to set the directed_read_options for all ReadRequests and ExecuteSqlRequests that indicates which replicas @@ -393,7 +393,7 @@ def execute_sql( ``partition_token``, the API will return an ``INVALID_ARGUMENT`` error. - :type directed_read_options: :class:`~google.cloud.spanner_v1.types.DirectedReadOptions` + :type directed_read_options: :class:`~google.cloud.spanner_v1.DirectedReadOptions` or :class:`dict` :param directed_read_options: (Optional) Request level option used to set the directed_read_options for all ReadRequests and ExecuteSqlRequests that indicates which replicas diff --git a/google/cloud/spanner_v1/types/spanner.py b/google/cloud/spanner_v1/types/spanner.py index 3dbacbe26b..e3628c0a77 100644 --- a/google/cloud/spanner_v1/types/spanner.py +++ b/google/cloud/spanner_v1/types/spanner.py @@ -398,7 +398,7 @@ class DirectedReadOptions(proto.Message): .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields Attributes: - include_replicas (google.cloud.spanner_v1.types.DirectedReadOptions.IncludeReplicas): + include_replicas (google.cloud.spanner_v1.DirectedReadOptions.IncludeReplicas): Include_replicas indicates the order of replicas (as they appear in this list) to process the request. If auto_failover_disabled is set to true and all replicas are @@ -407,7 +407,7 @@ class DirectedReadOptions(proto.Message): may fail due to ``DEADLINE_EXCEEDED`` errors. This field is a member of `oneof`_ ``replicas``. - exclude_replicas (google.cloud.spanner_v1.types.DirectedReadOptions.ExcludeReplicas): + exclude_replicas (google.cloud.spanner_v1.DirectedReadOptions.ExcludeReplicas): Exclude_replicas indicates that should be excluded from serving requests. Spanner will not route requests to the replicas in this list. @@ -437,7 +437,7 @@ class ReplicaSelection(proto.Message): location (str): The location or region of the serving requests, e.g. "us-east1". - type_ (google.cloud.spanner_v1.types.DirectedReadOptions.ReplicaSelection.Type): + type_ (google.cloud.spanner_v1.DirectedReadOptions.ReplicaSelection.Type): The type of replica. """ @@ -474,7 +474,7 @@ class IncludeReplicas(proto.Message): should be considered. Attributes: - replica_selections (MutableSequence[google.cloud.spanner_v1.types.DirectedReadOptions.ReplicaSelection]): + replica_selections (MutableSequence[google.cloud.spanner_v1.DirectedReadOptions.ReplicaSelection]): The directed read replica selector. auto_failover_disabled (bool): If true, Spanner will not route requests to a replica @@ -500,7 +500,7 @@ class ExcludeReplicas(proto.Message): ReplicaSelection that should be excluded from serving requests. Attributes: - replica_selections (MutableSequence[google.cloud.spanner_v1.types.DirectedReadOptions.ReplicaSelection]): + replica_selections (MutableSequence[google.cloud.spanner_v1.DirectedReadOptions.ReplicaSelection]): The directed read replica selector. """ @@ -626,7 +626,7 @@ class ExecuteSqlRequest(proto.Message): given query. request_options (google.cloud.spanner_v1.types.RequestOptions): Common options for this request. - directed_read_options (google.cloud.spanner_v1.types.DirectedReadOptions): + directed_read_options (google.cloud.spanner_v1.DirectedReadOptions): Directed read options for this request. data_boost_enabled (bool): If this is for a partitioned query and this field is set to @@ -1294,7 +1294,7 @@ class ReadRequest(proto.Message): create this partition_token. request_options (google.cloud.spanner_v1.types.RequestOptions): Common options for this request. - directed_read_options (google.cloud.spanner_v1.types.DirectedReadOptions): + directed_read_options (google.cloud.spanner_v1.DirectedReadOptions): Directed read options for this request. data_boost_enabled (bool): If this is for a partitioned read and this field is set to From 7162a427ddeda0444723c73e360df57c80414b01 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Mon, 8 Jan 2024 12:11:15 +0000 Subject: [PATCH 19/19] feat(spanner): remove comments --- google/cloud/spanner_v1/types/spanner.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/google/cloud/spanner_v1/types/spanner.py b/google/cloud/spanner_v1/types/spanner.py index e3628c0a77..3dbacbe26b 100644 --- a/google/cloud/spanner_v1/types/spanner.py +++ b/google/cloud/spanner_v1/types/spanner.py @@ -398,7 +398,7 @@ class DirectedReadOptions(proto.Message): .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields Attributes: - include_replicas (google.cloud.spanner_v1.DirectedReadOptions.IncludeReplicas): + include_replicas (google.cloud.spanner_v1.types.DirectedReadOptions.IncludeReplicas): Include_replicas indicates the order of replicas (as they appear in this list) to process the request. If auto_failover_disabled is set to true and all replicas are @@ -407,7 +407,7 @@ class DirectedReadOptions(proto.Message): may fail due to ``DEADLINE_EXCEEDED`` errors. This field is a member of `oneof`_ ``replicas``. - exclude_replicas (google.cloud.spanner_v1.DirectedReadOptions.ExcludeReplicas): + exclude_replicas (google.cloud.spanner_v1.types.DirectedReadOptions.ExcludeReplicas): Exclude_replicas indicates that should be excluded from serving requests. Spanner will not route requests to the replicas in this list. @@ -437,7 +437,7 @@ class ReplicaSelection(proto.Message): location (str): The location or region of the serving requests, e.g. "us-east1". - type_ (google.cloud.spanner_v1.DirectedReadOptions.ReplicaSelection.Type): + type_ (google.cloud.spanner_v1.types.DirectedReadOptions.ReplicaSelection.Type): The type of replica. """ @@ -474,7 +474,7 @@ class IncludeReplicas(proto.Message): should be considered. Attributes: - replica_selections (MutableSequence[google.cloud.spanner_v1.DirectedReadOptions.ReplicaSelection]): + replica_selections (MutableSequence[google.cloud.spanner_v1.types.DirectedReadOptions.ReplicaSelection]): The directed read replica selector. auto_failover_disabled (bool): If true, Spanner will not route requests to a replica @@ -500,7 +500,7 @@ class ExcludeReplicas(proto.Message): ReplicaSelection that should be excluded from serving requests. Attributes: - replica_selections (MutableSequence[google.cloud.spanner_v1.DirectedReadOptions.ReplicaSelection]): + replica_selections (MutableSequence[google.cloud.spanner_v1.types.DirectedReadOptions.ReplicaSelection]): The directed read replica selector. """ @@ -626,7 +626,7 @@ class ExecuteSqlRequest(proto.Message): given query. request_options (google.cloud.spanner_v1.types.RequestOptions): Common options for this request. - directed_read_options (google.cloud.spanner_v1.DirectedReadOptions): + directed_read_options (google.cloud.spanner_v1.types.DirectedReadOptions): Directed read options for this request. data_boost_enabled (bool): If this is for a partitioned query and this field is set to @@ -1294,7 +1294,7 @@ class ReadRequest(proto.Message): create this partition_token. request_options (google.cloud.spanner_v1.types.RequestOptions): Common options for this request. - directed_read_options (google.cloud.spanner_v1.DirectedReadOptions): + directed_read_options (google.cloud.spanner_v1.types.DirectedReadOptions): Directed read options for this request. data_boost_enabled (bool): If this is for a partitioned read and this field is set to