kodit 0.4.3__py3-none-any.whl → 0.5.1__py3-none-any.whl

kodit/infrastructure/sqlalchemy/unit_of_work.py

@@ -14,17 +14,10 @@ class SqlAlchemyUnitOfWork:
         self._session_factory = session_factory
         self._session: AsyncSession | None = None
 
-    @property
-    def session(self) -> AsyncSession:
-        """Get the current session."""
-        if self._session is None:
-            raise RuntimeError("UnitOfWork must be used within async context")
-        return self._session
-
-    async def __aenter__(self) -> "SqlAlchemyUnitOfWork":
+    async def __aenter__(self) -> "AsyncSession":
         """Enter the unit of work context."""
         self._session = self._session_factory()
-        return self
+        return self._session
 
     async def __aexit__(
         self,
@@ -34,11 +27,14 @@ class SqlAlchemyUnitOfWork:
     ) -> None:
         """Exit the unit of work context."""
         if self._session:
-            if exc_type is not None:
-                await self._session.rollback()
-            await self._session.commit()
-            await self._session.close()
-            self._session = None
+            try:
+                if exc_type is None:  # Only commit if no exception
+                    await self._session.commit()
+                else:
+                    await self._session.rollback()
+            finally:
+                await self._session.close()
+                self._session = None
 
     async def commit(self) -> None:
         """Commit the current transaction."""

kodit/mcp.py

@@ -1,6 +1,6 @@
 """MCP server for kodit."""
 
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
@@ -9,35 +9,27 @@ from typing import Annotated
 import structlog
 from fastmcp import Context, FastMCP
 from pydantic import Field
-from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit._version import version
-from kodit.application.factories.code_indexing_factory import (
-    create_code_indexing_application_service,
-)
-from kodit.application.factories.reporting_factory import create_server_operation
+from kodit.application.factories.server_factory import ServerFactory
+from kodit.application.services.code_search_application_service import MultiSearchResult
 from kodit.config import AppContext
 from kodit.database import Database
 from kodit.domain.value_objects import (
     MultiSearchRequest,
-    MultiSearchResult,
     SnippetSearchFilters,
 )
-from kodit.infrastructure.sqlalchemy.task_status_repository import (
-    create_task_status_repository,
-)
 
 # Global database connection for MCP server
 _mcp_db: Database | None = None
+_mcp_server_factory: ServerFactory | None = None
 
 
 @dataclass
 class MCPContext:
     """Context for the MCP server."""
 
-    session: AsyncSession
-    session_factory: Callable[[], AsyncSession]
-    app_context: AppContext
+    server_factory: ServerFactory
 
 
 @asynccontextmanager
