basic-memory 0.2.12__py3-none-any.whl → 0.16.1__py3-none-any.whl

basic_memory/repository/repository.py

@@ -1,6 +1,6 @@
 """Base repository implementation."""
 
-from typing import Type, Optional, Any, Sequence, TypeVar, List
+from typing import Type, Optional, Any, Sequence, TypeVar, List, Dict
 
 from loguru import logger
 from sqlalchemy import (
@@ -10,13 +10,13 @@ from sqlalchemy import (
     Executable,
     inspect,
     Result,
-    Column,
     and_,
     delete,
 )
 from sqlalchemy.exc import NoResultFound
 from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from sqlalchemy.orm.interfaces import LoaderOption
+from sqlalchemy.sql.elements import ColumnElement
 
 from basic_memory import db
 from basic_memory.models import Base
@@ -27,12 +27,30 @@ 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]):
+    def __init__(
+        self,
+        session_maker: async_sessionmaker[AsyncSession],
+        Model: Type[T],
+        project_id: Optional[int] = None,
+    ):
         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]
+        self.project_id = project_id
+        if Model:
+            self.Model = Model
+            self.mapper = inspect(self.Model).mapper
+            self.primary_key: ColumnElement[Any] = self.mapper.primary_key[0]
+            self.valid_columns = [column.key for column in self.mapper.columns]
+            # Check if this model has a project_id column
+            self.has_project_id = "project_id" in self.valid_columns
+
+    def _set_project_id_if_needed(self, model: T) -> None:
+        """Set project_id on model if needed and available."""
+        if (
+            self.has_project_id
+            and self.project_id is not None
+            and getattr(model, "project_id", None) is None
+        ):
+            setattr(model, "project_id", self.project_id)
 
     def get_model_data(self, entity_data):
         model_data = {
@@ -40,6 +58,19 @@ class Repository[T: Base]:
         }
         return model_data
 
+    def _add_project_filter(self, query: Select) -> Select:
+        """Add project_id filter to query if applicable.
+
+        Args:
+            query: The SQLAlchemy query to modify
+
+        Returns:
+            Updated query with project filter if applicable
+        """
+        if self.has_project_id and self.project_id is not None:
+            query = query.filter(getattr(self.Model, "project_id") == self.project_id)
+        return query
+
     async def select_by_id(self, session: AsyncSession, entity_id: int) -> Optional[T]:
         """Select an entity by ID using an existing session."""
         query = (
@@ -47,6 +78,9 @@ class Repository[T: Base]:
             .filter(self.primary_key == entity_id)
             .options(*self.get_load_options())
         )
+        # Add project filter if applicable
+        query = self._add_project_filter(query)
+
         result = await session.execute(query)
         return result.scalars().one_or_none()
 
@@ -55,6 +89,9 @@ class Repository[T: Base]:
         query = (
             select(self.Model).where(self.primary_key.in_(ids)).options(*self.get_load_options())
         )
+        # Add project filter if applicable
+        query = self._add_project_filter(query)
+
         result = await session.execute(query)
         return result.scalars().all()
 
@@ -65,12 +102,23 @@ class Repository[T: Base]:
         :return: the added model instance
         """
         async with db.scoped_session(self.session_maker) as session:
+            # Set project_id if applicable and not already set
+            self._set_project_id_if_needed(model)
+
             session.add(model)
             await session.flush()
 
             # Query within same session
             found = await self.select_by_id(session, model.id)  # pyright: ignore [reportAttributeAccessIssue]
-            assert found is not None, "can't find model after session.add"
+            if found is None:  # pragma: no cover
+                logger.error(
+                    "Failed to retrieve model after add",
+                    model_type=self.Model.__name__,
+                    model_id=model.id,  # pyright: ignore
+                )
+                raise ValueError(
+                    f"Can't find {self.Model.__name__} with ID {model.id} after session.add"  # pyright: ignore
+                )
             return found
 
     async def add_all(self, models: List[T]) -> Sequence[T]:
@@ -80,6 +128,10 @@ class Repository[T: Base]:
         :return: the added models instances
         """
         async with db.scoped_session(self.session_maker) as session:
+            # set the project id if not present in models
+            for model in models:
+                self._set_project_id_if_needed(model)
+
             session.add_all(models)
             await session.flush()
 
@@ -95,14 +147,33 @@ class Repository[T: Base]:
         """
         if not entities:
             entities = (self.Model,)
-        return select(*entities)
+        query = select(*entities)
+
+        # Add project filter if applicable
+        return self._add_project_filter(query)
+
+    async def find_all(
+        self, skip: int = 0, limit: Optional[int] = None, use_load_options: bool = True
+    ) -> Sequence[T]:
+        """Fetch records from the database with pagination.
 
