diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index bb21147e4c..fccaa8e844 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,4 +13,4 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:3abfa0f1886adaf0b83f07cb117b24a639ea1cb9cffe56d43280b977033563eb + digest: sha256:3bf87e47c2173d7eed42714589dc4da2c07c3268610f1e47f8e1a30decbfc7f1 diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index 9c1b9be34e..05dc4672ed 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -20,9 +20,9 @@ cachetools==5.2.0 \ --hash=sha256:6a94c6402995a99c3970cc7e4884bb60b4a8639938157eeed436098bf9831757 \ --hash=sha256:f9f17d2aec496a9aa6b76f53e3b614c965223c061982d434d160f930c698a9db # via google-auth -certifi==2022.9.24 \ - --hash=sha256:0d9c601124e5a6ba9712dbc60d9c53c21e34f5f641fe83002317394311bdce14 \ - --hash=sha256:90c1a32f1d68f940488354e36370f6cca89f0f106db09518524c88d6ed83f382 +certifi==2022.12.7 \ + --hash=sha256:35824b4c3a97115964b408844d64aa14db1cc518f6562e8d7261699d1350a9e3 \ + --hash=sha256:4ad3232f5e926d6718ec31cfc1fcadfde020920e278684144551c91769c7bc18 # via requests cffi==1.15.1 \ --hash=sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5 \ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 46d237160f..5405cc8ff1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -25,7 +25,7 @@ repos: rev: 22.3.0 hooks: - id: black -- repo: https://gitlab.com/pycqa/flake8 +- repo: https://github.com/pycqa/flake8 rev: 3.9.2 hooks: - id: flake8 diff --git a/CHANGELOG.md b/CHANGELOG.md index 99f346e89a..9561b0cc7d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,18 @@ [1]: https://pypi.org/project/google-cloud-spanner/#history +## [3.26.0](https://github.com/googleapis/python-spanner/compare/v3.25.0...v3.26.0) (2022-12-15) + + +### Features + +* Inline Begin transction for RW transactions ([#840](https://github.com/googleapis/python-spanner/issues/840)) ([c2456be](https://github.com/googleapis/python-spanner/commit/c2456bed513dc4ab8954e5227605fca12e776b63)) + + +### Bug Fixes + +* Fix for binding of pinging and bursty pool with database role ([#871](https://github.com/googleapis/python-spanner/issues/871)) ([89da17e](https://github.com/googleapis/python-spanner/commit/89da17efccdf4f686f73f87f997128a96c614839)) + ## [3.25.0](https://github.com/googleapis/python-spanner/compare/v3.24.0...v3.25.0) (2022-12-13) diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 0d27763432..f919fa2c5e 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -578,7 +578,6 @@ def execute_pdml(): request = ExecuteSqlRequest( session=session.name, sql=dml, - transaction=txn_selector, params=params_pb, param_types=param_types, query_options=query_options, @@ -589,7 +588,11 @@ def execute_pdml(): metadata=metadata, ) - iterator = _restart_on_unavailable(method, request) + iterator = _restart_on_unavailable( + method=method, + request=request, + transaction_selector=txn_selector, + ) result_set = StreamedResultSet(iterator) list(result_set) # consume all partials diff --git a/google/cloud/spanner_v1/pool.py b/google/cloud/spanner_v1/pool.py index 216ba5aeff..886e28d7f7 100644 --- a/google/cloud/spanner_v1/pool.py +++ b/google/cloud/spanner_v1/pool.py @@ -21,7 +21,7 @@ from google.cloud.spanner_v1 import BatchCreateSessionsRequest from google.cloud.spanner_v1 import Session from google.cloud.spanner_v1._helpers import _metadata_with_prefix - +from warnings import warn _NOW = datetime.datetime.utcnow # unit tests may replace @@ -193,14 +193,14 @@ def bind(self, database): metadata = _metadata_with_prefix(database.name) self._database_role = self._database_role or self._database.database_role request = BatchCreateSessionsRequest( + database=database.database_id, + session_count=self.size - self._sessions.qsize(), session_template=Session(creator_role=self.database_role), ) while not self._sessions.full(): resp = api.batch_create_sessions( request=request, - database=database.name, - session_count=self.size - self._sessions.qsize(), metadata=metadata, ) for session_pb in resp.session: @@ -406,14 +406,14 @@ def bind(self, database): self._database_role = self._database_role or self._database.database_role request = BatchCreateSessionsRequest( + database=database.database_id, + session_count=self.size - created_session_count, session_template=Session(creator_role=self.database_role), ) while created_session_count < self.size: resp = api.batch_create_sessions( request=request, - database=database.name, - session_count=self.size - created_session_count, metadata=metadata, ) for session_pb in resp.session: @@ -497,6 +497,10 @@ def ping(self): class TransactionPingingPool(PingingPool): """Concrete session pool implementation: + Deprecated: TransactionPingingPool no longer begins a transaction for each of its sessions at startup. + Hence the TransactionPingingPool is same as :class:`PingingPool` and maybe removed in the future. + + In addition to the features of :class:`PingingPool`, this class creates and begins a transaction for each of its sessions at startup. @@ -532,6 +536,12 @@ def __init__( labels=None, database_role=None, ): + """This throws a deprecation warning on initialization.""" + warn( + f"{self.__class__.__name__} is deprecated.", + DeprecationWarning, + stacklevel=2, + ) self._pending_sessions = queue.Queue() super(TransactionPingingPool, self).__init__( @@ -579,7 +589,6 @@ def begin_pending_transactions(self): """Begin all transactions for sessions added to the pool.""" while not self._pending_sessions.empty(): session = self._pending_sessions.get() - session._transaction.begin() super(TransactionPingingPool, self).put(session) diff --git a/google/cloud/spanner_v1/session.py b/google/cloud/spanner_v1/session.py index c210f8f61d..5b1ca6fbb8 100644 --- a/google/cloud/spanner_v1/session.py +++ b/google/cloud/spanner_v1/session.py @@ -366,8 +366,6 @@ def run_in_transaction(self, func, *args, **kw): txn.transaction_tag = transaction_tag else: txn = self._transaction - if txn._transaction_id is None: - txn.begin() try: attempts += 1 diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index a55c3994c4..f1fff8b533 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -15,7 +15,7 @@ """Model a set of read-only queries to a database as a snapshot.""" import functools - +import threading from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import ReadRequest @@ -27,6 +27,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 import gapic_v1 from google.cloud.spanner_v1._helpers import _make_value_pb from google.cloud.spanner_v1._helpers import _merge_query_options @@ -43,7 +44,13 @@ def _restart_on_unavailable( - method, request, trace_name=None, session=None, attributes=None + method, + request, + trace_name=None, + session=None, + attributes=None, + transaction=None, + transaction_selector=None, ): """Restart iteration after :exc:`.ServiceUnavailable`. @@ -52,15 +59,41 @@ def _restart_on_unavailable( :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 = b"" item_buffer = [] + + if transaction is not None: + transaction_selector = transaction._make_txn_selector() + elif transaction_selector is None: + raise InvalidArgument( + "Either transaction or transaction_selector should be set" + ) + + request.transaction = transaction_selector with trace_call(trace_name, session, attributes): iterator = method(request=request) while True: try: for item in iterator: item_buffer.append(item) + # Setting the transaction id because the transaction begin was inlined for first rpc. + if ( + transaction is not None + and transaction._transaction_id is None + and item.metadata is not None + and item.metadata.transaction is not None + and item.metadata.transaction.id is not None + ): + transaction._transaction_id = item.metadata.transaction.id if item.resume_token: resume_token = item.resume_token break @@ -68,6 +101,9 @@ def _restart_on_unavailable( del item_buffer[:] with trace_call(trace_name, session, attributes): request.resume_token = resume_token + if transaction is not None: + transaction_selector = transaction._make_txn_selector() + request.transaction = transaction_selector iterator = method(request=request) continue except InternalServerError as exc: @@ -80,6 +116,9 @@ def _restart_on_unavailable( del item_buffer[:] with trace_call(trace_name, session, attributes): request.resume_token = resume_token + if transaction is not None: + transaction_selector = transaction._make_txn_selector() + request.transaction = transaction_selector iterator = method(request=request) continue @@ -106,6 +145,7 @@ class _SnapshotBase(_SessionWrapper): _transaction_id = None _read_request_count = 0 _execute_sql_count = 0 + _lock = threading.Lock() def _make_txn_selector(self): """Helper for :meth:`read` / :meth:`execute_sql`. @@ -180,13 +220,12 @@ def read( 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: + if self._transaction_id is None and self._read_only: raise ValueError("Transaction ID pending.") database = self._session._database api = database.spanner_api metadata = _metadata_with_prefix(database.name) - transaction = self._make_txn_selector() if request_options is None: request_options = RequestOptions() @@ -204,7 +243,6 @@ def read( table=table, columns=columns, key_set=keyset._to_pb(), - transaction=transaction, index=index, limit=limit, partition_token=partition, @@ -219,13 +257,32 @@ def read( ) trace_attributes = {"table_id": table, "columns": columns} - iterator = _restart_on_unavailable( - restart, - request, - "CloudSpanner.ReadOnlyTransaction", - self._session, - trace_attributes, - ) + + if self._transaction_id is None: + # lock is added to handle the inline begin for first rpc + with self._lock: + iterator = _restart_on_unavailable( + restart, + request, + "CloudSpanner.ReadOnlyTransaction", + self._session, + trace_attributes, + transaction=self, + ) + self._read_request_count += 1 + if self._multi_use: + return StreamedResultSet(iterator, source=self) + else: + return StreamedResultSet(iterator) + else: + iterator = _restart_on_unavailable( + restart, + request, + "CloudSpanner.ReadOnlyTransaction", + self._session, + trace_attributes, + transaction=self, + ) self._read_request_count += 1 @@ -301,7 +358,7 @@ def execute_sql( 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: + if self._transaction_id is None and self._read_only: raise ValueError("Transaction ID pending.") if params is not None: @@ -315,7 +372,7 @@ def execute_sql( database = self._session._database metadata = _metadata_with_prefix(database.name) - transaction = self._make_txn_selector() + api = database.spanner_api # Query-level options have higher precedence than client-level and @@ -336,7 +393,6 @@ def execute_sql( request = ExecuteSqlRequest( session=self._session.name, sql=sql, - transaction=transaction, params=params_pb, param_types=param_types, query_mode=query_mode, @@ -354,13 +410,34 @@ def execute_sql( ) trace_attributes = {"db.statement": sql} - iterator = _restart_on_unavailable( - restart, - request, - "CloudSpanner.ReadWriteTransaction", - self._session, - trace_attributes, - ) + + if self._transaction_id is None: + # lock is added to handle the inline begin for first rpc + with self._lock: + iterator = _restart_on_unavailable( + restart, + request, + "CloudSpanner.ReadWriteTransaction", + self._session, + trace_attributes, + transaction=self, + ) + self._read_request_count += 1 + self._execute_sql_count += 1 + + if self._multi_use: + return StreamedResultSet(iterator, source=self) + else: + return StreamedResultSet(iterator) + else: + iterator = _restart_on_unavailable( + restart, + request, + "CloudSpanner.ReadWriteTransaction", + self._session, + trace_attributes, + transaction=self, + ) self._read_request_count += 1 self._execute_sql_count += 1 diff --git a/google/cloud/spanner_v1/transaction.py b/google/cloud/spanner_v1/transaction.py index d776b12469..ce34054ab9 100644 --- a/google/cloud/spanner_v1/transaction.py +++ b/google/cloud/spanner_v1/transaction.py @@ -13,7 +13,8 @@ # limitations under the License. """Spanner read-write transaction support.""" - +import functools +import threading from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1._helpers import ( @@ -48,6 +49,7 @@ class Transaction(_SnapshotBase, _BatchBase): commit_stats = None _multi_use = True _execute_sql_count = 0 + _lock = threading.Lock() def __init__(self, session): if session._transaction is not None: @@ -61,8 +63,6 @@ def _check_state(self): :raises: :exc:`ValueError` if the object's state is invalid for making API requests. """ - if self._transaction_id is None: - raise ValueError("Transaction is not begun") if self.committed is not None: raise ValueError("Transaction is already committed") @@ -78,7 +78,31 @@ def _make_txn_selector(self): :returns: a selector configured for read-write transaction semantics. """ self._check_state() - return TransactionSelector(id=self._transaction_id) + + if self._transaction_id is None: + return TransactionSelector( + begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()) + ) + else: + return TransactionSelector(id=self._transaction_id) + + def _execute_request( + self, method, request, trace_name=None, session=None, attributes=None + ): + """Helper method to execute request after fetching transaction selector. + + :type method: callable + :param method: function returning iterator + + :type request: proto + :param request: request proto to call the method with + """ + transaction = self._make_txn_selector() + request.transaction = transaction + with trace_call(trace_name, session, attributes): + response = method(request=request) + + return response def begin(self): """Begin a transaction on the database. @@ -111,15 +135,17 @@ def begin(self): def rollback(self): """Roll back a transaction on the database.""" self._check_state() - database = self._session._database - api = database.spanner_api - metadata = _metadata_with_prefix(database.name) - with trace_call("CloudSpanner.Rollback", self._session): - api.rollback( - session=self._session.name, - transaction_id=self._transaction_id, - metadata=metadata, - ) + + if self._transaction_id is not None: + database = self._session._database + api = database.spanner_api + metadata = _metadata_with_prefix(database.name) + with trace_call("CloudSpanner.Rollback", self._session): + api.rollback( + session=self._session.name, + transaction_id=self._transaction_id, + metadata=metadata, + ) self.rolled_back = True del self._session._transaction @@ -142,6 +168,10 @@ def commit(self, return_commit_stats=False, request_options=None): :raises ValueError: if there are no mutations to commit. """ self._check_state() + if self._transaction_id is None and len(self._mutations) > 0: + self.begin() + elif self._transaction_id is None and len(self._mutations) == 0: + raise ValueError("Transaction is not begun") database = self._session._database api = database.spanner_api @@ -264,7 +294,6 @@ def execute_update( params_pb = self._make_params_pb(params, param_types) database = self._session._database metadata = _metadata_with_prefix(database.name) - transaction = self._make_txn_selector() api = database.spanner_api seqno, self._execute_sql_count = ( @@ -288,7 +317,6 @@ def execute_update( request = ExecuteSqlRequest( session=self._session.name, sql=dml, - transaction=transaction, params=params_pb, param_types=param_types, query_mode=query_mode, @@ -296,12 +324,42 @@ def execute_update( seqno=seqno, request_options=request_options, ) - with trace_call( - "CloudSpanner.ReadWriteTransaction", self._session, trace_attributes - ): - response = api.execute_sql( - request=request, metadata=metadata, retry=retry, timeout=timeout + + method = functools.partial( + api.execute_sql, + request=request, + metadata=metadata, + retry=retry, + timeout=timeout, + ) + + if self._transaction_id is None: + # lock is added to handle the inline begin for first rpc + with self._lock: + response = self._execute_request( + method, + request, + "CloudSpanner.ReadWriteTransaction", + self._session, + trace_attributes, + ) + # Setting the transaction id because the transaction begin was inlined for first rpc. + if ( + self._transaction_id is None + and response is not None + and response.metadata is not None + and response.metadata.transaction is not None + ): + self._transaction_id = response.metadata.transaction.id + else: + response = self._execute_request( + method, + request, + "CloudSpanner.ReadWriteTransaction", + self._session, + trace_attributes, ) + return response.stats.row_count_exact def batch_update(self, statements, request_options=None): @@ -348,7 +406,6 @@ def batch_update(self, statements, request_options=None): database = self._session._database metadata = _metadata_with_prefix(database.name) - transaction = self._make_txn_selector() api = database.spanner_api seqno, self._execute_sql_count = ( @@ -368,21 +425,53 @@ def batch_update(self, statements, request_options=None): } request = ExecuteBatchDmlRequest( session=self._session.name, - transaction=transaction, statements=parsed, seqno=seqno, request_options=request_options, ) - with trace_call("CloudSpanner.DMLTransaction", self._session, trace_attributes): - response = api.execute_batch_dml(request=request, metadata=metadata) + + method = functools.partial( + api.execute_batch_dml, + request=request, + metadata=metadata, + ) + + if self._transaction_id is None: + # lock is added to handle the inline begin for first rpc + with self._lock: + response = self._execute_request( + method, + request, + "CloudSpanner.DMLTransaction", + self._session, + trace_attributes, + ) + # Setting the transaction id because the transaction begin was inlined for first rpc. + for result_set in response.result_sets: + if ( + self._transaction_id is None + and result_set.metadata is not None + and result_set.metadata.transaction is not None + ): + self._transaction_id = result_set.metadata.transaction.id + break + else: + response = self._execute_request( + method, + request, + "CloudSpanner.DMLTransaction", + self._session, + trace_attributes, + ) + row_counts = [ result_set.stats.row_count_exact for result_set in response.result_sets ] + return response.status, row_counts def __enter__(self): """Begin ``with`` block.""" - self.begin() return self def __exit__(self, exc_type, exc_val, exc_tb): diff --git a/samples/samples/requirements.txt b/samples/samples/requirements.txt index 689f92044c..3ece31cb72 100644 --- a/samples/samples/requirements.txt +++ b/samples/samples/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-spanner==3.23.0 +google-cloud-spanner==3.25.0 futures==3.4.0; python_version < "3" diff --git a/setup.py b/setup.py index ddb8ca503b..e75a858af1 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,7 @@ name = "google-cloud-spanner" description = "Cloud Spanner API client library" -version = "3.25.0" +version = "3.26.0" # Should be one of: # 'Development Status :: 3 - Alpha' # 'Development Status :: 4 - Beta' diff --git a/tests/system/test_database_api.py b/tests/system/test_database_api.py index 9fac10ed4d..699b3f4a69 100644 --- a/tests/system/test_database_api.py +++ b/tests/system/test_database_api.py @@ -20,6 +20,7 @@ from google.api_core import exceptions from google.iam.v1 import policy_pb2 from google.cloud import spanner_v1 +from google.cloud.spanner_v1.pool import FixedSizePool, PingingPool from google.type import expr_pb2 from . import _helpers from . import _sample_data @@ -73,6 +74,61 @@ def test_create_database(shared_instance, databases_to_delete, database_dialect) assert temp_db.name in database_ids +def test_database_binding_of_fixed_size_pool( + not_emulator, shared_instance, databases_to_delete, not_postgres +): + temp_db_id = _helpers.unique_id("fixed_size_db", separator="_") + temp_db = shared_instance.database(temp_db_id) + + create_op = temp_db.create() + databases_to_delete.append(temp_db) + create_op.result(DBAPI_OPERATION_TIMEOUT) # raises on failure / timeout. + + # Create role and grant select permission on table contacts for parent role. + ddl_statements = _helpers.DDL_STATEMENTS + [ + "CREATE ROLE parent", + "GRANT SELECT ON TABLE contacts TO ROLE parent", + ] + operation = temp_db.update_ddl(ddl_statements) + operation.result(DBAPI_OPERATION_TIMEOUT) # raises on failure / timeout. + + pool = FixedSizePool( + size=1, + default_timeout=500, + database_role="parent", + ) + database = shared_instance.database(temp_db.name, pool=pool) + assert database._pool.database_role == "parent" + + +def test_database_binding_of_pinging_pool( + not_emulator, shared_instance, databases_to_delete, not_postgres +): + temp_db_id = _helpers.unique_id("binding_db", separator="_") + temp_db = shared_instance.database(temp_db_id) + + create_op = temp_db.create() + databases_to_delete.append(temp_db) + create_op.result(DBAPI_OPERATION_TIMEOUT) # raises on failure / timeout. + + # Create role and grant select permission on table contacts for parent role. + ddl_statements = _helpers.DDL_STATEMENTS + [ + "CREATE ROLE parent", + "GRANT SELECT ON TABLE contacts TO ROLE parent", + ] + operation = temp_db.update_ddl(ddl_statements) + operation.result(DBAPI_OPERATION_TIMEOUT) # raises on failure / timeout. + + pool = PingingPool( + size=1, + default_timeout=500, + ping_interval=100, + database_role="parent", + ) + database = shared_instance.database(temp_db.name, pool=pool) + assert database._pool.database_role == "parent" + + def test_create_database_pitr_invalid_retention_period( not_emulator, # PITR-lite features are not supported by the emulator not_postgres, diff --git a/tests/system/test_dbapi.py b/tests/system/test_dbapi.py index 0b92d7a15d..6354f2091f 100644 --- a/tests/system/test_dbapi.py +++ b/tests/system/test_dbapi.py @@ -314,27 +314,35 @@ def test_execute_many(shared_instance, dbapi_database): def test_DDL_autocommit(shared_instance, dbapi_database): """Check that DDLs in autocommit mode are immediately executed.""" - conn = Connection(shared_instance, dbapi_database) - conn.autocommit = True - cur = conn.cursor() - cur.execute( + try: + conn = Connection(shared_instance, dbapi_database) + conn.autocommit = True + + cur = conn.cursor() + cur.execute( + """ + CREATE TABLE Singers ( + SingerId INT64 NOT NULL, + Name STRING(1024), + ) PRIMARY KEY (SingerId) """ - CREATE TABLE Singers ( - SingerId INT64 NOT NULL, - Name STRING(1024), - ) PRIMARY KEY (SingerId) - """ - ) - conn.close() + ) + conn.close() - # if previous DDL wasn't committed, the next DROP TABLE - # statement will fail with a ProgrammingError - conn = Connection(shared_instance, dbapi_database) - cur = conn.cursor() + # if previous DDL wasn't committed, the next DROP TABLE + # statement will fail with a ProgrammingError + conn = Connection(shared_instance, dbapi_database) + cur = conn.cursor() - cur.execute("DROP TABLE Singers") - conn.commit() + cur.execute("DROP TABLE Singers") + conn.commit() + finally: + # Delete table + table = dbapi_database.table("Singers") + if table.exists(): + op = dbapi_database.update_ddl(["DROP TABLE Singers"]) + op.result() @pytest.mark.skipif(_helpers.USE_EMULATOR, reason="Emulator does not support json.") @@ -343,93 +351,114 @@ def test_autocommit_with_json_data(shared_instance, dbapi_database): Check that DDLs in autocommit mode are immediately executed for json fields. """ - # Create table - conn = Connection(shared_instance, dbapi_database) - conn.autocommit = True + try: + # Create table + conn = Connection(shared_instance, dbapi_database) + conn.autocommit = True - cur = conn.cursor() - cur.execute( + cur = conn.cursor() + cur.execute( + """ + CREATE TABLE JsonDetails ( + DataId INT64 NOT NULL, + Details JSON, + ) PRIMARY KEY (DataId) """ - CREATE TABLE JsonDetails ( - DataId INT64 NOT NULL, - Details JSON, - ) PRIMARY KEY (DataId) - """ - ) + ) - # Insert data to table - cur.execute( - sql="INSERT INTO JsonDetails (DataId, Details) VALUES (%s, %s)", - args=(123, JsonObject({"name": "Jakob", "age": "26"})), - ) + # Insert data to table + cur.execute( + sql="INSERT INTO JsonDetails (DataId, Details) VALUES (%s, %s)", + args=(123, JsonObject({"name": "Jakob", "age": "26"})), + ) - # Read back the data. - cur.execute("""select * from JsonDetails;""") - got_rows = cur.fetchall() + # Read back the data. + cur.execute("""select * from JsonDetails;""") + got_rows = cur.fetchall() - # Assert the response - assert len(got_rows) == 1 - assert got_rows[0][0] == 123 - assert got_rows[0][1] == {"age": "26", "name": "Jakob"} + # Assert the response + assert len(got_rows) == 1 + assert got_rows[0][0] == 123 + assert got_rows[0][1] == {"age": "26", "name": "Jakob"} - # Drop the table - cur.execute("DROP TABLE JsonDetails") - conn.commit() - conn.close() + # Drop the table + cur.execute("DROP TABLE JsonDetails") + conn.commit() + conn.close() + finally: + # Delete table + table = dbapi_database.table("JsonDetails") + if table.exists(): + op = dbapi_database.update_ddl(["DROP TABLE JsonDetails"]) + op.result() @pytest.mark.skipif(_helpers.USE_EMULATOR, reason="Emulator does not support json.") def test_json_array(shared_instance, dbapi_database): - # Create table - conn = Connection(shared_instance, dbapi_database) - conn.autocommit = True + try: + # Create table + conn = Connection(shared_instance, dbapi_database) + conn.autocommit = True - cur = conn.cursor() - cur.execute( + cur = conn.cursor() + cur.execute( + """ + CREATE TABLE JsonDetails ( + DataId INT64 NOT NULL, + Details JSON, + ) PRIMARY KEY (DataId) """ - CREATE TABLE JsonDetails ( - DataId INT64 NOT NULL, - Details JSON, - ) PRIMARY KEY (DataId) - """ - ) - cur.execute( - "INSERT INTO JsonDetails (DataId, Details) VALUES (%s, %s)", - [1, JsonObject([1, 2, 3])], - ) + ) + cur.execute( + "INSERT INTO JsonDetails (DataId, Details) VALUES (%s, %s)", + [1, JsonObject([1, 2, 3])], + ) - cur.execute("SELECT * FROM JsonDetails WHERE DataId = 1") - row = cur.fetchone() - assert isinstance(row[1], JsonObject) - assert row[1].serialize() == "[1,2,3]" + cur.execute("SELECT * FROM JsonDetails WHERE DataId = 1") + row = cur.fetchone() + assert isinstance(row[1], JsonObject) + assert row[1].serialize() == "[1,2,3]" - cur.execute("DROP TABLE JsonDetails") - conn.close() + cur.execute("DROP TABLE JsonDetails") + conn.close() + finally: + # Delete table + table = dbapi_database.table("JsonDetails") + if table.exists(): + op = dbapi_database.update_ddl(["DROP TABLE JsonDetails"]) + op.result() def test_DDL_commit(shared_instance, dbapi_database): """Check that DDLs in commit mode are executed on calling `commit()`.""" - conn = Connection(shared_instance, dbapi_database) - cur = conn.cursor() + try: + conn = Connection(shared_instance, dbapi_database) + cur = conn.cursor() - cur.execute( + cur.execute( + """ + CREATE TABLE Singers ( + SingerId INT64 NOT NULL, + Name STRING(1024), + ) PRIMARY KEY (SingerId) """ - CREATE TABLE Singers ( - SingerId INT64 NOT NULL, - Name STRING(1024), - ) PRIMARY KEY (SingerId) - """ - ) - conn.commit() - conn.close() + ) + conn.commit() + conn.close() - # if previous DDL wasn't committed, the next DROP TABLE - # statement will fail with a ProgrammingError - conn = Connection(shared_instance, dbapi_database) - cur = conn.cursor() + # if previous DDL wasn't committed, the next DROP TABLE + # statement will fail with a ProgrammingError + conn = Connection(shared_instance, dbapi_database) + cur = conn.cursor() - cur.execute("DROP TABLE Singers") - conn.commit() + cur.execute("DROP TABLE Singers") + conn.commit() + finally: + # Delete table + table = dbapi_database.table("Singers") + if table.exists(): + op = dbapi_database.update_ddl(["DROP TABLE Singers"]) + op.result() def test_ping(shared_instance, dbapi_database): @@ -505,53 +534,62 @@ def test_staleness(shared_instance, dbapi_database): @pytest.mark.parametrize("autocommit", [False, True]) def test_rowcount(shared_instance, dbapi_database, autocommit): - conn = Connection(shared_instance, dbapi_database) - conn.autocommit = autocommit - cur = conn.cursor() + try: + conn = Connection(shared_instance, dbapi_database) + conn.autocommit = autocommit + cur = conn.cursor() - cur.execute( + cur.execute( + """ + CREATE TABLE Singers ( + SingerId INT64 NOT NULL, + Name STRING(1024), + ) PRIMARY KEY (SingerId) """ - CREATE TABLE Singers ( - SingerId INT64 NOT NULL, - Name STRING(1024), - ) PRIMARY KEY (SingerId) - """ - ) - conn.commit() - - # executemany sets rowcount to the total modified rows - rows = [(i, f"Singer {i}") for i in range(100)] - cur.executemany("INSERT INTO Singers (SingerId, Name) VALUES (%s, %s)", rows[:98]) - assert cur.rowcount == 98 - - # execute with INSERT - cur.execute( - "INSERT INTO Singers (SingerId, Name) VALUES (%s, %s), (%s, %s)", - [x for row in rows[98:] for x in row], - ) - assert cur.rowcount == 2 - - # execute with UPDATE - cur.execute("UPDATE Singers SET Name = 'Cher' WHERE SingerId < 25") - assert cur.rowcount == 25 - - # execute with SELECT - cur.execute("SELECT Name FROM Singers WHERE SingerId < 75") - assert len(cur.fetchall()) == 75 - # rowcount is not available for SELECT - assert cur.rowcount == -1 - - # execute with DELETE - cur.execute("DELETE FROM Singers") - assert cur.rowcount == 100 + ) + conn.commit() - # execute with UPDATE matching 0 rows - cur.execute("UPDATE Singers SET Name = 'Cher' WHERE SingerId < 25") - assert cur.rowcount == 0 + # executemany sets rowcount to the total modified rows + rows = [(i, f"Singer {i}") for i in range(100)] + cur.executemany( + "INSERT INTO Singers (SingerId, Name) VALUES (%s, %s)", rows[:98] + ) + assert cur.rowcount == 98 - conn.commit() - cur.execute("DROP TABLE Singers") - conn.commit() + # execute with INSERT + cur.execute( + "INSERT INTO Singers (SingerId, Name) VALUES (%s, %s), (%s, %s)", + [x for row in rows[98:] for x in row], + ) + assert cur.rowcount == 2 + + # execute with UPDATE + cur.execute("UPDATE Singers SET Name = 'Cher' WHERE SingerId < 25") + assert cur.rowcount == 25 + + # execute with SELECT + cur.execute("SELECT Name FROM Singers WHERE SingerId < 75") + assert len(cur.fetchall()) == 75 + # rowcount is not available for SELECT + assert cur.rowcount == -1 + + # execute with DELETE + cur.execute("DELETE FROM Singers") + assert cur.rowcount == 100 + + # execute with UPDATE matching 0 rows + cur.execute("UPDATE Singers SET Name = 'Cher' WHERE SingerId < 25") + assert cur.rowcount == 0 + + conn.commit() + cur.execute("DROP TABLE Singers") + conn.commit() + finally: + # Delete table + table = dbapi_database.table("Singers") + if table.exists(): + op = dbapi_database.update_ddl(["DROP TABLE Singers"]) + op.result() @pytest.mark.parametrize("autocommit", [False, True]) diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index aedcbcaa55..c9c5c8a959 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -1027,6 +1027,7 @@ def test_transaction_batch_update_wo_statements(sessions_database, sessions_to_d sessions_to_delete.append(session) with session.transaction() as transaction: + transaction.begin() with pytest.raises(exceptions.InvalidArgument): transaction.batch_update([]) @@ -1088,11 +1089,10 @@ def unit_of_work(transaction): session.run_in_transaction(unit_of_work) span_list = ot_exporter.get_finished_spans() - assert len(span_list) == 6 + assert len(span_list) == 5 expected_span_names = [ "CloudSpanner.CreateSession", "CloudSpanner.Commit", - "CloudSpanner.BeginTransaction", "CloudSpanner.DMLTransaction", "CloudSpanner.Commit", "Test Span", diff --git a/tests/unit/test_pool.py b/tests/unit/test_pool.py index 1a53aa1604..3a9d35bc92 100644 --- a/tests/unit/test_pool.py +++ b/tests/unit/test_pool.py @@ -724,7 +724,7 @@ def test_bind(self): for session in SESSIONS: session.create.assert_not_called() txn = session._transaction - txn.begin.assert_called_once_with() + txn.begin.assert_not_called() self.assertTrue(pool._pending_sessions.empty()) @@ -753,7 +753,7 @@ def test_bind_w_timestamp_race(self): for session in SESSIONS: session.create.assert_not_called() txn = session._transaction - txn.begin.assert_called_once_with() + txn.begin.assert_not_called() self.assertTrue(pool._pending_sessions.empty()) @@ -839,7 +839,7 @@ def test_begin_pending_transactions_non_empty(self): pool.begin_pending_transactions() # no raise for txn in TRANSACTIONS: - txn.begin.assert_called_once_with() + txn.begin.assert_not_called() self.assertTrue(pending.empty()) @@ -956,11 +956,10 @@ def __init__(self, name): self.name = name self._sessions = [] self._database_role = None + self.database_id = name def mock_batch_create_sessions( request=None, - database=None, - session_count=10, timeout=10, metadata=[], labels={}, @@ -969,7 +968,7 @@ def mock_batch_create_sessions( from google.cloud.spanner_v1 import Session database_role = request.session_template.creator_role if request else None - if session_count < 2: + if request.session_count < 2: response = BatchCreateSessionsResponse( session=[Session(creator_role=database_role, labels=labels)] ) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 005cd0cd1f..edad4ce777 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -758,7 +758,6 @@ def test_transaction_w_existing_txn(self): def test_run_in_transaction_callback_raises_non_gax_error(self): from google.cloud.spanner_v1 import ( Transaction as TransactionPB, - TransactionOptions, ) from google.cloud.spanner_v1.transaction import Transaction @@ -799,24 +798,16 @@ def unit_of_work(txn, *args, **kw): self.assertTrue(txn.rolled_back) self.assertEqual(args, ()) self.assertEqual(kw, {}) - - expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite()) - gax_api.begin_transaction.assert_called_once_with( - session=self.SESSION_NAME, - options=expected_options, - metadata=[("google-cloud-resource-prefix", database.name)], - ) - gax_api.rollback.assert_called_once_with( - session=self.SESSION_NAME, - transaction_id=TRANSACTION_ID, - metadata=[("google-cloud-resource-prefix", database.name)], - ) + # Transaction only has mutation operations. + # Exception was raised before commit, hence transaction did not begin. + # Therefore rollback and begin transaction were not called. + gax_api.rollback.assert_not_called() + gax_api.begin_transaction.assert_not_called() def test_run_in_transaction_callback_raises_non_abort_rpc_error(self): from google.api_core.exceptions import Cancelled from google.cloud.spanner_v1 import ( Transaction as TransactionPB, - TransactionOptions, ) from google.cloud.spanner_v1.transaction import Transaction @@ -855,12 +846,6 @@ def unit_of_work(txn, *args, **kw): self.assertEqual(args, ()) self.assertEqual(kw, {}) - expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite()) - gax_api.begin_transaction.assert_called_once_with( - session=self.SESSION_NAME, - options=expected_options, - metadata=[("google-cloud-resource-prefix", database.name)], - ) gax_api.rollback.assert_not_called() def test_run_in_transaction_w_args_w_kwargs_wo_abort(self): @@ -1216,16 +1201,12 @@ def unit_of_work(txn, *args, **kw): self.assertEqual(kw, {}) expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite()) - self.assertEqual( - gax_api.begin_transaction.call_args_list, - [ - mock.call( - session=self.SESSION_NAME, - options=expected_options, - metadata=[("google-cloud-resource-prefix", database.name)], - ) - ] - * 2, + + # First call was aborted before commit operation, therefore no begin rpc was made during first attempt. + gax_api.begin_transaction.assert_called_once_with( + session=self.SESSION_NAME, + options=expected_options, + metadata=[("google-cloud-resource-prefix", database.name)], ) request = CommitRequest( session=self.SESSION_NAME, diff --git a/tests/unit/test_snapshot.py b/tests/unit/test_snapshot.py index 5b515f1bbb..c3ea162f11 100644 --- a/tests/unit/test_snapshot.py +++ b/tests/unit/test_snapshot.py @@ -49,23 +49,65 @@ class Test_restart_on_unavailable(OpenTelemetryBase): + def _getTargetClass(self): + from google.cloud.spanner_v1.snapshot import _SnapshotBase + + return _SnapshotBase + + def _makeDerived(self, session): + class _Derived(self._getTargetClass()): + + _transaction_id = None + _multi_use = False + + def _make_txn_selector(self): + from google.cloud.spanner_v1 import ( + TransactionOptions, + TransactionSelector, + ) + + if self._transaction_id: + return TransactionSelector(id=self._transaction_id) + options = TransactionOptions( + read_only=TransactionOptions.ReadOnly(strong=True) + ) + if self._multi_use: + return TransactionSelector(begin=options) + return TransactionSelector(single_use=options) + + return _Derived(session) + + def _make_spanner_api(self): + from google.cloud.spanner_v1 import SpannerClient + + return mock.create_autospec(SpannerClient, instance=True) + def _call_fut( - self, restart, request, span_name=None, session=None, attributes=None + self, derived, restart, request, span_name=None, session=None, attributes=None ): from google.cloud.spanner_v1.snapshot import _restart_on_unavailable - return _restart_on_unavailable(restart, request, span_name, session, attributes) + return _restart_on_unavailable( + restart, request, span_name, session, attributes, transaction=derived + ) - def _make_item(self, value, resume_token=b""): + def _make_item(self, value, resume_token=b"", metadata=None): return mock.Mock( - value=value, resume_token=resume_token, spec=["value", "resume_token"] + value=value, + resume_token=resume_token, + metadata=metadata, + spec=["value", "resume_token", "metadata"], ) def test_iteration_w_empty_raw(self): raw = _MockIterator() request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], return_value=raw) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), []) restart.assert_called_once_with(request=request) self.assertNoSpans() @@ -75,7 +117,11 @@ def test_iteration_w_non_empty_raw(self): raw = _MockIterator(*ITEMS) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], return_value=raw) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(ITEMS)) restart.assert_called_once_with(request=request) self.assertNoSpans() @@ -90,7 +136,11 @@ def test_iteration_w_raw_w_resume_tken(self): raw = _MockIterator(*ITEMS) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], return_value=raw) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(ITEMS)) restart.assert_called_once_with(request=request) self.assertNoSpans() @@ -107,7 +157,11 @@ def test_iteration_w_raw_raising_unavailable_no_token(self): after = _MockIterator(*ITEMS) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(ITEMS)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, b"") @@ -130,7 +184,11 @@ def test_iteration_w_raw_raising_retryable_internal_error_no_token(self): after = _MockIterator(*ITEMS) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(ITEMS)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, b"") @@ -148,7 +206,11 @@ def test_iteration_w_raw_raising_non_retryable_internal_error_no_token(self): after = _MockIterator(*ITEMS) request = mock.Mock(spec=["resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) with self.assertRaises(InternalServerError): list(resumable) restart.assert_called_once_with(request=request) @@ -166,7 +228,11 @@ def test_iteration_w_raw_raising_unavailable(self): after = _MockIterator(*LAST) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(FIRST + LAST)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, RESUME_TOKEN) @@ -188,7 +254,11 @@ def test_iteration_w_raw_raising_retryable_internal_error(self): after = _MockIterator(*LAST) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(FIRST + LAST)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, RESUME_TOKEN) @@ -206,7 +276,11 @@ def test_iteration_w_raw_raising_non_retryable_internal_error(self): after = _MockIterator(*LAST) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) with self.assertRaises(InternalServerError): list(resumable) restart.assert_called_once_with(request=request) @@ -223,12 +297,120 @@ def test_iteration_w_raw_raising_unavailable_after_token(self): after = _MockIterator(*SECOND) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(FIRST + SECOND)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, RESUME_TOKEN) self.assertNoSpans() + def test_iteration_w_raw_w_multiuse(self): + from google.cloud.spanner_v1 import ( + ReadRequest, + ) + + FIRST = ( + self._make_item(0), + self._make_item(1), + ) + before = _MockIterator(*FIRST) + request = ReadRequest(transaction=None) + restart = mock.Mock(spec=[], return_value=before) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + derived._multi_use = True + resumable = self._call_fut(derived, restart, request) + self.assertEqual(list(resumable), list(FIRST)) + self.assertEqual(len(restart.mock_calls), 1) + begin_count = sum( + [1 for args in restart.call_args_list if "begin" in args.kwargs.__str__()] + ) + self.assertEqual(begin_count, 1) + self.assertNoSpans() + + def test_iteration_w_raw_raising_unavailable_w_multiuse(self): + from google.api_core.exceptions import ServiceUnavailable + from google.cloud.spanner_v1 import ( + ReadRequest, + ) + + FIRST = ( + self._make_item(0), + self._make_item(1), + ) + SECOND = (self._make_item(2), self._make_item(3)) + before = _MockIterator( + *FIRST, fail_after=True, error=ServiceUnavailable("testing") + ) + after = _MockIterator(*SECOND) + request = ReadRequest(transaction=None) + restart = mock.Mock(spec=[], side_effect=[before, after]) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + derived._multi_use = True + resumable = self._call_fut(derived, restart, request) + self.assertEqual(list(resumable), list(SECOND)) + self.assertEqual(len(restart.mock_calls), 2) + begin_count = sum( + [1 for args in restart.call_args_list if "begin" in args.kwargs.__str__()] + ) + + # Since the transaction id was not set before the Unavailable error, the statement will be retried with inline begin. + self.assertEqual(begin_count, 2) + self.assertNoSpans() + + def test_iteration_w_raw_raising_unavailable_after_token_w_multiuse(self): + from google.api_core.exceptions import ServiceUnavailable + + from google.cloud.spanner_v1 import ResultSetMetadata + from google.cloud.spanner_v1 import ( + Transaction as TransactionPB, + ReadRequest, + ) + + transaction_pb = TransactionPB(id=TXN_ID) + metadata_pb = ResultSetMetadata(transaction=transaction_pb) + FIRST = ( + self._make_item(0), + self._make_item(1, resume_token=RESUME_TOKEN, metadata=metadata_pb), + ) + SECOND = (self._make_item(2), self._make_item(3)) + before = _MockIterator( + *FIRST, fail_after=True, error=ServiceUnavailable("testing") + ) + after = _MockIterator(*SECOND) + request = ReadRequest(transaction=None) + restart = mock.Mock(spec=[], side_effect=[before, after]) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + derived._multi_use = True + + resumable = self._call_fut(derived, restart, request) + + self.assertEqual(list(resumable), list(FIRST + SECOND)) + self.assertEqual(len(restart.mock_calls), 2) + self.assertEqual(request.resume_token, RESUME_TOKEN) + transaction_id_selector_count = sum( + [ + 1 + for args in restart.call_args_list + if 'id: "DEAFBEAD"' in args.kwargs.__str__() + ] + ) + + # Statement will be retried with Transaction id. + self.assertEqual(transaction_id_selector_count, 2) + self.assertNoSpans() + def test_iteration_w_raw_raising_retryable_internal_error_after_token(self): from google.api_core.exceptions import InternalServerError @@ -244,7 +426,11 @@ def test_iteration_w_raw_raising_retryable_internal_error_after_token(self): after = _MockIterator(*SECOND) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) self.assertEqual(list(resumable), list(FIRST + SECOND)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, RESUME_TOKEN) @@ -261,7 +447,11 @@ def test_iteration_w_raw_raising_non_retryable_internal_error_after_token(self): after = _MockIterator(*SECOND) request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) - resumable = self._call_fut(restart, request) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut(derived, restart, request) with self.assertRaises(InternalServerError): list(resumable) restart.assert_called_once_with(request=request) @@ -273,8 +463,12 @@ def test_iteration_w_span_creation(self): raw = _MockIterator() request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], return_value=raw) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) resumable = self._call_fut( - restart, request, name, _Session(_Database()), extra_atts + derived, restart, request, name, _Session(_Database()), extra_atts ) self.assertEqual(list(resumable), []) self.assertSpanAttributes(name, attributes=dict(BASE_ATTRIBUTES, test_att=1)) @@ -293,7 +487,13 @@ def test_iteration_w_multiple_span_creation(self): request = mock.Mock(test="test", spec=["test", "resume_token"]) restart = mock.Mock(spec=[], side_effect=[before, after]) name = "TestSpan" - resumable = self._call_fut(restart, request, name, _Session(_Database())) + database = _Database() + database.spanner_api = self._make_spanner_api() + session = _Session(database) + derived = self._makeDerived(session) + resumable = self._call_fut( + derived, restart, request, name, _Session(_Database()) + ) self.assertEqual(list(resumable), list(FIRST + LAST)) self.assertEqual(len(restart.mock_calls), 2) self.assertEqual(request.resume_token, RESUME_TOKEN) @@ -876,7 +1076,6 @@ def _partition_read_helper( derived._multi_use = multi_use if w_txn: derived._transaction_id = TXN_ID - tokens = list( derived.partition_read( TABLE_NAME, diff --git a/tests/unit/test_spanner.py b/tests/unit/test_spanner.py new file mode 100644 index 0000000000..a7c41c5f4f --- /dev/null +++ b/tests/unit/test_spanner.py @@ -0,0 +1,873 @@ +# Copyright 2022 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. + + +import threading +from google.protobuf.struct_pb2 import Struct +from google.cloud.spanner_v1 import ( + PartialResultSet, + ResultSetMetadata, + ResultSetStats, + ResultSet, + RequestOptions, + Type, + TypeCode, + ExecuteSqlRequest, + ReadRequest, + StructType, + TransactionOptions, + TransactionSelector, + ExecuteBatchDmlRequest, + ExecuteBatchDmlResponse, + param_types, +) +from google.cloud.spanner_v1.types import transaction as transaction_type +from google.cloud.spanner_v1.keyset import KeySet + +from google.cloud.spanner_v1._helpers import ( + _make_value_pb, + _merge_query_options, +) + +import mock + +from google.api_core import gapic_v1 + +from tests._helpers import OpenTelemetryBase + +TABLE_NAME = "citizens" +COLUMNS = ["email", "first_name", "last_name", "age"] +VALUES = [ + ["phred@exammple.com", "Phred", "Phlyntstone", 32], + ["bharney@example.com", "Bharney", "Rhubble", 31], +] +DML_QUERY = """\ +INSERT INTO citizens(first_name, last_name, age) +VALUES ("Phred", "Phlyntstone", 32) +""" +DML_QUERY_WITH_PARAM = """ +INSERT INTO citizens(first_name, last_name, age) +VALUES ("Phred", "Phlyntstone", @age) +""" +SQL_QUERY = """\ +SELECT first_name, last_name, age FROM citizens ORDER BY age""" +SQL_QUERY_WITH_PARAM = """ +SELECT first_name, last_name, email FROM citizens WHERE age <= @max_age""" +PARAMS = {"age": 30} +PARAM_TYPES = {"age": Type(code=TypeCode.INT64)} +KEYS = [["bharney@example.com"], ["phred@example.com"]] +KEYSET = KeySet(keys=KEYS) +INDEX = "email-address-index" +LIMIT = 20 +MODE = 2 +RETRY = gapic_v1.method.DEFAULT +TIMEOUT = gapic_v1.method.DEFAULT +REQUEST_OPTIONS = RequestOptions() +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} +update_dml = 'UPDATE table SET desc = desc + "-amended"' +delete_dml = "DELETE FROM table WHERE desc IS NULL" + +dml_statements = [ + (insert_dml, insert_params, insert_param_types), + update_dml, + delete_dml, +] + + +class TestTransaction(OpenTelemetryBase): + + PROJECT_ID = "project-id" + INSTANCE_ID = "instance-id" + INSTANCE_NAME = "projects/" + PROJECT_ID + "/instances/" + INSTANCE_ID + DATABASE_ID = "database-id" + DATABASE_NAME = INSTANCE_NAME + "/databases/" + DATABASE_ID + SESSION_ID = "session-id" + SESSION_NAME = DATABASE_NAME + "/sessions/" + SESSION_ID + TRANSACTION_ID = b"DEADBEEF" + TRANSACTION_TAG = "transaction-tag" + + BASE_ATTRIBUTES = { + "db.type": "spanner", + "db.url": "spanner.googleapis.com", + "db.instance": "testing", + "net.host.name": "spanner.googleapis.com", + } + + def _getTargetClass(self): + from google.cloud.spanner_v1.transaction import Transaction + + return Transaction + + def _make_one(self, session, *args, **kwargs): + transaction = self._getTargetClass()(session, *args, **kwargs) + session._transaction = transaction + return transaction + + def _make_spanner_api(self): + from google.cloud.spanner_v1 import SpannerClient + + return mock.create_autospec(SpannerClient, instance=True) + + def _execute_update_helper( + self, + transaction, + api, + count=0, + query_options=None, + ): + stats_pb = ResultSetStats(row_count_exact=1) + + transaction_pb = transaction_type.Transaction(id=self.TRANSACTION_ID) + metadata_pb = ResultSetMetadata(transaction=transaction_pb) + api.execute_sql.return_value = ResultSet(stats=stats_pb, metadata=metadata_pb) + + transaction.transaction_tag = self.TRANSACTION_TAG + transaction._execute_sql_count = count + + row_count = transaction.execute_update( + DML_QUERY_WITH_PARAM, + PARAMS, + PARAM_TYPES, + query_mode=MODE, + query_options=query_options, + request_options=REQUEST_OPTIONS, + retry=RETRY, + timeout=TIMEOUT, + ) + self.assertEqual(row_count, count + 1) + + def _execute_update_expected_request( + self, database, query_options=None, begin=True, count=0 + ): + if begin is True: + expected_transaction = TransactionSelector( + begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()) + ) + else: + expected_transaction = TransactionSelector(id=self.TRANSACTION_ID) + + expected_params = Struct( + fields={key: _make_value_pb(value) for (key, value) in PARAMS.items()} + ) + + expected_query_options = database._instance._client._query_options + if query_options: + expected_query_options = _merge_query_options( + expected_query_options, query_options + ) + expected_request_options = REQUEST_OPTIONS + expected_request_options.transaction_tag = self.TRANSACTION_TAG + + expected_request = ExecuteSqlRequest( + session=self.SESSION_NAME, + sql=DML_QUERY_WITH_PARAM, + transaction=expected_transaction, + params=expected_params, + param_types=PARAM_TYPES, + query_mode=MODE, + query_options=expected_query_options, + request_options=expected_request_options, + seqno=count, + ) + + return expected_request + + def _execute_sql_helper( + self, + transaction, + api, + count=0, + partition=None, + sql_count=0, + query_options=None, + ): + VALUES = [["bharney", "rhubbyl", 31], ["phred", "phlyntstone", 32]] + VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES] + struct_type_pb = StructType( + fields=[ + StructType.Field(name="first_name", type_=Type(code=TypeCode.STRING)), + StructType.Field(name="last_name", type_=Type(code=TypeCode.STRING)), + StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), + ] + ) + transaction_pb = transaction_type.Transaction(id=self.TRANSACTION_ID) + metadata_pb = ResultSetMetadata( + row_type=struct_type_pb, transaction=transaction_pb + ) + stats_pb = ResultSetStats( + query_stats=Struct(fields={"rows_returned": _make_value_pb(2)}) + ) + result_sets = [ + PartialResultSet(metadata=metadata_pb), + PartialResultSet(stats=stats_pb), + ] + for i in range(len(result_sets)): + result_sets[i].values.extend(VALUE_PBS[i]) + iterator = _MockIterator(*result_sets) + api.execute_streaming_sql.return_value = iterator + transaction._execute_sql_count = sql_count + transaction._read_request_count = count + + result_set = transaction.execute_sql( + SQL_QUERY_WITH_PARAM, + PARAMS, + PARAM_TYPES, + query_mode=MODE, + query_options=query_options, + request_options=REQUEST_OPTIONS, + partition=partition, + retry=RETRY, + timeout=TIMEOUT, + ) + + self.assertEqual(transaction._read_request_count, count + 1) + + self.assertEqual(list(result_set), VALUES) + self.assertEqual(result_set.metadata, metadata_pb) + self.assertEqual(result_set.stats, stats_pb) + 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 + ): + if begin is True: + expected_transaction = TransactionSelector( + begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()) + ) + else: + expected_transaction = TransactionSelector(id=self.TRANSACTION_ID) + + expected_params = Struct( + fields={key: _make_value_pb(value) for (key, value) in PARAMS.items()} + ) + + expected_query_options = database._instance._client._query_options + if query_options: + expected_query_options = _merge_query_options( + expected_query_options, query_options + ) + + expected_request_options = REQUEST_OPTIONS + expected_request_options.transaction_tag = None + expected_request = ExecuteSqlRequest( + session=self.SESSION_NAME, + sql=SQL_QUERY_WITH_PARAM, + transaction=expected_transaction, + params=expected_params, + param_types=PARAM_TYPES, + query_mode=MODE, + query_options=expected_query_options, + request_options=expected_request_options, + partition_token=partition, + seqno=sql_count, + ) + + return expected_request + + def _read_helper( + self, + transaction, + api, + count=0, + partition=None, + ): + VALUES = [["bharney", 31], ["phred", 32]] + VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES] + struct_type_pb = StructType( + fields=[ + StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), + StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), + ] + ) + + transaction_pb = transaction_type.Transaction(id=self.TRANSACTION_ID) + metadata_pb = ResultSetMetadata( + row_type=struct_type_pb, transaction=transaction_pb + ) + + stats_pb = ResultSetStats( + query_stats=Struct(fields={"rows_returned": _make_value_pb(2)}) + ) + result_sets = [ + PartialResultSet(metadata=metadata_pb), + PartialResultSet(stats=stats_pb), + ] + for i in range(len(result_sets)): + result_sets[i].values.extend(VALUE_PBS[i]) + + api.streaming_read.return_value = _MockIterator(*result_sets) + transaction._read_request_count = count + + if partition is not None: # 'limit' and 'partition' incompatible + result_set = transaction.read( + TABLE_NAME, + COLUMNS, + KEYSET, + index=INDEX, + partition=partition, + retry=RETRY, + timeout=TIMEOUT, + request_options=REQUEST_OPTIONS, + ) + else: + result_set = transaction.read( + TABLE_NAME, + COLUMNS, + KEYSET, + index=INDEX, + limit=LIMIT, + retry=RETRY, + timeout=TIMEOUT, + request_options=REQUEST_OPTIONS, + ) + + self.assertEqual(transaction._read_request_count, count + 1) + + self.assertIs(result_set._source, transaction) + + self.assertEqual(list(result_set), VALUES) + 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): + + if begin is True: + expected_transaction = TransactionSelector( + begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()) + ) + else: + expected_transaction = TransactionSelector(id=self.TRANSACTION_ID) + + if partition is not None: + expected_limit = 0 + else: + expected_limit = LIMIT + + # Transaction tag is ignored for read request. + expected_request_options = REQUEST_OPTIONS + expected_request_options.transaction_tag = None + + expected_request = ReadRequest( + session=self.SESSION_NAME, + table=TABLE_NAME, + columns=COLUMNS, + key_set=KEYSET._to_pb(), + transaction=expected_transaction, + index=INDEX, + limit=expected_limit, + partition_token=partition, + request_options=expected_request_options, + ) + + return expected_request + + def _batch_update_helper( + self, + transaction, + database, + api, + error_after=None, + count=0, + ): + from google.rpc.status_pb2 import Status + + stats_pbs = [ + ResultSetStats(row_count_exact=1), + ResultSetStats(row_count_exact=2), + ResultSetStats(row_count_exact=3), + ] + if error_after is not None: + stats_pbs = stats_pbs[:error_after] + expected_status = Status(code=400) + else: + expected_status = Status(code=200) + expected_row_counts = [stats.row_count_exact for stats in stats_pbs] + transaction_pb = transaction_type.Transaction(id=self.TRANSACTION_ID) + metadata_pb = ResultSetMetadata(transaction=transaction_pb) + result_sets_pb = [ + ResultSet(stats=stats_pb, metadata=metadata_pb) for stats_pb in stats_pbs + ] + + response = ExecuteBatchDmlResponse( + status=expected_status, + result_sets=result_sets_pb, + ) + + api.execute_batch_dml.return_value = response + transaction.transaction_tag = self.TRANSACTION_TAG + transaction._execute_sql_count = count + + status, row_counts = transaction.batch_update( + dml_statements, request_options=REQUEST_OPTIONS + ) + + self.assertEqual(status, expected_status) + self.assertEqual(row_counts, expected_row_counts) + self.assertEqual(transaction._execute_sql_count, count + 1) + + def _batch_update_expected_request(self, begin=True, count=0): + if begin is True: + expected_transaction = TransactionSelector( + begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()) + ) + else: + expected_transaction = TransactionSelector(id=self.TRANSACTION_ID) + + expected_insert_params = Struct( + fields={ + key: _make_value_pb(value) for (key, value) in insert_params.items() + } + ) + expected_statements = [ + ExecuteBatchDmlRequest.Statement( + sql=insert_dml, + params=expected_insert_params, + param_types=insert_param_types, + ), + ExecuteBatchDmlRequest.Statement(sql=update_dml), + ExecuteBatchDmlRequest.Statement(sql=delete_dml), + ] + + expected_request_options = REQUEST_OPTIONS + expected_request_options.transaction_tag = self.TRANSACTION_TAG + + expected_request = ExecuteBatchDmlRequest( + session=self.SESSION_NAME, + transaction=expected_transaction, + statements=expected_statements, + seqno=count, + request_options=expected_request_options, + ) + + return expected_request + + def test_transaction_should_include_begin_with_first_update(self): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._execute_update_helper(transaction=transaction, api=api) + + api.execute_sql.assert_called_once_with( + request=self._execute_update_expected_request(database=database), + retry=RETRY, + timeout=TIMEOUT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_transaction_should_include_begin_with_first_query(self): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._execute_sql_helper(transaction=transaction, api=api) + + api.execute_streaming_sql.assert_called_once_with( + request=self._execute_sql_expected_request(database=database), + metadata=[("google-cloud-resource-prefix", database.name)], + timeout=TIMEOUT, + retry=RETRY, + ) + + def test_transaction_should_include_begin_with_first_read(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) + + api.streaming_read.assert_called_once_with( + request=self._read_helper_expected_request(), + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + def test_transaction_should_include_begin_with_first_batch_update(self): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._batch_update_helper(transaction=transaction, database=database, api=api) + api.execute_batch_dml.assert_called_once_with( + request=self._batch_update_expected_request(), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_transaction_should_use_transaction_id_if_error_with_first_batch_update( + self, + ): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._batch_update_helper( + transaction=transaction, database=database, api=api, error_after=2 + ) + api.execute_batch_dml.assert_called_once_with( + request=self._batch_update_expected_request(begin=True), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + self._execute_update_helper(transaction=transaction, api=api) + api.execute_sql.assert_called_once_with( + request=self._execute_update_expected_request( + database=database, begin=False + ), + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_transaction_should_use_transaction_id_returned_by_first_query(self): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._execute_sql_helper(transaction=transaction, api=api) + api.execute_streaming_sql.assert_called_once_with( + request=self._execute_sql_expected_request(database=database), + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + self._execute_update_helper(transaction=transaction, api=api) + api.execute_sql.assert_called_once_with( + request=self._execute_update_expected_request( + database=database, begin=False + ), + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_transaction_should_use_transaction_id_returned_by_first_update(self): + database = _Database() + session = _Session(database) + api = database.spanner_api = self._make_spanner_api() + transaction = self._make_one(session) + self._execute_update_helper(transaction=transaction, api=api) + api.execute_sql.assert_called_once_with( + request=self._execute_update_expected_request(database=database), + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + self._execute_sql_helper(transaction=transaction, api=api) + api.execute_streaming_sql.assert_called_once_with( + request=self._execute_sql_expected_request(database=database, begin=False), + retry=gapic_v1.method.DEFAULT, + timeout=gapic_v1.method.DEFAULT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_transaction_should_use_transaction_id_returned_by_first_read(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) + api.streaming_read.assert_called_once_with( + request=self._read_helper_expected_request(), + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + self._batch_update_helper(transaction=transaction, database=database, api=api) + api.execute_batch_dml.assert_called_once_with( + request=self._batch_update_expected_request(begin=False), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_transaction_should_use_transaction_id_returned_by_first_batch_update(self): + database = _Database() + api = database.spanner_api = self._make_spanner_api() + session = _Session(database) + transaction = self._make_one(session) + self._batch_update_helper(transaction=transaction, database=database, api=api) + api.execute_batch_dml.assert_called_once_with( + request=self._batch_update_expected_request(), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + self._read_helper(transaction=transaction, api=api) + api.streaming_read.assert_called_once_with( + request=self._read_helper_expected_request(begin=False), + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_execute_update( + self, + ): + database = _Database() + api = database.spanner_api = self._make_spanner_api() + session = _Session(database) + transaction = self._make_one(session) + threads = [] + threads.append( + threading.Thread( + target=self._execute_update_helper, + kwargs={"transaction": transaction, "api": api}, + ) + ) + threads.append( + threading.Thread( + target=self._execute_update_helper, + kwargs={"transaction": transaction, "api": api}, + ) + ) + for thread in threads: + thread.start() + + for thread in threads: + thread.join() + + self._batch_update_helper(transaction=transaction, database=database, api=api) + + api.execute_sql.assert_any_call( + request=self._execute_update_expected_request(database), + retry=RETRY, + timeout=TIMEOUT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + api.execute_sql.assert_any_call( + request=self._execute_update_expected_request(database, begin=False), + retry=RETRY, + timeout=TIMEOUT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + api.execute_batch_dml.assert_any_call( + request=self._batch_update_expected_request(begin=False), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + self.assertEqual(api.execute_sql.call_count, 2) + self.assertEqual(api.execute_batch_dml.call_count, 1) + + def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_batch_update( + self, + ): + database = _Database() + api = database.spanner_api = self._make_spanner_api() + session = _Session(database) + transaction = self._make_one(session) + threads = [] + threads.append( + threading.Thread( + target=self._batch_update_helper, + kwargs={"transaction": transaction, "database": database, "api": api}, + ) + ) + threads.append( + threading.Thread( + target=self._batch_update_helper, + kwargs={"transaction": transaction, "database": database, "api": api}, + ) + ) + for thread in threads: + thread.start() + + for thread in threads: + thread.join() + + self._execute_update_helper(transaction=transaction, api=api) + + api.execute_sql.assert_any_call( + request=self._execute_update_expected_request(database, begin=False), + retry=RETRY, + timeout=TIMEOUT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + api.execute_batch_dml.assert_any_call( + request=self._batch_update_expected_request(), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + api.execute_batch_dml.assert_any_call( + request=self._batch_update_expected_request(begin=False), + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + self.assertEqual(api.execute_sql.call_count, 1) + self.assertEqual(api.execute_batch_dml.call_count, 2) + + def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_read( + self, + ): + database = _Database() + api = database.spanner_api = self._make_spanner_api() + session = _Session(database) + transaction = self._make_one(session) + threads = [] + threads.append( + threading.Thread( + target=self._read_helper, + kwargs={"transaction": transaction, "api": api}, + ) + ) + threads.append( + threading.Thread( + target=self._read_helper, + kwargs={"transaction": transaction, "api": api}, + ) + ) + for thread in threads: + thread.start() + + for thread in threads: + thread.join() + + self._execute_update_helper(transaction=transaction, api=api) + + begin_read_write_count = sum( + [1 for call in api.mock_calls if "read_write" in call.kwargs.__str__()] + ) + + self.assertEqual(begin_read_write_count, 1) + api.execute_sql.assert_any_call( + request=self._execute_update_expected_request(database, begin=False), + retry=RETRY, + timeout=TIMEOUT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + api.streaming_read.assert_any_call( + request=self._read_helper_expected_request(), + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + api.streaming_read.assert_any_call( + request=self._read_helper_expected_request(begin=False), + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + self.assertEqual(api.execute_sql.call_count, 1) + self.assertEqual(api.streaming_read.call_count, 2) + + def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_query( + self, + ): + database = _Database() + api = database.spanner_api = self._make_spanner_api() + session = _Session(database) + transaction = self._make_one(session) + threads = [] + threads.append( + threading.Thread( + target=self._execute_sql_helper, + kwargs={"transaction": transaction, "api": api}, + ) + ) + threads.append( + threading.Thread( + target=self._execute_sql_helper, + kwargs={"transaction": transaction, "api": api}, + ) + ) + for thread in threads: + thread.start() + + for thread in threads: + thread.join() + + self._execute_update_helper(transaction=transaction, api=api) + + begin_read_write_count = sum( + [1 for call in api.mock_calls if "read_write" in call.kwargs.__str__()] + ) + + self.assertEqual(begin_read_write_count, 1) + api.execute_sql.assert_any_call( + request=self._execute_update_expected_request(database, begin=False), + retry=RETRY, + timeout=TIMEOUT, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + req = self._execute_sql_expected_request(database) + api.execute_streaming_sql.assert_any_call( + request=req, + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + api.execute_streaming_sql.assert_any_call( + request=self._execute_sql_expected_request(database, begin=False), + metadata=[("google-cloud-resource-prefix", database.name)], + retry=RETRY, + timeout=TIMEOUT, + ) + + self.assertEqual(api.execute_sql.call_count, 1) + self.assertEqual(api.execute_streaming_sql.call_count, 2) + + +class _Client(object): + def __init__(self): + from google.cloud.spanner_v1 import ExecuteSqlRequest + + self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1") + + +class _Instance(object): + def __init__(self): + self._client = _Client() + + +class _Database(object): + def __init__(self): + self.name = "testing" + self._instance = _Instance() + + +class _Session(object): + + _transaction = None + + def __init__(self, database=None, name=TestTransaction.SESSION_NAME): + self._database = database + self.name = name + + +class _MockIterator(object): + def __init__(self, *values, **kw): + self._iter_values = iter(values) + self._fail_after = kw.pop("fail_after", False) + self._error = kw.pop("error", Exception) + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._iter_values) + except StopIteration: + if self._fail_after: + raise self._error + raise + + next = __next__ diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py index d4d9c99c02..5fb69b4979 100644 --- a/tests/unit/test_transaction.py +++ b/tests/unit/test_transaction.py @@ -91,12 +91,6 @@ def test_ctor_defaults(self): self.assertTrue(transaction._multi_use) self.assertEqual(transaction._execute_sql_count, 0) - def test__check_state_not_begun(self): - session = _Session() - transaction = self._make_one(session) - with self.assertRaises(ValueError): - transaction._check_state() - def test__check_state_already_committed(self): session = _Session() transaction = self._make_one(session) @@ -195,10 +189,16 @@ def test_begin_ok(self): ) def test_rollback_not_begun(self): - session = _Session() + database = _Database() + api = database.spanner_api = self._make_spanner_api() + session = _Session(database) transaction = self._make_one(session) - with self.assertRaises(ValueError): - transaction.rollback() + + transaction.rollback() + self.assertTrue(transaction.rolled_back) + + # Since there was no transaction to be rolled back, rollbacl rpc is not called. + api.rollback.assert_not_called() self.assertNoSpans() @@ -835,16 +835,11 @@ def test_context_mgr_failure(self): raise Exception("bail out") self.assertEqual(transaction.committed, None) + # Rollback rpc will not be called as there is no transaction id to be rolled back, rolled_back flag will be marked as true. self.assertTrue(transaction.rolled_back) self.assertEqual(len(transaction._mutations), 1) - self.assertEqual(api._committed, None) - session_id, txn_id, metadata = api._rolled_back - self.assertEqual(session_id, session.name) - self.assertEqual(txn_id, self.TRANSACTION_ID) - self.assertEqual(metadata, [("google-cloud-resource-prefix", database.name)]) - class _Client(object): def __init__(self):