Skip to content

Commit f32b057

Browse files
finn-the-coderfinn-the-coder
committed
feat: add send and ack methods to Batch for Cloud Spanner Queues and include unit and system tests
1 parent 35e16da commit f32b057

6 files changed

Lines changed: 340 additions & 0 deletions

File tree

‎packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,12 @@
2323

2424
from google.cloud.aio._cross_sync import CrossSync
2525
from google.cloud.spanner_v1._async._helpers import _retry, _retry_on_aborted_exception
26+
from google.cloud._helpers import _datetime_to_pb_timestamp
2627
from google.cloud.spanner_v1._helpers import (
2728
AtomicCounter,
2829
_check_rst_stream_error,
30+
_make_value_pb,
31+
_make_list_value_pb,
2932
_make_list_value_pbs,
3033
_merge_client_context,
3134
_merge_request_options,
@@ -165,6 +168,55 @@ def delete(self, table, keyset):
165168
# TODO: Decide if we should add a span event per mutation:
166169
# https://github.com/googleapis/python-spanner/issues/1269
167170

171+
def send(self, queue, key, payload=None, deliver_time=None):
172+
"""Send a message to a Cloud Spanner queue.
173+
174+
:type queue: str
175+
:param queue: Name of the queue to which the message will be sent.
176+
177+
:type key: list
178+
:param key: The primary key of the message to be sent.
179+
180+
:type payload: object
181+
:param payload: (Optional) The payload of the message.
182+
183+
:type deliver_time: :class:`datetime.datetime`
184+
:param deliver_time: (Optional) The time at which Spanner will begin attempting to deliver the message.
185+
"""
186+
send_kwargs = {
187+
"queue": queue,
188+
"key": _make_list_value_pb(key)
189+
}
190+
if payload is not None:
191+
send_kwargs["payload"] = _make_value_pb(payload)
192+
if deliver_time is not None:
193+
send_kwargs["deliver_time"] = _datetime_to_pb_timestamp(deliver_time)
194+
195+
send = Mutation.Send(**send_kwargs)
196+
self._mutations.append(Mutation(send=send))
197+
198+
def ack(self, queue, key, ignore_not_found=None):
199+
"""Acknowledge a message in a Cloud Spanner queue.
200+
201+
:type queue: str
202+
:param queue: Name of the queue where the message to be acked is stored.
203+
204+
:type key: list
205+
:param key: The primary key of the message to be acked.
206+
207+
:type ignore_not_found: bool
208+
:param ignore_not_found: (Optional) Whether to ignore if the message does not exist.
209+
"""
210+
ack_kwargs = {
211+
"queue": queue,
212+
"key": _make_list_value_pb(key)
213+
}
214+
if ignore_not_found is not None:
215+
ack_kwargs["ignore_not_found"] = ignore_not_found
216+
217+
ack = Mutation.Ack(**ack_kwargs)
218+
self._mutations.append(Mutation(ack=ack))
219+
168220

169221
class Batch(_BatchBase):
170222
"""Accumulate mutations for transmission during :meth:`commit`."""

‎packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,12 @@
2323

2424
from google.api_core.exceptions import InternalServerError
2525

26+
from google.cloud._helpers import _datetime_to_pb_timestamp
2627
from google.cloud.spanner_v1._helpers import (
2728
AtomicCounter,
2829
_check_rst_stream_error,
30+
_make_value_pb,
31+
_make_list_value_pb,
2932
_make_list_value_pbs,
3033
_merge_client_context,
3134
_merge_request_options,
@@ -142,6 +145,55 @@ def delete(self, table, keyset):
142145
delete = Mutation.Delete(table=table, key_set=keyset._to_pb())
143146
self._mutations.append(Mutation(delete=delete))
144147

148+
def send(self, queue, key, payload=None, deliver_time=None):
149+
"""Send a message to a Cloud Spanner queue.
150+
151+
:type queue: str
152+
:param queue: Name of the queue to which the message will be sent.
153+
154+
:type key: list
155+
:param key: The primary key of the message to be sent.
156+
157+
:type payload: object
158+
:param payload: (Optional) The payload of the message.
159+
160+
:type deliver_time: :class:`datetime.datetime`
161+
:param deliver_time: (Optional) The time at which Spanner will begin attempting to deliver the message.
162+
"""
163+
send_kwargs = {
164+
"queue": queue,
165+
"key": _make_list_value_pb(key)
166+
}
167+
if payload is not None:
168+
send_kwargs["payload"] = _make_value_pb(payload)
169+
if deliver_time is not None:
170+
send_kwargs["deliver_time"] = _datetime_to_pb_timestamp(deliver_time)
171+
172+
send = Mutation.Send(**send_kwargs)
173+
self._mutations.append(Mutation(send=send))
174+
175+
def ack(self, queue, key, ignore_not_found=None):
176+
"""Acknowledge a message in a Cloud Spanner queue.
177+
178+
:type queue: str
179+
:param queue: Name of the queue where the message to be acked is stored.
180+
181+
:type key: list
182+
:param key: The primary key of the message to be acked.
183+
184+
:type ignore_not_found: bool
185+
:param ignore_not_found: (Optional) Whether to ignore if the message does not exist.
186+
"""
187+
ack_kwargs = {
188+
"queue": queue,
189+
"key": _make_list_value_pb(key)
190+
}
191+
if ignore_not_found is not None:
192+
ack_kwargs["ignore_not_found"] = ignore_not_found
193+
194+
ack = Mutation.Ack(**ack_kwargs)
195+
self._mutations.append(Mutation(ack=ack))
196+
145197

146198
class Batch(_BatchBase):
147199
"""Accumulate mutations for transmission during :meth:`commit`."""

‎packages/google-cloud-spanner/tests/system/_async/test_database_api.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,93 @@ async def test_db_batch_insert_then_db_snapshot_read(shared_database):
8787
sd._check_rows_data(from_snap)
8888

8989

90+
@pytest.mark.asyncio
91+
async def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, instance_config):
92+
import uuid
93+
from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin
94+
from google.cloud.spanner_admin_database_v1 import DatabaseDialect
95+
from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError
96+
97+
instance_id = f"test-instance-{uuid.uuid4().hex[:8]}"
98+
db_name = f"test-db-{uuid.uuid4().hex[:8]}"
99+
queue_name = f"test_queue_{uuid.uuid4().hex[:8]}"
100+
101+
config_name = instance_config.name
102+
request = spanner_instance_admin.CreateInstanceRequest(
103+
parent=spanner_client.project_name,
104+
instance_id=instance_id,
105+
instance=spanner_instance_admin.Instance(
106+
config=config_name,
107+
display_name=instance_id,
108+
node_count=1,
109+
edition=spanner_instance_admin.Instance.Edition.ENTERPRISE,
110+
),
111+
)
112+
print(f"instance creation request: {request}")
113+
operation = await spanner_client.instance_admin_api.create_instance(request=request)
114+
operation.result(600)
115+
116+
test_instance = spanner_client.instance(instance_id, configuration_name=config_name)
117+
118+
try:
119+
test_database = await test_instance.database(db_name, database_dialect=database_dialect)
120+
operation = await test_database.create()
121+
operation.result(300)
122+
print("Database created successfully!")
123+
124+
try:
125+
# 3. Create the Queue
126+
if database_dialect == DatabaseDialect.POSTGRESQL:
127+
queue_ddl = f"""CREATE QUEUE {queue_name} (
128+
id bigint NOT NULL,
129+
"Payload" varchar NOT NULL,
130+
PRIMARY KEY (id)
131+
)"""
132+
else:
133+
queue_ddl = f"""CREATE QUEUE {queue_name} (
134+
Id INT64 NOT NULL,
135+
Payload STRING(MAX) NOT NULL
136+
) PRIMARY KEY (Id)"""
137+
138+
try:
139+
operation = await test_database.update_ddl([queue_ddl])
140+
await operation.result(600)
141+
except MethodNotImplemented as e:
142+
print(f"MethodNotImplemented. Skipping test because Queues are not implemented yet: {e}")
143+
return
144+
except GoogleAPIError as e:
145+
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):
146+
print(f"Skipping test because Queues are not implemented yet: {e}")
147+
return
148+
raise
149+
print("Queue created successfully.")
150+
151+
# 4. Run mutations
152+
print("Sending message to queue...")
153+
async with test_database.batch() as batch:
154+
batch.send(
155+
queue=queue_name,
156+
key=(2,),
157+
payload="Hello, Queues!",
158+
)
159+
print("Send successful.")
160+
161+
print("Acking message in queue...")
162+
async with test_database.batch() as batch:
163+
batch.ack(
164+
queue=queue_name,
165+
key=(2,),
166+
)
167+
print("Ack successful.")
168+
169+
finally:
170+
print("Dropping database...")
171+
await test_database.drop()
172+
finally:
173+
print("Dropping instance...")
174+
await test_instance.delete()
175+
176+
90177
@pytest.mark.asyncio
91178
async def test_db_run_in_transaction_then_snapshot_execute_sql(shared_database):
92179
await shared_database.reload()