@@ -55,16 +47,12 @@ async def mcp_lifespan(_: FastMCP) -> AsyncIterator[MCPContext]:
     Since they don't provide a good way to handle global state, we must use a
     global variable to store the database connection.
     """
-    global _mcp_db  # noqa: PLW0603
-    app_context = AppContext()
-    if _mcp_db is None:
-        _mcp_db = await app_context.get_db()
-    async with _mcp_db.session_factory() as session:
-        yield MCPContext(
-            session=session,
-            app_context=app_context,
-            session_factory=_mcp_db.session_factory,
-        )
+    global _mcp_server_factory  # noqa: PLW0603
+    if _mcp_server_factory is None:
+        app_context = AppContext()
+        db = await app_context.get_db()
+        _mcp_server_factory = ServerFactory(app_context, db.session_factory)
+    yield MCPContext(_mcp_server_factory)
 
 
 def create_mcp_server(name: str, instructions: str | None = None) -> FastMCP:
@@ -180,13 +168,7 @@ def register_mcp_tools(mcp_server: FastMCP) -> None:
         mcp_context: MCPContext = ctx.request_context.lifespan_context
 
         # Use the unified application service
-        service = create_code_indexing_application_service(
-            app_context=mcp_context.app_context,
-            session_factory=mcp_context.session_factory,
-            operation=create_server_operation(
-                create_task_status_repository(mcp_context.session_factory)
-            ),
-        )
+        service = mcp_context.server_factory.code_search_application_service()
 
         log.debug("Searching for snippets")
 

kodit/migrations/env.py

@@ -41,6 +41,7 @@ def run_migrations_offline() -> None:
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        render_as_batch=True,
     )
 
     with context.begin_transaction():

kodit/migrations/versions/04b80f802e0c_foreign_key_review.py

@@ -0,0 +1,100 @@
+# ruff: noqa
+"""foreign key review
+
+Revision ID: 04b80f802e0c
+Revises: 7f15f878c3a1
+Create Date: 2025-09-22 11:21:43.432880
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "04b80f802e0c"
+down_revision: Union[str, None] = "7f15f878c3a1"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # SQLite doesn't support complex constraint alterations, so we'll drop and recreate tables
+
+    # Drop and recreate commit_indexes table with commit_sha as primary key
+    op.drop_table("commit_indexes")
+    op.create_table(
+        "commit_indexes",
+        sa.Column("commit_sha", sa.String(64), nullable=False),
+        sa.Column("status", sa.String(255), nullable=False),
+        sa.Column("indexed_at", sa.DateTime(), nullable=True),
+        sa.Column("error_message", sa.UnicodeText(), nullable=True),
+        sa.Column("files_processed", sa.Integer(), nullable=False, default=0),
+        sa.Column("processing_time_seconds", sa.Float(), nullable=False, default=0.0),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.PrimaryKeyConstraint("commit_sha", name="pk_commit_indexes"),
+    )
+    op.create_index("ix_commit_indexes_status", "commit_indexes", ["status"])
+
+    # Drop and recreate git_tracking_branches table with proper constraints
+    op.drop_table("git_tracking_branches")
+    op.create_table(
+        "git_tracking_branches",
+        sa.Column("repo_id", sa.Integer(), nullable=False),
+        sa.Column("name", sa.String(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_name", "git_tracking_branches", ["name"])
+    op.create_index(
+        "ix_git_tracking_branches_repo_id", "git_tracking_branches", ["repo_id"]
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # Recreate the original tables
+
+    # Recreate commit_indexes table with id-based primary key (original structure)
+    op.drop_table("commit_indexes")
+    op.create_table(
+        "commit_indexes",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("commit_sha", sa.String(64), nullable=False),
+        sa.Column("status", sa.String(255), nullable=False),
+        sa.Column("indexed_at", sa.DateTime(), nullable=True),
+        sa.Column("error_message", sa.UnicodeText(), nullable=True),
+        sa.Column("files_processed", sa.Integer(), nullable=False, default=0),
+        sa.Column("processing_time_seconds", sa.Float(), nullable=False, default=0.0),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index("ix_commit_indexes_status", "commit_indexes", ["status"])
+
+    # Recreate git_tracking_branches table with original structure
+    op.drop_table("git_tracking_branches")
+    op.create_table(
+        "git_tracking_branches",
+        sa.Column("repo_id", sa.Integer(), nullable=False),
+        sa.Column("name", sa.String(255), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["repo_id", "name"], ["git_branches.repo_id", "git_branches.name"]
+        ),
+        sa.PrimaryKeyConstraint("repo_id"),
+        sa.UniqueConstraint("repo_id", "name", name="uix_repo_tracking_branch"),
+    )
+    op.create_index("ix_git_tracking_branches_name", "git_tracking_branches", ["name"])
+    op.create_index(
+        "ix_git_tracking_branches_repo_id", "git_tracking_branches", ["repo_id"]
+    )

kodit/migrations/versions/19f8c7faf8b9_add_generic_enrichment_type.py

@@ -0,0 +1,260 @@
+# ruff: noqa
+"""add generic enrichment type
+
+Revision ID: 19f8c7faf8b9
+Revises: f9e5ef5e688f
+Create Date: 2025-09-29 21:38:19.093821
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from kodit.domain.enrichments.development.development import ENRICHMENT_TYPE_DEVELOPMENT
+from kodit.domain.enrichments.development.snippet.snippet import (
+    ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
+)
+
+
+# revision identifiers, used by Alembic.
+revision: str = "19f8c7faf8b9"
+down_revision: Union[str, None] = "f9e5ef5e688f"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "enrichments_v2",
+        sa.Column("content", sa.UnicodeText(), nullable=False),
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("type", sa.String(length=255), nullable=False),
+        sa.Column("subtype", sa.String(length=255), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        "idx_type_subtype", "enrichments_v2", ["type", "subtype"], unique=False
+    )
+    op.create_index(
+        op.f("ix_enrichments_v2_type"), "enrichments_v2", ["type"], unique=False
+    )
+    op.create_index(
+        op.f("ix_enrichments_v2_subtype"), "enrichments_v2", ["subtype"], unique=False
+    )
+    op.create_table(
+        "enrichment_associations",
+        sa.Column("enrichment_id", sa.Integer(), nullable=False),
+        sa.Column("entity_type", sa.String(length=50), nullable=False),
+        sa.Column("entity_id", sa.String(length=255), nullable=False),
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+        ),
+        sa.ForeignKeyConstraint(
+            ["enrichment_id"], ["enrichments_v2.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint(
+            "entity_type", "entity_id", "enrichment_id", name="uix_entity_enrichment"
+        ),
+        sqlite_autoincrement=True,
+    )
+    op.create_index(
+        "idx_entity_lookup",
+        "enrichment_associations",
+        ["entity_type", "entity_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_enrichment_associations_enrichment_id"),
+        "enrichment_associations",
+        ["enrichment_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_enrichment_associations_entity_id"),
+        "enrichment_associations",
+        ["entity_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_enrichment_associations_entity_type"),
+        "enrichment_associations",
+        ["entity_type"],
+        unique=False,
+    )
+
+    # Migrate existing enrichments from old table to new structure
+    # All old enrichments are snippet enrichments
+    connection = op.get_bind()
+
+    # Check if old enrichments table exists
+    inspector = sa.inspect(connection)
+    if "enrichments" in inspector.get_table_names():
+        # Copy enrichments to new table
+        old_enrichments = connection.execute(
+            sa.text(
+                "SELECT id, content, created_at, updated_at, snippet_sha FROM enrichments"
+            )
+        ).fetchall()
+
+        for old_enrichment in old_enrichments:
+            # Check if database supports RETURNING (PostgreSQL) or use lastrowid (SQLite)
+            dialect_name = connection.dialect.name
+
+            if dialect_name == "postgresql":
+                result = connection.execute(
+                    sa.text(
+                        "INSERT INTO enrichments_v2 (type, subtype, content, created_at, updated_at) "
+                        "VALUES (:type, :subtype, :content, :created_at, :updated_at) "
+                        "RETURNING id"
+                    ),
+                    {
+                        "type": ENRICHMENT_TYPE_DEVELOPMENT,
+                        "subtype": ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
+                        "content": old_enrichment[1],
+                        "created_at": old_enrichment[2],
+                        "updated_at": old_enrichment[3],
+                    },
+                )
+                row = result.fetchone()
+                if row is None:
+                    raise RuntimeError("Failed to insert enrichment")
+                new_enrichment_id = row[0]
+            else:
+                # SQLite and other databases
+                result = connection.execute(
+                    sa.text(
+                        "INSERT INTO enrichments_v2 (type, subtype, content, created_at, updated_at) "
+                        "VALUES (:type, :subtype, :content, :created_at, :updated_at)"
+                    ),
+                    {
+                        "type": ENRICHMENT_TYPE_DEVELOPMENT,
+                        "subtype": ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
+                        "content": old_enrichment[1],
+                        "created_at": old_enrichment[2],
+                        "updated_at": old_enrichment[3],
+                    },
+                )
+                new_enrichment_id = result.lastrowid
+
+            # Insert association (snippet_v2 is the entity type for snippets)
+            connection.execute(
+                sa.text(
+                    "INSERT INTO enrichment_associations "
+                    "(enrichment_id, entity_type, entity_id, created_at, updated_at) "
+                    "VALUES (:enrichment_id, :entity_type, :entity_id, :created_at, :updated_at)"
+                ),
+                {
+                    "enrichment_id": new_enrichment_id,
+                    "entity_type": "snippet_v2",
+                    "entity_id": old_enrichment[4],  # snippet_sha
+                    "created_at": old_enrichment[2],
+                    "updated_at": old_enrichment[3],
+                },
+            )
+
+        # Drop old enrichments table
+        op.drop_table("enrichments")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    # Recreate old enrichments table
+    op.create_table(
+        "enrichments",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+        ),
+        sa.Column("snippet_sha", sa.String(length=64), nullable=False),
+        sa.Column("type", sa.String(length=50), nullable=False),
+        sa.Column("content", sa.UnicodeText(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+        sa.ForeignKeyConstraint(["snippet_sha"], ["snippets_v2.sha"]),
+    )
+    op.create_index(
+        op.f("ix_enrichments_snippet_sha"), "enrichments", ["snippet_sha"], unique=False
+    )
+    op.create_index(op.f("ix_enrichments_type"), "enrichments", ["type"], unique=False)
+
+    # Migrate data back from new structure to old table
+    connection = op.get_bind()
+
+    # Get enrichments with their associations
+    enrichments_with_associations = connection.execute(
+        sa.text(
+            "SELECT e.id, e.content, e.created_at, e.updated_at, "
+            "a.entity_id, a.entity_type "
+            "FROM enrichments_v2 e "
+            "JOIN enrichment_associations a ON e.id = a.enrichment_id "
+            "WHERE a.entity_type = 'snippet_v2'"
+        )
+    ).fetchall()
+
+    for enrichment in enrichments_with_associations:
+        connection.execute(
+            sa.text(
+                "INSERT INTO enrichments "
+                "(content, created_at, updated_at, snippet_sha, type) "
+                "VALUES (:content, :created_at, :updated_at, :snippet_sha, :type)"
+            ),
+            {
+                "content": enrichment[1],
+                "created_at": enrichment[2],
+                "updated_at": enrichment[3],
+                "snippet_sha": enrichment[4],
+                "type": "summarization",  # Default to summarization type
+            },
+        )
+
+    op.drop_index(
+        op.f("ix_enrichment_associations_entity_type"),
+        table_name="enrichment_associations",
+    )
+    op.drop_index(
+        op.f("ix_enrichment_associations_entity_id"),
+        table_name="enrichment_associations",
+    )
+    op.drop_index(
+        op.f("ix_enrichment_associations_enrichment_id"),
+        table_name="enrichment_associations",
+    )
+    op.drop_index("idx_entity_lookup", table_name="enrichment_associations")
+    op.drop_table("enrichment_associations")
+    op.drop_index(op.f("ix_enrichments_v2_subtype"), table_name="enrichments_v2")
+    op.drop_index(op.f("ix_enrichments_v2_type"), table_name="enrichments_v2")
+    op.drop_index("idx_type_subtype", table_name="enrichments_v2")
+    op.drop_table("enrichments_v2")
+    # ### end Alembic commands ###