feat: Add project_id to Relation and Observation for efficient project-scoped queries

phernandez · claude · phernandez · commit a920a9ff2938 · 2025-12-01T21:54:56.000-06:00
Denormalizes project_id onto Relation and Observation tables to enable efficient project-scoped queries without joins. Migration backfills from associated entity and adds pg_trgm extension with GIN indexes for fuzzy link resolution on PostgreSQL. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: phernandez <paul@basicmachines.co>
diff --git a/src/basic_memory/alembic/versions/f8a9b2c3d4e5_add_pg_trgm_for_fuzzy_link_resolution.py b/src/basic_memory/alembic/versions/f8a9b2c3d4e5_add_pg_trgm_for_fuzzy_link_resolution.py
@@ -0,0 +1,166 @@
+"""Add project_id to relation/observation and pg_trgm for fuzzy link resolution
+
+Revision ID: f8a9b2c3d4e5
+Revises: 314f1ea54dc4
+Create Date: 2025-12-01 12:00:00.000000
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision: str = "f8a9b2c3d4e5"
+down_revision: Union[str, None] = "314f1ea54dc4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Add project_id to relation and observation tables, plus pg_trgm indexes.
+
+    This migration:
+    1. Adds project_id column to relation and observation tables (denormalization)
+    2. Backfills project_id from the associated entity
+    3. Enables pg_trgm extension for trigram-based fuzzy matching (Postgres only)
+    4. Creates GIN indexes on entity title and permalink for fast similarity searches
+    5. Creates partial index on unresolved relations for efficient bulk resolution
+    """
+    connection = op.get_bind()
+    dialect = connection.dialect.name
+
+    # -------------------------------------------------------------------------
+    # Add project_id to relation table
+    # -------------------------------------------------------------------------
+
+    # Step 1: Add project_id column as nullable first
+    op.add_column("relation", sa.Column("project_id", sa.Integer(), nullable=True))
+
+    # Step 2: Backfill project_id from entity.project_id via from_id
+    if dialect == "postgresql":
+        op.execute("""
+            UPDATE relation
+            SET project_id = entity.project_id
+            FROM entity
+            WHERE relation.from_id = entity.id
+        """)
+    else:
+        # SQLite syntax
+        op.execute("""
+            UPDATE relation
+            SET project_id = (
+                SELECT entity.project_id
+                FROM entity
+                WHERE entity.id = relation.from_id
+            )
+        """)
+
+    # Step 3: Make project_id NOT NULL and add foreign key
+    op.alter_column("relation", "project_id", nullable=False)
+    op.create_foreign_key(
+        "fk_relation_project_id",
+        "relation",
+        "project",
+        ["project_id"],
+        ["id"],
+    )
+
+    # Step 4: Create index on relation.project_id
+    op.create_index("ix_relation_project_id", "relation", ["project_id"])
+
+    # -------------------------------------------------------------------------
+    # Add project_id to observation table
+    # -------------------------------------------------------------------------
+
+    # Step 1: Add project_id column as nullable first
+    op.add_column("observation", sa.Column("project_id", sa.Integer(), nullable=True))
+
+    # Step 2: Backfill project_id from entity.project_id via entity_id
+    if dialect == "postgresql":
+        op.execute("""
+            UPDATE observation
+            SET project_id = entity.project_id
+            FROM entity
+            WHERE observation.entity_id = entity.id
+        """)
+    else:
+        # SQLite syntax
+        op.execute("""
+            UPDATE observation
+            SET project_id = (
+                SELECT entity.project_id
+                FROM entity
+                WHERE entity.id = observation.entity_id
+            )
+        """)
+
+    # Step 3: Make project_id NOT NULL and add foreign key
+    op.alter_column("observation", "project_id", nullable=False)
+    op.create_foreign_key(
+        "fk_observation_project_id",
+        "observation",
+        "project",
+        ["project_id"],
+        ["id"],
+    )
+
+    # Step 4: Create index on observation.project_id
+    op.create_index("ix_observation_project_id", "observation", ["project_id"])
+
+    # Postgres-specific: pg_trgm and GIN indexes
+    if dialect == "postgresql":
+        # Enable pg_trgm extension for fuzzy string matching
+        op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
+
+        # Create trigram indexes on entity table for fuzzy matching
+        # GIN indexes with gin_trgm_ops support similarity searches
+        op.execute("""
+            CREATE INDEX IF NOT EXISTS idx_entity_title_trgm
+            ON entity USING gin (title gin_trgm_ops)
+        """)
+
+        op.execute("""
+            CREATE INDEX IF NOT EXISTS idx_entity_permalink_trgm
+            ON entity USING gin (permalink gin_trgm_ops)
+        """)
+
+        # Create partial index on unresolved relations for efficient bulk resolution
+        # This makes "WHERE to_id IS NULL AND project_id = X" queries very fast
+        op.execute("""
+            CREATE INDEX IF NOT EXISTS idx_relation_unresolved
+            ON relation (project_id, to_name)
+            WHERE to_id IS NULL
+        """)
+
+        # Create index on relation.to_name for join performance in bulk resolution
+        op.execute("""
+            CREATE INDEX IF NOT EXISTS idx_relation_to_name
+            ON relation (to_name)
+        """)
+
+
+def downgrade() -> None:
+    """Remove project_id from relation/observation and pg_trgm indexes."""
+    connection = op.get_bind()
+    dialect = connection.dialect.name
+
+    if dialect == "postgresql":
+        # Drop Postgres-specific indexes
+        op.execute("DROP INDEX IF EXISTS idx_relation_to_name")
+        op.execute("DROP INDEX IF EXISTS idx_relation_unresolved")
+        op.execute("DROP INDEX IF EXISTS idx_entity_permalink_trgm")
+        op.execute("DROP INDEX IF EXISTS idx_entity_title_trgm")
+        # Note: We don't drop the pg_trgm extension as other code may depend on it
+
+    # Drop project_id from observation
+    op.drop_index("ix_observation_project_id", table_name="observation")
+    op.drop_constraint("fk_observation_project_id", "observation", type_="foreignkey")
+    op.drop_column("observation", "project_id")
+
+    # Drop project_id from relation
+    op.drop_index("ix_relation_project_id", table_name="relation")
+    op.drop_constraint("fk_relation_project_id", "relation", type_="foreignkey")
+    op.drop_column("relation", "project_id")
diff --git a/src/basic_memory/markdown/utils.py b/src/basic_memory/markdown/utils.py
@@ -14,7 +14,10 @@
 
 @logfire.instrument()
 def entity_model_from_markdown(
-    file_path: Path, markdown: EntityMarkdown, entity: Optional[Entity] = None
+    file_path: Path,
+    markdown: EntityMarkdown,
+    entity: Optional[Entity] = None,
+    project_id: Optional[int] = None,
 ) -> Entity:
     """
     Convert markdown entity to model. Does not include relations.
@@ -23,6 +26,7 @@ def entity_model_from_markdown(
         file_path: Path to the markdown file
         markdown: Parsed markdown entity
         entity: Optional existing entity to update
+        project_id: Project ID for new observations (uses entity.project_id if not provided)
 
     Returns:
         Entity model populated from markdown
@@ -52,9 +56,13 @@ def entity_model_from_markdown(
     metadata = markdown.frontmatter.metadata or {}
     model.entity_metadata = {k: str(v) for k, v in metadata.items() if v is not None}
 
+    # Get project_id from entity if not provided
+    obs_project_id = project_id or (model.project_id if hasattr(model, "project_id") else None)
+
     # Convert observations
     model.observations = [
         ObservationModel(
+            project_id=obs_project_id,
             content=obs.content,
             category=obs.category,
             context=obs.context,
diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py
@@ -145,6 +145,7 @@ class Observation(Base):
     )
 
     id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    project_id: Mapped[int] = mapped_column(Integer, ForeignKey("project.id"), index=True)
     entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id", ondelete="CASCADE"))
     content: Mapped[str] = mapped_column(Text)
     category: Mapped[str] = mapped_column(String, nullable=False, default="note")
@@ -191,6 +192,7 @@ class Relation(Base):
     )
 
     id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    project_id: Mapped[int] = mapped_column(Integer, ForeignKey("project.id"), index=True)
     from_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id", ondelete="CASCADE"))
     to_id: Mapped[Optional[int]] = mapped_column(
         Integer, ForeignKey("entity.id", ondelete="CASCADE"), nullable=True
diff --git a/src/basic_memory/repository/relation_repository.py b/src/basic_memory/repository/relation_repository.py
@@ -117,6 +117,7 @@ async def add_all_ignore_duplicates(self, relations: List[Relation]) -> int:
         # Convert Relation objects to dicts for insert
         values = [
             {
+                "project_id": r.project_id if r.project_id else self.project_id,
                 "from_id": r.from_id,
                 "to_id": r.to_id,
                 "to_name": r.to_name,
diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py
@@ -398,7 +398,9 @@ async def create_entity_from_markdown(
         Uses UPSERT approach to handle permalink/file_path conflicts cleanly.
         """
         logger.debug(f"Creating entity: {markdown.frontmatter.title} file_path: {file_path}")