‎packages/google-cloud-spanner/tests/system/test_database_api.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -567,6 +567,90 @@ def test_db_batch_insert_then_db_snapshot_read(shared_database):
567567
sd._check_rows_data(from_snap)
568568

569569

570+
def test_db_batch_send_and_ack(not_emulator, spanner_client, database_dialect, instance_config):
571+
import uuid
572+
from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin
573+
from google.cloud.spanner_admin_database_v1 import DatabaseDialect
574+
from google.api_core.exceptions import MethodNotImplemented, GoogleAPIError
575+
576+
instance_id = f"test-instance-{uuid.uuid4().hex[:8]}"
577+
db_id = f"test-db-{uuid.uuid4().hex[:8]}"
578+
queue_name = f"test_queue_{uuid.uuid4().hex[:8]}"
579+
580+
config_name = instance_config.name
581+
request = spanner_instance_admin.CreateInstanceRequest(
582+
parent=spanner_client.project_name,
583+
instance_id=instance_id,
584+
instance=spanner_instance_admin.Instance(
585+
config=config_name,
586+
display_name=instance_id,
587+
node_count=1,
588+
edition=spanner_instance_admin.Instance.Edition.ENTERPRISE,
589+
),
590+
)
591+
print(f"instance creation request: {request}")
592+
operation = spanner_client.instance_admin_api.create_instance(request=request)
593+
operation.result(600)
594+
test_instance = spanner_client.instance(instance_id, configuration_name=config_name)
595+
596+
try:
597+
test_database = test_instance.database(db_id, database_dialect=database_dialect)
598+
operation = test_database.create()
599+
operation.result(300)
600+
print("Database created successfully!")
601+
602+
try:
603+
# 3. Create the Queue
604+
if database_dialect == DatabaseDialect.POSTGRESQL:
605+
queue_ddl = f"""CREATE QUEUE {queue_name} (
606+
id bigint NOT NULL,
607+
"Payload" varchar NOT NULL,
608+
PRIMARY KEY (id)
609+
)"""
610+
else:
611+
queue_ddl = f"""CREATE QUEUE {queue_name} (
612+
Id INT64 NOT NULL,
613+
Payload STRING(MAX) NOT NULL
614+
) PRIMARY KEY (Id)"""
615+
try:
616+
operation = test_database.update_ddl([queue_ddl])
617+
operation.result(600)
618+
except MethodNotImplemented as e:
619+
print(f"MethodNotImplemented. Skipping test because Queues are not implemented yet: {e}")
620+
return
621+
except GoogleAPIError as e:
622+
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):
623+
print(f"Skipping test because Queues are not implemented yet: {e}")
624+
return
625+
raise
626+
print("Queue created successfully.")
627+
628+
# 4. Run mutations
629+
print("Sending message to queue...")
630+
with test_database.batch() as batch:
631+
batch.send(
632+
queue=queue_name,
633+
key=(2,),
634+
payload="Hello, Queues!",
635+
)
636+
print("Send successful.")
637+
638+
print("Acking message in queue...")
639+
with test_database.batch() as batch:
640+
batch.ack(
641+
queue=queue_name,
642+
key=(2,),
643+
)
644+
print("Ack successful.")
645+
646+
finally:
647+
print("Dropping database...")
648+
test_database.drop()
649+
finally:
650+
print("Dropping instance...")
651+
test_instance.delete()
652+
653+
570654
def test_db_run_in_transaction_then_snapshot_execute_sql(shared_database):
571655
_helpers.retry_has_all_dll(shared_database.reload)()
572656
sd = _sample_data

