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 new file mode 100644 index 00000000..99d9502a --- /dev/null +++ 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 index f34d1f39..b8b0ed9a 100644 --- a/src/basic_memory/markdown/utils.py +++ b/src/basic_memory/markdown/utils.py @@ -14,7 +14,10 @@ from basic_memory.models import Observation as ObservationModel @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 index a7b6c778..e3275102 100644 --- 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 index 93d1d0b9..bc1494f9 100644 --- a/src/basic_memory/repository/relation_repository.py +++ b/src/basic_memory/repository/relation_repository.py @@ -117,6 +117,7 @@ class RelationRepository(Repository[Relation]): # 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 index c60fb860..8f16a1df 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -398,7 +398,9 @@ class EntityService(BaseService[EntityModel]): 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 @@ class EntityService(BaseService[EntityModel]): # 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 @@ class EntityService(BaseService[EntityModel]): # 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 index 393b69da..35e4548c 100644 --- 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 index 62fa939b..4b57c2f5 100644 --- 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 index 5e513c44..aa9e990f 100644 --- 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 index ad030a80..47772d00 100644 --- 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 index 57331870..3bf1b499 100644 --- a/tests/repository/test_observation_repository.py +++ b/tests/repository/test_observation_repository.py @@ -22,6 +22,7 @@ async def repo(observation_repository): async def sample_observation(repo, sample_entity: Entity): """Create a sample observation for testing""" observation_data = { + "project_id": sample_entity.project_id, "entity_id": sample_entity.id, "content": "Test observation", "context": "test-context", @@ -35,6 +36,7 @@ async def test_create_observation( ): """Test creating a new observation""" observation_data = { + "project_id": sample_entity.project_id, "entity_id": sample_entity.id, "content": "Test content", "context": "test-context", @@ -52,6 +54,7 @@ async def test_create_observation_entity_does_not_exist( ): """Test creating a new observation""" observation_data = { + "project_id": sample_entity.project_id, "entity_id": 99999, # Non-existent entity ID (integer for Postgres compatibility) "content": "Test content", "context": "test-context", @@ -104,10 +107,12 @@ async def test_delete_observations(session_maker: async_sessionmaker, repo, test # Create test observations obs1 = Observation( + project_id=test_project.id, entity_id=entity.id, content="Test observation 1", ) obs2 = Observation( + project_id=test_project.id, entity_id=entity.id, content="Test observation 2", ) @@ -144,6 +149,7 @@ async def test_delete_observation_by_id( # Create test observation obs = Observation( + project_id=test_project.id, entity_id=entity.id, content="Test observation", ) @@ -180,10 +186,12 @@ async def test_delete_observation_by_content( # Create test observations obs1 = Observation( + project_id=test_project.id, entity_id=entity.id, content="Delete this observation", ) obs2 = Observation( + project_id=test_project.id, entity_id=entity.id, content="Keep this observation", ) @@ -220,16 +228,19 @@ async def test_find_by_category(session_maker: async_sessionmaker, repo, test_pr # Create test observations with different categories observations = [ Observation( + project_id=test_project.id, entity_id=entity.id, content="Tech observation", category="tech", ), Observation( + project_id=test_project.id, entity_id=entity.id, content="Design observation", category="design", ), Observation( + project_id=test_project.id, entity_id=entity.id, content="Another tech observation", category="tech", @@ -278,21 +289,25 @@ async def test_observation_categories( # Create observations with various categories observations = [ Observation( + project_id=test_project.id, entity_id=entity.id, content="First tech note", category="tech", ), Observation( + project_id=test_project.id, entity_id=entity.id, content="Second tech note", category="tech", # Duplicate category ), Observation( + project_id=test_project.id, entity_id=entity.id, content="Design note", category="design", ), Observation( + project_id=test_project.id, entity_id=entity.id, content="Feature note", category="feature", @@ -341,6 +356,7 @@ async def test_find_by_category_case_sensitivity( # Create a test observation obs = Observation( + project_id=test_project.id, entity_id=entity.id, content="Tech note", category="tech", # lowercase in database @@ -386,6 +402,7 @@ async def test_observation_permalink_truncates_long_content( # Create observation with very long content (5000+ chars to simulate transcript) long_content = "A" * 5000 # Well over the 200 char limit obs = Observation( + project_id=test_project.id, entity_id=entity.id, content=long_content, category="transcript", @@ -436,6 +453,7 @@ async def test_observation_permalink_short_content_unchanged( # Create observation with short content short_content = "Short observation content" obs = Observation( + project_id=test_project.id, entity_id=entity.id, content=short_content, category="note", diff --git a/tests/repository/test_relation_repository.py b/tests/repository/test_relation_repository.py index a73df956..01e08ccc 100644 --- a/tests/repository/test_relation_repository.py +++ b/tests/repository/test_relation_repository.py @@ -50,16 +50,18 @@ async def target_entity(session_maker, test_project: Project): @pytest_asyncio.fixture -async def test_relations(session_maker, source_entity, target_entity): +async def test_relations(session_maker, source_entity, target_entity, test_project: Project): """Create test relations.""" relations = [ Relation( + project_id=test_project.id, from_id=source_entity.id, to_id=target_entity.id, to_name=target_entity.title, relation_type="connects_to", ), Relation( + project_id=test_project.id, from_id=source_entity.id, to_id=target_entity.id, to_name=target_entity.title, diff --git a/tests/services/test_context_service.py b/tests/services/test_context_service.py index 58449d5f..2fda3404 100644 --- a/tests/services/test_context_service.py +++ b/tests/services/test_context_service.py @@ -286,6 +286,7 @@ async def test_project_isolation_in_find_related(session_maker, app_config): # Create relation in project1 (between entities of project1) relation_p1 = Relation( + project_id=project1.id, from_id=entity1_p1.id, to_id=entity2_p1.id, to_name="Entity2_P1",