-        model = entity_model_from_markdown(file_path, markdown)
+        model = entity_model_from_markdown(
+            file_path, markdown, project_id=self.repository.project_id
+        )
 
         # Mark as incomplete because we still need to add relations
         model.checksum = None
@@ -429,6 +431,7 @@ async def update_entity_and_observations(
         # add new observations
         observations = [
             Observation(
+                project_id=self.observation_repository.project_id,
                 entity_id=db_entity.id,
                 content=obs.content,
                 category=obs.category,
@@ -496,6 +499,7 @@ async def update_entity_relations(
 
                 # Create the relation
                 relation = Relation(
+                    project_id=self.relation_repository.project_id,
                     from_id=db_entity.id,
                     to_id=target_id,
                     to_name=target_name,
diff --git a/test-int/test_db_wal_mode.py b/test-int/test_db_wal_mode.py
@@ -142,21 +142,6 @@ async def test_null_pool_on_windows(tmp_path, monkeypatch):
         assert isinstance(engine.pool, NullPool)
 
 
-@pytest.mark.asyncio
-@pytest.mark.skipif(
-    __import__("os").name == "nt", reason="Non-Windows test - cannot mock POSIX paths on Windows"
-)
-async def test_regular_pool_on_non_windows(tmp_path):
-    """Test that regular pooling is used on non-Windows platforms."""
-    from basic_memory.db import engine_session_factory, DatabaseType
-    from sqlalchemy.pool import NullPool
-
-    db_path = tmp_path / "test_posix_pool.db"
-
-    with patch("basic_memory.db.os.name", "posix"):
-        async with engine_session_factory(db_path, DatabaseType.FILESYSTEM) as (engine, _):
-            # Engine should NOT be using NullPool on non-Windows
-            assert not isinstance(engine.pool, NullPool)
 
 
 @pytest.mark.asyncio
diff --git a/tests/markdown/test_entity_parser.py b/tests/markdown/test_entity_parser.py
@@ -85,11 +85,11 @@ async def test_parse_complete_file(project_config, entity_parser, valid_entity_c
     ), "missing [[Auth API Spec]]"
 
     # inline links in content
-    assert Relation(type="links to", target="Random Link", context=None) in entity.relations, (
+    assert Relation(type="links_to", target="Random Link", context=None) in entity.relations, (
         "missing [[Random Link]]"
     )
     assert (
-        Relation(type="links to", target="Random Link with Title|Titled Link", context=None)
+        Relation(type="links_to", target="Random Link with Title|Titled Link", context=None)
         in entity.relations
     ), "missing [[Random Link with Title|Titled Link]]"
 
diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py
@@ -18,10 +18,12 @@ async def entity_with_observations(session_maker, sample_entity):
     async with db.scoped_session(session_maker) as session:
         observations = [
             Observation(
+                project_id=sample_entity.project_id,
                 entity_id=sample_entity.id,
                 content="First observation",
             ),
             Observation(
+                project_id=sample_entity.project_id,
                 entity_id=sample_entity.id,
                 content="Second observation",
             ),
@@ -59,6 +61,7 @@ async def related_results(session_maker, test_project: Project):
         await session.flush()
 
         relation = Relation(
+            project_id=test_project.id,
             from_id=source.id,
             to_id=target.id,
             to_name=target.title,
@@ -199,6 +202,7 @@ async def test_update_entity_returns_with_relations_and_observations(
         await session.flush()
 
         relation = Relation(
+            project_id=test_project.id,
             from_id=entity.id,
             to_id=target.id,
             to_name=target.title,
@@ -785,6 +789,7 @@ async def test_get_all_file_paths_performance(entity_repository: EntityRepositor
 
         # Add observations to entity1
         observation = Observation(
+            project_id=entity_repository.project_id,
             entity_id=entity1.id,
             content="Test observation",
             category="note",
@@ -793,6 +798,7 @@ async def test_get_all_file_paths_performance(entity_repository: EntityRepositor
 
         # Add relation between entities
         relation = Relation(
+            project_id=entity_repository.project_id,
             from_id=entity1.id,
             to_id=entity2.id,
             to_name=entity2.title,
diff --git a/tests/repository/test_entity_upsert_issue_187.py b/tests/repository/test_entity_upsert_issue_187.py
@@ -28,6 +28,7 @@ async def test_upsert_entity_with_observations_conflict(entity_repository: Entit
 
     # Add observations to the entity
     obs1 = Observation(
+        project_id=entity_repository.project_id,
         content="This is a test observation",
         category="testing",
         tags=["test"],
@@ -56,11 +57,13 @@ async def test_upsert_entity_with_observations_conflict(entity_repository: Entit
 
     # Add different observations
     obs2 = Observation(
+        project_id=entity_repository.project_id,
         content="This is an updated observation",
         category="updated",
         tags=["updated"],
     )
     obs3 = Observation(
+        project_id=entity_repository.project_id,
         content="This is a second observation",
         category="second",
         tags=["second"],
diff --git a/tests/repository/test_observation_repository.py b/tests/repository/test_observation_repository.py
diff --git a/tests/repository/test_relation_repository.py b/tests/repository/test_relation_repository.py
diff --git a/tests/services/test_context_service.py b/tests/services/test_context_service.py

Original file line number	Diff line number	Diff line change
`@@ -145,6 +145,7 @@ class Observation(Base):`
`145`	`145`	`)`
`146`	`146`
`147`	`147`	`id: Mapped[int] = mapped_column(Integer, primary_key=True)`
	`148`	`+ project_id: Mapped[int] = mapped_column(Integer, ForeignKey("project.id"), index=True)`
`148`	`149`	`entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id", ondelete="CASCADE"))`
`149`	`150`	`content: Mapped[str] = mapped_column(Text)`
`150`	`151`	`category: Mapped[str] = mapped_column(String, nullable=False, default="note")`
`@@ -191,6 +192,7 @@ class Relation(Base):`
`191`	`192`	`)`
`192`	`193`
`193`	`194`	`id: Mapped[int] = mapped_column(Integer, primary_key=True)`
	`195`	`+ project_id: Mapped[int] = mapped_column(Integer, ForeignKey("project.id"), index=True)`
`194`	`196`	`from_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id", ondelete="CASCADE"))`
`195`	`197`	`to_id: Mapped[Optional[int]] = mapped_column(`
`196`	`198`	`Integer, ForeignKey("entity.id", ondelete="CASCADE"), nullable=True`
Original file line number	Diff line number	Diff line change
`@@ -117,6 +117,7 @@ async def add_all_ignore_duplicates(self, relations: List[Relation]) -> int:`
`117`	`117`	`# Convert Relation objects to dicts for insert`
`118`	`118`	`values = [`
`119`	`119`	`{`
	`120`	`+ "project_id": r.project_id if r.project_id else self.project_id,`
`120`	`121`	`"from_id": r.from_id,`
`121`	`122`	`"to_id": r.to_id,`
`122`	`123`	`"to_name": r.to_name,`