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

kodit/infrastructure/sqlalchemy/task_repository.py

@@ -1,16 +1,17 @@
 """Task repository for the task queue."""
 
 from collections.abc import Callable
+from typing import Any
 
 import structlog
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.domain.entities import Task
+import kodit.domain.entities as domain_entities
 from kodit.domain.protocols import TaskRepository
-from kodit.domain.value_objects import TaskOperation
 from kodit.infrastructure.mappers.task_mapper import TaskMapper
 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
 
 
@@ -21,7 +22,12 @@ def create_task_repository(
     return SqlAlchemyTaskRepository(session_factory=session_factory)
 
 
-class SqlAlchemyTaskRepository(TaskRepository):
+# TODO(Phil): Stop using dedup_key as the primary key. Add some DDD to this instead.
+
+
+class SqlAlchemyTaskRepository(
+    SqlAlchemyRepository[domain_entities.Task, db_entities.Task], TaskRepository
+):
     """Repository for task persistence using the existing Task entity."""
 
     def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
@@ -29,25 +35,81 @@ class SqlAlchemyTaskRepository(TaskRepository):
         self.session_factory = session_factory
         self.log = structlog.get_logger(__name__)
 
-    async def add(
-        self,
-        task: Task,
-    ) -> None:
-        """Create a new task in the database."""
+    @property
+    def db_entity_type(self) -> type[db_entities.Task]:
+        """The SQLAlchemy model type."""
+        return db_entities.Task
+
+    def _get_id(self, entity: domain_entities.Task) -> Any:
+        """Extract ID from domain entity."""
+        return entity.id
+
+    @staticmethod
+    def to_domain(db_entity: db_entities.Task) -> domain_entities.Task:
+        """Map database entity to domain entity."""
+        return TaskMapper.to_domain_task(db_entity)
+
+    @staticmethod
+    def to_db(domain_entity: domain_entities.Task) -> db_entities.Task:
+        """Map domain entity to database entity."""
+        return TaskMapper.from_domain_task(domain_entity)
+
+    async def get(self, entity_id: Any) -> domain_entities.Task:
+        """Get entity by dedup_key."""
         async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            session.add(TaskMapper.from_domain_task(task))
+            stmt = select(db_entities.Task).where(
+                db_entities.Task.dedup_key == entity_id
+            )
+            result = await session.execute(stmt)
+            db_entity = result.scalar_one_or_none()
+            if not db_entity:
+                raise ValueError(f"Entity with id {entity_id} not found")
+            return self.to_domain(db_entity)
 
-    async def get(self, task_id: str) -> Task | None:
-        """Get a task by ID."""
+    async def exists(self, entity_id: Any) -> bool:
+        """Check if entity exists by dedup_key."""
         async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            stmt = select(db_entities.Task).where(db_entities.Task.dedup_key == task_id)
+            stmt = select(db_entities.Task).where(
+                db_entities.Task.dedup_key == entity_id
+            )
             result = await session.execute(stmt)
-            db_task = result.scalar_one_or_none()
-            if not db_task:
-                return None
-            return TaskMapper.to_domain_task(db_task)
+            return result.scalar_one_or_none() is not None
+
+    async def save(self, entity: domain_entities.Task) -> domain_entities.Task:
+        """Save entity (create new or update existing)."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Query by dedup_key instead of primary key
+            stmt = select(db_entities.Task).where(
+                db_entities.Task.dedup_key == entity.id
+            )
+            result = await session.execute(stmt)
+            existing_db_entity = result.scalar_one_or_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 delete(self, entity: domain_entities.Task) -> None:
+        """Remove entity by dedup_key."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(db_entities.Task).where(
+                db_entities.Task.dedup_key == entity.id
+            )
+            result = await session.execute(stmt)
+            db_entity = result.scalar_one_or_none()
+            if db_entity:
+                await session.delete(db_entity)
 
