From 049ee4b8d05b06a7157147bbc20b8ec87136da36 Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Wed, 15 Jul 2026 21:43:36 +0000 Subject: [PATCH 1/4] 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 384a1855bb14..5398e9163567 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 9adb62eb4ebda58e8352ae794a58d72acaaf2604 Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Thu, 16 Jul 2026 21:21:59 +0000 Subject: [PATCH 2/4] 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 5398e9163567..e17687092242 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): From ef3168ae3c2d1c95837635dc834628c6c16adbc9 Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Mon, 3 Aug 2026 20:06:20 +0000 Subject: [PATCH 3/4] fix linting & format issues --- .../google/cloud/spanner_v1/_async/batch.py | 18 +++++--------- .../google/cloud/spanner_v1/batch.py | 18 +++++--------- .../tests/system/_async/conftest.py | 8 ++++--- .../tests/system/_async/test_database_api.py | 24 ++++++++++++++----- .../tests/system/conftest.py | 5 +++- .../tests/system/test_database_api.py | 18 ++++++++++---- .../tests/unit/_async/test_batch.py | 2 +- .../tests/unit/test_batch.py | 4 ++-- 8 files changed, 55 insertions(+), 42 deletions(-) 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 519122622276..966519265d50 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 @@ -20,16 +20,16 @@ from typing import List, Optional from google.api_core.exceptions import InternalServerError +from google.cloud._helpers import _datetime_to_pb_timestamp 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, + _make_value_pb, _merge_client_context, _merge_request_options, _merge_Transaction_Options, @@ -183,15 +183,12 @@ def send(self, queue, key, payload=None, deliver_time=None): :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) - } + 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)) @@ -207,13 +204,10 @@ def ack(self, queue, key, ignore_not_found=None): :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) - } + ack_kwargs = {"queue": queue, "key": _make_list_value_pb(key)} if ignore_not_found is not None: ack_kwargs["ignore_not_found"] = ignore_not_found - + ack = Mutation.Ack(**ack_kwargs) self._mutations.append(Mutation(ack=ack)) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py index af26874b5a0f..c32b65812db4 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py @@ -22,14 +22,14 @@ from typing import List, Optional 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, + _make_value_pb, _merge_client_context, _merge_request_options, _merge_Transaction_Options, @@ -160,15 +160,12 @@ def send(self, queue, key, payload=None, deliver_time=None): :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) - } + 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)) @@ -184,13 +181,10 @@ def ack(self, queue, key, ignore_not_found=None): :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) - } + ack_kwargs = {"queue": queue, "key": _make_list_value_pb(key)} if ignore_not_found is not None: ack_kwargs["ignore_not_found"] = ignore_not_found - + ack = Mutation.Ack(**ack_kwargs) self._mutations.append(Mutation(ack=ack)) diff --git a/packages/google-cloud-spanner/tests/system/_async/conftest.py b/packages/google-cloud-spanner/tests/system/_async/conftest.py index 997e701983b7..a4be2250a065 100644 --- a/packages/google-cloud-spanner/tests/system/_async/conftest.py +++ b/packages/google-cloud-spanner/tests/system/_async/conftest.py @@ -140,9 +140,10 @@ async def shared_instance( spanner_client._instance_admin_api = None if _helpers.CREATE_INSTANCE: - from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin import time + 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} @@ -157,7 +158,9 @@ async def shared_instance( edition=spanner_instance_admin.Instance.Edition.ENTERPRISE_PLUS, ), ) - created_op = await spanner_client.instance_admin_api.create_instance(request=request) + created_op = await spanner_client.instance_admin_api.create_instance( + request=request + ) await created_op.result(instance_operation_timeout) instance = spanner_client.instance( @@ -232,4 +235,3 @@ def not_postgres(database_dialect): 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 1beef4e4f457..fff280f2dd76 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,16 +88,21 @@ 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, shared_instance): +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.api_core.exceptions import GoogleAPIError, MethodNotImplemented + from google.cloud.spanner_admin_database_v1 import DatabaseDialect - from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError db_name = f"test-db-{uuid.uuid4().hex[:8]}" queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" - test_database = await shared_instance.database(db_name, database_dialect=database_dialect) + test_database = await shared_instance.database( + db_name, database_dialect=database_dialect + ) operation = await test_database.create() operation.result(300) @@ -121,7 +126,14 @@ async def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dial 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): + 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 @@ -132,7 +144,7 @@ async def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dial key=(2,), payload="Hello, Queues!", ) - + print("Acking message in queue...") async with test_database.batch() as batch: batch.ack( diff --git a/packages/google-cloud-spanner/tests/system/conftest.py b/packages/google-cloud-spanner/tests/system/conftest.py index f3ff0c38aac9..7adcd9c81137 100644 --- a/packages/google-cloud-spanner/tests/system/conftest.py +++ b/packages/google-cloud-spanner/tests/system/conftest.py @@ -206,6 +206,7 @@ def shared_instance( 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} @@ -220,7 +221,9 @@ def shared_instance( 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 = _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( 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 9e5219c75243..6939f411ef82 100644 --- a/packages/google-cloud-spanner/tests/system/test_database_api.py +++ b/packages/google-cloud-spanner/tests/system/test_database_api.py @@ -567,11 +567,14 @@ 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, shared_instance): +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.api_core.exceptions import GoogleAPIError, MethodNotImplemented + from google.cloud.spanner_admin_database_v1 import DatabaseDialect - from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError db_id = f"test-db-{uuid.uuid4().hex[:8]}" queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" @@ -600,7 +603,12 @@ def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, s 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): + 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.") @@ -614,7 +622,7 @@ def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, s payload="Hello, Queues!", ) print("Send successful.") - + print("Acking message in queue...") with test_database.batch() as batch: batch.ack( 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 77ac30ae6ba2..12b465be877e 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py @@ -25,8 +25,8 @@ 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._helpers import _make_list_value_pb, _make_value_pb from google.cloud.spanner_v1.keyset import KeySet TABLE_NAME = "citizens" diff --git a/packages/google-cloud-spanner/tests/unit/test_batch.py b/packages/google-cloud-spanner/tests/unit/test_batch.py index e17687092242..933c47bfa92d 100644 --- a/packages/google-cloud-spanner/tests/unit/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/test_batch.py @@ -35,10 +35,10 @@ from google.cloud.spanner_v1._helpers import ( AtomicCounter, _augment_errors_with_request_id, + _make_list_value_pb, + _make_value_pb, _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 From 9a5eae68b8b0f24bf0afe91b36de6872341c05ca Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Mon, 3 Aug 2026 20:39:49 +0000 Subject: [PATCH 4/4] Trigger CI workflows