ai-memory-backend / memory /oracle_adapter.py
Baida07's picture
sync: 171 files from Baida98/AI [deploy-all]
24480a0 verified
Raw
History Blame
8.97 kB
"""
backend/memory/oracle_adapter.py — Oracle Database MemoryBackend adapter.
Implements MemoryBackend for Oracle using python-oracledb (thin mode, pooled).
Bug fixes vs the original prototype:
BUG-1 add_memory accepts an external memory_id so IDs stay consistent with
the Supabase primary backend (no more key divergence on dual-write).
BUG-2 uuid4() instead of hash() — hash() is non-deterministic across
processes (PYTHONHASHSEED), causing silent duplicate-key collisions
on Railway restarts.
BUG-3 is_connected flag guards every operation; the flag is only set True
after a successful pool creation, so callers can check it cheaply.
BUG-5 oracledb.create_pool() instead of a single connection — safe under
concurrent asyncio tasks and survives individual connection drops.
BUG-6 Removed isinstance(self.connection, str) branch — that was a
test-artifact that leaked mock awareness into production code.
BUG-7 time.time() instead of asyncio.get_event_loop().time() — the latter
is deprecated since Python 3.10 and raises DeprecationWarning on 3.12.
"""
from __future__ import annotations
import asyncio
import json
import uuid
from typing import Any, Dict, List, Optional
import oracledb
from .memory_backend import MemoryBackend
class OracleMemoryAdapter(MemoryBackend):
"""MemoryBackend implementation for Oracle Database (thin mode, connection pool)."""
def __init__(self) -> None:
self._pool: Optional[oracledb.ConnectionPool] = None
self.is_connected: bool = False
# ── Lifecycle ────────────────────────────────────────────────────────────────
async def connect(self, config: Dict[str, Any]) -> None:
"""Create a connection pool. config keys: user, password, dsn."""
try:
self._pool = await asyncio.to_thread(
oracledb.create_pool,
user=config["user"],
password=config["password"],
dsn=config["dsn"],
min=1,
max=5,
increment=1,
)
self.is_connected = True
except oracledb.Error as exc:
self.is_connected = False
raise ConnectionError(f"OracleMemoryAdapter: pool creation failed — {exc}") from exc
async def disconnect(self) -> None:
"""Close the connection pool gracefully."""
if self._pool and self.is_connected:
await asyncio.to_thread(self._pool.close)
self.is_connected = False
self._pool = None
def _require_pool(self) -> None:
if not self.is_connected or self._pool is None:
raise ConnectionError(
"OracleMemoryAdapter: not connected — call connect() first."
)
# ── Helpers ──────────────────────────────────────────────────────────────────
async def _acquire(self): # type: ignore[return]
return await asyncio.to_thread(self._pool.acquire)
async def _release(self, conn) -> None: # type: ignore[type-arg]
await asyncio.to_thread(self._pool.release, conn)
# ── CRUD ─────────────────────────────────────────────────────────────────────
async def add_memory(
self,
user_id: str,
memory_data: Dict[str, Any],
memory_id: Optional[str] = None, # BUG-1 fix
) -> str:
self._require_pool()
if memory_id is None:
memory_id = str(uuid.uuid4()) # BUG-2 fix
data_json = json.dumps(memory_data)
conn = cursor = None
try:
conn = await self._acquire()
cursor = await asyncio.to_thread(conn.cursor)
await asyncio.to_thread(
cursor.execute,
"INSERT INTO memories (user_id, memory_id, data)"
" VALUES (:user_id, :memory_id, :data)",
user_id=user_id,
memory_id=memory_id,
data=data_json,
)
await asyncio.to_thread(conn.commit)
return memory_id
except oracledb.Error as exc:
if conn:
await asyncio.to_thread(conn.rollback)
raise RuntimeError(f"OracleMemoryAdapter.add_memory failed: {exc}") from exc
finally:
if cursor:
await asyncio.to_thread(cursor.close)
if conn:
await self._release(conn)
async def get_memory(
self, user_id: str, memory_id: str
) -> Optional[Dict[str, Any]]:
self._require_pool()
conn = cursor = None
try:
conn = await self._acquire()
cursor = await asyncio.to_thread(conn.cursor)
await asyncio.to_thread(
cursor.execute,
"SELECT data FROM memories"
" WHERE user_id = :user_id AND memory_id = :memory_id",
user_id=user_id,
memory_id=memory_id,
)
row = await asyncio.to_thread(cursor.fetchone)
return json.loads(row[0]) if row else None
except oracledb.Error as exc:
raise RuntimeError(f"OracleMemoryAdapter.get_memory failed: {exc}") from exc
finally:
if cursor:
await asyncio.to_thread(cursor.close)
if conn:
await self._release(conn)
async def update_memory(
self, user_id: str, memory_id: str, new_data: Dict[str, Any]
) -> bool:
self._require_pool()
data_json = json.dumps(new_data)
conn = cursor = None
try:
conn = await self._acquire()
cursor = await asyncio.to_thread(conn.cursor)
await asyncio.to_thread(
cursor.execute,
"UPDATE memories SET data = :data"
" WHERE user_id = :user_id AND memory_id = :memory_id",
data=data_json,
user_id=user_id,
memory_id=memory_id,
)
rowcount = cursor.rowcount
await asyncio.to_thread(conn.commit)
return rowcount > 0
except oracledb.Error as exc:
if conn:
await asyncio.to_thread(conn.rollback)
raise RuntimeError(f"OracleMemoryAdapter.update_memory failed: {exc}") from exc
finally:
if cursor:
await asyncio.to_thread(cursor.close)
if conn:
await self._release(conn)
async def delete_memory(self, user_id: str, memory_id: str) -> bool:
self._require_pool()
conn = cursor = None
try:
conn = await self._acquire()
cursor = await asyncio.to_thread(conn.cursor)
await asyncio.to_thread(
cursor.execute,
"DELETE FROM memories"
" WHERE user_id = :user_id AND memory_id = :memory_id",
user_id=user_id,
memory_id=memory_id,
)
rowcount = cursor.rowcount
await asyncio.to_thread(conn.commit)
return rowcount > 0
except oracledb.Error as exc:
if conn:
await asyncio.to_thread(conn.rollback)
raise RuntimeError(f"OracleMemoryAdapter.delete_memory failed: {exc}") from exc
finally:
if cursor:
await asyncio.to_thread(cursor.close)
if conn:
await self._release(conn)
async def list_memories(
self, user_id: str, limit: int = 100, offset: int = 0
) -> List[Dict[str, Any]]:
self._require_pool()
conn = cursor = None
try:
conn = await self._acquire()
cursor = await asyncio.to_thread(conn.cursor)
await asyncio.to_thread(
cursor.execute,
"""SELECT data FROM memories
WHERE user_id = :user_id
ORDER BY memory_id
OFFSET :offset ROWS FETCH NEXT :limit ROWS ONLY""",
user_id=user_id,
offset=offset,
limit=limit,
)
rows = await asyncio.to_thread(cursor.fetchall)
return [json.loads(row[0]) for row in rows]
except oracledb.Error as exc:
raise RuntimeError(f"OracleMemoryAdapter.list_memories failed: {exc}") from exc
finally:
if cursor:
await asyncio.to_thread(cursor.close)
if conn:
await self._release(conn)