Append all event data in a single transaction

Author: ScottMansfield

src/google/adk/integrations/firestore/firestore_session_service.py

@@ -138,13 +138,6 @@ async def create_session(
 
     session_ref = self._get_sessions_ref(app_name, user_id).document(session_id)
 
-    # Check if session already exists
-    doc = await session_ref.get()
-    if doc.exists:
-      from ...errors.already_exists_error import AlreadyExistsError
-
-      raise AlreadyExistsError(f"Session {session_id} already exists.")
-
     session_data = {
         "id": session_id,
         "appName": app_name,
@@ -154,7 +147,16 @@ async def create_session(
         "updateTime": now,
     }
 
-    await session_ref.set(session_data)
+    @firestore.async_transactional  # type: ignore[untyped-decorator]
+    async def _create_txn(transaction: firestore.AsyncTransaction) -> None:
+      snap = await session_ref.get(transaction=transaction)
+      if snap.exists:
+        from ...errors.already_exists_error import AlreadyExistsError
+        raise AlreadyExistsError(f"Session {session_id} already exists.")
+      transaction.set(session_ref, session_data)
+
+    transaction_obj = self.client.transaction()
+    await _create_txn(transaction_obj)
 
     # We need a timestamp for the Session object. Since SERVER_TIMESTAMP is
     # evaluated on the server, we might want to use local time for the object
@@ -318,6 +320,7 @@ async def delete_session(
 
     await session_ref.delete()
 
+
   async def _update_app_state_transactional(
       self, app_name: str, delta: dict[str, Any]
   ) -> dict[str, Any]:
@@ -389,44 +392,57 @@ async def append_event(self, session: Session, event: Event) -> Event:
         else:
           session_updates[key] = value
 
-      if app_updates:
-        await self._update_app_state_transactional(
-            session.app_name, app_updates
-        )
+      app_ref = self.client.collection(self.app_state_collection).document(session.app_name)
+      user_ref = (
+          self.client.collection(self.user_state_collection)
+          .document(session.app_name)
+          .collection("users")
+          .document(session.user_id)
+      )
 
-      if user_updates:
-        await self._update_user_state_transactional(
-            session.app_name, session.user_id, user_updates
+      @firestore.async_transactional  # type: ignore[untyped-decorator]
+      async def _append_txn(transaction: firestore.AsyncTransaction) -> None:
+        # 1. Reads
+        app_snap = await app_ref.get(transaction=transaction) if app_updates else None
+        user_snap = await user_ref.get(transaction=transaction) if user_updates else None
+
+        # 2. Writes
+        if app_updates and app_snap is not None:
+          current_app = app_snap.to_dict() if app_snap.exists else {}
+          current_app.update(app_updates)
+          transaction.set(app_ref, current_app, merge=True)
+
+        if user_updates and user_snap is not None:
+          current_user = user_snap.to_dict() if user_snap.exists else {}
+          current_user.update(user_updates)
+          transaction.set(user_ref, current_user, merge=True)
+
+        for k, v in session_updates.items():
+          session.state[k] = v
+
+        transaction.update(
+            session_ref,
+            {
+                "state": session.state,
+                "updateTime": firestore.SERVER_TIMESTAMP,
+            },
         )
 
-      for k, v in session_updates.items():
-        session.state[k] = v
-
-      batch = self.client.batch()
-      batch.update(
-          session_ref,
-          {
-              "state": session.state,
-              "updateTime": firestore.SERVER_TIMESTAMP,
-          },
-      )
-
-      event_id = event.id
-      event_ref = session_ref.collection(self.events_collection).document(
-          event_id
-      )
-      event_data = event.model_dump(exclude_none=True, mode="json")
-      batch.set(
-          event_ref,
-          {
-              "event_data": event_data,
-              "timestamp": firestore.SERVER_TIMESTAMP,
-              "appName": session.app_name,
-              "userId": session.user_id,
-          },
-      )
+        event_id = event.id
+        event_ref = session_ref.collection(self.events_collection).document(event_id)
+        event_data = event.model_dump(exclude_none=True, mode="json")
+        transaction.set(
+            event_ref,
+            {
+                "event_data": event_data,
+                "timestamp": firestore.SERVER_TIMESTAMP,
+                "appName": session.app_name,
+                "userId": session.user_id,
+            },
+        )
 
-      await batch.commit()
+      transaction_obj = self.client.transaction()
+      await _append_txn(transaction_obj)
     else:
       batch = self.client.batch()
       event_id = event.id

tests/unittests/integrations/firestore/test_firestore_session_service.py

@@ -89,7 +89,8 @@ async def test_create_session(mock_firestore_client):
   app_name = "test_app"
   user_id = "test_user"
 
-  session = await service.create_session(app_name=app_name, user_id=user_id)
+  with mock.patch("google.cloud.firestore.async_transactional", lambda x: x):
+    session = await service.create_session(app_name=app_name, user_id=user_id)
 
   assert session.app_name == app_name
   assert session.user_id == user_id
@@ -109,14 +110,16 @@ async def test_create_session(mock_firestore_client):
   session_doc_ref = sessions_ref.document.return_value
   from google.cloud import firestore
 
-  session_doc_ref.set.assert_called_once()
-  args, kwargs = session_doc_ref.set.call_args
-  assert args[0]["id"] == session.id
-  assert args[0]["appName"] == app_name
-  assert args[0]["userId"] == user_id
-  assert args[0]["state"] == {}
-  assert args[0]["createTime"] == firestore.SERVER_TIMESTAMP
-  assert args[0]["updateTime"] == firestore.SERVER_TIMESTAMP
+  transaction = mock_firestore_client.transaction.return_value
+  transaction.set.assert_called_once()
+  args, kwargs = transaction.set.call_args
+  assert args[0] == session_doc_ref
+  assert args[1]["id"] == session.id
+  assert args[1]["appName"] == app_name
+  assert args[1]["userId"] == user_id
+  assert args[1]["state"] == {}
+  assert args[1]["createTime"] == firestore.SERVER_TIMESTAMP
+  assert args[1]["updateTime"] == firestore.SERVER_TIMESTAMP
 
 
 @pytest.mark.asyncio
@@ -228,7 +231,8 @@ async def test_append_event(mock_firestore_client):
   session = Session(id="test_session", app_name=app_name, user_id=user_id)
   event = Event(invocation_id="test_inv", author="user")
 
-  await service.append_event(session, event)
+  with mock.patch("google.cloud.firestore.async_transactional", lambda x: x):
+    await service.append_event(session, event)
 
   from google.cloud import firestore
 
@@ -273,31 +277,22 @@ async def test_append_event_with_state_delta(mock_firestore_client):
   service._update_app_state_transactional = mock.AsyncMock()
   service._update_user_state_transactional = mock.AsyncMock()
 
-  await service.append_event(session, event)
-
-  mock_firestore_client.batch.assert_called_once()
-  service._update_app_state_transactional.assert_called_once_with(
-      "test_app", {"my_key": "app_val"}
-  )
-  service._update_user_state_transactional.assert_called_once_with(
-      "test_app", "test_user", {"my_key": "user_val"}
-  )
-
-  batch = mock_firestore_client.batch.return_value
+  with mock.patch("google.cloud.firestore.async_transactional", lambda x: x):
+    await service.append_event(session, event)
 
-  batch.set.assert_called()
+  transaction = mock_firestore_client.transaction.return_value
+  transaction.set.assert_called()
 
   assert session.state["session_key"] == "session_val"
 
   from google.cloud import firestore
 
-  batch.update.assert_called_once()
-  args, kwargs = batch.update.call_args
+  transaction.update.assert_called_once()
+  args, kwargs = transaction.update.call_args
+  assert args[0] == session_ref
   assert args[1]["state"] == session.state
   assert args[1]["updateTime"] == firestore.SERVER_TIMESTAMP
 
-  batch.commit.assert_called_once()
-
 
 @pytest.mark.asyncio
 async def test_list_sessions_with_user_id(mock_firestore_client):
@@ -497,10 +492,11 @@ async def test_create_session_already_exists(mock_firestore_client):
 
   from google.adk.errors.already_exists_error import AlreadyExistsError
 
-  with pytest.raises(AlreadyExistsError):
-    await service.create_session(
-        app_name=app_name, user_id=user_id, session_id="existing_id"
-    )
+  with mock.patch("google.cloud.firestore.async_transactional", lambda x: x):
+    with pytest.raises(AlreadyExistsError):
+      await service.create_session(
+          app_name=app_name, user_id=user_id, session_id="existing_id"
+      )
 
 
 @pytest.mark.asyncio