-    async def next(self) -> Task | None:
+    async def next(self) -> domain_entities.Task | None:
         """Take a task for processing and remove it from the database."""
         async with SqlAlchemyUnitOfWork(self.session_factory) as session:
             stmt = (
@@ -60,44 +122,3 @@ class SqlAlchemyTaskRepository(TaskRepository):
             if not db_task:
                 return None
             return TaskMapper.to_domain_task(db_task)
-
-    async def remove(self, task: Task) -> None:
-        """Remove a task from the database."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            db_task = await session.scalar(
-                select(db_entities.Task).where(db_entities.Task.dedup_key == task.id)
-            )
-            if not db_task:
-                raise ValueError(f"Task not found: {task.id}")
-            await session.delete(db_task)
-
-    async def update(self, task: Task) -> None:
-        """Update a task in the database."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            stmt = select(db_entities.Task).where(db_entities.Task.dedup_key == task.id)
-            result = await session.execute(stmt)
-            db_task = result.scalar_one_or_none()
-
-            if not db_task:
-                raise ValueError(f"Task not found: {task.id}")
-
-            db_task.priority = task.priority
-            db_task.payload = task.payload
-
-    async def list(self, task_operation: TaskOperation | None = None) -> list[Task]:
-        """List tasks with optional status filter."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            stmt = select(db_entities.Task)
-
-            if task_operation:
-                stmt = stmt.where(db_entities.Task.type == task_operation.value)
-
-            stmt = stmt.order_by(
-                db_entities.Task.priority.desc(), db_entities.Task.created_at
-            )
-
-            result = await session.execute(stmt)
-            records = result.scalars().all()
-
-            # Convert to domain entities
-            return [TaskMapper.to_domain_task(record) for record in records]

kodit/infrastructure/sqlalchemy/task_status_repository.py

@@ -1,15 +1,16 @@
 """Task repository for the task queue."""
 
 from collections.abc import Callable
+from typing import Any, override
 
-import structlog
-from sqlalchemy import delete, select
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.domain import entities as domain_entities
 from kodit.domain.protocols import TaskStatusRepository
 from kodit.infrastructure.mappers.task_status_mapper import TaskStatusMapper
 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
 
 
@@ -20,52 +21,52 @@ def create_task_status_repository(
     return SqlAlchemyTaskStatusRepository(session_factory=session_factory)
 
 
-class SqlAlchemyTaskStatusRepository(TaskStatusRepository):
+class SqlAlchemyTaskStatusRepository(
+    SqlAlchemyRepository[domain_entities.TaskStatus, db_entities.TaskStatus],
+    TaskStatusRepository,
+):
     """Repository for persisting TaskStatus entities."""
 
-    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
-        """Initialize the repository."""
-        self.session_factory = session_factory
-        self.mapper = TaskStatusMapper()
-        self.log = structlog.get_logger(__name__)
+    @property
+    def db_entity_type(self) -> type[db_entities.TaskStatus]:
+        """The SQLAlchemy model type."""
+        return db_entities.TaskStatus
 
-    async def save(self, status: domain_entities.TaskStatus) -> None:
+    def _get_id(self, entity: domain_entities.TaskStatus) -> Any:
+        """Extract ID from domain entity."""
+        return entity.id
+
+    @staticmethod
+    def to_domain(
+        db_entity: db_entities.TaskStatus,
+    ) -> domain_entities.TaskStatus:
+        """Map database entity to domain entity."""
+        return TaskStatusMapper.to_domain_task_status(db_entity)
+
+    @staticmethod
+    def to_db(
+        domain_entity: domain_entities.TaskStatus,
+    ) -> db_entities.TaskStatus:
+        """Map domain entity to database entity."""
+        return TaskStatusMapper.from_domain_task_status(domain_entity)
+
+    @override
+    async def save(
+        self, entity: domain_entities.TaskStatus
+    ) -> domain_entities.TaskStatus:
         """Save a TaskStatus to database."""
-        # If this task has a parent, ensure the parent exists in the database first
-        if status.parent is not None:
-            async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-                parent_stmt = select(db_entities.TaskStatus).where(
-                    db_entities.TaskStatus.id == status.parent.id,
-                )
-                parent_result = await session.execute(parent_stmt)
-                existing_parent = parent_result.scalar_one_or_none()
-
-            if not existing_parent:
-                # Recursively save the parent first
-                await self.save(status.parent)
+        # Recursively convert parents to a list of domain entities, parents first
+        parents: list[domain_entities.TaskStatus] = []
+        current = entity
+        while current.parent is not None:
+            parents.insert(0, current.parent)
+            current = current.parent
 
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            # Convert domain entity to database entity
-            db_status = self.mapper.from_domain_task_status(status)
-            stmt = select(db_entities.TaskStatus).where(
-                db_entities.TaskStatus.id == db_status.id,
-            )
-            result = await session.execute(stmt)
-            existing = result.scalar_one_or_none()
-
-            if not existing:
-                session.add(db_status)
-            else:
-                # Update existing record with new values
-                existing.operation = db_status.operation
-                existing.state = db_status.state
-                existing.error = db_status.error
-                existing.total = db_status.total
-                existing.current = db_status.current
-                existing.updated_at = db_status.updated_at
-                existing.parent = db_status.parent
-                existing.trackable_id = db_status.trackable_id
-                existing.trackable_type = db_status.trackable_type
+        # Add current entity to the end of the list
+        parents.append(entity)
+
+        await self.save_bulk(parents)
+        return entity
 
     async def load_with_hierarchy(
         self, trackable_type: str, trackable_id: int
@@ -80,12 +81,4 @@ class SqlAlchemyTaskStatusRepository(TaskStatusRepository):
             db_statuses = list(result.scalars().all())
 
             # Use mapper to convert and reconstruct hierarchy
-            return self.mapper.to_domain_task_status_with_hierarchy(db_statuses)
-
-    async def delete(self, status: domain_entities.TaskStatus) -> None:
-        """Delete a TaskStatus."""
-        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
-            stmt = delete(db_entities.TaskStatus).where(
-                db_entities.TaskStatus.id == status.id,
-            )
-            await session.execute(stmt)
+            return TaskStatusMapper.to_domain_task_status_with_hierarchy(db_statuses)

kodit/migrations/versions/4b1a3b2c8fa5_refactor_git_tracking.py

@@ -0,0 +1,190 @@
+# ruff: noqa
+"""refactorings
+
+Revision ID: 4b1a3b2c8fa5
+Revises: 19f8c7faf8b9
+Create Date: 2025-10-29 13:38:10.737704
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from kodit.domain.tracking.trackable import TrackableReferenceType
+
+
+# revision identifiers, used by Alembic.
+revision: str = "4b1a3b2c8fa5"
+down_revision: Union[str, None] = "19f8c7faf8b9"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    op.drop_index("ix_commit_snippets_v2_commit_sha", table_name="commit_snippets_v2")
+    op.drop_index("ix_commit_snippets_v2_snippet_sha", table_name="commit_snippets_v2")
+    op.drop_table("commit_snippets_v2")
+    op.drop_index("ix_snippet_v2_files_blob_sha", table_name="snippet_v2_files")
+    op.drop_index("ix_snippet_v2_files_commit_sha", table_name="snippet_v2_files")
+    op.drop_index("ix_snippet_v2_files_file_path", table_name="snippet_v2_files")
+    op.drop_index("ix_snippet_v2_files_snippet_sha", table_name="snippet_v2_files")
+    op.drop_table("snippet_v2_files")
+    op.drop_index("ix_snippets_v2_extension", table_name="snippets_v2")
+    op.drop_table("snippets_v2")
+    op.drop_index("ix_git_tracking_branches_name", table_name="git_tracking_branches")
+    op.drop_index(
+        "ix_git_tracking_branches_repo_id", table_name="git_tracking_branches"
+    )
+    op.drop_table("git_tracking_branches")
+
+    # Use batch_alter_table for SQLite compatibility
+    with op.batch_alter_table("git_repos", schema=None) as batch_op:
+        # Add columns as nullable first
+        batch_op.add_column(
+            sa.Column("tracking_type", sa.String(length=255), nullable=True)
+        )
+        batch_op.add_column(
+            sa.Column("tracking_name", sa.String(length=255), nullable=True)
+        )
+
+    # Set default values for existing rows
+    op.execute(
+        f"UPDATE git_repos SET tracking_type = '{TrackableReferenceType.BRANCH}' WHERE tracking_type IS NULL"
+    )
+    op.execute(
+        f"UPDATE git_repos SET tracking_name = 'main' WHERE tracking_name IS NULL"
+    )
+
+    # Make columns non-nullable using batch_alter_table for SQLite compatibility
+    with op.batch_alter_table("git_repos", schema=None) as batch_op:
+        batch_op.alter_column("tracking_type", nullable=False)
+        batch_op.alter_column("tracking_name", nullable=False)
+        batch_op.create_index(
+            op.f("ix_git_repos_tracking_name"),
+            ["tracking_name"],
+            unique=False,
+        )
+        batch_op.create_index(
+            op.f("ix_git_repos_tracking_type"),
+            ["tracking_type"],
+            unique=False,
+        )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.create_table(
+        "snippets_v2",
+        sa.Column("sha", sa.VARCHAR(length=64), nullable=False),
+        sa.Column("created_at", sa.DATETIME(), nullable=False),
+        sa.Column("updated_at", sa.DATETIME(), nullable=False),
+        sa.Column("content", sa.TEXT(), nullable=False),
+        sa.Column("extension", sa.VARCHAR(length=255), nullable=False),
+        sa.PrimaryKeyConstraint("sha"),
+    )
+    op.create_index(
+        "ix_snippets_v2_extension", "snippets_v2", ["extension"], unique=False
+    )
+    op.create_table(
+        "snippet_v2_files",
+        sa.Column("id", sa.INTEGER(), nullable=False),
+        sa.Column("snippet_sha", sa.VARCHAR(length=64), nullable=False),
+        sa.Column("blob_sha", sa.VARCHAR(length=64), nullable=False),
+        sa.Column("commit_sha", sa.VARCHAR(length=64), nullable=False),
+        sa.Column("file_path", sa.VARCHAR(length=1024), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["commit_sha", "file_path"],
+            ["git_commit_files.commit_sha", "git_commit_files.path"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["snippet_sha"],
+            ["snippets_v2.sha"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint(
+            "snippet_sha",
+            "blob_sha",
+            "commit_sha",
+            "file_path",
+            name="uix_snippet_file",
+        ),
+    )
+    op.create_index(
+        "ix_snippet_v2_files_snippet_sha",
+        "snippet_v2_files",
+        ["snippet_sha"],
+        unique=False,
+    )
+    op.create_index(
+        "ix_snippet_v2_files_file_path", "snippet_v2_files", ["file_path"], unique=False
+    )
+    op.create_index(
+        "ix_snippet_v2_files_commit_sha",
+        "snippet_v2_files",
+        ["commit_sha"],
+        unique=False,
+    )
+    op.create_index(
+        "ix_snippet_v2_files_blob_sha", "snippet_v2_files", ["blob_sha"], unique=False
+    )
+    op.create_table(
+        "commit_snippets_v2",
+        sa.Column("id", sa.INTEGER(), nullable=False),
+        sa.Column("commit_sha", sa.VARCHAR(length=64), nullable=False),
+        sa.Column("snippet_sha", sa.VARCHAR(length=64), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["commit_sha"],
+            ["git_commits.commit_sha"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["snippet_sha"],
+            ["snippets_v2.sha"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("commit_sha", "snippet_sha", name="uix_commit_snippet"),
+    )
+    op.create_index(
+        "ix_commit_snippets_v2_snippet_sha",
+        "commit_snippets_v2",
+        ["snippet_sha"],
+        unique=False,
+    )
+    op.create_index(
+        "ix_commit_snippets_v2_commit_sha",
+        "commit_snippets_v2",
+        ["commit_sha"],
+        unique=False,
+    )
+    # Use batch_alter_table for SQLite compatibility
+    with op.batch_alter_table("git_repos", schema=None) as batch_op:
+        batch_op.drop_index(op.f("ix_git_repos_tracking_type"))
+        batch_op.drop_index(op.f("ix_git_repos_tracking_name"))
+        batch_op.drop_column("tracking_name")
+        batch_op.drop_column("tracking_type")
+
+    op.create_table(
+        "git_tracking_branches",
+        sa.Column("repo_id", sa.INTEGER(), nullable=False),
+        sa.Column("name", sa.VARCHAR(length=255), nullable=False),
+        sa.Column("created_at", sa.DATETIME(), nullable=False),
+        sa.Column("updated_at", sa.DATETIME(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["repo_id"], ["git_repos.id"], name="fk_tracking_branch_repo"
+        ),
+        sa.PrimaryKeyConstraint("repo_id", "name", name="pk_git_tracking_branches"),
+    )
+    op.create_index(
+        "ix_git_tracking_branches_repo_id",
+        "git_tracking_branches",
+        ["repo_id"],
+        unique=False,
+    )
+    op.create_index(
+        "ix_git_tracking_branches_name",
+        "git_tracking_branches",
+        ["name"],
+        unique=False,
+    )

{kodit-0.5.4.dist-info → kodit-0.5.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.5.4
+Version: 0.5.5
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/