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/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 fa0ffaeab623..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 @@ -87,6 +87,64 @@ 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, 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 + + 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) + operation = await test_database.create() + operation.result(300) + + try: + # 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: + 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 database...") + await test_database.drop() + + @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/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 793d72f18374..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,6 +567,67 @@ 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): + 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 + + db_id = f"test-db-{uuid.uuid4().hex[:8]}" + queue_name = f"test_queue_{uuid.uuid4().hex[:8]}" + + test_database = shared_instance.database(db_id, database_dialect=database_dialect) + operation = test_database.create() + operation.result(300) + print("Database created successfully!") + + try: + # 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: + 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 database...") + test_database.drop() + + 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..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 @@ -85,6 +86,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._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" + 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.assertEqual(base._mutations[0].ack._pb.key, _make_list_value_pb(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..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 @@ -182,6 +184,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._pb.payload, _make_value_pb(payload)) + self.assertEqual(send._pb.key, _make_list_value_pb(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.assertEqual(ack._pb.key, _make_list_value_pb(key)) + class TestBatch(_BaseTest, OpenTelemetryBase): def _getTargetClass(self):