From f32b057cce6e1106ef54621b285d447ed939be6e Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Wed, 15 Jul 2026 21:43:36 +0000 Subject: [PATCH 1/2] feat: add send and ack methods to Batch for Cloud Spanner Queues and include unit and system tests --- .../google/cloud/spanner_v1/_async/batch.py | 52 +++++++++++ .../google/cloud/spanner_v1/batch.py | 52 +++++++++++ .../tests/system/_async/test_database_api.py | 87 +++++++++++++++++++ .../tests/system/test_database_api.py | 84 ++++++++++++++++++ .../tests/unit/_async/test_batch.py | 31 +++++++ .../tests/unit/test_batch.py | 34 ++++++++ 6 files changed, 340 insertions(+) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py index 8fda14c3f1a2..519122622276 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py @@ -23,9 +23,12 @@ from google.cloud.aio._cross_sync import CrossSync from google.cloud.spanner_v1._async._helpers import _retry, _retry_on_aborted_exception +from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.spanner_v1._helpers import ( AtomicCounter, _check_rst_stream_error, + _make_value_pb, + _make_list_value_pb, _make_list_value_pbs, _merge_client_context, _merge_request_options, @@ -165,6 +168,55 @@ def delete(self, table, keyset): # TODO: Decide if we should add a span event per mutation: # https://github.com/googleapis/python-spanner/issues/1269 + def send(self, queue, key, payload=None, deliver_time=None): + """Send a message to a Cloud Spanner queue. + + :type queue: str + :param queue: Name of the queue to which the message will be sent. + + :type key: list + :param key: The primary key of the message to be sent. + + :type payload: object + :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. + """ + 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)) + + def ack(self, queue, key, ignore_not_found=None): + """Acknowledge a message in a Cloud Spanner queue. + + :type queue: str + :param queue: Name of the queue where the message to be acked is stored. + + :type key: list + :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. + """ + 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)) + class Batch(_BatchBase): """Accumulate mutations for transmission during :meth:`commit`.""" 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 6bb79e5e2955..af26874b5a0f 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py @@ -23,9 +23,12 @@ from google.api_core.exceptions import InternalServerError +from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.spanner_v1._helpers import ( AtomicCounter, _check_rst_stream_error, + _make_value_pb, + _make_list_value_pb, _make_list_value_pbs, _merge_client_context, _merge_request_options, @@ -142,6 +145,55 @@ def delete(self, table, keyset): delete = Mutation.Delete(table=table, key_set=keyset._to_pb()) self._mutations.append(Mutation(delete=delete)) + def send(self, queue, key, payload=None, deliver_time=None): + """Send a message to a Cloud Spanner queue. + + :type queue: str + :param queue: Name of the queue to which the message will be sent. + + :type key: list + :param key: The primary key of the message to be sent. + + :type payload: object + :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. + """ + 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)) + + def ack(self, queue, key, ignore_not_found=None): + """Acknowledge a message in a Cloud Spanner queue. + + :type queue: str + :param queue: Name of the queue where the message to be acked is stored. + + :type key: list + :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. + """ + 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)) + class Batch(_BatchBase): """Accumulate mutations for transmission during :meth:`commit`.""" diff --git a/packages/google-cloud-spanner/tests/system/_async/test_database_api.py b/packages/google-cloud-spanner/tests/system/_async/test_database_api.py index fa0ffaeab623..376627a5f62d 100644 --- a/packages/google-cloud-spanner/tests/system/_async/test_database_api.py +++ b/packages/google-cloud-spanner/tests/system/_async/test_database_api.py @@ -87,6 +87,93 @@ async def test_db_batch_insert_then_db_snapshot_read(shared_database): sd._check_rows_data(from_snap) +@pytest.mark.asyncio +async def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, instance_config): + import uuid + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin + from google.cloud.spanner_admin_database_v1 import DatabaseDialect + from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError + + instance_id = f"test-instance-{uuid.uuid4().hex[:8]}" + db_name = f"test-db-{uuid.uuid4().hex[:8]}" + queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" + + config_name = instance_config.name + request = spanner_instance_admin.CreateInstanceRequest( + parent=spanner_client.project_name, + instance_id=instance_id, + instance=spanner_instance_admin.Instance( + config=config_name, + display_name=instance_id, + node_count=1, + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE, + ), + ) + print(f"instance creation request: {request}") + operation = await spanner_client.instance_admin_api.create_instance(request=request) + operation.result(600) + + test_instance = spanner_client.instance(instance_id, configuration_name=config_name) + + try: + test_database = await test_instance.database(db_name, database_dialect=database_dialect) + operation = await test_database.create() + operation.result(300) + print("Database created successfully!") + + try: + # 3. Create the Queue + if database_dialect == DatabaseDialect.POSTGRESQL: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + id bigint NOT NULL, + "Payload" varchar NOT NULL, + PRIMARY KEY (id) + )""" + else: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + Id INT64 NOT NULL, + Payload STRING(MAX) NOT NULL + ) PRIMARY KEY (Id)""" + + try: + operation = await test_database.update_ddl([queue_ddl]) + await operation.result(600) + except MethodNotImplemented as e: + print(f"MethodNotImplemented. Skipping test because Queues are not implemented yet: {e}") + return + except GoogleAPIError as e: + if getattr(e, 'code', None) == 501 or (getattr(e, 'grpc_status_code', None) and e.grpc_status_code.name == 'UNIMPLEMENTED') or "UNIMPLEMENTED" in str(e): + print(f"Skipping test because Queues are not implemented yet: {e}") + return + raise + print("Queue created successfully.") + + # 4. Run mutations + print("Sending message to queue...") + async with test_database.batch() as batch: + batch.send( + queue=queue_name, + key=(2,), + payload="Hello, Queues!", + ) + print("Send successful.") + + print("Acking message in queue...") + async with test_database.batch() as batch: + batch.ack( + queue=queue_name, + key=(2,), + ) + print("Ack successful.") + + finally: + print("Dropping database...") + await test_database.drop() + finally: + print("Dropping instance...") + await test_instance.delete() + + @pytest.mark.asyncio async def test_db_run_in_transaction_then_snapshot_execute_sql(shared_database): await shared_database.reload() diff --git a/packages/google-cloud-spanner/tests/system/test_database_api.py b/packages/google-cloud-spanner/tests/system/test_database_api.py index 793d72f18374..517282612e99 100644 --- a/packages/google-cloud-spanner/tests/system/test_database_api.py +++ b/packages/google-cloud-spanner/tests/system/test_database_api.py @@ -567,6 +567,90 @@ def test_db_batch_insert_then_db_snapshot_read(shared_database): sd._check_rows_data(from_snap) +def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, instance_config): + import uuid + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin + from google.cloud.spanner_admin_database_v1 import DatabaseDialect + from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError + + instance_id = f"test-instance-{uuid.uuid4().hex[:8]}" + db_id = f"test-db-{uuid.uuid4().hex[:8]}" + queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" + + config_name = instance_config.name + request = spanner_instance_admin.CreateInstanceRequest( + parent=spanner_client.project_name, + instance_id=instance_id, + instance=spanner_instance_admin.Instance( + config=config_name, + display_name=instance_id, + node_count=1, + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE, + ), + ) + print(f"instance creation request: {request}") + operation = spanner_client.instance_admin_api.create_instance(request=request) + operation.result(600) + test_instance = spanner_client.instance(instance_id, configuration_name=config_name) + + try: + test_database = test_instance.database(db_id, database_dialect=database_dialect) + operation = test_database.create() + operation.result(300) + print("Database created successfully!") + + try: + # 3. Create the Queue + if database_dialect == DatabaseDialect.POSTGRESQL: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + id bigint NOT NULL, + "Payload" varchar NOT NULL, + PRIMARY KEY (id) + )""" + else: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + Id INT64 NOT NULL, + Payload STRING(MAX) NOT NULL + ) PRIMARY KEY (Id)""" + try: + operation = test_database.update_ddl([queue_ddl]) + operation.result(600) + except MethodNotImplemented as e: + print(f"MethodNotImplemented. Skipping test because Queues are not implemented yet: {e}") + return + except GoogleAPIError as e: + if getattr(e, 'code', None) == 501 or getattr(e, 'grpc_status_code', None) and e.grpc_status_code.name == 'UNIMPLEMENTED' or "UNIMPLEMENTED" in str(e): + print(f"Skipping test because Queues are not implemented yet: {e}") + return + raise + print("Queue created successfully.") + + # 4. Run mutations + print("Sending message to queue...") + with test_database.batch() as batch: + batch.send( + queue=queue_name, + key=(2,), + payload="Hello, Queues!", + ) + print("Send successful.") + + print("Acking message in queue...") + with test_database.batch() as batch: + batch.ack( + queue=queue_name, + key=(2,), + ) + print("Ack successful.") + + finally: + print("Dropping database...") + test_database.drop() + finally: + print("Dropping instance...") + test_instance.delete() + + def test_db_run_in_transaction_then_snapshot_execute_sql(shared_database): _helpers.retry_has_all_dll(shared_database.reload)() sd = _sample_data diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py index 9cd49742b289..94d623caf285 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py @@ -43,6 +43,15 @@ def _getTargetClass(self): def _make_one(self, *args, **kwargs): return self._getTargetClass()(*args, **kwargs) + def _compare_values(self, result, source): + for found, expected in zip(result, source): + self.assertEqual(len(found), len(expected)) + for found_cell, expected_cell in zip(found, expected): + if isinstance(expected_cell, int): + self.assertEqual(int(found_cell), expected_cell) + else: + self.assertEqual(found_cell, expected_cell) + def test_ctor(self): session = mock.Mock() base = self._make_one(session) @@ -85,6 +94,28 @@ def test_delete(self): self.assertEqual(len(base._mutations), 1) self.assertEqual(base._mutations[0].delete.table, TABLE_NAME) + def test_send(self): + queue = "TestQueue" + key = [2] + payload = "Hello, Queues!" + session = mock.Mock() + base = self._make_one(session) + base.send(queue=queue, key=key, payload=payload) + self.assertEqual(len(base._mutations), 1) + self.assertEqual(base._mutations[0].send.queue, queue) + self.assertEqual(base._mutations[0].send.payload, payload) + self._compare_values([base._mutations[0].send.key], [key]) + + def test_ack(self): + queue = "TestQueue" + key = [2] + session = mock.Mock() + base = self._make_one(session) + base.ack(queue=queue, key=key) + self.assertEqual(len(base._mutations), 1) + self.assertEqual(base._mutations[0].ack.queue, queue) + self._compare_values([base._mutations[0].ack.key], [key]) + class TestBatch(unittest.IsolatedAsyncioTestCase): def _getTargetClass(self): diff --git a/packages/google-cloud-spanner/tests/unit/test_batch.py b/packages/google-cloud-spanner/tests/unit/test_batch.py index 228c2025c14c..db4ebba0b7f4 100644 --- a/packages/google-cloud-spanner/tests/unit/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/test_batch.py @@ -182,6 +182,40 @@ def test_delete(self): for found, expected in zip(key_set_pb.keys, keys): self.assertEqual([int(value) for value in found], expected) + def test_send(self): + queue = "TestQueue" + key = [2] + payload = "Hello, Queues!" + session = _Session() + base = self._make_one(session) + + base.send(queue, key=key, payload=payload) + + self.assertEqual(len(base._mutations), 1) + mutation = base._mutations[0] + self.assertIsInstance(mutation, Mutation) + send = mutation.send + self.assertIsInstance(send, Mutation.Send) + self.assertEqual(send.queue, queue) + self.assertEqual(send.payload, payload) + self._compare_values([send.key], [key]) + + def test_ack(self): + queue = "TestQueue" + key = [2] + session = _Session() + base = self._make_one(session) + + base.ack(queue, key=key) + + self.assertEqual(len(base._mutations), 1) + mutation = base._mutations[0] + self.assertIsInstance(mutation, Mutation) + ack = mutation.ack + self.assertIsInstance(ack, Mutation.Ack) + self.assertEqual(ack.queue, queue) + self._compare_values([ack.key], [key]) + class TestBatch(_BaseTest, OpenTelemetryBase): def _getTargetClass(self): From 19d5083b6f3c8284e0025bbe89ec660b611ce9f7 Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Thu, 16 Jul 2026 21:21:59 +0000 Subject: [PATCH 2/2] fix comments: - Use shared instance fixture - Update the shared instance to use ENTERPRISE_PLUS edition - Use pytest skip when Queues not supported --- packages/google-cloud-spanner/test_spanner.py | 9 ++ .../tests/system/_async/conftest.py | 34 ++++- .../tests/system/_async/test_database_api.py | 113 +++++++---------- .../tests/system/conftest.py | 17 ++- .../tests/system/test_database_api.py | 117 +++++++----------- .../tests/unit/_async/test_batch.py | 16 +-- .../tests/unit/test_batch.py | 8 +- 7 files changed, 153 insertions(+), 161 deletions(-) create mode 100644 packages/google-cloud-spanner/test_spanner.py diff --git a/packages/google-cloud-spanner/test_spanner.py b/packages/google-cloud-spanner/test_spanner.py new file mode 100644 index 000000000000..1444183ddd7a --- /dev/null +++ b/packages/google-cloud-spanner/test_spanner.py @@ -0,0 +1,9 @@ +from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin +request = spanner_instance_admin.UpdateInstanceRequest( + instance=spanner_instance_admin.Instance( + name="projects/my-project/instances/my-instance", + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE, + ), + field_mask={"paths": ["edition"]}, +) +print("SUCCESS") diff --git a/packages/google-cloud-spanner/tests/system/_async/conftest.py b/packages/google-cloud-spanner/tests/system/_async/conftest.py index ce103488552f..997e701983b7 100644 --- a/packages/google-cloud-spanner/tests/system/_async/conftest.py +++ b/packages/google-cloud-spanner/tests/system/_async/conftest.py @@ -138,12 +138,33 @@ async def shared_instance( instance_config, ): spanner_client._instance_admin_api = None - instance = spanner_client.instance(shared_instance_id, instance_config.name) if _helpers.CREATE_INSTANCE: - op = await instance.create() - await op.result(instance_operation_timeout) + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin + import time + + create_time = str(int(time.time())) + labels = {"python-spanner-systests": "true", "created": create_time} + + request = spanner_instance_admin.CreateInstanceRequest( + parent=spanner_client.project_name, + instance_id=shared_instance_id, + instance=spanner_instance_admin.Instance( + config=instance_config.name, + display_name=shared_instance_id, + node_count=1, + labels=labels, + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE_PLUS, + ), + ) + created_op = await spanner_client.instance_admin_api.create_instance(request=request) + await created_op.result(instance_operation_timeout) + + instance = spanner_client.instance( + shared_instance_id, instance_config.name, labels=labels + ) else: + instance = spanner_client.instance(shared_instance_id, instance_config.name) await instance.reload() yield instance @@ -205,3 +226,10 @@ async def databases_to_delete(): def not_postgres(database_dialect): if database_dialect == DatabaseDialect.POSTGRESQL: pytest.skip("Skip for Postgres") + + +@pytest.fixture(scope="function") +def not_emulator(): + if _helpers.USE_EMULATOR: + pytest.skip(f"{_helpers.USE_EMULATOR_ENVVAR} set in environment.") + diff --git a/packages/google-cloud-spanner/tests/system/_async/test_database_api.py b/packages/google-cloud-spanner/tests/system/_async/test_database_api.py index 376627a5f62d..1beef4e4f457 100644 --- a/packages/google-cloud-spanner/tests/system/_async/test_database_api.py +++ b/packages/google-cloud-spanner/tests/system/_async/test_database_api.py @@ -88,90 +88,61 @@ async def test_db_batch_insert_then_db_snapshot_read(shared_database): @pytest.mark.asyncio -async def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, instance_config): +async def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, shared_instance): import uuid from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin from google.cloud.spanner_admin_database_v1 import DatabaseDialect from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError - instance_id = f"test-instance-{uuid.uuid4().hex[:8]}" db_name = f"test-db-{uuid.uuid4().hex[:8]}" queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" - config_name = instance_config.name - request = spanner_instance_admin.CreateInstanceRequest( - parent=spanner_client.project_name, - instance_id=instance_id, - instance=spanner_instance_admin.Instance( - config=config_name, - display_name=instance_id, - node_count=1, - edition=spanner_instance_admin.Instance.Edition.ENTERPRISE, - ), - ) - print(f"instance creation request: {request}") - operation = await spanner_client.instance_admin_api.create_instance(request=request) - operation.result(600) - - test_instance = spanner_client.instance(instance_id, configuration_name=config_name) + test_database = await shared_instance.database(db_name, database_dialect=database_dialect) + operation = await test_database.create() + operation.result(300) try: - test_database = await test_instance.database(db_name, database_dialect=database_dialect) - operation = await test_database.create() - operation.result(300) - print("Database created successfully!") + # Create the Queue + if database_dialect == DatabaseDialect.POSTGRESQL: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + id bigint NOT NULL, + "Payload" varchar NOT NULL, + PRIMARY KEY (id) + )""" + else: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + Id INT64 NOT NULL, + Payload STRING(MAX) NOT NULL + ) PRIMARY KEY (Id)""" try: - # 3. Create the Queue - if database_dialect == DatabaseDialect.POSTGRESQL: - queue_ddl = f"""CREATE QUEUE {queue_name} ( - id bigint NOT NULL, - "Payload" varchar NOT NULL, - PRIMARY KEY (id) - )""" - else: - queue_ddl = f"""CREATE QUEUE {queue_name} ( - Id INT64 NOT NULL, - Payload STRING(MAX) NOT NULL - ) PRIMARY KEY (Id)""" - - try: - operation = await test_database.update_ddl([queue_ddl]) - await operation.result(600) - except MethodNotImplemented as e: - print(f"MethodNotImplemented. Skipping test because Queues are not implemented yet: {e}") - return - except GoogleAPIError as e: - if getattr(e, 'code', None) == 501 or (getattr(e, 'grpc_status_code', None) and e.grpc_status_code.name == 'UNIMPLEMENTED') or "UNIMPLEMENTED" in str(e): - print(f"Skipping test because Queues are not implemented yet: {e}") - return - raise - print("Queue created successfully.") - - # 4. Run mutations - print("Sending message to queue...") - async with test_database.batch() as batch: - batch.send( - queue=queue_name, - key=(2,), - payload="Hello, Queues!", - ) - print("Send successful.") - - print("Acking message in queue...") - async with test_database.batch() as batch: - batch.ack( - queue=queue_name, - key=(2,), - ) - print("Ack successful.") - - finally: - print("Dropping database...") - await test_database.drop() + operation = await test_database.update_ddl([queue_ddl]) + await operation.result(600) + except MethodNotImplemented as e: + pytest.skip(f"Queues are not implemented yet: {e}") + except GoogleAPIError as e: + if getattr(e, 'code', None) == 501 or (getattr(e, 'grpc_status_code', None) and e.grpc_status_code.name == 'UNIMPLEMENTED') or "UNIMPLEMENTED" in str(e): + pytest.skip(f"Queues are not implemented yet: {e}") + raise + + # Run mutations + async with test_database.batch() as batch: + batch.send( + queue=queue_name, + key=(2,), + payload="Hello, Queues!", + ) + + print("Acking message in queue...") + async with test_database.batch() as batch: + batch.ack( + queue=queue_name, + key=(2,), + ) + finally: - print("Dropping instance...") - await test_instance.delete() + print("Dropping database...") + await test_database.drop() @pytest.mark.asyncio diff --git a/packages/google-cloud-spanner/tests/system/conftest.py b/packages/google-cloud-spanner/tests/system/conftest.py index 0ebb8bce0b11..f3ff0c38aac9 100644 --- a/packages/google-cloud-spanner/tests/system/conftest.py +++ b/packages/google-cloud-spanner/tests/system/conftest.py @@ -205,14 +205,27 @@ def shared_instance( _helpers.cleanup_old_instances(spanner_client) if _helpers.CREATE_INSTANCE: + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin create_time = str(int(time.time())) labels = {"python-spanner-systests": "true", "created": create_time} + request = spanner_instance_admin.CreateInstanceRequest( + parent=spanner_client.project_name, + instance_id=shared_instance_id, + instance=spanner_instance_admin.Instance( + config=instance_config.name, + display_name=shared_instance_id, + node_count=1, + labels=labels, + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE_PLUS, + ), + ) + created_op = _helpers.retry_429_503(spanner_client.instance_admin_api.create_instance)(request=request) + created_op.result(instance_operation_timeout) # block until completion + instance = spanner_client.instance( shared_instance_id, instance_config.name, labels=labels ) - created_op = _helpers.retry_429_503(instance.create)() - created_op.result(instance_operation_timeout) # block until completion else: # reuse existing instance instance = spanner_client.instance(shared_instance_id) diff --git a/packages/google-cloud-spanner/tests/system/test_database_api.py b/packages/google-cloud-spanner/tests/system/test_database_api.py index 517282612e99..9e5219c75243 100644 --- a/packages/google-cloud-spanner/tests/system/test_database_api.py +++ b/packages/google-cloud-spanner/tests/system/test_database_api.py @@ -567,88 +567,65 @@ def test_db_batch_insert_then_db_snapshot_read(shared_database): sd._check_rows_data(from_snap) -def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, instance_config): +def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, shared_instance): import uuid from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin from google.cloud.spanner_admin_database_v1 import DatabaseDialect from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError - instance_id = f"test-instance-{uuid.uuid4().hex[:8]}" db_id = f"test-db-{uuid.uuid4().hex[:8]}" queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" - config_name = instance_config.name - request = spanner_instance_admin.CreateInstanceRequest( - parent=spanner_client.project_name, - instance_id=instance_id, - instance=spanner_instance_admin.Instance( - config=config_name, - display_name=instance_id, - node_count=1, - edition=spanner_instance_admin.Instance.Edition.ENTERPRISE, - ), - ) - print(f"instance creation request: {request}") - operation = spanner_client.instance_admin_api.create_instance(request=request) - operation.result(600) - test_instance = spanner_client.instance(instance_id, configuration_name=config_name) + test_database = shared_instance.database(db_id, database_dialect=database_dialect) + operation = test_database.create() + operation.result(300) + print("Database created successfully!") try: - test_database = test_instance.database(db_id, database_dialect=database_dialect) - operation = test_database.create() - operation.result(300) - print("Database created successfully!") - + # Create the Queue + if database_dialect == DatabaseDialect.POSTGRESQL: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + id bigint NOT NULL, + "Payload" varchar NOT NULL, + PRIMARY KEY (id) + )""" + else: + queue_ddl = f"""CREATE QUEUE {queue_name} ( + Id INT64 NOT NULL, + Payload STRING(MAX) NOT NULL + ) PRIMARY KEY (Id)""" try: - # 3. Create the Queue - if database_dialect == DatabaseDialect.POSTGRESQL: - queue_ddl = f"""CREATE QUEUE {queue_name} ( - id bigint NOT NULL, - "Payload" varchar NOT NULL, - PRIMARY KEY (id) - )""" - else: - queue_ddl = f"""CREATE QUEUE {queue_name} ( - Id INT64 NOT NULL, - Payload STRING(MAX) NOT NULL - ) PRIMARY KEY (Id)""" - try: - operation = test_database.update_ddl([queue_ddl]) - operation.result(600) - except MethodNotImplemented as e: - print(f"MethodNotImplemented. Skipping test because Queues are not implemented yet: {e}") - return - except GoogleAPIError as e: - if getattr(e, 'code', None) == 501 or getattr(e, 'grpc_status_code', None) and e.grpc_status_code.name == 'UNIMPLEMENTED' or "UNIMPLEMENTED" in str(e): - print(f"Skipping test because Queues are not implemented yet: {e}") - return - raise - print("Queue created successfully.") - - # 4. Run mutations - print("Sending message to queue...") - with test_database.batch() as batch: - batch.send( - queue=queue_name, - key=(2,), - payload="Hello, Queues!", - ) - print("Send successful.") - - print("Acking message in queue...") - with test_database.batch() as batch: - batch.ack( - queue=queue_name, - key=(2,), - ) - print("Ack successful.") - - finally: - print("Dropping database...") - test_database.drop() + operation = test_database.update_ddl([queue_ddl]) + operation.result(600) + except MethodNotImplemented as e: + pytest.skip(f"Queues are not implemented yet: {e}") + except GoogleAPIError as e: + if getattr(e, 'code', None) == 501 or getattr(e, 'grpc_status_code', None) and e.grpc_status_code.name == 'UNIMPLEMENTED' or "UNIMPLEMENTED" in str(e): + pytest.skip(f"Queues are not implemented yet: {e}") + raise + print("Queue created successfully.") + + # Run mutations + print("Sending message to queue...") + with test_database.batch() as batch: + batch.send( + queue=queue_name, + key=(2,), + payload="Hello, Queues!", + ) + print("Send successful.") + + print("Acking message in queue...") + with test_database.batch() as batch: + batch.ack( + queue=queue_name, + key=(2,), + ) + print("Ack successful.") + finally: - print("Dropping instance...") - test_instance.delete() + print("Dropping database...") + test_database.drop() def test_db_run_in_transaction_then_snapshot_execute_sql(shared_database): diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py index 94d623caf285..77ac30ae6ba2 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py @@ -25,6 +25,7 @@ RequestOptions, TransactionOptions, ) +from google.cloud.spanner_v1._helpers import _make_value_pb, _make_list_value_pb from google.cloud.spanner_v1._async.batch import Batch, MutationGroups, _BatchBase from google.cloud.spanner_v1.keyset import KeySet @@ -43,15 +44,6 @@ def _getTargetClass(self): def _make_one(self, *args, **kwargs): return self._getTargetClass()(*args, **kwargs) - def _compare_values(self, result, source): - for found, expected in zip(result, source): - self.assertEqual(len(found), len(expected)) - for found_cell, expected_cell in zip(found, expected): - if isinstance(expected_cell, int): - self.assertEqual(int(found_cell), expected_cell) - else: - self.assertEqual(found_cell, expected_cell) - def test_ctor(self): session = mock.Mock() base = self._make_one(session) @@ -103,8 +95,8 @@ def test_send(self): base.send(queue=queue, key=key, payload=payload) self.assertEqual(len(base._mutations), 1) self.assertEqual(base._mutations[0].send.queue, queue) - self.assertEqual(base._mutations[0].send.payload, payload) - self._compare_values([base._mutations[0].send.key], [key]) + self.assertEqual(base._mutations[0].send._pb.payload, _make_value_pb(payload)) + self.assertEqual(base._mutations[0].send._pb.key, _make_list_value_pb(key)) def test_ack(self): queue = "TestQueue" @@ -114,7 +106,7 @@ def test_ack(self): base.ack(queue=queue, key=key) self.assertEqual(len(base._mutations), 1) self.assertEqual(base._mutations[0].ack.queue, queue) - self._compare_values([base._mutations[0].ack.key], [key]) + self.assertEqual(base._mutations[0].ack._pb.key, _make_list_value_pb(key)) class TestBatch(unittest.IsolatedAsyncioTestCase): diff --git a/packages/google-cloud-spanner/tests/unit/test_batch.py b/packages/google-cloud-spanner/tests/unit/test_batch.py index db4ebba0b7f4..cf43fccbc777 100644 --- a/packages/google-cloud-spanner/tests/unit/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/test_batch.py @@ -37,6 +37,8 @@ _augment_errors_with_request_id, _metadata_with_request_id, _metadata_with_request_id_and_req_id, + _make_value_pb, + _make_list_value_pb, ) from google.cloud.spanner_v1.batch import Batch, MutationGroups, _BatchBase from google.cloud.spanner_v1.keyset import KeySet @@ -197,8 +199,8 @@ def test_send(self): send = mutation.send self.assertIsInstance(send, Mutation.Send) self.assertEqual(send.queue, queue) - self.assertEqual(send.payload, payload) - self._compare_values([send.key], [key]) + self.assertEqual(send._pb.payload, _make_value_pb(payload)) + self.assertEqual(send._pb.key, _make_list_value_pb(key)) def test_ack(self): queue = "TestQueue" @@ -214,7 +216,7 @@ def test_ack(self): ack = mutation.ack self.assertIsInstance(ack, Mutation.Ack) self.assertEqual(ack.queue, queue) - self._compare_values([ack.key], [key]) + self.assertEqual(ack._pb.key, _make_list_value_pb(key)) class TestBatch(_BaseTest, OpenTelemetryBase):