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/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")

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.")
+

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

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)

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

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