fix(memory): address database memory review feedback

Author: Raman369AI

src/google/adk_community/memory/database_memory_service.py

@@ -33,7 +33,11 @@
 from google.adk.memory.memory_entry import MemoryEntry
 from google.genai import types
 from sqlalchemy import delete
+from sqlalchemy import func
 from sqlalchemy import select
+from sqlalchemy.dialects.mysql import insert as mysql_insert
+from sqlalchemy.dialects.postgresql import insert as postgresql_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import ArgumentError
 from sqlalchemy.ext.asyncio import async_sessionmaker
@@ -57,6 +61,7 @@
 logger = logging.getLogger('google_adk.' + __name__)
 
 _SQLITE_DIALECT = 'sqlite'
+_MYSQL_DIALECTS = frozenset({'mysql', 'mariadb'})
 
 
 def _format_timestamp(timestamp: float) -> str:
@@ -378,21 +383,43 @@ async def set_scratchpad(
     """
     await self._prepare_tables()
     async with self._session() as sql:
-      existing = await sql.get(
-          StorageScratchpadKV, (app_name, user_id, session_id, key)
-      )
-      if existing is not None:
-        existing.value_json = value
-      else:
-        sql.add(
-            StorageScratchpadKV(
-                app_name=app_name,
-                user_id=user_id,
-                session_id=session_id,
-                key=key,
-                value_json=value,
-            )
+      values = {
+          'app_name': app_name,
+          'user_id': user_id,
+          'session_id': session_id,
+          'key': key,
+          'value_json': value,
+      }
+      dialect_name = sql.get_bind().dialect.name
+      if dialect_name == 'postgresql':
+        stmt = postgresql_insert(StorageScratchpadKV).values(**values)
+        stmt = stmt.on_conflict_do_update(
+            index_elements=['app_name', 'user_id', 'session_id', 'key'],
+            set_={
+                'value_json': stmt.excluded.value_json,
+                'updated_at': func.now(),
+            },
+        )
+        await sql.execute(stmt)
+      elif dialect_name == _SQLITE_DIALECT:
+        stmt = sqlite_insert(StorageScratchpadKV).values(**values)
+        stmt = stmt.on_conflict_do_update(
+            index_elements=['app_name', 'user_id', 'session_id', 'key'],
+            set_={
+                'value_json': stmt.excluded.value_json,
+                'updated_at': func.now(),
+            },
         )
+        await sql.execute(stmt)
+      elif dialect_name in _MYSQL_DIALECTS:
+        stmt = mysql_insert(StorageScratchpadKV).values(**values)
+        stmt = stmt.on_duplicate_key_update(
+            value_json=stmt.inserted.value_json,
+            updated_at=func.now(),
+        )
+        await sql.execute(stmt)
+      else:
+        await sql.merge(StorageScratchpadKV(**values))
 
   async def get_scratchpad(
       self,

src/google/adk_community/memory/memory_search_backend.py

@@ -30,7 +30,7 @@
 if TYPE_CHECKING:
   from sqlalchemy.ext.asyncio import AsyncSession
 
-_ILIKE_DIALECTS = frozenset({'postgresql', 'mysql', 'mariadb'})
+_ILIKE_DIALECTS = frozenset({'postgresql'})
 
 
 class MemorySearchBackend(ABC):
@@ -68,8 +68,8 @@ class KeywordSearchBackend(MemorySearchBackend):
     2. Try an AND predicate (all tokens must appear) — return if found.
     3. Fall back to OR (any token matches) if AND yields nothing.
 
-  Uses ILIKE on PostgreSQL/MySQL/MariaDB and LIKE on SQLite
-  (case-insensitive by default collation).
+  Uses ILIKE on PostgreSQL and LIKE on other dialects
+  (case-insensitive by default for common SQLite/MySQL/MariaDB collations).
   """
 
   async def search(
@@ -88,9 +88,7 @@ async def search(
     tokens = [
         cleaned
         for raw in query.split()
-        if raw.strip()
-        for cleaned in [re.sub(r'[^\w]', '', raw).lower()]
-        if cleaned
+        if (cleaned := re.sub(r'[^\w]', '', raw).lower())
     ]
     if not tokens:
       return []

tests/unittests/memory/test_database_memory_service.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Sequence
 import time
 from typing import Any
@@ -278,6 +279,28 @@ async def test_scratchpad_kv_overwrite(svc):
   assert val == 'new'
 
 
+@pytest.mark.asyncio
+async def test_scratchpad_kv_concurrent_set_same_new_key(tmp_path):
+  db_path = tmp_path / 'scratchpad.db'
+  svc = DatabaseMemoryService(f'sqlite+aiosqlite:///{db_path}')
+
+  await asyncio.gather(*[
+      svc.set_scratchpad(
+          app_name=_APP,
+          user_id=_USER,
+          session_id=_SESSION,
+          key='shared',
+          value=i,
+      )
+      for i in range(10)
+  ])
+
+  val = await svc.get_scratchpad(
+      app_name=_APP, user_id=_USER, session_id=_SESSION, key='shared'
+  )
+  assert val in range(10)
+
+
 @pytest.mark.asyncio
 async def test_scratchpad_kv_missing_returns_none(svc):
   val = await svc.get_scratchpad(