feat: add send and ack mutations Cloud Spanner Queues

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`."""

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`."""

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()

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

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):