‎packages/google-cloud-spanner/tests/unit/_async/test_batch.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,15 @@ def _getTargetClass(self):
4343
def _make_one(self, *args, **kwargs):
4444
return self._getTargetClass()(*args, **kwargs)
4545

46+
def _compare_values(self, result, source):
47+
for found, expected in zip(result, source):
48+
self.assertEqual(len(found), len(expected))
49+
for found_cell, expected_cell in zip(found, expected):
50+
if isinstance(expected_cell, int):
51+
self.assertEqual(int(found_cell), expected_cell)
52+
else:
53+
self.assertEqual(found_cell, expected_cell)
54+
4655
def test_ctor(self):
4756
session = mock.Mock()
4857
base = self._make_one(session)
@@ -85,6 +94,28 @@ def test_delete(self):
8594
self.assertEqual(len(base._mutations), 1)
8695
self.assertEqual(base._mutations[0].delete.table, TABLE_NAME)
8796

97+
def test_send(self):
98+
queue = "TestQueue"
99+
key = [2]
100+
payload = "Hello, Queues!"
101+
session = mock.Mock()
102+
base = self._make_one(session)
103+
base.send(queue=queue, key=key, payload=payload)
104+
self.assertEqual(len(base._mutations), 1)
105+
self.assertEqual(base._mutations[0].send.queue, queue)
106+
self.assertEqual(base._mutations[0].send.payload, payload)
107+
self._compare_values([base._mutations[0].send.key], [key])
108+
109+
def test_ack(self):
110+
queue = "TestQueue"
111+
key = [2]
112+
session = mock.Mock()
113+
base = self._make_one(session)
114+
base.ack(queue=queue, key=key)
115+
self.assertEqual(len(base._mutations), 1)
116+
self.assertEqual(base._mutations[0].ack.queue, queue)
117+
self._compare_values([base._mutations[0].ack.key], [key])
118+
88119

89120
class TestBatch(unittest.IsolatedAsyncioTestCase):
90121
def _getTargetClass(self):

0 commit comments

Comments
 (0)