-    async def find_all(self, skip: int = 0, limit: Optional[int] = 0) -> Sequence[T]:
-        """Fetch records from the database with pagination."""
+        Args:
+            skip: Number of records to skip
+            limit: Maximum number of records to return
+            use_load_options: Whether to apply eager loading options (default: True)
+        """
         logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})")
 
         async with db.scoped_session(self.session_maker) as session:
-            query = select(self.Model).offset(skip).options(*self.get_load_options())
+            query = select(self.Model).offset(skip)
+
+            # Only apply load options if requested
+            if use_load_options:
+                query = query.options(*self.get_load_options())
+
+            # Add project filter if applicable
+            query = self._add_project_filter(query)
+
             if limit:
                 query = query.limit(limit)
 
@@ -128,17 +199,15 @@ class Repository[T: Base]:
 
     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}")
-
         # add in load options
         query = query.options(*self.get_load_options())
         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)}")
+            logger.trace(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}")
         else:
-            logger.debug(f"No {self.Model.__name__} found")
+            logger.trace(f"No {self.Model.__name__} found")
         return entity
 
     async def create(self, data: dict) -> T:
@@ -147,12 +216,29 @@ class Repository[T: Base]:
         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)
+
+            # Add project_id if applicable and not already provided
+            if (
+                self.has_project_id
+                and self.project_id is not None
+                and "project_id" not in model_data
+            ):
+                model_data["project_id"] = self.project_id
+
             model = self.Model(**model_data)
             session.add(model)
             await session.flush()
 
             return_instance = await self.select_by_id(session, model.id)  # pyright: ignore [reportAttributeAccessIssue]
-            assert return_instance is not None, "can't find model after session.add"
+            if return_instance is None:  # pragma: no cover
+                logger.error(
+                    "Failed to retrieve model after create",
+                    model_type=self.Model.__name__,
+                    model_id=model.id,  # pyright: ignore
+                )
+                raise ValueError(
+                    f"Can't find {self.Model.__name__} with ID {model.id} after session.add"  # pyright: ignore
+                )
             return return_instance
 
     async def create_all(self, data_list: List[dict]) -> Sequence[T]:
@@ -161,12 +247,20 @@ class Repository[T: Base]:
 
         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
-            ]
+            model_list = []
+            for d in data_list:
+                model_data = self.get_model_data(d)
+
+                # Add project_id if applicable and not already provided
+                if (
+                    self.has_project_id
+                    and self.project_id is not None
+                    and "project_id" not in model_data
+                ):
+                    model_data["project_id"] = self.project_id  # pragma: no cover
+
+                model_list.append(self.Model(**model_data))
+
             session.add_all(model_list)
             await session.flush()
 
@@ -222,7 +316,13 @@ class Repository[T: Base]:
         """Delete records matching given IDs."""
         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))
+            conditions = [self.primary_key.in_(ids)]
+
+            # Add project_id filter if applicable
+            if self.has_project_id and self.project_id is not None:  # pragma: no cover
+                conditions.append(getattr(self.Model, "project_id") == self.project_id)
+
+            query = delete(self.Model).where(and_(*conditions))
             result = await session.execute(query)
             logger.debug(f"Deleted {result.rowcount} records")
             return result.rowcount
@@ -232,6 +332,11 @@ class Repository[T: Base]:
         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()]
+
+            # Add project_id filter if applicable
+            if self.has_project_id and self.project_id is not None:
+                conditions.append(getattr(self.Model, "project_id") == self.project_id)
+
             query = delete(self.Model).where(and_(*conditions))
             result = await session.execute(query)
             deleted = result.rowcount > 0
@@ -243,21 +348,34 @@ class Repository[T: Base]:
         async with db.scoped_session(self.session_maker) as session:
             if query is None:
                 query = select(func.count()).select_from(self.Model)
+                # Add project filter if applicable
+                if (
+                    isinstance(query, Select)
+                    and self.has_project_id
+                    and self.project_id is not None
+                ):
+                    query = query.where(
+                        getattr(self.Model, "project_id") == self.project_id
+                    )  # pragma: no cover
+
             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, use_query_options: bool = True) -> Result[Any]:
+    async def execute_query(
+        self,
+        query: Executable,
+        params: Optional[Dict[str, Any]] = None,
+        use_query_options: bool = True,
+    ) -> Result[Any]:
         """Execute a query asynchronously."""
 
         query = query.options(*self.get_load_options()) if use_query_options else query
-
-        logger.debug(f"Executing query: {query}")
+        logger.trace(f"Executing query: {query}, params: {params}")
         async with db.scoped_session(self.session_maker) as session:
-            result = await session.execute(query)
-            logger.debug("Query executed successfully")
+            result = await session.execute(query, params)
             return result
 
     def get_load_options(self) -> List[LoaderOption]: