kodit 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

kodit/infrastructure/sqlalchemy/entities.py

@@ -10,6 +10,7 @@ from sqlalchemy import (
     Float,
     ForeignKey,
     ForeignKeyConstraint,
+    Index,
     Integer,
     String,
     TypeDecorator,
@@ -483,44 +484,6 @@ class CommitSnippetV2(Base):
         self.snippet_sha = snippet_sha
 
 
-# Enrichment model for SnippetV2
-
-
-class EnrichmentType(Enum):
-    """Enrichment type enum."""
-
-    UNKNOWN = "unknown"
-    SUMMARIZATION = "summarization"
-
-
-class Enrichment(Base, CommonMixin):
-    """Enrichment model for snippet enrichments."""
-
-    __tablename__ = "enrichments"
-
-    snippet_sha: Mapped[str] = mapped_column(ForeignKey("snippets_v2.sha"), index=True)
-    type: Mapped[EnrichmentType] = mapped_column(
-        SQLAlchemyEnum(EnrichmentType), index=True
-    )
-    content: Mapped[str] = mapped_column(UnicodeText)
-
-    __table_args__ = (
-        UniqueConstraint("snippet_sha", "type", name="uix_snippet_enrichment"),
-    )
-
-    def __init__(
-        self,
-        snippet_sha: str,
-        type: EnrichmentType,  # noqa: A002
-        content: str,
-    ) -> None:
-        """Initialize enrichment."""
-        super().__init__()
-        self.snippet_sha = snippet_sha
-        self.type = type
-        self.content = content
-
-
 class CommitIndex(Base):
     """Commit index model."""
 
@@ -559,3 +522,48 @@ class CommitIndex(Base):
         self.error_message = error_message
         self.files_processed = files_processed
         self.processing_time_seconds = processing_time_seconds
+
+
+class EnrichmentV2(Base, CommonMixin):
+    """Generic enrichment entity."""
+
+    __tablename__ = "enrichments_v2"
+
+    type: Mapped[str] = mapped_column(String, nullable=False, index=True)
+    subtype: Mapped[str] = mapped_column(String, nullable=False, index=True)
+    content: Mapped[str] = mapped_column(UnicodeText, nullable=False)
+
+    __table_args__ = (Index("idx_type_subtype", "type", "subtype"),)
+
+
+class EnrichmentAssociation(Base, CommonMixin):
+    """Polymorphic association between enrichments and entities."""
+
+    __tablename__ = "enrichment_associations"
+
+    enrichment_id: Mapped[int] = mapped_column(
+        ForeignKey("enrichments_v2.id", ondelete="CASCADE"),
+        nullable=False,
+        index=True,
+    )
+    entity_type: Mapped[str] = mapped_column(
+        String(50),
+        nullable=False,
+        index=True,
+    )
+    entity_id: Mapped[str] = mapped_column(
+        String(255),
+        nullable=False,
+        index=True,
+    )
+
+    __table_args__ = (
+        UniqueConstraint(
+            "entity_type",
+            "entity_id",
+            "enrichment_id",
+            name="uix_entity_enrichment",
+        ),
+        Index("idx_entity_lookup", "entity_type", "entity_id"),
+        {"sqlite_autoincrement": True},
+    )

kodit/infrastructure/sqlalchemy/git_branch_repository.py

@@ -98,17 +98,28 @@ class SqlAlchemyGitBranchRepository(GitBranchRepository):
 
             commit_shas = [branch.head_commit_sha for branch in db_branches]
 
-            # Get all head commits for these branches
-            commits_stmt = select(db_entities.GitCommit).where(
-                db_entities.GitCommit.commit_sha.in_(commit_shas)
-            )
-            db_commits = (await session.scalars(commits_stmt)).all()
-
-            # Get all files for these commits
-            files_stmt = select(db_entities.GitCommitFile).where(
-                db_entities.GitCommitFile.commit_sha.in_(commit_shas)
-            )
-            db_files = (await session.scalars(files_stmt)).all()
+            # Get all head commits for these branches in chunks
+            # to avoid parameter limits
+            db_commits: list[db_entities.GitCommit] = []
+            chunk_size = 1000
+            for i in range(0, len(commit_shas), chunk_size):
+                chunk = commit_shas[i : i + chunk_size]
+                commits_stmt = select(db_entities.GitCommit).where(
+                    db_entities.GitCommit.commit_sha.in_(chunk)
+                )
+                chunk_commits = (await session.scalars(commits_stmt)).all()
+                db_commits.extend(chunk_commits)
+
+            # Get all files for these commits in chunks
+            # to avoid parameter limits
+            db_files: list[db_entities.GitCommitFile] = []
+            for i in range(0, len(commit_shas), chunk_size):
+                chunk = commit_shas[i : i + chunk_size]
+                files_stmt = select(db_entities.GitCommitFile).where(
+                    db_entities.GitCommitFile.commit_sha.in_(chunk)
+                )
+                chunk_files = (await session.scalars(files_stmt)).all()
+                db_files.extend(chunk_files)
 
             # Group files by commit SHA
             from kodit.domain.entities.git import GitFile

kodit/infrastructure/sqlalchemy/git_commit_repository.py

@@ -77,11 +77,17 @@ class SqlAlchemyGitCommitRepository(GitCommitRepository):
 
             commit_shas = [commit.commit_sha for commit in db_commits]
 
-            # Get all files for these commits
-            files_stmt = select(db_entities.GitCommitFile).where(
-                db_entities.GitCommitFile.commit_sha.in_(commit_shas)
-            )
-            db_files = (await session.scalars(files_stmt)).all()
+            # Get all files for these commits in chunks
+            # to avoid parameter limits
+            db_files: list[db_entities.GitCommitFile] = []
+            chunk_size = 1000
+            for i in range(0, len(commit_shas), chunk_size):
+                chunk = commit_shas[i : i + chunk_size]
+                files_stmt = select(db_entities.GitCommitFile).where(
+                    db_entities.GitCommitFile.commit_sha.in_(chunk)
+                )
+                chunk_files = (await session.scalars(files_stmt)).all()
+                db_files.extend(chunk_files)
 
             # Group files by commit SHA
             files_by_commit: dict[str, list[GitFile]] = {}
@@ -149,13 +155,16 @@ class SqlAlchemyGitCommitRepository(GitCommitRepository):
         async with SqlAlchemyUnitOfWork(self.session_factory) as session:
             commit_shas = [commit.commit_sha for commit in commits]
 
-            # Get existing commits in bulk
-            existing_commits_stmt = select(db_entities.GitCommit.commit_sha).where(
-                db_entities.GitCommit.commit_sha.in_(commit_shas)
-            )
-            existing_commit_shas = set(
-                (await session.scalars(existing_commits_stmt)).all()
-            )
+            # Get existing commits in bulk (chunked to avoid parameter limits)
+            existing_commit_shas: set[str] = set()
+            chunk_size = 1000
+            for i in range(0, len(commit_shas), chunk_size):
+                chunk = commit_shas[i : i + chunk_size]
+                existing_commits_stmt = select(db_entities.GitCommit.commit_sha).where(
+                    db_entities.GitCommit.commit_sha.in_(chunk)
+                )
+                chunk_existing = (await session.scalars(existing_commits_stmt)).all()
+                existing_commit_shas.update(chunk_existing)
 
             # Prepare new commits for bulk insert
             new_commits_data = []
@@ -176,8 +185,8 @@ class SqlAlchemyGitCommitRepository(GitCommitRepository):
             if new_commits_data:
                 chunk_size = 1000  # Conservative chunk size for parameter limits
                 for i in range(0, len(new_commits_data), chunk_size):
-                    chunk = new_commits_data[i : i + chunk_size]
-                    stmt = insert(db_entities.GitCommit).values(chunk)
+                    data_chunk = new_commits_data[i : i + chunk_size]
+                    stmt = insert(db_entities.GitCommit).values(data_chunk)
                     await session.execute(stmt)
 
                 # Bulk save files for new commits

kodit/infrastructure/sqlalchemy/git_repository.py

@@ -219,28 +219,38 @@ class SqlAlchemyGitRepoRepository(GitRepoRepository):
         if db_tracking_branch_entity:
             referenced_commit_shas.add(db_tracking_branch_entity.head_commit_sha)
 
-        # Load only the referenced commits
+        # Load only the referenced commits in chunks to avoid parameter limits
         referenced_commits = []
         referenced_files = []
         if referenced_commit_shas:
-            referenced_commits = list(
-                (
-                    await session.scalars(
-                        select(db_entities.GitCommit).where(
-                            db_entities.GitCommit.commit_sha.in_(referenced_commit_shas)
+            commit_shas_list = list(referenced_commit_shas)
+            chunk_size = 1000
+
+            for i in range(0, len(commit_shas_list), chunk_size):
+                chunk = commit_shas_list[i : i + chunk_size]
+                chunk_commits = list(
+                    (
+                        await session.scalars(
+                            select(db_entities.GitCommit).where(
+                                db_entities.GitCommit.commit_sha.in_(chunk)
+                            )
                         )
-                    )
-                ).all()
-            )
-            referenced_files = list(
-                (
-                    await session.scalars(
-                        select(db_entities.GitCommitFile).where(
-                            db_entities.GitCommitFile.commit_sha.in_(referenced_commit_shas)
+                    ).all()
+                )
+                referenced_commits.extend(chunk_commits)
+
+            for i in range(0, len(commit_shas_list), chunk_size):
+                chunk = commit_shas_list[i : i + chunk_size]
+                chunk_files = list(
+                    (
+                        await session.scalars(
+                            select(db_entities.GitCommitFile).where(
+                                db_entities.GitCommitFile.commit_sha.in_(chunk)
+                            )
                         )
-                    )
-                ).all()
-            )
+                    ).all()
+                )
+                referenced_files.extend(chunk_files)
 
         return self._mapper.to_domain_git_repo(
             db_repo=db_repo,

kodit/infrastructure/sqlalchemy/git_tag_repository.py

@@ -96,17 +96,28 @@ class SqlAlchemyGitTagRepository(GitTagRepository):
 
             commit_shas = [tag.target_commit_sha for tag in db_tags]
 
-            # Get all target commits for these tags
-            commits_stmt = select(db_entities.GitCommit).where(
-                db_entities.GitCommit.commit_sha.in_(commit_shas)
-            )
-            db_commits = (await session.scalars(commits_stmt)).all()
-
-            # Get all files for these commits
-            files_stmt = select(db_entities.GitCommitFile).where(
-                db_entities.GitCommitFile.commit_sha.in_(commit_shas)
-            )
-            db_files = (await session.scalars(files_stmt)).all()
+            # Get all target commits for these tags in chunks
+            # to avoid parameter limits
+            db_commits: list[db_entities.GitCommit] = []
+            chunk_size = 1000
+            for i in range(0, len(commit_shas), chunk_size):
+                chunk = commit_shas[i : i + chunk_size]
+                commits_stmt = select(db_entities.GitCommit).where(
+                    db_entities.GitCommit.commit_sha.in_(chunk)
+                )
+                chunk_commits = (await session.scalars(commits_stmt)).all()
+                db_commits.extend(chunk_commits)
+
+            # Get all files for these commits in chunks
+            # to avoid parameter limits
+            db_files: list[db_entities.GitCommitFile] = []
+            for i in range(0, len(commit_shas), chunk_size):
+                chunk = commit_shas[i : i + chunk_size]
+                files_stmt = select(db_entities.GitCommitFile).where(
+                    db_entities.GitCommitFile.commit_sha.in_(chunk)
+                )
+                chunk_files = (await session.scalars(files_stmt)).all()
+                db_files.extend(chunk_files)
 
             # Group files by commit SHA
             files_by_commit: dict[str, list[GitFile]] = {}

kodit/infrastructure/sqlalchemy/snippet_v2_repository.py

@@ -2,18 +2,36 @@
 
 import zlib
 from collections.abc import Callable
+from datetime import datetime
+from typing import TypedDict
 
 from sqlalchemy import delete, insert, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from kodit.domain.enrichments.development.snippet.snippet import SnippetEnrichment
 from kodit.domain.entities.git import SnippetV2
 from kodit.domain.protocols import SnippetRepositoryV2
 from kodit.domain.value_objects import MultiSearchRequest
 from kodit.infrastructure.mappers.snippet_mapper import SnippetMapper
 from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.enrichment_v2_repository import (
+    EnrichmentV2Repository,
+)
 from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
 
+class _GitFileData(TypedDict):
+    """Type for GitCommitFile creation data."""
+
+    commit_sha: str
+    path: str
+    blob_sha: str
+    mime_type: str
+    size: int
+    extension: str
+    created_at: datetime
+
+
 def create_snippet_v2_repository(
     session_factory: Callable[[], AsyncSession],
 ) -> SnippetRepositoryV2:
@@ -27,6 +45,7 @@ class SqlAlchemySnippetRepositoryV2(SnippetRepositoryV2):
     def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
         """Initialize the repository."""
         self.session_factory = session_factory
+        self._enrichment_repo = EnrichmentV2Repository(session_factory)
 
     @property
     def _mapper(self) -> SnippetMapper:
@@ -112,10 +131,15 @@ class SqlAlchemySnippetRepositoryV2(SnippetRepositoryV2):
                 stmt = insert(db_entities.CommitSnippetV2).values(chunk)
                 await session.execute(stmt)
 
-    async def _bulk_create_file_associations(
+    async def _bulk_create_file_associations(  # noqa: C901
         self, session: AsyncSession, commit_sha: str, snippets: list[SnippetV2]
     ) -> None:
-        """Bulk create snippet-file associations."""
+        """Bulk create snippet-file associations.
+
+        Creates SnippetV2File records linking snippets to GitCommitFile records.
+        If a GitCommitFile doesn't exist, it creates it automatically to prevent
+        losing file associations during enrichment cycles.
+        """
         # Collect all file paths from all snippets
         file_paths = set()
         for snippet in snippets:
@@ -150,18 +174,55 @@ class SqlAlchemySnippetRepositoryV2(SnippetRepositoryV2):
         existing_snippet_files = set(await session.execute(existing_snippet_files_stmt))
 
         # Prepare new file associations
-        new_file_associations = []
+        new_file_associations: list[dict[str, str]] = []
+        missing_git_files: list[_GitFileData] = []
+
         for snippet in snippets:
             for file in snippet.derives_from:
                 association_key = (snippet.sha, file.path)
-                if (association_key not in existing_snippet_files
-                    and file.path in existing_files_map):
-                    new_file_associations.append({
-                        "snippet_sha": snippet.sha,
-                        "blob_sha": existing_files_map[file.path],
-                        "commit_sha": commit_sha,
-                        "file_path": file.path,
-                    })
+                if association_key not in existing_snippet_files:
+                    if file.path in existing_files_map:
+                        # GitCommitFile exists, use its blob_sha
+                        new_file_associations.append({
+                            "snippet_sha": snippet.sha,
+                            "blob_sha": existing_files_map[file.path],
+                            "commit_sha": commit_sha,
+                            "file_path": file.path,
+                        })
+                    else:
+                        # GitCommitFile doesn't exist - create it and the association
+                        missing_git_files.append({
+                            "commit_sha": commit_sha,
+                            "path": file.path,
+                            "blob_sha": file.blob_sha,
+                            "mime_type": file.mime_type,
+                            "size": file.size,
+                            "extension": file.extension,
+                            "created_at": file.created_at,
+                        })
+                        new_file_associations.append({
+                            "snippet_sha": snippet.sha,
+                            "blob_sha": file.blob_sha,
+                            "commit_sha": commit_sha,
+                            "file_path": file.path,
+                        })
+                        # Add to map so subsequent snippets can find it
+                        existing_files_map[file.path] = file.blob_sha
+
+        # Create missing GitCommitFile records
+        if missing_git_files:
+            for git_file_data in missing_git_files:
+                git_file = db_entities.GitCommitFile(
+                    commit_sha=git_file_data["commit_sha"],
+                    path=git_file_data["path"],
+                    blob_sha=git_file_data["blob_sha"],
+                    mime_type=git_file_data["mime_type"],
+                    size=git_file_data["size"],
+                    extension=git_file_data["extension"],
+                    created_at=git_file_data["created_at"],
+                )
+                session.add(git_file)
+            await session.flush()
 
         # Bulk insert new file associations in chunks to avoid parameter limits
         if new_file_associations:
@@ -172,70 +233,29 @@ class SqlAlchemySnippetRepositoryV2(SnippetRepositoryV2):
                 await session.execute(stmt)
 
     async def _bulk_update_enrichments(
-        self, session: AsyncSession, snippets: list[SnippetV2]
+        self, session: AsyncSession, snippets: list[SnippetV2]  # noqa: ARG002
     ) -> None:
-        """Bulk update enrichments for snippets."""
-        snippet_shas = [snippet.sha for snippet in snippets]
-
-        # Get all existing enrichments for these snippets
-        existing_enrichments_stmt = select(
-            db_entities.Enrichment.snippet_sha,
-            db_entities.Enrichment.type,
-            db_entities.Enrichment.content
-        ).where(
-            db_entities.Enrichment.snippet_sha.in_(snippet_shas)
-        )
-        existing_enrichments = await session.execute(existing_enrichments_stmt)
-
-        # Create lookup for existing enrichment hashes
-        existing_enrichment_map = {}
-        for snippet_sha, enrichment_type, content in existing_enrichments:
-            content_hash = self._hash_string(content)
-            key = (snippet_sha, enrichment_type)
-            existing_enrichment_map[key] = content_hash
-
-        # Collect enrichments to delete and add
-        enrichments_to_delete = []
-        enrichments_to_add = []
+        """Bulk update enrichments for snippets using new enrichment_v2."""
+        # Collect all enrichments from snippets using list comprehension
+        snippet_enrichments = [
+            SnippetEnrichment(
+                entity_id=snippet.sha,
+                content=enrichment.content,
+            )
+            for snippet in snippets
+            for enrichment in snippet.enrichments
+        ]
 
-        for snippet in snippets:
-            for enrichment in snippet.enrichments:
-                key = (snippet.sha, db_entities.EnrichmentType(enrichment.type.value))
-                new_hash = self._hash_string(enrichment.content)
-
-                if key in existing_enrichment_map:
-                    if existing_enrichment_map[key] != new_hash:
-                        # Content changed, mark for deletion and re-addition
-                        enrichments_to_delete.append(key)
-                        enrichments_to_add.append({
-                            "snippet_sha": snippet.sha,
-                            "type": db_entities.EnrichmentType(enrichment.type.value),
-                            "content": enrichment.content,
-                        })
-                else:
-                    # New enrichment
-                    enrichments_to_add.append({
-                        "snippet_sha": snippet.sha,
-                        "type": db_entities.EnrichmentType(enrichment.type.value),
-                        "content": enrichment.content,
-                    })
-
-        # Bulk delete changed enrichments
-        if enrichments_to_delete:
-            for snippet_sha, enrichment_type in enrichments_to_delete:
-                stmt = delete(db_entities.Enrichment).where(
-                    db_entities.Enrichment.snippet_sha == snippet_sha,
-                    db_entities.Enrichment.type == enrichment_type,
-                )
-                await session.execute(stmt)
+        if snippet_enrichments:
+            # First delete existing enrichments for these snippets
+            snippet_shas = [snippet.sha for snippet in snippets]
+            await self._enrichment_repo.bulk_delete_enrichments(
+                entity_type="snippet_v2",
+                entity_ids=snippet_shas,
+            )
 
-        # Bulk insert new/updated enrichments in chunks to avoid parameter limits
-        if enrichments_to_add:
-            chunk_size = 1000  # Conservative chunk size for parameter limits
-            for i in range(0, len(enrichments_to_add), chunk_size):
-                chunk = enrichments_to_add[i : i + chunk_size]
-                insert_stmt = insert(db_entities.Enrichment).values(chunk)
-                await session.execute(insert_stmt)
+            # Then save the new enrichments
+            await self._enrichment_repo.bulk_save_enrichments(snippet_enrichments)
 
     async def _get_or_create_raw_snippet(
         self, session: AsyncSession, commit_sha: str, domain_snippet: SnippetV2
@@ -281,33 +301,8 @@ class SqlAlchemySnippetRepositoryV2(SnippetRepositoryV2):
         domain_snippet: SnippetV2,
     ) -> None:
         """Update enrichments if they have changed."""
-        current_enrichments = await session.scalars(
-            select(db_entities.Enrichment).where(
-                db_entities.Enrichment.snippet_sha == db_snippet.sha
-            )
-        )
-        current_enrichment_shas = {
-            self._hash_string(enrichment.content)
-            for enrichment in list(current_enrichments)
-        }
-        for enrichment in domain_snippet.enrichments:
-            if self._hash_string(enrichment.content) in current_enrichment_shas:
-                continue
-
-            # If not present, delete the existing enrichment for this type if it exists
-            stmt = delete(db_entities.Enrichment).where(
-                db_entities.Enrichment.snippet_sha == db_snippet.sha,
-                db_entities.Enrichment.type
-                == db_entities.EnrichmentType(enrichment.type.value),
-            )
-            await session.execute(stmt)
-
-            db_enrichment = db_entities.Enrichment(
-                snippet_sha=db_snippet.sha,
-                type=db_entities.EnrichmentType(enrichment.type.value),
-                content=enrichment.content,
-            )
-            session.add(db_enrichment)
+        # For now, enrichments are not yet implemented with the new schema
+        # This method will need to be updated once we migrate to EnrichmentV2
 
     async def get_snippets_for_commit(self, commit_sha: str) -> list[SnippetV2]:
         """Get all snippets for a specific commit."""
@@ -469,16 +464,16 @@ class SqlAlchemySnippetRepositoryV2(SnippetRepositoryV2):
             )
             .where(db_entities.SnippetV2File.snippet_sha == db_snippet.sha)
         )
+        db_files_list = list(db_files)
 
-        # Enrichments related to this snippet
-        db_enrichments = await session.scalars(
-            select(db_entities.Enrichment).where(
-                db_entities.Enrichment.snippet_sha == db_snippet.sha
-            )
+        # Get enrichments for this snippet
+        db_enrichments = await self._enrichment_repo.enrichments_for_entity_type(
+            entity_type="snippet_v2",
+            entity_ids=[db_snippet.sha],
         )
 
         return self._mapper.to_domain_snippet_v2(
             db_snippet=db_snippet,
-            db_files=list(db_files),
-            db_enrichments=list(db_enrichments),
+            db_files=db_files_list,
+            db_enrichments=db_enrichments,
         )