kodit 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

kodit/infrastructure/sqlalchemy/embedding_repository.py

@@ -1,12 +1,15 @@
 """SQLAlchemy implementation of embedding repository."""
 
 from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.infrastructure.sqlalchemy.entities import Embedding, EmbeddingType
+from kodit.infrastructure.sqlalchemy.query import FilterOperator, QueryBuilder
+from kodit.infrastructure.sqlalchemy.repository import SqlAlchemyRepository
 from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
 
@@ -17,68 +20,75 @@ def create_embedding_repository(
     return SqlAlchemyEmbeddingRepository(session_factory=session_factory)
 
 
-class SqlAlchemyEmbeddingRepository:
+class SqlAlchemyEmbeddingRepository(SqlAlchemyRepository[Embedding, Embedding]):
     """SQLAlchemy implementation of embedding repository."""
 
-    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
-        """Initialize the SQLAlchemy embedding repository."""
-        self.session_factory = session_factory
+    @property
+    def db_entity_type(self) -> type[Embedding]:
+        """The SQLAlchemy model type."""
+        return Embedding
+
+    @staticmethod
+    def to_domain(db_entity: Embedding) -> Embedding:
+        """Map database entity to domain entity."""
+        return db_entity
+
+    @staticmethod
+    def to_db(domain_entity: Embedding) -> Embedding:
+        """Map domain entity to database entity."""
+        return domain_entity
+
+    def _get_id(self, entity: Embedding) -> Any:
+        """Extract ID from domain entity."""
+        return entity.id
 
     async def create_embedding(self, embedding: Embedding) -> None:
         """Create a new embedding record in the database."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            session.add(embedding)
+        await self.save(embedding)
 
     async def get_embedding_by_snippet_id_and_type(
         self, snippet_id: int, embedding_type: EmbeddingType
     ) -> Embedding | None:
         """Get an embedding by its snippet ID and type."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            query = select(Embedding).where(
-                Embedding.snippet_id == snippet_id,
-                Embedding.type == embedding_type,
-            )
-            result = await session.execute(query)
-            return result.scalar_one_or_none()
+        query = (
+            QueryBuilder()
+            .filter("snippet_id", FilterOperator.EQ, snippet_id)
+            .filter("type", FilterOperator.EQ, embedding_type)
+        )
+        results = await self.find(query)
+        return results[0] if results else None
 
     async def list_embeddings_by_type(
         self, embedding_type: EmbeddingType
     ) -> list[Embedding]:
         """List all embeddings of a given type."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            query = select(Embedding).where(Embedding.type == embedding_type)
-            result = await session.execute(query)
-            return list(result.scalars())
+        query = QueryBuilder().filter("type", FilterOperator.EQ, embedding_type)
+        return await self.find(query)
 
     async def delete_embeddings_by_snippet_id(self, snippet_id: str) -> None:
         """Delete all embeddings for a snippet."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            query = select(Embedding).where(Embedding.snippet_id == snippet_id)
-            result = await session.execute(query)
-            embeddings = result.scalars().all()
-            for embedding in embeddings:
-                await session.delete(embedding)
+        query = QueryBuilder().filter("snippet_id", FilterOperator.EQ, snippet_id)
+        embeddings = await self.find(query)
+        for embedding in embeddings:
+            await self.delete(embedding)
 
     async def list_embeddings_by_snippet_ids_and_type(
         self, snippet_ids: list[str], embedding_type: EmbeddingType
     ) -> list[Embedding]:
         """Get all embeddings for the given snippet IDs."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            query = select(Embedding).where(
-                Embedding.snippet_id.in_(snippet_ids),
-                Embedding.type == embedding_type,
-            )
-            result = await session.execute(query)
-            return list(result.scalars())
+        query = (
+            QueryBuilder()
+            .filter("snippet_id", FilterOperator.IN, snippet_ids)
+            .filter("type", FilterOperator.EQ, embedding_type)
+        )
+        return await self.find(query)
 
     async def get_embeddings_by_snippet_ids(
         self, snippet_ids: list[str]
     ) -> list[Embedding]:
         """Get all embeddings for the given snippet IDs."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            query = select(Embedding).where(Embedding.snippet_id.in_(snippet_ids))
