mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
update main with changes
This commit is contained in:
phernandez
@@ -1,235 +0,0 @@
|
||||
"""Base repository implementation."""
|
||||
from typing import Type, Optional, Any, Sequence, TypeVar, List, Dict
|
||||
from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, insert, and_, delete
|
||||
from sqlalchemy.exc import NoResultFound
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import Mapped
|
||||
from loguru import logger
|
||||
|
||||
from basic_memory.models import Base, Entity
|
||||
|
||||
T = TypeVar('T', bound=Base)
|
||||
|
||||
class Repository[T: Base]:
|
||||
"""Base repository implementation with generic CRUD operations."""
|
||||
|
||||
def __init__(self, session: AsyncSession, Model: Type[T]):
|
||||
self.session = session
|
||||
self.Model = Model
|
||||
self.mapper = inspect(self.Model).mapper
|
||||
self.primary_key: Column[Any] = self.mapper.primary_key[0]
|
||||
self.valid_columns = [column.key for column in self.mapper.columns]
|
||||
|
||||
def select(self, *entities: Any) -> Select:
|
||||
"""Create a new SELECT statement.
|
||||
|
||||
Returns:
|
||||
A SQLAlchemy Select object configured with the provided entities
|
||||
or this repository's model if no entities provided.
|
||||
"""
|
||||
if not entities:
|
||||
entities = (self.Model,)
|
||||
return select(*entities)
|
||||
|
||||
async def refresh(self, instance: T, relationships: list[str] | None = None) -> None:
|
||||
"""Refresh instance and optionally specified relationships."""
|
||||
logger.debug(f"Refreshing {self.Model.__name__} instance: {getattr(instance, 'id', None)}")
|
||||
try:
|
||||
await self.session.refresh(instance, relationships or [])
|
||||
logger.debug(f"Refreshed relationships: {relationships}")
|
||||
except Exception:
|
||||
logger.exception(f"Failed to refresh {self.Model.__name__} instance")
|
||||
raise
|
||||
|
||||
async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]:
|
||||
"""Fetch records from the database with pagination."""
|
||||
logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})")
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).offset(skip).limit(limit)
|
||||
)
|
||||
items = result.scalars().all()
|
||||
logger.debug(f"Found {len(items)} {self.Model.__name__} records")
|
||||
return items
|
||||
except Exception:
|
||||
logger.exception(f"Failed to find all {self.Model.__name__}")
|
||||
raise
|
||||
|
||||
async def find_by_id(self, entity_id: str) -> Optional[T]:
|
||||
"""Fetch an entity by its unique identifier."""
|
||||
logger.debug(f"Finding {self.Model.__name__} by ID: {entity_id}")
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
logger.debug(f"Found {self.Model.__name__}: {entity_id}")
|
||||
return entity
|
||||
except NoResultFound:
|
||||
logger.debug(f"No {self.Model.__name__} found with ID: {entity_id}")
|
||||
return None
|
||||
except Exception:
|
||||
logger.exception(f"Failed to find {self.Model.__name__} by ID: {entity_id}")
|
||||
raise
|
||||
|
||||
async def create(self, entity_data: dict, model: Type[Base] | None = None) -> T:
|
||||
"""Create a new entity in the database from the provided data."""
|
||||
model = model or self.Model
|
||||
logger.debug(f"Creating {model.__name__} with data: {entity_data}")
|
||||
try:
|
||||
# Only include valid columns that are provided in entity_data
|
||||
model_data = {
|
||||
k: v for k, v in entity_data.items()
|
||||
if k in self.valid_columns and v is not None
|
||||
}
|
||||
|
||||
# Generate ID if this is an Entity model and no ID provided
|
||||
if model is Entity and 'id' not in model_data:
|
||||
model_data['id'] = Entity.generate_id(
|
||||
model_data['entity_type'],
|
||||
model_data['name']
|
||||
)
|
||||
|
||||
logger.debug(f"Filtered data for valid columns: {model_data}")
|
||||
|
||||
# Create insert statement with only provided data
|
||||
stmt = insert(model).values(**model_data).returning(model)
|
||||
result = await self.session.execute(stmt)
|
||||
entity: T = result.scalar_one() # pyright: ignore [reportAssignmentType]
|
||||
|
||||
logger.debug(f"Created {model.__name__}: {getattr(entity, 'id', None)}")
|
||||
return entity
|
||||
|
||||
except Exception:
|
||||
logger.exception(f"Failed to create {model.__name__}")
|
||||
raise
|
||||
|
||||
async def instance_create(self, instance: T) -> T:
|
||||
"""Create a new record from a model instance."""
|
||||
logger.debug(f"Creating {self.Model.__name__} from instance: {instance}")
|
||||
try:
|
||||
self.session.add(instance)
|
||||
await self.session.flush()
|
||||
return instance
|
||||
except Exception:
|
||||
logger.exception(f"Failed to create {self.Model.__name__}")
|
||||
raise
|
||||
|
||||
async def bulk_create(self, instances: List[T]) -> List[T]:
|
||||
"""Create multiple records in a single transaction."""
|
||||
logger.debug(f"Bulk creating {len(instances)} {self.Model.__name__} instances")
|
||||
try:
|
||||
for instance in instances:
|
||||
self.session.add(instance)
|
||||
await self.session.flush()
|
||||
return instances
|
||||
except Exception:
|
||||
logger.exception(f"Failed to bulk create {self.Model.__name__}")
|
||||
raise
|
||||
|
||||
async def update(self, entity_id: str, entity_data: dict) -> Optional[T]:
|
||||
"""Update an entity with the given data."""
|
||||
logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}")
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
|
||||
for key, value in entity_data.items():
|
||||
if key in self.valid_columns:
|
||||
setattr(entity, key, value)
|
||||
await self.session.flush()
|
||||
|
||||
logger.debug(f"Updated {self.Model.__name__}: {entity_id}")
|
||||
return entity
|
||||
except NoResultFound:
|
||||
logger.debug(f"No {self.Model.__name__} found to update: {entity_id}")
|
||||
return None
|
||||
except Exception:
|
||||
logger.exception(f"Failed to update {self.Model.__name__}: {entity_id}")
|
||||
raise
|
||||
|
||||
async def delete(self, entity_id: str) -> bool:
|
||||
"""Delete an entity from the database."""
|
||||
logger.debug(f"Deleting {self.Model.__name__}: {entity_id}")
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
await self.session.delete(entity)
|
||||
await self.session.flush()
|
||||
|
||||
logger.debug(f"Deleted {self.Model.__name__}: {entity_id}")
|
||||
return True
|
||||
except NoResultFound:
|
||||
logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}")
|
||||
return False
|
||||
except Exception:
|
||||
logger.exception(f"Failed to delete {self.Model.__name__}: {entity_id}")
|
||||
raise
|
||||
|
||||
async def delete_by_fields(self, **filters: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Delete records matching given field values.
|
||||
|
||||
Args:
|
||||
**filters: Field names and values to filter by
|
||||
|
||||
Returns:
|
||||
bool: True if any records were deleted
|
||||
"""
|
||||
logger.debug(f"Deleting {self.Model.__name__} by fields: {filters}")
|
||||
try:
|
||||
conditions = [getattr(self.Model, field) == value for field, value in filters.items()]
|
||||
query = delete(self.Model).where(and_(*conditions))
|
||||
result = await self.execute_query(query)
|
||||
await self.session.flush()
|
||||
deleted = result.rowcount > 0 # pyright: ignore [reportAttributeAccessIssue]
|
||||
logger.debug(f"Deleted {result.rowcount} records") # pyright: ignore [reportAttributeAccessIssue]
|
||||
return deleted # pyright: ignore [reportAttributeAccessIssue]
|
||||
except Exception:
|
||||
logger.exception(f"Failed to delete {self.Model.__name__} by fields")
|
||||
raise
|
||||
|
||||
async def count(self, query: Executable | None = None) -> int:
|
||||
"""Count entities in the database table."""
|
||||
try:
|
||||
if query is None:
|
||||
query = select(func.count()).select_from(self.Model)
|
||||
result = await self.session.execute(query)
|
||||
scalar = result.scalar()
|
||||
count = scalar if scalar is not None else 0
|
||||
|
||||
logger.debug(f"Counted {count} {self.Model.__name__} records")
|
||||
return count
|
||||
except Exception:
|
||||
logger.exception(f"Failed to count {self.Model.__name__}")
|
||||
raise
|
||||
|
||||
async def execute_query(self, query: Executable) -> Result[Any]:
|
||||
"""Execute a query asynchronously."""
|
||||
logger.debug(f"Executing query: {query}")
|
||||
try:
|
||||
result = await self.session.execute(query)
|
||||
logger.debug("Query executed successfully")
|
||||
return result
|
||||
except Exception:
|
||||
logger.exception("Failed to execute query")
|
||||
raise
|
||||
|
||||
async def find_one(self, query: Select[tuple[T]]) -> Optional[T]:
|
||||
"""Execute a query and retrieve a single record."""
|
||||
logger.debug(f"Finding one {self.Model.__name__} with query: {query}")
|
||||
try:
|
||||
result = await self.execute_query(query)
|
||||
entity = result.scalars().one_or_none()
|
||||
if entity:
|
||||
logger.debug(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}")
|
||||
else:
|
||||
logger.debug(f"No {self.Model.__name__} found")
|
||||
return entity
|
||||
except Exception:
|
||||
logger.exception(f"Failed to find one {self.Model.__name__}")
|
||||
raise
|
||||
@@ -1,132 +1,165 @@
|
||||
"""Repository for managing Entity objects."""
|
||||
from typing import Optional, Sequence, Type
|
||||
|
||||
from typing import Optional, Sequence, List
|
||||
|
||||
from loguru import logger
|
||||
from sqlalchemy import select, or_
|
||||
from sqlalchemy.exc import NoResultFound
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from basic_memory import db
|
||||
from basic_memory.models import Entity, Observation
|
||||
from basic_memory.repository import Repository
|
||||
from loguru import logger
|
||||
from basic_memory.repository.repository import Repository
|
||||
|
||||
|
||||
class EntityRepository(Repository[Entity]):
|
||||
"""Repository for Entity model with memory-specific operations."""
|
||||
|
||||
def __init__(self, session):
|
||||
super().__init__(session, Entity)
|
||||
|
||||
def __init__(self, session_maker: async_sessionmaker):
|
||||
super().__init__(session_maker, Entity)
|
||||
logger.debug("Initialized EntityRepository")
|
||||
|
||||
async def create(self, entity_data: dict, model: Type[Entity] | None = None) -> Entity:
|
||||
async def create(self, data: dict) -> Entity: # pyright: ignore [reportIncompatibleMethodOverride]
|
||||
"""Create a new entity in the database from the provided data."""
|
||||
entity_id = Entity.generate_id(
|
||||
entity_data['entity_type'],
|
||||
entity_data['name']
|
||||
)
|
||||
return await super().create({**entity_data, 'id': entity_id})
|
||||
entity_id = Entity.generate_id(data["entity_type"], data["name"])
|
||||
await super().create({**data, "id": entity_id})
|
||||
|
||||
# we have to find to get relations
|
||||
created = await self.find_by_id(entity_id)
|
||||
assert created is not None, f"Created entity {entity_id} should not be None"
|
||||
return created
|
||||
|
||||
async def create_all(self, data_list: List[dict]) -> Sequence[Entity]: # pyright: ignore [reportIncompatibleMethodOverride]
|
||||
"""Create a new entity in the database from the provided data."""
|
||||
for data in data_list:
|
||||
entity_id = Entity.generate_id(data["entity_type"], data["name"])
|
||||
data["id"] = entity_id
|
||||
created = await super().create_all(data_list)
|
||||
# we have to find to get relations
|
||||
return await self.find_by_ids([e.id for e in created])
|
||||
|
||||
async def find_by_id(self, entity_id: str) -> Optional[Entity]:
|
||||
"""Find entity by ID with all relationships eagerly loaded."""
|
||||
logger.debug(f"Finding entity by ID: {entity_id}")
|
||||
try:
|
||||
# First load base entity
|
||||
result = await self.session.execute(
|
||||
select(Entity).filter(Entity.id == entity_id)
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
try:
|
||||
result = await session.execute(
|
||||
select(Entity)
|
||||
.filter(Entity.id == entity_id)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations),
|
||||
)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
logger.debug(f"Found entity: {entity.id}")
|
||||
return entity
|
||||
except NoResultFound:
|
||||
logger.debug(f"No entity found with ID: {entity_id}")
|
||||
return None
|
||||
|
||||
async def find_by_ids(self, ids: List[str]) -> Sequence[Entity]:
|
||||
"""Search for entities of a specific type."""
|
||||
logger.debug(f"Find entities by ids: {ids}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(
|
||||
select(Entity)
|
||||
.where(self.primary_key.in_(ids))
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations),
|
||||
)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
logger.debug(f"Found base entity: {entity.id}")
|
||||
|
||||
# Force refresh of all relationships
|
||||
await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations'])
|
||||
logger.debug(f"Refreshed entity relationships: {entity.id}")
|
||||
|
||||
return entity
|
||||
except NoResultFound:
|
||||
logger.debug(f"No entity found with ID: {entity_id}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.exception(f"Error finding entity by ID: {entity_id}")
|
||||
raise
|
||||
entities = result.scalars().all()
|
||||
logger.debug(f"Found {len(entities)}")
|
||||
return entities
|
||||
|
||||
async def find_by_name(self, name: str) -> Optional[Entity]:
|
||||
"""Find an entity by its unique name."""
|
||||
logger.debug(f"Finding entity by name: {name}")
|
||||
try:
|
||||
query = (
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(
|
||||
select(Entity)
|
||||
.filter(Entity.name == name)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations),
|
||||
)
|
||||
)
|
||||
result = await self.session.execute(query)
|
||||
entity = result.scalars().one_or_none()
|
||||
if entity:
|
||||
logger.debug(f"Found entity: {entity.id}")
|
||||
await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations'])
|
||||
logger.debug(f"Refreshed entity relationships: {entity.id}")
|
||||
else:
|
||||
logger.debug(f"No entity found with name: {name}")
|
||||
return entity
|
||||
except Exception as e:
|
||||
logger.exception(f"Error finding entity by name: {name}")
|
||||
raise
|
||||
|
||||
async def find_by_type_and_name(self, entity_type: str, name: str) -> Optional[Entity]:
|
||||
"""Find an entity by its type and name combination."""
|
||||
logger.debug(f"Finding entity by type and name: {entity_type}/{name}")
|
||||
try:
|
||||
query = (
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(
|
||||
select(Entity)
|
||||
.filter(Entity.entity_type == entity_type)
|
||||
.filter(Entity.name == name)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations)
|
||||
selectinload(Entity.incoming_relations),
|
||||
)
|
||||
)
|
||||
result = await self.session.execute(query)
|
||||
entity = result.scalars().one_or_none()
|
||||
if entity:
|
||||
logger.debug(f"Found entity: {entity.id}")
|
||||
else:
|
||||
logger.debug(f"No entity found with type/name: {entity_type}/{name}")
|
||||
return entity
|
||||
except Exception as e:
|
||||
logger.exception(f"Error finding entity by type/name: {entity_type}/{name}")
|
||||
raise
|
||||
|
||||
async def search_by_type(self, entity_type: str, skip: int = 0, limit: int = 100) -> Sequence[Entity]:
|
||||
|
||||
async def search_by_type(
|
||||
self, entity_type: str, skip: int = 0, limit: int = 100
|
||||
) -> Sequence[Entity]:
|
||||
"""Search for entities of a specific type."""
|
||||
logger.debug(f"Searching entities by type: {entity_type} (skip={skip}, limit={limit})")
|
||||
try:
|
||||
query = select(Entity).filter(Entity.entity_type == entity_type).offset(skip).limit(limit)
|
||||
result = await self.execute_query(query)
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(
|
||||
select(Entity)
|
||||
.filter(Entity.entity_type == entity_type)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations),
|
||||
)
|
||||
.offset(skip)
|
||||
.limit(limit)
|
||||
)
|
||||
entities = result.scalars().all()
|
||||
logger.debug(f"Found {len(entities)} entities of type {entity_type}")
|
||||
return entities
|
||||
except Exception as e:
|
||||
logger.exception(f"Error searching entities by type: {entity_type}")
|
||||
raise
|
||||
|
||||
async def search(self, query: str) -> Sequence[Entity]:
|
||||
"""Search entities using LIKE pattern matching."""
|
||||
logger.debug(f"Searching entities with query: {query}")
|
||||
try:
|
||||
stmt = select(Entity).distinct().where(
|
||||
or_(
|
||||
Entity.name.ilike(f"%{query}%"),
|
||||
Entity.entity_type.ilike(f"%{query}%"),
|
||||
Entity.observations.any(
|
||||
Observation.content.ilike(f"%{query}%")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
stmt = (
|
||||
select(Entity)
|
||||
.distinct()
|
||||
.where(
|
||||
or_(
|
||||
Entity.name.ilike(f"%{query}%"),
|
||||
Entity.entity_type.ilike(f"%{query}%"),
|
||||
Entity.observations.any(Observation.content.ilike(f"%{query}%")),
|
||||
)
|
||||
)
|
||||
).options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations),
|
||||
)
|
||||
)
|
||||
result = await self.session.execute(stmt)
|
||||
result = await session.execute(stmt)
|
||||
entities = list(result.scalars())
|
||||
logger.debug(f"Found {len(entities)} matching entities")
|
||||
return entities
|
||||
except Exception as e:
|
||||
logger.exception(f"Error searching entities: {query}")
|
||||
raise
|
||||
@@ -1,25 +1,30 @@
|
||||
"""Repository for managing Observation objects."""
|
||||
|
||||
from typing import Sequence
|
||||
|
||||
from loguru import logger
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker
|
||||
|
||||
from basic_memory.models import Observation
|
||||
from basic_memory.repository import Repository
|
||||
from basic_memory.repository.repository import Repository
|
||||
|
||||
|
||||
class ObservationRepository(Repository[Observation]):
|
||||
"""Repository for Observation model with memory-specific operations."""
|
||||
|
||||
def __init__(self, session):
|
||||
super().__init__(session, Observation)
|
||||
|
||||
|
||||
def __init__(self, session_maker: async_sessionmaker):
|
||||
super().__init__(session_maker, Observation)
|
||||
logger.debug("Initialized ObservationRepository")
|
||||
|
||||
async def find_by_entity(self, entity_id: str) -> Sequence[Observation]:
|
||||
"""Find all observations for a specific entity."""
|
||||
query = select(Observation).filter(Observation.entity_id == entity_id)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def find_by_context(self, context: str) -> Sequence[Observation]:
|
||||
"""Find observations with a specific context."""
|
||||
query = select(Observation).filter(Observation.context == context)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
return result.scalars().all()
|
||||
|
||||
@@ -1,16 +1,19 @@
|
||||
"""Repository for managing Relation objects."""
|
||||
|
||||
from typing import Sequence
|
||||
|
||||
from sqlalchemy import select, and_
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker
|
||||
|
||||
from basic_memory.models import Relation
|
||||
from basic_memory.repository import Repository
|
||||
from basic_memory.repository.repository import Repository
|
||||
|
||||
|
||||
class RelationRepository(Repository[Relation]):
|
||||
"""Repository for Relation model with memory-specific operations."""
|
||||
|
||||
def __init__(self, session):
|
||||
super().__init__(session, Relation)
|
||||
|
||||
def __init__(self, session_maker: async_sessionmaker):
|
||||
super().__init__(session_maker, Relation)
|
||||
|
||||
async def find_by_entity(self, from_entity_id: str) -> Sequence[Relation]:
|
||||
"""Find all relations from a specific entity."""
|
||||
@@ -20,17 +23,12 @@ class RelationRepository(Repository[Relation]):
|
||||
|
||||
async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]:
|
||||
"""Find all relations between two entities."""
|
||||
query = select(Relation).filter(
|
||||
and_(
|
||||
Relation.from_id == from_id,
|
||||
Relation.to_id == to_id
|
||||
)
|
||||
)
|
||||
query = select(Relation).filter(and_(Relation.from_id == from_id, Relation.to_id == to_id))
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def find_by_type(self, relation_type: str) -> Sequence[Relation]:
|
||||
"""Find all relations of a specific type."""
|
||||
query = select(Relation).filter(Relation.relation_type == relation_type)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
return result.scalars().all()
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
"""Base repository implementation."""
|
||||
|
||||
from typing import Type, Optional, Any, Sequence, TypeVar, List
|
||||
|
||||
from loguru import logger
|
||||
from sqlalchemy import (
|
||||
select,
|
||||
func,
|
||||
Select,
|
||||
Executable,
|
||||
inspect,
|
||||
Result,
|
||||
Column,
|
||||
and_,
|
||||
delete,
|
||||
)
|
||||
from sqlalchemy.exc import NoResultFound
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
|
||||
|
||||
from basic_memory import db
|
||||
from basic_memory.models import Base
|
||||
|
||||
T = TypeVar("T", bound=Base)
|
||||
|
||||
|
||||
class Repository[T: Base]:
|
||||
"""Base repository implementation with generic CRUD operations."""
|
||||
|
||||
def __init__(self, session_maker: async_sessionmaker[AsyncSession], Model: Type[T]):
|
||||
self.session_maker = session_maker
|
||||
self.Model = Model
|
||||
self.mapper = inspect(self.Model).mapper
|
||||
self.primary_key: Column[Any] = self.mapper.primary_key[0]
|
||||
self.valid_columns = [column.key for column in self.mapper.columns]
|
||||
|
||||
def get_model_data(self, entity_data):
|
||||
model_data = {
|
||||
k: v for k, v in entity_data.items() if k in self.valid_columns and v is not None
|
||||
}
|
||||
return model_data
|
||||
|
||||
async def add(self, model: T) -> T:
|
||||
"""
|
||||
Add a model to the repository. This will also add related objects
|
||||
:param model: the model to add
|
||||
:return: the added model instance
|
||||
"""
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
session.add(model)
|
||||
await session.flush()
|
||||
found = await self.find_by_id(model.id) # pyright: ignore [reportAttributeAccessIssue]
|
||||
assert found is not None, "can't find model after session.add"
|
||||
return found
|
||||
|
||||
async def add_all(self, models: List[T]) -> Sequence[T]:
|
||||
"""
|
||||
Add a list of models to the repository. This will also add related objects
|
||||
:param model: the models to add
|
||||
:return: the added models instances
|
||||
"""
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
session.add_all(models)
|
||||
await session.flush()
|
||||
# we have to find to get relations
|
||||
return await self.find_by_ids([m.id for m in models]) # pyright: ignore [reportAttributeAccessIssue]
|
||||
|
||||
def select(self, *entities: Any) -> Select:
|
||||
"""Create a new SELECT statement.
|
||||
|
||||
Returns:
|
||||
A SQLAlchemy Select object configured with the provided entities
|
||||
or this repository's model if no entities provided.
|
||||
"""
|
||||
if not entities:
|
||||
entities = (self.Model,)
|
||||
return select(*entities)
|
||||
|
||||
async def refresh(self, instance: T, relationships: list[str] | None = None) -> None:
|
||||
"""Refresh instance and optionally specified relationships."""
|
||||
logger.debug(f"Refreshing {self.Model.__name__} instance: {getattr(instance, 'id', None)}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
await session.refresh(instance, relationships or [])
|
||||
logger.debug(f"Refreshed relationships: {relationships}")
|
||||
|
||||
async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]:
|
||||
"""Fetch records from the database with pagination."""
|
||||
logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(select(self.Model).offset(skip).limit(limit))
|
||||
items = result.scalars().all()
|
||||
logger.debug(f"Found {len(items)} {self.Model.__name__} records")
|
||||
return items
|
||||
|
||||
async def find_by_id(self, entity_id: str) -> Optional[T]:
|
||||
"""Fetch an entity by its unique identifier."""
|
||||
logger.debug(f"Finding {self.Model.__name__} by ID: {entity_id}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
try:
|
||||
result = await session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
logger.debug(f"Found {self.Model.__name__}: {entity_id}")
|
||||
return entity
|
||||
except NoResultFound:
|
||||
logger.debug(f"No {self.Model.__name__} found with ID: {entity_id}")
|
||||
return None
|
||||
|
||||
async def find_one(self, query: Select[tuple[T]]) -> Optional[T]:
|
||||
"""Execute a query and retrieve a single record."""
|
||||
logger.debug(f"Finding one {self.Model.__name__} with query: {query}")
|
||||
result = await self.execute_query(query)
|
||||
entity = result.scalars().one_or_none()
|
||||
if entity:
|
||||
logger.debug(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}")
|
||||
else:
|
||||
logger.debug(f"No {self.Model.__name__} found")
|
||||
return entity
|
||||
|
||||
async def find_by_ids(self, ids: List[str]) -> Sequence[T]:
|
||||
"""Fetch multiple entities by their identifiers in a single query."""
|
||||
logger.debug(f"Finding {self.Model.__name__} by IDs: {ids}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(select(self.Model).where(self.primary_key.in_(ids)))
|
||||
entities = result.scalars().all()
|
||||
logger.debug(f"Found {len(entities)} {self.Model.__name__} records")
|
||||
return entities
|
||||
|
||||
async def create(self, data: dict) -> T:
|
||||
"""Create a new record from a model instance."""
|
||||
logger.debug(f"Creating {self.Model.__name__} from entity_data: {data}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
# Only include valid columns that are provided in entity_data
|
||||
model_data = self.get_model_data(data)
|
||||
model = self.Model(**model_data)
|
||||
session.add(model)
|
||||
await session.flush()
|
||||
return model
|
||||
|
||||
async def create_all(self, data_list: List[dict]) -> List[T]:
|
||||
"""Create multiple records in a single transaction."""
|
||||
logger.debug(f"Bulk creating {len(data_list)} {self.Model.__name__} instances")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
# Only include valid columns that are provided in entity_data
|
||||
model_list = [self.Model(**self.get_model_data(d)) for d in data_list]
|
||||
session.add_all(model_list)
|
||||
return model_list
|
||||
|
||||
async def update(self, entity_id: str, entity_data: dict) -> Optional[T]:
|
||||
"""Update an entity with the given data."""
|
||||
logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
try:
|
||||
result = await session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
|
||||
for key, value in entity_data.items():
|
||||
if key in self.valid_columns:
|
||||
setattr(entity, key, value)
|
||||
|
||||
logger.debug(f"Updated {self.Model.__name__}: {entity_id}")
|
||||
return entity
|
||||
except NoResultFound:
|
||||
logger.debug(f"No {self.Model.__name__} found to update: {entity_id}")
|
||||
return None
|
||||
|
||||
async def delete(self, entity_id: str) -> bool:
|
||||
"""Delete an entity from the database."""
|
||||
logger.debug(f"Deleting {self.Model.__name__}: {entity_id}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
try:
|
||||
result = await session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
)
|
||||
entity = result.scalars().one()
|
||||
await session.delete(entity)
|
||||
|
||||
logger.debug(f"Deleted {self.Model.__name__}: {entity_id}")
|
||||
return True
|
||||
except NoResultFound:
|
||||
logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}")
|
||||
return False
|
||||
|
||||
async def delete_by_ids(self, ids: List[str]) -> int:
|
||||
"""Delete records matching given field values."""
|
||||
logger.debug(f"Deleting {self.Model.__name__} by ids: {ids}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
query = delete(self.Model).where(self.primary_key.in_(ids))
|
||||
result = await session.execute(query)
|
||||
logger.debug(f"Deleted {result.rowcount} records")
|
||||
return result.rowcount
|
||||
|
||||
async def delete_by_fields(self, **filters: Any) -> bool:
|
||||
"""Delete records matching given field values."""
|
||||
logger.debug(f"Deleting {self.Model.__name__} by fields: {filters}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
conditions = [getattr(self.Model, field) == value for field, value in filters.items()]
|
||||
query = delete(self.Model).where(and_(*conditions))
|
||||
result = await session.execute(query)
|
||||
deleted = result.rowcount > 0
|
||||
logger.debug(f"Deleted {result.rowcount} records")
|
||||
return deleted
|
||||
|
||||
async def count(self, query: Executable | None = None) -> int:
|
||||
"""Count entities in the database table."""
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
if query is None:
|
||||
query = select(func.count()).select_from(self.Model)
|
||||
result = await session.execute(query)
|
||||
scalar = result.scalar()
|
||||
count = scalar if scalar is not None else 0
|
||||
logger.debug(f"Counted {count} {self.Model.__name__} records")
|
||||
return count
|
||||
|
||||
async def execute_query(self, query: Executable) -> Result[Any]:
|
||||
"""Execute a query asynchronously."""
|
||||
logger.debug(f"Executing query: {query}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
result = await session.execute(query)
|
||||
logger.debug("Query executed successfully")
|
||||
return result
|
||||
Reference in New Issue
Block a user