From e4797f5b4c4cee426ecb6c437b5b55aa29845e3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Tue, 15 Sep 2026 17:00:42 +0200 Subject: [PATCH] chore(spanner): restore CrossSync parity for snapshot.py snapshot.py is generated from _async/snapshot.py, but the generated file had picked up code the async source never had. Running the generator deleted it. The drift came in via 25aba4f4f45 (PR #16488): a branch authored before the CrossSync migration but merged after it, so it edited snapshot.py as if that file were still hand-maintained. Fixes: - Port the _transaction_begin_event guard and the read()/execute_sql() docstrings back into _async/snapshot.py, so the generator reproduces them instead of dropping them. - Use CrossSync.Lock / CrossSync.Event instead of raw threading, so the primitives survive the async -> sync conversion. - De-duplicate the guard: read() and execute_sql() had identical copies, now one _wait_for_transaction_begin() helper. - Async fix: read()/execute_sql() raised "Transaction has not begun." when a concurrent inline begin was in flight. They now wait for it, matching the sync client. - generate.py: format with ruff instead of black. The repo uses ruff, so every run produced a formatting-only diff. - generate.py: exit non-zero when given a file instead of a directory. It used to print "Generated 0 artifacts" and do nothing, which is how the drift went unnoticed. - Delete snapshot_helpers.py, an orphaned generated file with a stale copy of execute_sql that nothing imports. - Add tests for the guard, which had none. - Speed up 4 tests that each waited out the real 30s begin timeout. They passed, just slowly; two were already slow on main. They now patch the timeout constant and assert the specific error message. Unit suite: 204s -> 57s. Regenerating now leaves the tree unchanged: PYTHONPATH=.cross_sync python3 .cross_sync/generate.py \ google/cloud/spanner_v1/_async/ git diff --exit-code google/cloud/spanner_v1/ --- .../.agents/workflows/verify-asyncio.md | 9 +- .../.cross_sync/generate.py | 85 +- .../cloud/spanner_v1/_async/snapshot.py | 246 +++++- .../cloud/spanner_v1/_async/streamed.py | 5 + .../google/cloud/spanner_v1/batch.py | 11 +- .../google/cloud/spanner_v1/client.py | 8 +- .../google/cloud/spanner_v1/database.py | 13 +- .../spanner_v1/database_sessions_manager.py | 11 +- .../google/cloud/spanner_v1/instance.py | 3 +- .../google/cloud/spanner_v1/pool.py | 15 +- .../google/cloud/spanner_v1/session.py | 3 +- .../google/cloud/spanner_v1/snapshot.py | 110 +-- .../cloud/spanner_v1/snapshot_helpers.py | 730 ------------------ .../google/cloud/spanner_v1/streamed.py | 12 +- .../cloud/spanner_v1/testing/database_test.py | 7 +- .../google/cloud/spanner_v1/transaction.py | 6 +- .../tests/unit/_async/test_snapshot.py | 52 +- .../tests/unit/test_snapshot.py | 77 +- 18 files changed, 507 insertions(+), 896 deletions(-) delete mode 100644 packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot_helpers.py diff --git a/packages/google-cloud-spanner/.agents/workflows/verify-asyncio.md b/packages/google-cloud-spanner/.agents/workflows/verify-asyncio.md index f5887a178773..2287597792bb 100644 --- a/packages/google-cloud-spanner/.agents/workflows/verify-asyncio.md +++ b/packages/google-cloud-spanner/.agents/workflows/verify-asyncio.md @@ -25,11 +25,18 @@ nox -s system -- tests/system/_async ## 3. Verify Sync/Async Parity Run the cross-sync generation tool and ensure no regressions in the synchronous codebase. + +`generate.py` must be pointed at a **directory**; it only rewrites files reachable from +that directory that carry a `__CROSS_SYNC_OUTPUT__` annotation. ```bash -python3 .cross_sync/generate.py +PYTHONPATH=.cross_sync python3 .cross_sync/generate.py google/cloud/spanner_v1/_async/ +git diff --exit-code google/cloud/spanner_v1/ nox -s unit-3.14 nox -s system-3.14 ``` +A non-empty `git diff` here means the generated sync code has drifted from +`google/cloud/spanner_v1/_async/`. Fix it in the `_async/` source, never in the +generated artifact. ## 4. Check for Coroutine Leaks Ensure all asynchronous GAPIC calls are properly awaited. Search for any unawaited coroutines in the `_async` directory. diff --git a/packages/google-cloud-spanner/.cross_sync/generate.py b/packages/google-cloud-spanner/.cross_sync/generate.py index 3c96b7469e61..07a619e1105b 100644 --- a/packages/google-cloud-spanner/.cross_sync/generate.py +++ b/packages/google-cloud-spanner/.cross_sync/generate.py @@ -39,6 +39,56 @@ def extract_header_comments(file_path) -> str: return "".join(header) +# Keep these in sync with the `format` / `lint` sessions in noxfile.py, otherwise +# regenerating the artifacts will produce a spurious formatting-only diff. +RUFF_TARGET_VERSION = "py310" +RUFF_LINE_LENGTH = "88" + + +def format_with_ruff(source: str, filename: str) -> str: + """ + Format generated source with ruff, the formatter used by this repository. + + Runs two passes over stdin, mirroring `nox -s format`: + 1. `ruff check --select I,F401 --fix` to sort imports and drop the + imports that became unused during the async -> sync conversion. + 2. `ruff format` to apply the code style. + + Args: + source: the generated python source + filename: the path the source will be written to. Only used to give + ruff a sensible filename for diagnostics. + Returns: + the formatted source + """ + import shutil + import subprocess + import sys + + ruff = shutil.which("ruff") + base_command = [ruff] if ruff else [sys.executable, "-m", "ruff"] + shared_args = [ + f"--target-version={RUFF_TARGET_VERSION}", + "--line-length", + RUFF_LINE_LENGTH, + "--stdin-filename", + filename, + "-", + ] + passes = [ + base_command + ["check", "--select", "I,F401", "--fix", "--quiet", *shared_args], + base_command + ["format", *shared_args], + ] + for command in passes: + result = subprocess.run(command, input=source, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"ruff failed for {filename}: {' '.join(command)}\n{result.stderr}" + ) + source = result.stdout + return source + + class CrossSyncOutputFile: def __init__(self, output_path: str, ast_tree, header: str | None = None): @@ -51,18 +101,12 @@ def render(self, with_formatter=True, save_to_disk: bool = True) -> str: Render the file to a string, and optionally save to disk Args: - with_formatter: whether to run the output through black before returning + with_formatter: whether to run the output through ruff before returning save_to_disk: whether to write the output to the file path """ full_str = self.header + ast.unparse(self.tree) if with_formatter: - import black # type: ignore - import autoflake # type: ignore - - full_str = black.format_str( - autoflake.fix_code(full_str, remove_all_unused_imports=True), - mode=black.FileMode(), - ) + full_str = format_with_ruff(full_str, self.output_path) if save_to_disk: import os os.makedirs(os.path.dirname(self.output_path), exist_ok=True) @@ -71,12 +115,18 @@ def render(self, with_formatter=True, save_to_disk: bool = True) -> str: return full_str -def convert_files_in_dir(directory: str) -> set[CrossSyncOutputFile]: +def convert_path(search_path: str) -> set[CrossSyncOutputFile]: import glob from transformers import CrossSyncFileProcessor - # find all python files in the directory - files = glob.glob(directory + "/**/*.py", recursive=True) + if os.path.isfile(search_path): + files = [search_path] + elif os.path.isdir(search_path): + files = glob.glob(search_path + "/**/*.py", recursive=True) + else: + print(f"Path does not exist: {search_path}") + sys.exit(1) + # keep track of the output files pointed to by the annotated classes artifacts: set[CrossSyncOutputFile] = set() file_transformer = CrossSyncFileProcessor() @@ -100,13 +150,22 @@ def save_artifacts(artifacts: Sequence[CrossSyncOutputFile]): if __name__ == "__main__": + import os import sys if len(sys.argv) < 2: - print("Usage: python .cross_sync/generate.py ") + print("Usage: python .cross_sync/generate.py ") sys.exit(1) search_root = sys.argv[1] - outputs = convert_files_in_dir(search_root) + if not os.path.exists(search_root): + print(f"Path does not exist: {search_root}") + sys.exit(1) + + outputs = convert_path(search_root) + if not outputs: + print(f"No __CROSS_SYNC_OUTPUT__ annotated files found under {search_root}") + sys.exit(1) + print(f"Generated {len(outputs)} artifacts: {[a.output_path for a in outputs]}") save_artifacts(outputs) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py index b54b5d314e1a..5e64e106db99 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py @@ -68,6 +68,10 @@ "Received unexpected EOS on DATA frame from server", ) +# Maximum time that a request will wait for a concurrent request to begin the +# transaction before giving up. +_TRANSACTION_BEGIN_TIMEOUT_SECONDS = 30.0 + @CrossSync.convert async def _restart_on_unavailable( @@ -221,10 +225,20 @@ def __init__(self, session, client_context=None): self._client_context = _validate_client_context(client_context) self._execute_sql_request_count: int = 0 self._read_request_count: int = 0 + self._begin_request_sent: bool = False + + # Identifier for the transaction. self._transaction_id: Optional[bytes] = None self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None + + # Operations within a transaction can be performed concurrently, so we + # need to use a lock when updating the transaction. self._lock: CrossSync.Lock = CrossSync.Lock() + # Event to coordinate concurrent requests beginning the transaction. + # This is used to prevent the "Transaction has not begun" race condition. + self._transaction_begin_event: CrossSync.Event = CrossSync.Event() + @property def _resource_info(self): """Resource information for metrics labels.""" @@ -235,6 +249,39 @@ def _resource_info(self): "database": database.database_id, } + @CrossSync.convert + async def _wait_for_transaction_begin(self) -> None: + """Claims the inline-begin for this request, or waits for it to complete. + + The first request against the transaction is the one that begins it + inline. Requests that are issued concurrently, before the transaction + id is available, must wait for that first request to complete instead + of assuming that the transaction has not begun. + + :raises ValueError: if the transaction has already been used to execute + a request, but is not a multi-use transaction, or if the concurrent + request that began the transaction did not complete in time. + """ + + async with self._lock: + if self._begin_request_sent or self._read_request_count > 0: + if not self._multi_use: + raise ValueError("Cannot re-use single-use snapshot.") + wait_needed = self._transaction_id is None + else: + wait_needed = False + self._begin_request_sent = True + + if not wait_needed: + return + + await CrossSync.event_wait( + self._transaction_begin_event, + timeout=_TRANSACTION_BEGIN_TIMEOUT_SECONDS, + ) + if not self._transaction_begin_event.is_set(): + raise ValueError("Timed out waiting for transaction to begin.") + @CrossSync.convert async def begin(self) -> bytes: """Begins a transaction on the database. @@ -246,7 +293,6 @@ async def begin(self) -> bytes: """ return await self._begin_transaction() - @CrossSync.convert @CrossSync.convert async def read( self, @@ -265,12 +311,88 @@ async def read( column_info=None, lazy_decode=False, ): - """Perform a ``StreamingRead`` API request for rows in a table.""" - if self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") - if self._transaction_id is None: - raise ValueError("Transaction has not begun.") + """Perform a ``StreamingRead`` API request for rows in a table. + + :type table: str + :param table: name of the table from which to fetch data + + :type columns: list of str + :param columns: names of columns to be retrieved + + :type keyset: :class:`~google.cloud.spanner_v1.keyset.KeySet` + :param keyset: keys / ranges identifying rows to be retrieved + + :type index: str + :param index: (Optional) name of index to use, rather than the + table's primary key + + :type limit: int + :param limit: (Optional) maximum number of rows to return. + Incompatible with ``partition``. + + :type partition: bytes + :param partition: (Optional) one of the partition tokens returned + from :meth:`partition_read`. Incompatible with + ``limit``. + + :type request_options: + :class:`google.cloud.spanner_v1.types.RequestOptions` + :param request_options: + (Optional) Common options for this request. + If a dict is provided, it must be of the same form as the protobuf + message :class:`~google.cloud.spanner_v1.types.RequestOptions`. + Please note, the `transactionTag` setting will be ignored for + snapshot as it's not supported for read-only transactions. + + :type retry: :class:`~google.api_core.retry.Retry` + :param retry: (Optional) The retry settings for this request. + + :type timeout: float + :param timeout: (Optional) The timeout for this request. + + :type data_boost_enabled: + :param data_boost_enabled: + (Optional) If this is for a partitioned read and this field is + set ``true``, the request will be executed via offline access. + If the field is set to ``true`` but the request does not set + ``partition_token``, the API will return an + ``INVALID_ARGUMENT`` error. + + :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 + or regions should be used for non-transactional reads or queries. + + :type column_info: dict + :param column_info: (Optional) dict of mapping between column names and additional column information. + An object where column names as keys and custom objects as corresponding + values for deserialization. It's specifically useful for data types like + protobuf where deserialization logic is on user-specific code. When provided, + the custom object enables deserialization of backend-received column data. + If not provided, data remains serialized as bytes for Proto Messages and + integer for Proto Enums. + + :type lazy_decode: bool + :param lazy_decode: + (Optional) If this argument is set to ``true``, the iterator + returns the underlying protobuf values instead of decoded Python + objects. This reduces the time that is needed to iterate through + large result sets. The application is responsible for decoding + the data that is needed. The returned row iterator contains two + functions that can be used for this. ``iterator.decode_row(row)`` + decodes all the columns in the given row to an array of Python + objects. ``iterator.decode_column(row, column_index)`` decodes one + specific column in the given row. + + :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet` + :returns: a result set instance which can be used to consume rows. + + :raises ValueError: if the Transaction already used to execute a + read request, but is not a multi-use transaction or has not begun. + """ + + await self._wait_for_transaction_begin() session = self._session database = session._database @@ -336,7 +458,6 @@ async def read( lazy_decode=lazy_decode, ) - @CrossSync.convert @CrossSync.convert async def execute_sql( self, @@ -355,12 +476,107 @@ async def execute_sql( column_info=None, lazy_decode=False, ): - """Perform an ``ExecuteStreamingSql`` API request.""" - if self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") - if self._transaction_id is None: - raise ValueError("Transaction has not begun.") + """Perform an ``ExecuteStreamingSql`` API request. + + :type sql: str + :param sql: SQL query statement + + :type params: dict, {str -> column value} + :param params: values for parameter replacement. Keys must match + the names used in ``sql``. + + :type param_types: dict[str -> Union[dict, .types.Type]] + :param param_types: + (Optional) maps explicit types for one or more param values; + required if parameters are passed. + + :type query_mode: + :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryMode` + :param query_mode: Mode governing return of results / query plan. + See: + `QueryMode `_. + + :type query_options: + :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions` + or :class:`dict` + :param query_options: + (Optional) Query optimizer configuration to use for the given query. + If a dict is provided, it must be of the same form as the protobuf + message :class:`~google.cloud.spanner_v1.types.QueryOptions` + + :type request_options: + :class:`google.cloud.spanner_v1.types.RequestOptions` + :param request_options: + (Optional) Common options for this request. + If a dict is provided, it must be of the same form as the protobuf + message :class:`~google.cloud.spanner_v1.types.RequestOptions`. + + :type last_statement: bool + :param last_statement: + If set to true, this option marks the end of the transaction. The + transaction should be committed or aborted after this statement + executes, and attempts to execute any other requests against this + transaction (including reads and queries) will be rejected. Mixing + mutations with statements that are marked as the last statement is + not allowed. + For DML statements, setting this option may cause some error + reporting to be deferred until commit time (e.g. validation of + unique constraints). Given this, successful execution of a DML + statement should not be assumed until the transaction commits. + + :type partition: bytes + :param partition: (Optional) one of the partition tokens returned + from :meth:`partition_query`. + + :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet` + :returns: a result set instance which can be used to consume rows. + + :type retry: :class:`~google.api_core.retry.Retry` + :param retry: (Optional) The retry settings for this request. + + :type timeout: float + :param timeout: (Optional) The timeout for this request. + + :type data_boost_enabled: + :param data_boost_enabled: + (Optional) If this is for a partitioned query and this field is + set ``true``, the request will be executed via offline access. + If the field is set to ``true`` but the request does not set + ``partition_token``, the API will return an + ``INVALID_ARGUMENT`` error. + + :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 + or regions should be used for non-transactional reads or queries. + + :type column_info: dict + :param column_info: (Optional) dict of mapping between column names and additional column information. + An object where column names as keys and custom objects as corresponding + values for deserialization. It's specifically useful for data types like + protobuf where deserialization logic is on user-specific code. When provided, + the custom object enables deserialization of backend-received column data. + If not provided, data remains serialized as bytes for Proto Messages and + integer for Proto Enums. + + :type lazy_decode: bool + :param lazy_decode: + (Optional) If this argument is set to ``true``, the iterator + returns the underlying protobuf values instead of decoded Python + objects. This reduces the time that is needed to iterate through + large result sets. The application is responsible for decoding + the data that is needed. The returned row iterator contains two + functions that can be used for this. ``iterator.decode_row(row)`` + decodes all the columns in the given row to an array of Python + objects. ``iterator.decode_column(row, column_index)`` decodes one + specific column in the given row. + + :raises ValueError: if the Transaction already used to execute a + read request, but is not a multi-use transaction or has not begun. + """ + + await self._wait_for_transaction_begin() if params is not None: params_pb = Struct( @@ -766,6 +982,8 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None: """Updates the snapshot for the given transaction.""" if self._transaction_id is None and transaction_pb.id: self._transaction_id = transaction_pb.id + # Release any request waiting for this transaction to begin. + self._transaction_begin_event.set() if transaction_pb._pb.HasField("precommit_token"): self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py index c47cc0ef0a17..95fda33f7b12 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py @@ -123,6 +123,11 @@ def _merge_chunk(self, value): def _merge_values(self, values): """Merge values into rows. + Note: We manually check value.HasField("null_value") here instead of + wrapping every decoder in _parse_nullable to avoid the overhead of + an extra Python function call layer for every cell value decoded in this loop. + If the nullable check logic is updated in _parse_nullable, update this check. + :type values: list of :class:`~google.protobuf.struct_pb2.Value` :param values: non-chunked values from partial result set. """ diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py index c32b65812db4..c0ff1cc1a613 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py @@ -58,8 +58,7 @@ class _BatchBase(_SessionWrapper): """Accumulate mutations for transmission during :meth:`commit`. :type session: :class:`~google.cloud.spanner_v1.session.Session` - :param session: the session used to perform the commit - """ + :param session: the session used to perform the commit""" def __init__(self, session, client_context=None): super(_BatchBase, self).__init__(session) @@ -158,14 +157,12 @@ def send(self, queue, key, payload=None, deliver_time=None): :param payload: (Optional) The payload of the message. :type deliver_time: :class:`datetime.datetime` - :param deliver_time: (Optional) The time at which Spanner will begin attempting to deliver the message. - """ + :param deliver_time: (Optional) The time at which Spanner will begin attempting to deliver the message.""" send_kwargs = {"queue": queue, "key": _make_list_value_pb(key)} if payload is not None: send_kwargs["payload"] = _make_value_pb(payload) if deliver_time is not None: send_kwargs["deliver_time"] = _datetime_to_pb_timestamp(deliver_time) - send = Mutation.Send(**send_kwargs) self._mutations.append(Mutation(send=send)) @@ -179,12 +176,10 @@ def ack(self, queue, key, ignore_not_found=None): :param key: The primary key of the message to be acked. :type ignore_not_found: bool - :param ignore_not_found: (Optional) Whether to ignore if the message does not exist. - """ + :param ignore_not_found: (Optional) Whether to ignore if the message does not exist.""" ack_kwargs = {"queue": queue, "key": _make_list_value_pb(key)} if ignore_not_found is not None: ack_kwargs["ignore_not_found"] = ignore_not_found - ack = Mutation.Ack(**ack_kwargs) self._mutations.append(Mutation(ack=ack)) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py index 83bde6d89ccc..a6b509b619e3 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py @@ -279,7 +279,6 @@ def __init__( ) else: self._client_options = client_options - host_endpoint = None if experimental_host is not None: warnings.warn( @@ -289,13 +288,11 @@ def __init__( ) instance_type = "omni" host_endpoint = experimental_host - if instance_type is not None: instance_type = instance_type.lower() if instance_type not in ("cloud", "omni"): raise ValueError("instance_type must be one of 'cloud' or 'omni'") self._instance_type = instance_type - if self._emulator_host: credentials = AnonymousCredentials() elif self._instance_type == "omni": @@ -305,12 +302,10 @@ def __init__( host_endpoint = self._client_options.api_endpoint elif isinstance(self._client_options, dict): host_endpoint = self._client_options.get("api_endpoint") - if not host_endpoint: raise ValueError( "Host must be set for connecting to Spanner Omni instances" ) - project = "default" self._use_plain_text = use_plain_text self._ca_certificate = ca_certificate @@ -518,8 +513,7 @@ def default_transaction_options(self): :rtype: :class:`~google.cloud.spanner_v1.DefaultTransactionOptions` or :class:`dict` - :returns: The default transaction options that are used by this client for all transactions. - """ + :returns: The default transaction options that are used by this client for all transactions.""" return self._default_transaction_options @property diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py index f98d6d09f70d..302b12f1cb46 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py @@ -42,9 +42,7 @@ RestoreDatabaseRequest, UpdateDatabaseDdlRequest, ) -from google.cloud.spanner_admin_database_v1 import ( - Database as DatabasePB, -) +from google.cloud.spanner_admin_database_v1 import Database as DatabasePB from google.cloud.spanner_admin_database_v1.types import DatabaseDialect from google.cloud.spanner_v1._helpers import ( _augment_errors_with_request_id, @@ -149,8 +147,7 @@ class Database(object): has drop protection enabled or not. :type proto_descriptors: bytes :param proto_descriptors: (Optional) Proto descriptors used by CREATE/ALTER PROTO BUNDLE - statements in 'ddl_statements' above. - """ + statements in 'ddl_statements' above.""" _spanner_api: SpannerClient = None __transport_lock = threading.Lock() @@ -720,8 +717,7 @@ def drop(self): """Drop this database. See - https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.database.v1#google.spanner.admin.database.v1.DatabaseAdmin.DropDatabase - """ + https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.database.v1#google.spanner.admin.database.v1.DatabaseAdmin.DropDatabase""" api = self._instance._client.database_admin_api metadata = _metadata_with_prefix(self.name) api.drop_database( @@ -1114,8 +1110,7 @@ def is_ready(self): """Test whether this database is ready for use. :rtype: bool - :returns: True if the database state is READY_OPTIMIZING or READY, else False. - """ + :returns: True if the database state is READY_OPTIMIZING or READY, else False.""" return ( self.state == DatabasePB.State.READY_OPTIMIZING or self.state == DatabasePB.State.READY diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py index 1b2c6231f46e..79862d90e44d 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py @@ -130,7 +130,6 @@ def _get_multiplexed_session(self) -> Session: session = self._multiplexed_session if session is not None: return session - with self._init_lock: if self._multiplexed_session_lock is None: self._multiplexed_session_lock = CrossSync._Sync_Impl.Lock() @@ -188,23 +187,19 @@ def _rotate_multiplexed_session(self) -> bool: """Rotates the multiplexed session by building and swapping in a new session. :rtype: bool - :returns: True if the session was successfully refreshed, False otherwise. - """ + :returns: True if the session was successfully refreshed, False otherwise.""" try: new_session = self._build_multiplexed_session() except Exception: return False - with self._multiplexed_session_lock: old_session = self._multiplexed_session self._multiplexed_session = new_session - if old_session is not None: try: CrossSync._Sync_Impl.run_if_async(old_session.delete) except Exception: pass - return True @staticmethod @@ -239,11 +234,9 @@ def _maintain_multiplexed_session(session_manager_ref) -> None: session_created_time = time.monotonic() manager = None continue - manager = None CrossSync._Sync_Impl.event_wait( - terminate_event, - timeout=polling_interval_seconds, + terminate_event, timeout=polling_interval_seconds ) @classmethod diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/instance.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/instance.py index 5fb824886d82..8e5a6b2bac91 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/instance.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/instance.py @@ -114,8 +114,7 @@ class Instance(object): :param labels: (Optional) User-assigned labels for this instance. :type experimental_host: str - :param experimental_host: (Deprecated) The instance type and host are now managed by the Client. - """ + :param experimental_host: (Deprecated) The instance type and host are now managed by the Client.""" def __init__( self, diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py index 80cc7ed2d7aa..74dbfa9e60ae 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py @@ -81,8 +81,7 @@ class AbstractSessionPool(object): by the pool. :type database_role: str - :param database_role: (Optional) user-assigned database_role for the session. - """ + :param database_role: (Optional) user-assigned database_role for the session.""" _database = None @@ -224,8 +223,7 @@ class FixedSizePool(AbstractSessionPool): by the pool. :type database_role: str - :param database_role: (Optional) user-assigned database_role for the session. - """ + :param database_role: (Optional) user-assigned database_role for the session.""" DEFAULT_SIZE = 10 DEFAULT_TIMEOUT = 10 @@ -441,8 +439,7 @@ class BurstyPool(AbstractSessionPool): by the pool. :type database_role: str - :param database_role: (Optional) user-assigned database_role for the session. - """ + :param database_role: (Optional) user-assigned database_role for the session.""" def __init__(self, target_size=10, labels=None, database_role=None): super(BurstyPool, self).__init__(labels=labels, database_role=database_role) @@ -556,8 +553,7 @@ class PingingPool(FixedSizePool): by the pool. :type database_role: str - :param database_role: (Optional) user-assigned database_role for the session. - """ + :param database_role: (Optional) user-assigned database_role for the session.""" def __init__( self, @@ -767,8 +763,7 @@ class TransactionPingingPool(PingingPool): by the pool. :type database_role: str - :param database_role: (Optional) user-assigned database_role for the session. - """ + :param database_role: (Optional) user-assigned database_role for the session.""" def __init__( self, diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py index 308c5d323624..9c1cb6ea1984 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py @@ -70,8 +70,7 @@ class Session(object): :param database_role: (Optional) user-assigned database_role for the session. :type is_multiplexed: bool - :param is_multiplexed: (Optional) whether this session is a multiplexed session. - """ + :param is_multiplexed: (Optional) whether this session is a multiplexed session.""" def __init__(self, database, labels=None, database_role=None, is_multiplexed=False): self._database = database diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py index 3d30e308c72a..ce2bd07306d8 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py @@ -18,7 +18,6 @@ """Model a set of read-only queries to a database as a snapshot.""" import functools -import threading from typing import List, Optional, Union from google.api_core import gapic_v1 @@ -70,6 +69,7 @@ "RST_STREAM", "Received unexpected EOS on DATA frame from server", ) +_TRANSACTION_BEGIN_TIMEOUT_SECONDS = 30.0 def _restart_on_unavailable( @@ -98,8 +98,7 @@ def _restart_on_unavailable( :type transaction_selector: :class:`transaction_pb2.TransactionSelector` :param transaction_selector: Transaction selector object to be used in request if transaction is not passed, - if both transaction_selector and transaction are passed, then transaction is given priority. - """ + if both transaction_selector and transaction are passed, then transaction is given priority.""" resume_token: bytes = b"" item_buffer: List[PartialResultSet] = [] if transaction is not None: @@ -126,11 +125,10 @@ def _restart_on_unavailable( ) as span, MetricsCapture(resource_info), ): - ( - call_metadata, - current_request_id, - ) = request_id_manager.metadata_and_request_id( - nth_request, attempt, metadata, span + call_metadata, current_request_id = ( + request_id_manager.metadata_and_request_id( + nth_request, attempt, metadata, span + ) ) iterator = CrossSync._Sync_Impl.run_if_async( method, request=request, metadata=call_metadata @@ -202,19 +200,12 @@ def __init__(self, session, client_context=None): self._execute_sql_request_count: int = 0 self._read_request_count: int = 0 self._begin_request_sent: bool = False - - # Identifier for the transaction. self._transaction_id: Optional[bytes] = None self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None self._lock: CrossSync._Sync_Impl.Lock = CrossSync._Sync_Impl.Lock() - - # Operation within a transaction can be performed using multiple - # threads, so we need to use a lock when updating the transaction. - self._lock: threading.Lock = threading.Lock() - - # Event to coordinate concurrent requests beginning the transaction. - # This is used to prevent the "Transaction has not begun" race condition. - self._transaction_begin_event = threading.Event() + self._transaction_begin_event: CrossSync._Sync_Impl.Event = ( + CrossSync._Sync_Impl.Event() + ) @property def _resource_info(self): @@ -226,6 +217,33 @@ def _resource_info(self): "database": database.database_id, } + def _wait_for_transaction_begin(self) -> None: + """Claims the inline-begin for this request, or waits for it to complete. + + The first request against the transaction is the one that begins it + inline. Requests that are issued concurrently, before the transaction + id is available, must wait for that first request to complete instead + of assuming that the transaction has not begun. + + :raises ValueError: if the transaction has already been used to execute + a request, but is not a multi-use transaction, or if the concurrent + request that began the transaction did not complete in time.""" + with self._lock: + if self._begin_request_sent or self._read_request_count > 0: + if not self._multi_use: + raise ValueError("Cannot re-use single-use snapshot.") + wait_needed = self._transaction_id is None + else: + wait_needed = False + self._begin_request_sent = True + if not wait_needed: + return + CrossSync._Sync_Impl.event_wait( + self._transaction_begin_event, timeout=_TRANSACTION_BEGIN_TIMEOUT_SECONDS + ) + if not self._transaction_begin_event.is_set(): + raise ValueError("Timed out waiting for transaction to begin.") + def begin(self) -> bytes: """Begins a transaction on the database. @@ -330,31 +348,8 @@ def read( :returns: a result set instance which can be used to consume rows. :raises ValueError: if the Transaction already used to execute a - read request, but is not a multi-use transaction or has not begun. - """ - - with self._lock: - # Check if this request is beginning the transaction. - # If a request is already in progress, other requests must wait - # until the transaction ID is available. - if self._begin_request_sent or self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") - if self._transaction_id is None: - wait_needed = True - else: - wait_needed = False - else: - wait_needed = False - self._begin_request_sent = True - - if wait_needed: - # Wait for the transaction to begin (set by another concurrent request). - # This prevents the race condition where concurrent requests think - # the transaction hasn't begun. - if not self._transaction_begin_event.wait(timeout=30.0): - raise ValueError("Timed out waiting for transaction to begin.") - + read request, but is not a multi-use transaction or has not begun.""" + self._wait_for_transaction_begin() session = self._session database = session._database api = database.spanner_api @@ -526,31 +521,8 @@ def execute_sql( specific column in the given row. :raises ValueError: if the Transaction already used to execute a - read request, but is not a multi-use transaction or has not begun. - """ - - with self._lock: - # Check if this request is beginning the transaction. - # If a request is already in progress, other requests must wait - # until the transaction ID is available. - if self._begin_request_sent or self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") - if self._transaction_id is None: - wait_needed = True - else: - wait_needed = False - else: - wait_needed = False - self._begin_request_sent = True - - if wait_needed: - # Wait for the transaction to begin (set by another concurrent request). - # This prevents the race condition where concurrent requests think - # the transaction hasn't begun. - if not self._transaction_begin_event.wait(timeout=30.0): - raise ValueError("Timed out waiting for transaction to begin.") - + read request, but is not a multi-use transaction or has not begun.""" + self._wait_for_transaction_begin() if params is not None: params_pb = Struct( fields={key: _make_value_pb(value) for key, value in params.items()} @@ -910,9 +882,7 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None: """Updates the snapshot for the given transaction.""" if self._transaction_id is None and transaction_pb.id: self._transaction_id = transaction_pb.id - # Notify waiting threads that the transaction has begun. self._transaction_begin_event.set() - if transaction_pb._pb.HasField("precommit_token"): self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot_helpers.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot_helpers.py deleted file mode 100644 index 5e1d6840665a..000000000000 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot_helpers.py +++ /dev/null @@ -1,730 +0,0 @@ -# Copyright 2016 Google LLC All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -# This file is automatically generated by CrossSync. Do not edit manually. - -"""Model a set of read-only queries to a database as a snapshot.""" - -import functools -from typing import List, Optional, Union - -from google.api_core import gapic_v1 -from google.api_core.exceptions import ( - Aborted, - InternalServerError, - InvalidArgument, - ServiceUnavailable, -) -from google.protobuf.struct_pb2 import Struct - -from google.cloud.aio._cross_sync import CrossSync -from google.cloud.spanner_v1._helpers import ( - AtomicCounter, - _augment_error_with_request_id, - _check_rst_stream_error, - _make_value_pb, - _merge_query_options, - _metadata_with_leader_aware_routing, - _metadata_with_prefix, - _retry, - _SessionWrapper, -) -from google.cloud.spanner_v1._opentelemetry_tracing import add_span_event, trace_call -from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture -from google.cloud.spanner_v1.streamed import StreamedResultSet -from google.cloud.spanner_v1.types import MultiplexedSessionPrecommitToken -from google.cloud.spanner_v1.types.mutation import Mutation -from google.cloud.spanner_v1.types.result_set import PartialResultSet, ResultSet -from google.cloud.spanner_v1.types.spanner import ( - BeginTransactionRequest, - ExecuteSqlRequest, - PartitionOptions, - PartitionQueryRequest, - PartitionReadRequest, - ReadRequest, - RequestOptions, -) -from google.cloud.spanner_v1.types.transaction import ( - Transaction, - TransactionOptions, - TransactionSelector, -) - -_STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES = ( - "RST_STREAM", - "Received unexpected EOS on DATA frame from server", -) - - -def _restart_on_unavailable( - method, - request, - metadata=None, - trace_name=None, - session=None, - attributes=None, - transaction=None, - transaction_selector=None, - observability_options=None, - request_id_manager=None, -): - """Restart iteration after :exc:`.ServiceUnavailable`. - - :type method: callable - :param method: function returning iterator - - :type request: proto - :param request: request proto to call the method with - - :type transaction: :class:`google.cloud.spanner_v1.snapshot._SnapshotBase` - :param transaction: Snapshot or Transaction class object based on the type of transaction - - :type transaction_selector: :class:`transaction_pb2.TransactionSelector` - :param transaction_selector: Transaction selector object to be used in request if transaction is not passed, - if both transaction_selector and transaction are passed, then transaction is given priority. - """ - resume_token: bytes = b"" - item_buffer: List[PartialResultSet] = [] - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - elif transaction_selector is None: - raise InvalidArgument( - "Either transaction or transaction_selector should be set" - ) - request.transaction = transaction_selector - iterator = None - attempt = 1 - nth_request = getattr(request_id_manager, "_next_nth_request", 0) - current_request_id = None - while True: - try: - if iterator is None: - with ( - trace_call( - trace_name, - session, - attributes, - observability_options=observability_options, - metadata=metadata, - ) as span, - MetricsCapture(), - ): - ( - call_metadata, - current_request_id, - ) = request_id_manager.metadata_and_request_id( - nth_request, attempt, metadata, span - ) - iterator = CrossSync._Sync_Impl.run_if_async( - method, request=request, metadata=call_metadata - ) - item: PartialResultSet - for item in iterator: - item_buffer.append(item) - if transaction is not None: - transaction._update_for_result_set_pb(item) - if ( - item._pb is not None - and item._pb.HasField("precommit_token") - and (transaction is not None) - ): - transaction._update_for_precommit_token_pb(item.precommit_token) - if item.resume_token: - resume_token = item.resume_token - break - except ServiceUnavailable: - del item_buffer[:] - request.resume_token = resume_token - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - request.transaction = transaction_selector - attempt += 1 - iterator = None - continue - except InternalServerError as exc: - resumable_error = any( - ( - resumable_message in exc.message - for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES - ) - ) - if not resumable_error: - raise _augment_error_with_request_id(exc, current_request_id) - del item_buffer[:] - request.resume_token = resume_token - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - attempt += 1 - request.transaction = transaction_selector - iterator = None - continue - except Exception as exc: - raise _augment_error_with_request_id(exc, current_request_id) - if len(item_buffer) == 0: - break - for item in item_buffer: - yield item - del item_buffer[:] - - -class _SnapshotBase(_SessionWrapper): - """Base class for Snapshot. - - Allows reuse of API request methods with different transaction selector. - - :type session: :class:`~google.cloud.spanner_v1.session.Session` - :param session: the session used to perform transaction operations. - """ - - _read_only: bool = True - _multi_use: bool = False - - def __init__(self, session): - super().__init__(session) - self._execute_sql_request_count: int = 0 - self._read_request_count: int = 0 - self._transaction_id: Optional[bytes] = None - self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None - self._lock: CrossSync._Sync_Impl.Lock = CrossSync._Sync_Impl.Lock() - - def begin(self) -> bytes: - """Begins a transaction on the database. - - :rtype: bytes - :returns: identifier for the transaction. - - :raises ValueError: if the transaction has already begun.""" - return self._begin_transaction() - - def read( - self, - table, - columns, - keyset, - index="", - limit=0, - partition=None, - request_options=None, - data_boost_enabled=False, - directed_read_options=None, - *, - retry=gapic_v1.method.DEFAULT, - timeout=gapic_v1.method.DEFAULT, - column_info=None, - lazy_decode=False, - ): - """Perform a ``StreamingRead`` API request for rows in a table.""" - if self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") - if self._transaction_id is None: - raise ValueError("Transaction has not begun.") - session = self._session - database = session._database - api = database.spanner_api - metadata = _metadata_with_prefix(database.name) - if not self._read_only and database._route_to_leader_enabled: - metadata.append( - _metadata_with_leader_aware_routing(database._route_to_leader_enabled) - ) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - if self._read_only: - request_options.transaction_tag = None - if ( - directed_read_options is None - and database._directed_read_options is not None - ): - directed_read_options = database._directed_read_options - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag - read_request = ReadRequest( - session=session.name, - table=table, - columns=columns, - key_set=keyset._to_pb(), - index=index, - limit=limit, - partition_token=partition, - request_options=request_options, - data_boost_enabled=data_boost_enabled, - directed_read_options=directed_read_options, - ) - streaming_read_method = functools.partial( - api.streaming_read, - request=read_request, - metadata=metadata, - retry=retry, - timeout=timeout, - ) - return self._get_streamed_result_set( - method=streaming_read_method, - request=read_request, - metadata=metadata, - trace_attributes={ - "table_id": table, - "columns": columns, - "request_options": request_options, - }, - column_info=column_info, - lazy_decode=lazy_decode, - ) - - def execute_sql( - self, - sql, - params=None, - param_types=None, - query_mode=None, - query_options=None, - request_options=None, - last_statement=False, - partition=None, - retry=gapic_v1.method.DEFAULT, - timeout=gapic_v1.method.DEFAULT, - data_boost_enabled=False, - directed_read_options=None, - column_info=None, - lazy_decode=False, - ): - """Perform an ``ExecuteStreamingSql`` API request.""" - if self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") - if self._transaction_id is None: - raise ValueError("Transaction has not begun.") - if params is not None: - params_pb = Struct( - fields={key: _make_value_pb(value) for key, value in params.items()} - ) - else: - params_pb = {} - session = self._session - database = session._database - api = database.spanner_api - metadata = _metadata_with_prefix(database.name) - if not self._read_only and database._route_to_leader_enabled: - metadata.append( - _metadata_with_leader_aware_routing(database._route_to_leader_enabled) - ) - default_query_options = database._instance._client._query_options - query_options = _merge_query_options(default_query_options, query_options) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - if self._read_only: - request_options.transaction_tag = None - if ( - directed_read_options is None - and database._directed_read_options is not None - ): - directed_read_options = database._directed_read_options - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag - execute_sql_request = ExecuteSqlRequest( - session=session.name, - sql=sql, - params=params_pb, - param_types=param_types, - query_mode=query_mode, - partition_token=partition, - seqno=self._execute_sql_request_count, - query_options=query_options, - request_options=request_options, - last_statement=last_statement, - data_boost_enabled=data_boost_enabled, - directed_read_options=directed_read_options, - ) - execute_streaming_sql_method = functools.partial( - api.execute_streaming_sql, - request=execute_sql_request, - metadata=metadata, - retry=retry, - timeout=timeout, - ) - return self._get_streamed_result_set( - method=execute_streaming_sql_method, - request=execute_sql_request, - metadata=metadata, - trace_attributes={"db.statement": sql, "request_options": request_options}, - column_info=column_info, - lazy_decode=lazy_decode, - ) - - def _get_streamed_result_set( - self, method, request, metadata, trace_attributes, column_info, lazy_decode - ): - """Returns the streamed result set for a read or execute SQL request.""" - session = self._session - database = session._database - is_execute_sql_request = isinstance(request, ExecuteSqlRequest) - trace_method_name = "execute_sql" if is_execute_sql_request else "read" - trace_name = f"CloudSpanner.{type(self).__name__}.{trace_method_name}" - is_inline_begin = False - if self._transaction_id is None: - is_inline_begin = True - self._lock.acquire() - try: - iterator = _restart_on_unavailable( - method=method, - request=request, - session=session, - metadata=metadata, - trace_name=trace_name, - attributes=trace_attributes, - transaction=self, - observability_options=getattr(database, "observability_options", None), - request_id_manager=database, - ) - if is_execute_sql_request: - self._execute_sql_request_count += 1 - self._read_request_count += 1 - streamed_result_set_args = { - "response_iterator": iterator, - "column_info": column_info, - "lazy_decode": lazy_decode, - } - if self._multi_use: - streamed_result_set_args["source"] = self - return StreamedResultSet(**streamed_result_set_args) - finally: - if is_inline_begin: - self._lock.release() - - def partition_read( - self, - table, - columns, - keyset, - index="", - partition_size_bytes=None, - max_partitions=None, - *, - retry=gapic_v1.method.DEFAULT, - timeout=gapic_v1.method.DEFAULT, - ): - """Perform a ``PartitionRead`` API request for rows in a table.""" - if self._transaction_id is None: - raise ValueError("Transaction has not begun.") - if not self._multi_use: - raise ValueError("Cannot partition a single-use transaction.") - session = self._session - database = session._database - api = database.spanner_api - metadata = _metadata_with_prefix(database.name) - if database._route_to_leader_enabled: - metadata.append( - _metadata_with_leader_aware_routing(database._route_to_leader_enabled) - ) - transaction = self._build_transaction_selector_pb() - partition_options = PartitionOptions( - partition_size_bytes=partition_size_bytes, max_partitions=max_partitions - ) - partition_read_request = PartitionReadRequest( - session=session.name, - table=table, - columns=columns, - key_set=keyset._to_pb(), - transaction=transaction, - index=index, - partition_options=partition_options, - ) - trace_attributes = {"table_id": table, "columns": columns} - can_include_index = index != "" and index is not None - if can_include_index: - trace_attributes["index"] = index - with ( - trace_call( - f"CloudSpanner.{type(self).__name__}.partition_read", - session, - extra_attributes=trace_attributes, - observability_options=getattr(database, "observability_options", None), - metadata=metadata, - ) as span, - MetricsCapture(), - ): - nth_request = getattr(database, "_next_nth_request", 0) - attempt = AtomicCounter() - - def attempt_tracking_method(): - all_metadata = database.metadata_with_request_id( - nth_request, attempt.increment(), metadata, span - ) - partition_read_method = functools.partial( - api.partition_read, - request=partition_read_request, - metadata=all_metadata, - retry=retry, - timeout=timeout, - ) - return partition_read_method() - - response = _retry( - attempt_tracking_method, - allowed_exceptions={InternalServerError: _check_rst_stream_error}, - ) - return [partition.partition_token for partition in response.partitions] - - def partition_query( - self, - sql, - params=None, - param_types=None, - partition_size_bytes=None, - max_partitions=None, - *, - retry=gapic_v1.method.DEFAULT, - timeout=gapic_v1.method.DEFAULT, - ): - """Perform a ``PartitionQuery`` API request.""" - if self._transaction_id is None: - raise ValueError("Transaction has not begun.") - if not self._multi_use: - raise ValueError("Cannot partition a single-use transaction.") - if params is not None: - params_pb = Struct( - fields={key: _make_value_pb(value) for key, value in params.items()} - ) - else: - params_pb = Struct() - session = self._session - database = session._database - api = database.spanner_api - metadata = _metadata_with_prefix(database.name) - if database._route_to_leader_enabled: - metadata.append( - _metadata_with_leader_aware_routing(database._route_to_leader_enabled) - ) - transaction = self._build_transaction_selector_pb() - partition_options = PartitionOptions( - partition_size_bytes=partition_size_bytes, max_partitions=max_partitions - ) - partition_query_request = PartitionQueryRequest( - session=session.name, - sql=sql, - transaction=transaction, - params=params_pb, - param_types=param_types, - partition_options=partition_options, - ) - trace_attributes = {"db.statement": sql} - with ( - trace_call( - f"CloudSpanner.{type(self).__name__}.partition_query", - session, - trace_attributes, - observability_options=getattr(database, "observability_options", None), - metadata=metadata, - ) as span, - MetricsCapture(), - ): - nth_request = getattr(database, "_next_nth_request", 0) - attempt = AtomicCounter() - - def attempt_tracking_method(): - all_metadata = database.metadata_with_request_id( - nth_request, attempt.increment(), metadata, span - ) - partition_query_method = functools.partial( - api.partition_query, - request=partition_query_request, - metadata=all_metadata, - retry=retry, - timeout=timeout, - ) - return partition_query_method() - - response = _retry( - attempt_tracking_method, - allowed_exceptions={InternalServerError: _check_rst_stream_error}, - ) - return [partition.partition_token for partition in response.partitions] - - def _begin_transaction( - self, mutation: Mutation = None, transaction_tag: str = None - ) -> bytes: - """Begins a transaction on the database.""" - if self._transaction_id is not None: - raise ValueError("Transaction has already begun.") - if not self._multi_use: - raise ValueError("Cannot begin a single-use transaction.") - if self._read_request_count > 0: - raise ValueError("Read-only transaction already pending") - session = self._session - database = session._database - api = database.spanner_api - metadata = _metadata_with_prefix(database.name) - if not self._read_only and database._route_to_leader_enabled: - metadata.append( - _metadata_with_leader_aware_routing(database._route_to_leader_enabled) - ) - begin_request_kwargs = { - "session": session.name, - "options": self._build_transaction_selector_pb().begin, - "mutation_key": mutation, - } - if transaction_tag: - begin_request_kwargs["request_options"] = RequestOptions( - transaction_tag=transaction_tag - ) - with ( - trace_call( - name=f"CloudSpanner.{type(self).__name__}.begin", - session=session, - observability_options=getattr(database, "observability_options", None), - metadata=metadata, - ) as span, - MetricsCapture(), - ): - nth_request = getattr(database, "_next_nth_request", 0) - attempt = AtomicCounter() - - def wrapped_method(): - begin_transaction_request = BeginTransactionRequest( - **begin_request_kwargs - ) - call_metadata, error_augmenter = database.with_error_augmentation( - nth_request, attempt.increment(), metadata, span - ) - begin_transaction_method = functools.partial( - api.begin_transaction, - request=begin_transaction_request, - metadata=call_metadata, - ) - with error_augmenter: - return begin_transaction_method() - - def before_next_retry(nth_retry, delay_in_seconds): - add_span_event( - span=span, - event_name="Transaction Begin Attempt Failed. Retrying", - event_attributes={ - "attempt": nth_retry, - "sleep_seconds": delay_in_seconds, - }, - ) - - transaction_pb: Transaction = _retry( - wrapped_method, - before_next_retry=before_next_retry, - allowed_exceptions={ - InternalServerError: _check_rst_stream_error, - Aborted: None, - }, - ) - self._update_for_transaction_pb(transaction_pb) - return self._transaction_id - - def _build_transaction_options_pb(self) -> TransactionOptions: - """Builds and returns the transaction options for this snapshot.""" - raise NotImplementedError - - def _build_transaction_selector_pb(self) -> TransactionSelector: - """Builds and returns a transaction selector for this snapshot.""" - if self._transaction_id is not None: - return TransactionSelector(id=self._transaction_id) - options = self._build_transaction_options_pb() - if not self._multi_use: - return TransactionSelector(single_use=options) - return TransactionSelector(begin=options) - - def _update_for_result_set_pb( - self, result_set_pb: Union[ResultSet, PartialResultSet] - ) -> None: - """Updates the snapshot for the given result set.""" - if result_set_pb.metadata and result_set_pb.metadata.transaction: - self._update_for_transaction_pb(result_set_pb.metadata.transaction) - - def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None: - """Updates the snapshot for the given transaction.""" - if self._transaction_id is None and transaction_pb.id: - self._transaction_id = transaction_pb.id - if transaction_pb._pb.HasField("precommit_token"): - self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token) - - def _update_for_precommit_token_pb( - self, precommit_token_pb: MultiplexedSessionPrecommitToken - ) -> None: - """Updates the snapshot for the given multiplexed session precommit token.""" - with self._lock: - self._update_for_precommit_token_pb_unsafe(precommit_token_pb) - - def _update_for_precommit_token_pb_unsafe( - self, precommit_token_pb: MultiplexedSessionPrecommitToken - ) -> None: - """Updates the snapshot for the given multiplexed session precommit token.""" - if ( - self._precommit_token is None - or precommit_token_pb.seq_num > self._precommit_token.seq_num - ): - self._precommit_token = precommit_token_pb - - -class Snapshot(_SnapshotBase): - """Allow a set of reads / SQL statements with shared staleness.""" - - def __init__( - self, - session, - read_timestamp=None, - min_read_timestamp=None, - max_staleness=None, - exact_staleness=None, - multi_use=False, - transaction_id=None, - ): - super(Snapshot, self).__init__(session) - opts = [read_timestamp, min_read_timestamp, max_staleness, exact_staleness] - flagged = [opt for opt in opts if opt is not None] - if len(flagged) > 1: - raise ValueError("Supply zero or one options.") - if multi_use: - if min_read_timestamp is not None or max_staleness is not None: - raise ValueError( - "'multi_use' is incompatible with 'min_read_timestamp' / 'max_staleness'" - ) - self._transaction_read_timestamp = None - self._strong = len(flagged) == 0 - self._read_timestamp = read_timestamp - self._min_read_timestamp = min_read_timestamp - self._max_staleness = max_staleness - self._exact_staleness = exact_staleness - self._multi_use = multi_use - self._transaction_id = transaction_id - - def _build_transaction_options_pb(self) -> TransactionOptions: - """Builds and returns transaction options for this snapshot.""" - read_only_pb_args = dict(return_read_timestamp=True) - if self._read_timestamp: - read_only_pb_args["read_timestamp"] = self._read_timestamp - elif self._min_read_timestamp: - read_only_pb_args["min_read_timestamp"] = self._min_read_timestamp - elif self._max_staleness: - read_only_pb_args["max_staleness"] = self._max_staleness - elif self._exact_staleness: - read_only_pb_args["exact_staleness"] = self._exact_staleness - else: - read_only_pb_args["strong"] = True - read_only_pb = TransactionOptions.ReadOnly(**read_only_pb_args) - return TransactionOptions(read_only=read_only_pb) - - def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None: - """Updates the snapshot for the given transaction.""" - super(Snapshot, self)._update_for_transaction_pb(transaction_pb) - if transaction_pb.read_timestamp is not None: - self._transaction_read_timestamp = transaction_pb.read_timestamp diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py index a92f008f5e32..d95705af038a 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py @@ -111,6 +111,11 @@ def _merge_chunk(self, value): def _merge_values(self, values): """Merge values into rows. + Note: We manually check value.HasField("null_value") here instead of + wrapping every decoder in _parse_nullable to avoid the overhead of + an extra Python function call layer for every cell value decoded in this loop. + If the nullable check logic is updated in _parse_nullable, update this check. + :type values: list of :class:`~google.protobuf.struct_pb2.Value` :param values: non-chunked values from partial result set.""" decoders = self._decoders @@ -131,10 +136,6 @@ def _merge_values(self, values): index = 0 else: for value in values: - # Note: We manually check value.HasField("null_value") here instead of - # wrapping every decoder in _parse_nullable to avoid the overhead of - # an extra Python function call layer for every cell value decoded in this loop. - # If the nullable check logic is updated in _parse_nullable, update this check. if value.HasField("null_value"): current_row_append(None) else: @@ -184,8 +185,7 @@ def decode_row(self, row: []) -> []: The array that is returned by this function is the same as the array that would have been returned by the rows iterator if ``lazy_decoding=False``. - :returns: an array containing the decoded values of all the columns in the given row - """ + :returns: an array containing the decoded values of all the columns in the given row""" if not hasattr(row, "__len__"): raise TypeError("row", "row must be an array of protobuf values") decoders = self._decoders diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/database_test.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/database_test.py index 523946ab3545..ab66bdb030c6 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/database_test.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/database_test.py @@ -119,10 +119,9 @@ def spanner_api(self): return self._spanner_api def _create_spanner_client_for_tests(self, client_options, credentials): - ( - api_endpoint, - client_cert_source_func, - ) = SpannerClient.get_mtls_endpoint_and_cert_source(client_options) + api_endpoint, client_cert_source_func = ( + SpannerClient.get_mtls_endpoint_and_cert_source(client_options) + ) channel = grpc_helpers.create_channel( api_endpoint, credentials=credentials, diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py index 5ba207f07d49..43c58400d1ad 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py @@ -60,8 +60,7 @@ class Transaction(_SnapshotBase, _BatchBase): :type session: :class:`~google.cloud.spanner_v1.session.Session` :param session: the session used to perform the commit - :raises ValueError: if session has an existing transaction - """ + :raises ValueError: if session has an existing transaction""" exclude_txn_from_change_streams: bool = False isolation_level: TransactionOptions.IsolationLevel = ( @@ -686,8 +685,7 @@ def _update_for_execute_batch_dml_response_pb( """Update the transaction for the given execute batch DML response. :type response_pb: :class:`~google.cloud.spanner_v1.types.ExecuteBatchDmlResponse` - :param response_pb: The execute batch DML response to update the transaction with. - """ + :param response_pb: The execute batch DML response to update the transaction with.""" if len(response_pb.result_sets) > 0: self._update_for_result_set_pb(response_pb.result_sets[0]) diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py index bc902a6f63d1..e4e8785f2c43 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import datetime import unittest from datetime import timedelta @@ -71,21 +72,56 @@ def _make_snapshot(self, *args, **kwargs): async def test_read_errors(self): snapshot = self._make_snapshot(_Session(), multi_use=False) snapshot._read_request_count = 1 - with self.assertRaises(ValueError): + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): await snapshot.read(TABLE_NAME, COLUMNS, None) + @mock.patch( + "google.cloud.spanner_v1._async.snapshot._TRANSACTION_BEGIN_TIMEOUT_SECONDS", + 0.01, + ) + async def test_read_w_multi_use_not_begun_times_out(self): + # No concurrent request ever begins the transaction, so this request + # waits for the begin timeout to expire before giving up. snapshot = self._make_snapshot(_Session(), multi_use=True) snapshot._read_request_count = 1 snapshot._transaction_id = None - with self.assertRaises(ValueError): + with self.assertRaisesRegex( + ValueError, "Timed out waiting for transaction to begin." + ): await snapshot.read(TABLE_NAME, COLUMNS, None) async def test_execute_sql_errors(self): snapshot = self._make_snapshot(_Session(), multi_use=False) snapshot._read_request_count = 1 - with self.assertRaises(ValueError): + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): await snapshot.execute_sql(SQL_QUERY) + async def test_wait_for_transaction_begin_waits_for_concurrent_begin(self): + """A concurrent request must wait for the in-flight inline begin. + + Regression test: without the begin event, the second request observes + ``_transaction_id is None`` and wrongly raises "Transaction has not begun." + """ + + snapshot = self._make_snapshot(_Session(), multi_use=True) + + # First request claims the inline begin, but has not completed yet, so + # no transaction id is available. + await snapshot._wait_for_transaction_begin() + self.assertTrue(snapshot._begin_request_sent) + + concurrent_request = asyncio.ensure_future( + snapshot._wait_for_transaction_begin() + ) + + # The concurrent request blocks while the transaction id is unknown. + done, _ = await asyncio.wait([concurrent_request], timeout=0.1) + self.assertEqual(done, set()) + + # Completing the inline begin releases it. + snapshot._update_for_transaction_pb(TransactionPB(id=TXN_ID)) + await asyncio.wait_for(concurrent_request, timeout=10) + async def test_partition_read_ok(self): token_1 = b"TOKEN1" response = PartitionResponse( @@ -695,12 +731,20 @@ async def test_execute_sql_w_partition(self): call_args = api.execute_streaming_sql.call_args self.assertEqual(call_args.kwargs["request"].partition_token, b"token") + @mock.patch( + "google.cloud.spanner_v1._async.snapshot._TRANSACTION_BEGIN_TIMEOUT_SECONDS", + 0.01, + ) async def test_execute_sql_not_begun_error(self): + # No concurrent request ever begins the transaction, so this request + # waits for the begin timeout to expire before giving up. session = _Session() snapshot = self._make_snapshot(session, multi_use=True) snapshot._read_request_count = 1 snapshot._transaction_id = None - with self.assertRaises(ValueError): + with self.assertRaisesRegex( + ValueError, "Timed out waiting for transaction to begin." + ): await snapshot.execute_sql(SQL_QUERY) async def test_execute_sql_w_params(self): diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index a5082e5b8aa1..4a65d46aea8c 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from datetime import datetime, timedelta -from threading import Lock +from threading import Event, Lock, Thread from typing import Mapping import mock @@ -695,12 +695,69 @@ def test_ctor(self): self.assertFalse(derived._multi_use) self.assertEqual(derived._execute_sql_request_count, 0) self.assertEqual(derived._read_request_count, 0) + self.assertFalse(derived._begin_request_sent) self.assertIsNone(derived._transaction_id) self.assertIsNone(derived._precommit_token) self.assertIsInstance(derived._lock, type(Lock())) + self.assertFalse(derived._transaction_begin_event.is_set()) self.assertNoSpans() + def test__wait_for_transaction_begin_claims_inline_begin(self): + derived = _build_snapshot_derived(multi_use=True) + + derived._wait_for_transaction_begin() + + # The first request is the one that begins the transaction inline. + self.assertTrue(derived._begin_request_sent) + + def test__wait_for_transaction_begin_wo_multi_use(self): + derived = _build_snapshot_derived(multi_use=False) + derived._read_request_count = 1 + + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + derived._wait_for_transaction_begin() + + def test__wait_for_transaction_begin_waits_for_concurrent_begin(self): + """A concurrent request must wait for the in-flight inline begin. + + Regression test: without the begin event, the second request observes + ``_transaction_id is None`` and wrongly raises "Transaction has not begun." + """ + + from google.cloud.spanner_v1 import Transaction as TransactionPB + + derived = _build_snapshot_derived(multi_use=True) + + # First request claims the inline begin, but has not completed yet, so + # no transaction id is available. + derived._wait_for_transaction_begin() + + errors = [] + released = Event() + + def concurrent_request(): + try: + derived._wait_for_transaction_begin() + except Exception as exc: # pragma: no cover - only on regression + errors.append(exc) + finally: + released.set() + + thread = Thread(target=concurrent_request, daemon=True) + thread.start() + + # The concurrent request blocks while the transaction id is unknown. + self.assertFalse(released.wait(timeout=0.1)) + + # Completing the inline begin releases it. + derived._update_for_transaction_pb(TransactionPB(id=TXN_ID)) + self.assertTrue(released.wait(timeout=10)) + + thread.join(timeout=10) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + def test__build_transaction_selector_pb_single_use(self): derived = _build_snapshot_derived(multi_use=False) @@ -1203,8 +1260,15 @@ def test_read_w_multi_use_w_first_w_partition(self, mock_region): "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", ) + @mock.patch( + "google.cloud.spanner_v1.snapshot._TRANSACTION_BEGIN_TIMEOUT_SECONDS", 0.01 + ) def test_read_w_multi_use_w_first_w_count_gt_0(self, mock_region): - with self.assertRaises(ValueError): + # No concurrent request ever begins the transaction, so this request + # waits for the begin timeout to expire before giving up. + with self.assertRaisesRegex( + ValueError, "Timed out waiting for transaction to begin." + ): self._execute_read(multi_use=True, first=True, count=1) @mock.patch( @@ -1510,8 +1574,15 @@ def test_execute_sql_w_multi_use_wo_first_w_count_gt_0(self, mock_region): def test_execute_sql_w_multi_use_w_first(self, mock_region): self._execute_sql_helper(multi_use=True, first=True) + @mock.patch( + "google.cloud.spanner_v1.snapshot._TRANSACTION_BEGIN_TIMEOUT_SECONDS", 0.01 + ) def test_execute_sql_w_multi_use_w_first_w_count_gt_0(self): - with self.assertRaises(ValueError): + # No concurrent request ever begins the transaction, so this request + # waits for the begin timeout to expire before giving up. + with self.assertRaisesRegex( + ValueError, "Timed out waiting for transaction to begin." + ): self._execute_sql_helper(multi_use=True, first=True, count=1) @mock.patch(