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

kodit/infrastructure/sqlalchemy/query.py

@@ -0,0 +1,331 @@
+"""Base query for SQLAlchemy repositories."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Self
+
+from sqlalchemy import Select
+
+from kodit.domain.enrichments.enrichment import EnrichmentV2
+from kodit.infrastructure.api.v1.query_params import PaginationParams
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+
+
+class Query(ABC):
+    """Base query/specification object for encapsulating query logic."""
+
+    @abstractmethod
+    def apply(self, stmt: Select, model_type: type) -> Select:
+        """Apply this query's criteria to a SQLAlchemy Select statement."""
+
+
+class FilterOperator(Enum):
+    """SQL filter operators."""
+
+    EQ = "eq"
+    NE = "ne"
+    GT = "gt"
+    GTE = "ge"
+    LT = "lt"
+    LTE = "le"
+    IN = "in_"
+    LIKE = "like"
+    ILIKE = "ilike"
+
+
+@dataclass
+class FilterCriteria:
+    """Filter criteria for a query."""
+
+    field: str
+    operator: FilterOperator
+    value: Any
+
+    def apply(self, model_type: type, stmt: Select) -> Select:  # noqa: C901
+        """Apply filter to statement."""
+        column = getattr(model_type, self.field)
+
+        # Convert AnyUrl to string for SQLAlchemy comparison
+        value = self.value
+        if hasattr(value, "__str__") and type(value).__module__ == "pydantic.networks":
+            value = str(value)
+
+        # Use column comparison methods instead of operators module
+        condition = None
+        match self.operator:
+            case FilterOperator.EQ:
+                condition = column == value
+            case FilterOperator.NE:
+                condition = column != value
+            case FilterOperator.GT:
+                condition = column > value
+            case FilterOperator.GTE:
+                condition = column >= value
+            case FilterOperator.LT:
+                condition = column < value
+            case FilterOperator.LTE:
+                condition = column <= value
+            case FilterOperator.IN:
+                condition = column.in_(value)
+            case FilterOperator.LIKE:
+                condition = column.like(value)
+            case FilterOperator.ILIKE:
+                condition = column.ilike(value)
+
+        return stmt.where(condition)
+
+
+@dataclass
+class SortCriteria:
+    """Sort criteria for a query."""
+
+    field: str
+    descending: bool = False
+
+    def apply(self, model_type: type, stmt: Select) -> Select:
+        """Apply sort to statement."""
+        column = getattr(model_type, self.field)
+        return stmt.order_by(column.desc() if self.descending else column.asc())
+
+
+@dataclass
+class PaginationCriteria:
+    """Pagination criteria for a query."""
+
+    limit: int | None = None
+    offset: int = 0
+
+    def apply(self, stmt: Select) -> Select:
+        """Apply pagination to statement."""
+        stmt = stmt.offset(self.offset)
+        if self.limit is not None:
+            stmt = stmt.limit(self.limit)
+        return stmt
+
+
+class QueryBuilder(Query):
+    """Composable query builder for constructing database queries."""
+
+    DEFAULT_SORT_FIELD = "created_at"
+    DEFAULT_SORT_DESCENDING = True
+
+    def __init__(self) -> None:
+        """Initialize query builder."""
+        self._filters: list[FilterCriteria] = []
+        self._sorts: list[SortCriteria] = []
+        self._pagination: PaginationCriteria | None = None
+
+    def filter(self, field: str, operator: FilterOperator, value: Any) -> Self:
+        """Add a filter criterion."""
+        self._filters.append(FilterCriteria(field, operator, value))
+        return self
+
+    def sort(self, field: str, *, descending: bool = False) -> Self:
+        """Add a sort criterion."""
+        self._sorts.append(SortCriteria(field, descending))
+        return self
+
+    def paginate(self, pagination: PaginationParams) -> Self:
+        """Add pagination."""
+        self._pagination = PaginationCriteria(
+            limit=pagination.limit, offset=pagination.offset
+        )
+        return self
+
+    def apply_filters_only(self, stmt: Select, model_type: type) -> Select:
+        """Apply only filter criteria to the statement."""
+        for filter_criteria in self._filters:
+            stmt = filter_criteria.apply(model_type, stmt)
+        return stmt
+
+    def apply(self, stmt: Select, model_type: type) -> Select:
+        """Apply all criteria to the statement."""
+        for filter_criteria in self._filters:
+            stmt = filter_criteria.apply(model_type, stmt)
+
+        if not self._sorts:
+            self._sorts = [
+                SortCriteria(
+                    field=self.DEFAULT_SORT_FIELD,
+                    descending=self.DEFAULT_SORT_DESCENDING,
+                )
+            ]
+        for sort_criteria in self._sorts:
+            stmt = sort_criteria.apply(model_type, stmt)
+
+        if self._pagination:
+            stmt = self._pagination.apply(stmt)
+
+        return stmt
+
+
+class EnrichmentAssociationQueryBuilder(QueryBuilder):
+    """Query builder for enrichment association entities."""
+
+    def for_enrichment_ids(self, enrichment_ids: list[int]) -> Self:
+        """Build a query for associations by enrichment IDs."""
+        self.filter(
+            db_entities.EnrichmentAssociation.enrichment_id.key,
+            FilterOperator.IN,
+            enrichment_ids,
+        )
+        return self
+
+    def for_enrichments(self, enrichments: list[EnrichmentV2]) -> Self:
+        """Build a query for enrichment associations by entity IDs."""
+        self.filter(
+            db_entities.EnrichmentAssociation.enrichment_id.key,
+            FilterOperator.IN,
+            [enrichment.id for enrichment in enrichments if enrichment.id is not None],
+        )
+        return self
+
+    def for_enrichment_type(self) -> Self:
+        """Build a query for enrichment types."""
+        return self.for_entity_type(db_entities.EnrichmentV2.__tablename__)
+
+    def for_entity_type(self, entity_type: str) -> Self:
+        """Build a query for enrichment associations by entity type."""
+        self.filter(
+            db_entities.EnrichmentAssociation.entity_type.key,
+            FilterOperator.EQ,
+            entity_type,
+        )
+        return self
+
+    def for_entity_ids(self, entity_ids: list[str]) -> Self:
+        """Build a query for enrichment associations by entity IDs."""
+        self.filter(
+            db_entities.EnrichmentAssociation.entity_id.key,
+            FilterOperator.IN,
+            entity_ids,
+        )
+        return self
+
+    @staticmethod
+    def for_enrichment_association(
+        entity_type: str,
+        entity_id: str,
+    ) -> QueryBuilder:
+        """Build a query for a specific enrichment association."""
+        return EnrichmentAssociationQueryBuilder.for_enrichment_associations(
+            entity_type,
+            [entity_id],
+        )
+
+    @staticmethod
+    def for_enrichment_associations(
+        entity_type: str, entity_ids: list[str]
+    ) -> QueryBuilder:
+        """Build a query for enrichment associations by entity type and IDs."""
+        return (
+            QueryBuilder()
+            .filter(
+                db_entities.EnrichmentAssociation.entity_type.key,
+                FilterOperator.EQ,
+                entity_type,
+            )
+            .filter(
+                db_entities.EnrichmentAssociation.entity_id.key,
+                FilterOperator.IN,
+                entity_ids,
+            )
+        )
+
+    @staticmethod
+    def type_and_ids(
+        entity_type: str,
+        enrichment_ids: list[int],
+    ) -> QueryBuilder:
+        """Build a query for enrichment associations by enrichment IDs."""
+        return (
+            QueryBuilder()
+            .filter(
+                db_entities.EnrichmentAssociation.entity_type.key,
+                FilterOperator.EQ,
+                entity_type,
+            )
+            .filter(
+                db_entities.EnrichmentAssociation.enrichment_id.key,
+                FilterOperator.IN,
+                enrichment_ids,
+            )
+        )
+
+    @staticmethod
+    def associations_pointing_to_these_enrichments(
+        enrichment_ids: list[int],
+    ) -> QueryBuilder:
+        """Build a query for enrichment associations pointing to these enrichments."""
+        return EnrichmentAssociationQueryBuilder.type_and_ids(
+            entity_type=db_entities.EnrichmentV2.__tablename__,
+            enrichment_ids=enrichment_ids,
+        )
+
+    def for_commit(self, commit_sha: str) -> Self:
+        """Build a query for enrichment associations for a commit."""
+        self.filter(
+            db_entities.EnrichmentAssociation.entity_type.key,
+            FilterOperator.EQ,
+            db_entities.GitCommit.__tablename__,
+        )
+        self.filter(
+            db_entities.EnrichmentAssociation.entity_id.key,
+            FilterOperator.EQ,
+            commit_sha,
+        )
+        return self
+
+
+class EnrichmentQueryBuilder(QueryBuilder):
+    """Query builder for enrichment entities."""
+
+    def for_ids(self, enrichment_ids: list[int]) -> Self:
+        """Build a query for enrichments by their IDs."""
+        self.filter(
+            db_entities.EnrichmentV2.id.key,
+            FilterOperator.IN,
+            enrichment_ids,
+        )
+        return self
+
+    def for_type(self, enrichment_type: str) -> Self:
+        """Build a query for enrichments by their type."""
+        self.filter(
+            db_entities.EnrichmentV2.type.key,
+            FilterOperator.EQ,
+            enrichment_type,
+        )
+        return self
+
+    def for_subtype(self, enrichment_subtype: str) -> Self:
+        """Build a query for enrichments by their subtype."""
+        self.filter(
+            db_entities.EnrichmentV2.subtype.key,
+            FilterOperator.EQ,
+            enrichment_subtype,
+        )
+        return self
+
+
+class GitFileQueryBuilder(QueryBuilder):
+    """Query builder for git file entities."""
+
+    def for_commit_sha(self, commit_sha: str) -> Self:
+        """Build a query for git files by their commit SHA."""
+        self.filter(
+            db_entities.GitCommitFile.commit_sha.key,
+            FilterOperator.EQ,
+            commit_sha,
+        )
+        return self
+
+    def for_blob_sha(self, blob_sha: str) -> Self:
+        """Build a query for git files by their blob SHA."""
+        self.filter(
+            db_entities.GitCommitFile.blob_sha.key,
+            FilterOperator.EQ,
+            blob_sha,
+        )
+        return self