-            result = await session.execute(query)
-            return list(result.scalars())
+        query = QueryBuilder().filter("snippet_id", FilterOperator.IN, snippet_ids)
+        return await self.find(query)
 
     async def list_semantic_results(
         self,

kodit/infrastructure/sqlalchemy/enrichment_association_repository.py

@@ -0,0 +1,73 @@
+"""Enrichment association repository."""
+
+from collections.abc import Callable
+
+import structlog
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.enrichments.enrichment import (
+    EnrichmentAssociation,
+)
+from kodit.domain.protocols import EnrichmentAssociationRepository
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.repository import SqlAlchemyRepository
+
+
+def create_enrichment_association_repository(
+    session_factory: Callable[[], AsyncSession],
+) -> EnrichmentAssociationRepository:
+    """Create a enrichment association repository."""
+    return SQLAlchemyEnrichmentAssociationRepository(session_factory=session_factory)
+
+
+class SQLAlchemyEnrichmentAssociationRepository(
+    SqlAlchemyRepository[EnrichmentAssociation, db_entities.EnrichmentAssociation],
+    EnrichmentAssociationRepository,
+):
+    """Repository for managing enrichment associations."""
+
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
+        """Initialize the repository."""
+        super().__init__(session_factory=session_factory)
+        self._log = structlog.get_logger(__name__)
+
+    def _get_id(self, entity: EnrichmentAssociation) -> int | None:
+        """Get the ID of an enrichment association."""
+        return entity.id
+
+    @property
+    def db_entity_type(self) -> type[db_entities.EnrichmentAssociation]:
+        """The SQLAlchemy model type."""
+        return db_entities.EnrichmentAssociation
+
+    @staticmethod
+    def to_domain(
+        db_entity: db_entities.EnrichmentAssociation,
+    ) -> EnrichmentAssociation:
+        """Map database entity to domain entity."""
+        return EnrichmentAssociation(
+            enrichment_id=db_entity.enrichment_id,
+            entity_type=db_entity.entity_type,
+            entity_id=db_entity.entity_id,
+            id=db_entity.id,
+        )
+
+    @staticmethod
+    def to_db(
+        domain_entity: EnrichmentAssociation,
+    ) -> db_entities.EnrichmentAssociation:
+        """Map domain entity to database entity."""
+        from datetime import UTC, datetime
+
+        now = datetime.now(UTC)
+        db_entity = db_entities.EnrichmentAssociation(
+            enrichment_id=domain_entity.enrichment_id,
+            entity_type=domain_entity.entity_type,
+            entity_id=domain_entity.entity_id,
+        )
+        if domain_entity.id is not None:
+            db_entity.id = domain_entity.id
+        # Always set timestamps since domain entity doesn't track them
+        db_entity.created_at = now
+        db_entity.updated_at = now
+        return db_entity

kodit/infrastructure/sqlalchemy/enrichment_v2_repository.py

@@ -1,118 +1,137 @@
 """EnrichmentV2 repository."""
 
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 
-import structlog
-from sqlalchemy import delete, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from kodit.domain.enrichments.architecture.architecture import (
+    ENRICHMENT_TYPE_ARCHITECTURE,
+)
+from kodit.domain.enrichments.architecture.physical.physical import (
+    ENRICHMENT_SUBTYPE_PHYSICAL,
+    PhysicalArchitectureEnrichment,
+)
+from kodit.domain.enrichments.development.development import ENRICHMENT_TYPE_DEVELOPMENT
+from kodit.domain.enrichments.development.snippet.snippet import (
+    ENRICHMENT_SUBTYPE_SNIPPET,
+    ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
+    SnippetEnrichment,
+    SnippetEnrichmentSummary,
+)
 from kodit.domain.enrichments.enrichment import EnrichmentV2
-from kodit.infrastructure.mappers.enrichment_mapper import EnrichmentMapper
+from kodit.domain.enrichments.usage.api_docs import (
+    ENRICHMENT_SUBTYPE_API_DOCS,
+    APIDocEnrichment,
+)
+from kodit.domain.enrichments.usage.usage import ENRICHMENT_TYPE_USAGE
+from kodit.domain.protocols import EnrichmentV2Repository
 from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.repository import SqlAlchemyRepository
 from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
 
-class EnrichmentV2Repository:
+def create_enrichment_v2_repository(
+    session_factory: Callable[[], AsyncSession],
+) -> EnrichmentV2Repository:
+    """Create a enrichment v2 repository."""
+    return SQLAlchemyEnrichmentV2Repository(session_factory=session_factory)
+
+
+class SQLAlchemyEnrichmentV2Repository(
+    SqlAlchemyRepository[EnrichmentV2, db_entities.EnrichmentV2], EnrichmentV2Repository
+):
     """Repository for managing enrichments and their associations."""
 
-    def __init__(
-        self,
-        session_factory: Callable[[], AsyncSession],
-    ) -> None:
-        """Initialize the repository."""
-        self.session_factory = session_factory
-        self.mapper = EnrichmentMapper()
-        self.log = structlog.get_logger(__name__)
+    def _get_id(self, entity: EnrichmentV2) -> int | None:
+        """Extract ID from domain entity."""
+        return entity.id
 
-    async def enrichments_for_entity_type(
-        self,
-        entity_type: str,
-        entity_ids: list[str],
-    ) -> list[EnrichmentV2]:
-        """Get all enrichments for multiple entities of the same type."""
-        if not entity_ids:
-            return []
+    @property
+    def db_entity_type(self) -> type[db_entities.EnrichmentV2]:
+        """The SQLAlchemy model type."""
+        return db_entities.EnrichmentV2
 
+    async def save(self, entity: EnrichmentV2) -> EnrichmentV2:
+        """Save entity (create new or update existing)."""
         async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            stmt = (
-                select(
-                    db_entities.EnrichmentV2,
-                    db_entities.EnrichmentAssociation.entity_id,
-                )
-                .join(db_entities.EnrichmentAssociation)
-                .where(
-                    db_entities.EnrichmentAssociation.entity_type == entity_type,
-                    db_entities.EnrichmentAssociation.entity_id.in_(entity_ids),
-                )
+            entity_id = self._get_id(entity)
+            # Skip session.get if entity_id is None (new entity not yet persisted)
+            existing_db_entity = (
+                await session.get(self.db_entity_type, entity_id)
+                if entity_id is not None
+                else None
             )
 
-            result = await session.execute(stmt)
-            rows = result.all()
-
-            return [
-                self.mapper.to_domain(db_enrichment, entity_type, entity_id)
-                for db_enrichment, entity_id in rows
-            ]
-
-    async def bulk_save_enrichments(
-        self,
-        enrichments: Sequence[EnrichmentV2],
-    ) -> None:
-        """Bulk save enrichments with their associations."""
-        if not enrichments:
-            return
-
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            enrichment_records = []
-            for enrichment in enrichments:
-                db_enrichment = db_entities.EnrichmentV2(
-                    type=enrichment.type,
-                    subtype=enrichment.subtype,
-                    content=enrichment.content,
-                )
-                session.add(db_enrichment)
-                enrichment_records.append((enrichment, db_enrichment))
+            if existing_db_entity:
+                # Update existing entity
+                new_db_entity = self.to_db(entity)
+                self._update_db_entity(existing_db_entity, new_db_entity)
+                db_entity = existing_db_entity
+            else:
+                # Create new entity
+                db_entity = self.to_db(entity)
+                session.add(db_entity)
 
             await session.flush()
-
-            for enrichment, db_enrichment in enrichment_records:
-                db_association = db_entities.EnrichmentAssociation(
-                    enrichment_id=db_enrichment.id,
-                    entity_type=enrichment.entity_type_key(),
-                    entity_id=enrichment.entity_id,
-                )
-                session.add(db_association)
-
-    async def bulk_delete_enrichments(
-        self,
-        entity_type: str,
-        entity_ids: list[str],
-    ) -> None:
-        """Bulk delete enrichments for multiple entities of the same type."""
-        if not entity_ids:
-            return
-
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            stmt = select(db_entities.EnrichmentAssociation.enrichment_id).where(
-                db_entities.EnrichmentAssociation.entity_type == entity_type,
-                db_entities.EnrichmentAssociation.entity_id.in_(entity_ids),
+            return self.to_domain(db_entity)
+
+    @staticmethod
+    def to_db(domain_entity: EnrichmentV2) -> db_entities.EnrichmentV2:
+        """Convert domain enrichment to database entity."""
+        enrichment = db_entities.EnrichmentV2(
+            type=domain_entity.type,
+            subtype=domain_entity.subtype,
+            content=domain_entity.content,
+        )
+        if domain_entity.id is not None:
+            enrichment.id = domain_entity.id
+        return enrichment
+
+    @staticmethod
+    def to_domain(db_entity: db_entities.EnrichmentV2) -> EnrichmentV2:
+        """Convert database enrichment to domain entity."""
+        # Use the stored type and subtype to determine the correct domain class
+        if (
+            db_entity.type == ENRICHMENT_TYPE_DEVELOPMENT
+            and db_entity.subtype == ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY
+        ):
+            return SnippetEnrichmentSummary(
+                id=db_entity.id,
+                content=db_entity.content,
+                created_at=db_entity.created_at,
+                updated_at=db_entity.updated_at,
             )
-            result = await session.execute(stmt)
-            enrichment_ids = result.scalars().all()
-
-            if enrichment_ids:
-                await session.execute(
-                    delete(db_entities.EnrichmentV2).where(
-                        db_entities.EnrichmentV2.id.in_(enrichment_ids)
-                    )
-                )
-
-    async def delete_enrichment(self, enrichment_id: int) -> bool:
-        """Delete a specific enrichment by ID."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            result = await session.execute(
-                delete(db_entities.EnrichmentV2).where(
-                    db_entities.EnrichmentV2.id == enrichment_id
-                )
+        if (
+            db_entity.type == ENRICHMENT_TYPE_DEVELOPMENT
+            and db_entity.subtype == ENRICHMENT_SUBTYPE_SNIPPET
+        ):
+            return SnippetEnrichment(
+                id=db_entity.id,
+                content=db_entity.content,
+                created_at=db_entity.created_at,
+                updated_at=db_entity.updated_at,
+            )
+        if (
+            db_entity.type == ENRICHMENT_TYPE_USAGE
+            and db_entity.subtype == ENRICHMENT_SUBTYPE_API_DOCS
+        ):
+            return APIDocEnrichment(
+                id=db_entity.id,
+                content=db_entity.content,
+                created_at=db_entity.created_at,
+                updated_at=db_entity.updated_at,
             )
-            return result.rowcount > 0
+        if (
+            db_entity.type == ENRICHMENT_TYPE_ARCHITECTURE
+            and db_entity.subtype == ENRICHMENT_SUBTYPE_PHYSICAL
+        ):
+            return PhysicalArchitectureEnrichment(
+                id=db_entity.id,
+                content=db_entity.content,
+                created_at=db_entity.created_at,
+                updated_at=db_entity.updated_at,
+            )
+
+        raise ValueError(
+            f"Unknown enrichment type: {db_entity.type}/{db_entity.subtype}"
+        )

kodit/infrastructure/sqlalchemy/entities.py

@@ -9,7 +9,6 @@ from sqlalchemy import (
     DateTime,
     Float,
     ForeignKey,
-    ForeignKeyConstraint,
     Index,
     Integer,
     String,
@@ -210,11 +209,15 @@ class GitRepo(Base, CommonMixin):
     num_commits: Mapped[int] = mapped_column(Integer, default=0)
     num_branches: Mapped[int] = mapped_column(Integer, default=0)
     num_tags: Mapped[int] = mapped_column(Integer, default=0)
+    tracking_type: Mapped[str] = mapped_column(String(255), index=True)
+    tracking_name: Mapped[str] = mapped_column(String(255), index=True)
 
     def __init__(  # noqa: PLR0913
         self,
         sanitized_remote_uri: str,
         remote_uri: str,
+        tracking_type: str,
+        tracking_name: str,
         cloned_path: Path | None,
         last_scanned_at: datetime | None = None,
         num_commits: int = 0,
@@ -225,6 +228,8 @@ class GitRepo(Base, CommonMixin):
         super().__init__()
         self.sanitized_remote_uri = sanitized_remote_uri
         self.remote_uri = remote_uri
+        self.tracking_type = tracking_type
+        self.tracking_name = tracking_name
         self.cloned_path = cloned_path
         self.last_scanned_at = last_scanned_at
         self.num_commits = num_commits
@@ -293,7 +298,12 @@ class GitBranch(Base):
 
     __table_args__ = (UniqueConstraint("repo_id", "name", name="uix_repo_branch"),)
 
-    def __init__(self, repo_id: int, name: str, head_commit_sha: str) -> None:
+    def __init__(
+        self,
+        repo_id: int,
+        name: str,
+        head_commit_sha: str,
+    ) -> None:
         """Initialize Git branch."""
         super().__init__()
         self.repo_id = repo_id
@@ -301,31 +311,6 @@ class GitBranch(Base):
         self.head_commit_sha = head_commit_sha
 
 
-class GitTrackingBranch(Base):
-    """Git tracking branch model."""
-
-    __tablename__ = "git_tracking_branches"
-    repo_id: Mapped[int] = mapped_column(
-        ForeignKey("git_repos.id"), index=True, primary_key=True
-    )
-    name: Mapped[str] = mapped_column(String(255), index=True, primary_key=True)
-    created_at: Mapped[datetime] = mapped_column(
-        TZDateTime, nullable=False, default=lambda: datetime.now(UTC)
-    )
-    updated_at: Mapped[datetime] = mapped_column(
-        TZDateTime,
-        nullable=False,
-        default=lambda: datetime.now(UTC),
-        onupdate=lambda: datetime.now(UTC),
-    )
-
-    def __init__(self, repo_id: int, name: str) -> None:
-        """Initialize Git tracking branch."""
-        super().__init__()
-        self.repo_id = repo_id
-        self.name = name
-
-
 class GitTag(Base):
     """Git tag model."""
 
@@ -395,95 +380,6 @@ class GitCommitFile(Base):
         self.extension = extension
 
 
-class SnippetV2(Base):
-    """SnippetV2 model for commit-based snippets."""
-
-    __tablename__ = "snippets_v2"
-
-    sha: Mapped[str] = mapped_column(String(64), primary_key=True)
-    created_at: Mapped[datetime] = mapped_column(
-        TZDateTime, nullable=False, default=lambda: datetime.now(UTC)
-    )
-    updated_at: Mapped[datetime] = mapped_column(
-        TZDateTime,
-        nullable=False,
-        default=lambda: datetime.now(UTC),
-        onupdate=lambda: datetime.now(UTC),
-    )
-    content: Mapped[str] = mapped_column(UnicodeText)
-    extension: Mapped[str] = mapped_column(String(255), index=True)
-
-    def __init__(
-        self,
-        sha: str,
-        content: str,
-        extension: str,
-    ) -> None:
-        """Initialize snippet."""
-        super().__init__()
-        self.sha = sha
-        self.content = content
-        self.extension = extension
-
-
-class SnippetV2File(Base):
-    """Association between snippets and files."""
-
-    __tablename__ = "snippet_v2_files"
-
-    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
-    snippet_sha: Mapped[str] = mapped_column(ForeignKey("snippets_v2.sha"), index=True)
-    blob_sha: Mapped[str] = mapped_column(String(64), index=True)
-    commit_sha: Mapped[str] = mapped_column(String(64), index=True)
-    file_path: Mapped[str] = mapped_column(String(1024), index=True)
-
-    __table_args__ = (
-        ForeignKeyConstraint(
-            ["commit_sha", "file_path"],
-            ["git_commit_files.commit_sha", "git_commit_files.path"],
-        ),
-        UniqueConstraint(
-            "snippet_sha",
-            "blob_sha",
-            "commit_sha",
-            "file_path",
-            name="uix_snippet_file",
-        ),
-    )
-
-    def __init__(
-        self, snippet_sha: str, blob_sha: str, commit_sha: str, file_path: str
-    ) -> None:
-        """Initialize snippet file association."""
-        super().__init__()
-        self.snippet_sha = snippet_sha
-        self.blob_sha = blob_sha
-        self.commit_sha = commit_sha
-        self.file_path = file_path
-
-
-class CommitSnippetV2(Base):
-    """Association table for commits and snippets v2."""
-
-    __tablename__ = "commit_snippets_v2"
-
-    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
-    commit_sha: Mapped[str] = mapped_column(
-        ForeignKey("git_commits.commit_sha"), index=True
-    )
-    snippet_sha: Mapped[str] = mapped_column(ForeignKey("snippets_v2.sha"), index=True)
-
-    __table_args__ = (
-        UniqueConstraint("commit_sha", "snippet_sha", name="uix_commit_snippet"),
-    )
-
-    def __init__(self, commit_sha: str, snippet_sha: str) -> None:
-        """Initialize commit snippet association."""
-        super().__init__()
-        self.commit_sha = commit_sha
-        self.snippet_sha = snippet_sha
-
-
 class CommitIndex(Base):
     """Commit index model."""