kodit/infrastructure/sqlalchemy/repository.py

@@ -0,0 +1,203 @@
+"""Abstract base classes for repositories."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Generator
+from typing import Any, Generic, TypeVar
+
+from sqlalchemy import func, inspect, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.infrastructure.sqlalchemy.query import Query
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
+
+DomainEntityType = TypeVar("DomainEntityType")
+DatabaseEntityType = TypeVar("DatabaseEntityType")
+
+
+class SqlAlchemyRepository(ABC, Generic[DomainEntityType, DatabaseEntityType]):
+    """Base repository with common SQLAlchemy patterns."""
+
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
+        """Initialize the repository."""
+        self.session_factory = session_factory
+        self._chunk_size = 1000
+
+    @abstractmethod
+    def _get_id(self, entity: DomainEntityType) -> Any:
+        """Extract ID from domain entity."""
+
+    @property
+    @abstractmethod
+    def db_entity_type(self) -> type[DatabaseEntityType]:
+        """The SQLAlchemy model type."""
+
+    @staticmethod
+    @abstractmethod
+    def to_domain(db_entity: DatabaseEntityType) -> DomainEntityType:
+        """Map database entity to domain entity."""
+
+    @staticmethod
+    @abstractmethod
+    def to_db(domain_entity: DomainEntityType) -> DatabaseEntityType:
+        """Map domain entity to database entity."""
+
+    def _update_db_entity(
+        self, existing: DatabaseEntityType, new: DatabaseEntityType
+    ) -> None:
+        """Update existing database entity with values from new entity."""
+        mapper = inspect(type(existing))
+        if mapper is None:
+            return
+        # Skip auto-managed columns
+        skip_columns = {"created_at", "updated_at", "id"}
+        for column in mapper.columns:
+            if not column.primary_key and column.key not in skip_columns:
+                setattr(existing, column.key, getattr(new, column.key))
+
+    async def get(self, entity_id: Any) -> DomainEntityType:
+        """Get entity by primary key."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            db_entity = await session.get(self.db_entity_type, entity_id)
+            if not db_entity:
+                raise ValueError(f"Entity with id {entity_id} not found")
+            return self.to_domain(db_entity)
+
+    async def find(self, query: Query) -> list[DomainEntityType]:
+        """Find all entities matching query."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(self.db_entity_type)
+            stmt = query.apply(stmt, self.db_entity_type)
+            db_entities = (await session.scalars(stmt)).all()
+            return [self.to_domain(db) for db in db_entities]
+
+    async def count(self, query: Query) -> int:
+        """Count the number of entities matching query."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(self.db_entity_type).with_only_columns(func.count())
+            # For count queries, only apply filters, not sorting or pagination
+            from kodit.infrastructure.sqlalchemy.query import QueryBuilder
+
+            if isinstance(query, QueryBuilder):
+                # Apply only filters, skip sorting and pagination for count queries
+                stmt = query.apply_filters_only(stmt, self.db_entity_type)
+            else:
+                stmt = query.apply(stmt, self.db_entity_type)
+            result = await session.scalar(stmt)
+            return result or 0
+
+    async def save(self, entity: DomainEntityType) -> DomainEntityType:
+        """Save entity (create new or update existing)."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            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
+            )
+
+            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()
+            return self.to_domain(db_entity)
+
+    async def save_bulk(
+        self, entities: list[DomainEntityType]
+    ) -> list[DomainEntityType]:
+        """Save multiple entities in bulk (create new or update existing)."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            all_saved_db_entities = []
+
+            for chunk in self._chunked_domain(entities):
+                # Get IDs for all entities in chunk
+                entity_ids = [self._get_id(entity) for entity in chunk]
+
+                # Fetch all existing entities in one query
+                existing_entities = {}
+                for entity_id in entity_ids:
+                    # Skip None IDs (new entities not yet persisted)
+                    if entity_id is None:
+                        continue
+                    existing = await session.get(self.db_entity_type, entity_id)
+                    if existing:
+                        existing_entities[entity_id] = existing
+
+                # Process each entity
+                new_entities = []
+                chunk_db_entities = []
+                for entity in chunk:
+                    entity_id = self._get_id(entity)
+                    new_db_entity = self.to_db(entity)
+
+                    if entity_id in existing_entities:
+                        # Update existing entity
+                        existing = existing_entities[entity_id]
+                        self._update_db_entity(existing, new_db_entity)
+                        chunk_db_entities.append(existing)
+                    else:
+                        # Collect new entities to add
+                        new_entities.append(new_db_entity)
+                        chunk_db_entities.append(new_db_entity)
+
+                # Add all new entities at once
+                if new_entities:
+                    session.add_all(new_entities)
+
+                await session.flush()
+                all_saved_db_entities.extend(chunk_db_entities)
+
+            return [self.to_domain(db) for db in all_saved_db_entities]
+
+    async def exists(self, entity_id: Any) -> bool:
+        """Check if entity exists by primary key."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            db_entity = await session.get(self.db_entity_type, entity_id)
+            return db_entity is not None
+
+    async def delete(self, entity: DomainEntityType) -> None:
+        """Remove entity."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            db_entity = await session.get(self.db_entity_type, self._get_id(entity))
+            if db_entity:
+                await session.delete(db_entity)
+
+    async def delete_by_query(self, query: Query) -> None:
+        """Remove entities by query."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(self.db_entity_type)
+            stmt = query.apply(stmt, self.db_entity_type)
+            db_entities = list((await session.scalars(stmt)).all())
+            if not db_entities:
+                return
+            for chunk in self._chunked_db(db_entities):
+                for db_entity in chunk:
+                    await session.delete(db_entity)
+            await session.flush()
+
+    def _chunked_domain(
+        self,
+        items: list[DomainEntityType],
+        chunk_size: int | None = None,
+    ) -> Generator[list[DomainEntityType], None, None]:
+        """Yield chunks of items."""
+        chunk_size = chunk_size or self._chunk_size
+        for i in range(0, len(items), chunk_size):
+            yield items[i : i + chunk_size]
+
+    def _chunked_db(
+        self,
+        items: list[DatabaseEntityType],
+        chunk_size: int | None = None,
+    ) -> Generator[list[DatabaseEntityType], None, None]:
+        """Yield chunks of items."""
+        chunk_size = chunk_size or self._chunk_size
+        for i in range(0, len(items), chunk_size):
+            yield items[i : i + chunk_size]