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

kodit/infrastructure/sqlalchemy/git_branch_repository.py

@@ -0,0 +1,274 @@
+"""SQLAlchemy implementation of GitBranchRepository."""
+
+from collections.abc import Callable
+
+from sqlalchemy import delete, func, insert, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities.git import GitBranch, GitCommit
+from kodit.domain.protocols import GitBranchRepository
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
+
+
+def create_git_branch_repository(
+    session_factory: Callable[[], AsyncSession],
+) -> GitBranchRepository:
+    """Create a git branch repository."""
+    return SqlAlchemyGitBranchRepository(session_factory=session_factory)
+
+
+class SqlAlchemyGitBranchRepository(GitBranchRepository):
+    """SQLAlchemy implementation of GitBranchRepository."""
+
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
+        """Initialize the repository."""
+        self.session_factory = session_factory
+
+    async def get_by_name(self, branch_name: str, repo_id: int) -> GitBranch:
+        """Get a branch by name and repository ID."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Get the branch
+            stmt = select(db_entities.GitBranch).where(
+                db_entities.GitBranch.name == branch_name,
+                db_entities.GitBranch.repo_id == repo_id,
+            )
+            db_branch = await session.scalar(stmt)
+            if not db_branch:
+                raise ValueError(f"Branch {branch_name} not found in repo {repo_id}")
+
+            # Get the head commit
+            commit_stmt = select(db_entities.GitCommit).where(
+                db_entities.GitCommit.commit_sha == db_branch.head_commit_sha
+            )
+            db_commit = await session.scalar(commit_stmt)
+            if not db_commit:
+                raise ValueError(f"Head commit {db_branch.head_commit_sha} not found")
+
+            # Get files for the head commit
+            files_stmt = select(db_entities.GitCommitFile).where(
+                db_entities.GitCommitFile.commit_sha == db_branch.head_commit_sha
+            )
+            db_files = (await session.scalars(files_stmt)).all()
+
+            from kodit.domain.entities.git import GitFile
+
+            domain_files = []
+            for db_file in db_files:
+                domain_file = GitFile(
+                    blob_sha=db_file.blob_sha,
+                    path=db_file.path,
+                    mime_type=db_file.mime_type,
+                    size=db_file.size,
+                    extension=db_file.extension,
+                    created_at=db_file.created_at,
+                )
+                domain_files.append(domain_file)
+
+            head_commit = GitCommit(
+                commit_sha=db_commit.commit_sha,
+                date=db_commit.date,
+                message=db_commit.message,
+                parent_commit_sha=db_commit.parent_commit_sha,
+                files=domain_files,
+                author=db_commit.author,
+                created_at=db_commit.created_at,
+                updated_at=db_commit.updated_at,
+            )
+
+            return GitBranch(
+                repo_id=db_branch.repo_id,
+                name=db_branch.name,
+                head_commit=head_commit,
+                created_at=db_branch.created_at,
+                updated_at=db_branch.updated_at,
+            )
+
+    async def get_by_repo_id(self, repo_id: int) -> list[GitBranch]:
+        """Get all branches for a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Get all branches for the repo
+            branches_stmt = select(db_entities.GitBranch).where(
+                db_entities.GitBranch.repo_id == repo_id
+            )
+            db_branches = (await session.scalars(branches_stmt)).all()
+
+            if not db_branches:
+                return []
+
+            commit_shas = [branch.head_commit_sha for branch in db_branches]
+
+            # 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
+
+            files_by_commit: dict[str, list[GitFile]] = {}
+            for db_file in db_files:
+                if db_file.commit_sha not in files_by_commit:
+                    files_by_commit[db_file.commit_sha] = []
+
+                domain_file = GitFile(
+                    blob_sha=db_file.blob_sha,
+                    path=db_file.path,
+                    mime_type=db_file.mime_type,
+                    size=db_file.size,
+                    extension=db_file.extension,
+                    created_at=db_file.created_at,
+                )
+                files_by_commit[db_file.commit_sha].append(domain_file)
+
+            # Create commit lookup
+            commits_by_sha = {commit.commit_sha: commit for commit in db_commits}
+
+            # Create domain branches
+            domain_branches = []
+            for db_branch in db_branches:
+                db_commit = commits_by_sha.get(db_branch.head_commit_sha)
+                if not db_commit:
+                    continue
+
+                commit_files = files_by_commit.get(db_branch.head_commit_sha, [])
+                head_commit = GitCommit(
+                    commit_sha=db_commit.commit_sha,
+                    date=db_commit.date,
+                    message=db_commit.message,
+                    parent_commit_sha=db_commit.parent_commit_sha,
+                    files=commit_files,
+                    author=db_commit.author,
+                    created_at=db_commit.created_at,
+                    updated_at=db_commit.updated_at,
+                )
+
+                domain_branch = GitBranch(
+                    repo_id=db_branch.repo_id,
+                    name=db_branch.name,
+                    head_commit=head_commit,
+                    created_at=db_branch.created_at,
+                    updated_at=db_branch.updated_at,
+                )
+                domain_branches.append(domain_branch)
+
+            return domain_branches
+
+    async def save(self, branch: GitBranch, repo_id: int) -> GitBranch:
+        """Save a branch to a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Set repo_id on the branch
+            branch.repo_id = repo_id
+
+            # Check if branch already exists
+            existing_branch = await session.get(
+                db_entities.GitBranch, (repo_id, branch.name)
+            )
+
+            if existing_branch:
+                # Update existing branch
+                existing_branch.head_commit_sha = branch.head_commit.commit_sha
+                if branch.updated_at:
+                    existing_branch.updated_at = branch.updated_at
+            else:
+                # Create new branch
+                db_branch = db_entities.GitBranch(
+                    repo_id=repo_id,
+                    name=branch.name,
+                    head_commit_sha=branch.head_commit.commit_sha,
+                )
+                session.add(db_branch)
+
+            return branch
+
+    async def save_bulk(self, branches: list[GitBranch], repo_id: int) -> None:
+        """Bulk save branches to a repository."""
+        if not branches:
+            return
+
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            [(repo_id, branch.name) for branch in branches]
+
+            # Get existing branches in bulk
+            existing_branches_stmt = select(db_entities.GitBranch).where(
+                db_entities.GitBranch.repo_id == repo_id,
+                db_entities.GitBranch.name.in_([branch.name for branch in branches]),
+            )
+            existing_branches = (await session.scalars(existing_branches_stmt)).all()
+            existing_branch_names = {branch.name for branch in existing_branches}
+
+            # Update existing branches
+            for existing_branch in existing_branches:
+                for branch in branches:
+                    if (
+                        branch.name == existing_branch.name
+                        and existing_branch.head_commit_sha
+                        != branch.head_commit.commit_sha
+                    ):
+                        existing_branch.head_commit_sha = branch.head_commit.commit_sha
+                        break
+
+            # Prepare new branches for bulk insert
+            new_branches_data = [
+                {
+                    "repo_id": repo_id,
+                    "name": branch.name,
+                    "head_commit_sha": branch.head_commit.commit_sha,
+                }
+                for branch in branches
+                if branch.name not in existing_branch_names
+            ]
+
+            # Bulk insert new branches in chunks to avoid parameter limits
+            if new_branches_data:
+                chunk_size = 1000  # Conservative chunk size for parameter limits
+                for i in range(0, len(new_branches_data), chunk_size):
+                    chunk = new_branches_data[i : i + chunk_size]
+                    stmt = insert(db_entities.GitBranch).values(chunk)
+                    await session.execute(stmt)
+
+    async def exists(self, branch_name: str, repo_id: int) -> bool:
+        """Check if a branch exists."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(db_entities.GitBranch.name).where(
+                db_entities.GitBranch.name == branch_name,
+                db_entities.GitBranch.repo_id == repo_id,
+            )
+            result = await session.scalar(stmt)
+            return result is not None
+
+    async def delete_by_repo_id(self, repo_id: int) -> None:
+        """Delete all branches for a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Delete branches
+            del_branches_stmt = delete(db_entities.GitBranch).where(
+                db_entities.GitBranch.repo_id == repo_id
+            )
+            await session.execute(del_branches_stmt)
+
+    async def count_by_repo_id(self, repo_id: int) -> int:
+        """Count the number of branches for a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(func.count()).select_from(db_entities.GitBranch).where(
+                db_entities.GitBranch.repo_id == repo_id
+            )
+            result = await session.scalar(stmt)
+            return result or 0

kodit/infrastructure/sqlalchemy/git_commit_repository.py

@@ -0,0 +1,346 @@
+"""SQLAlchemy implementation of GitCommitRepository."""
+
+from collections.abc import Callable
+
+from sqlalchemy import delete, func, insert, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities.git import GitCommit, GitFile
+from kodit.domain.protocols import GitCommitRepository
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
+
+
+def create_git_commit_repository(
+    session_factory: Callable[[], AsyncSession],
+) -> GitCommitRepository:
+    """Create a git commit repository."""
+    return SqlAlchemyGitCommitRepository(session_factory=session_factory)
+
+
+class SqlAlchemyGitCommitRepository(GitCommitRepository):
+    """SQLAlchemy implementation of GitCommitRepository."""
+
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
+        """Initialize the repository."""
+        self.session_factory = session_factory
+
+    async def get_by_sha(self, commit_sha: str) -> GitCommit:
+        """Get a commit by its SHA."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Get the commit
+            stmt = select(db_entities.GitCommit).where(
+                db_entities.GitCommit.commit_sha == commit_sha
+            )
+            db_commit = await session.scalar(stmt)
+            if not db_commit:
+                raise ValueError(f"Commit with SHA {commit_sha} not found")
+
+            # Get associated files
+            files_stmt = select(db_entities.GitCommitFile).where(
+                db_entities.GitCommitFile.commit_sha == commit_sha
+            )
+            db_files = (await session.scalars(files_stmt)).all()
+
+            domain_files = []
+            for db_file in db_files:
+                domain_file = GitFile(
+                    blob_sha=db_file.blob_sha,
+                    path=db_file.path,
+                    mime_type=db_file.mime_type,
+                    size=db_file.size,
+                    extension=db_file.extension,
+                    created_at=db_file.created_at,
+                )
+                domain_files.append(domain_file)
+
+            return GitCommit(
+                commit_sha=db_commit.commit_sha,
+                date=db_commit.date,
+                message=db_commit.message,
+                parent_commit_sha=db_commit.parent_commit_sha,
+                files=domain_files,
+                author=db_commit.author,
+            )
+
+    async def get_by_repo_id(self, repo_id: int) -> list[GitCommit]:
+        """Get all commits for a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Get all commits for the repo
+            commits_stmt = select(db_entities.GitCommit).where(
+                db_entities.GitCommit.repo_id == repo_id
+            )
+            db_commits = (await session.scalars(commits_stmt)).all()
+
+            if not db_commits:
+                return []
+
+            commit_shas = [commit.commit_sha for commit in db_commits]
+
+            # 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]] = {}
+            for db_file in db_files:
+                if db_file.commit_sha not in files_by_commit:
+                    files_by_commit[db_file.commit_sha] = []
+
+                domain_file = GitFile(
+                    blob_sha=db_file.blob_sha,
+                    path=db_file.path,
+                    mime_type=db_file.mime_type,
+                    size=db_file.size,
+                    extension=db_file.extension,
+                    created_at=db_file.created_at,
+                )
+                files_by_commit[db_file.commit_sha].append(domain_file)
+
+            # Create domain commits
+            domain_commits = []
+            for db_commit in db_commits:
+                commit_files = files_by_commit.get(db_commit.commit_sha, [])
+                domain_commit = GitCommit(
+                    commit_sha=db_commit.commit_sha,
+                    date=db_commit.date,
+                    message=db_commit.message,
+                    parent_commit_sha=db_commit.parent_commit_sha,
+                    files=commit_files,
+                    author=db_commit.author,
+                )
+                domain_commits.append(domain_commit)
+
+            return domain_commits
+
+    async def save(self, commit: GitCommit, repo_id: int) -> GitCommit:
+        """Save a commit to a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Check if commit already exists
+            existing_commit = await session.get(
+                db_entities.GitCommit, commit.commit_sha
+            )
+
+            if not existing_commit:
+                # Create new commit
+                db_commit = db_entities.GitCommit(
+                    commit_sha=commit.commit_sha,
+                    repo_id=repo_id,
+                    date=commit.date,
+                    message=commit.message,
+                    parent_commit_sha=commit.parent_commit_sha,
+                    author=commit.author,
+                )
+                session.add(db_commit)
+                await session.flush()
+
+                # Save associated files
+                await self._save_commit_files(session, commit)
+
+            return commit
+
+    async def save_bulk(self, commits: list[GitCommit], repo_id: int) -> None:
+        """Bulk save commits to a repository."""
+        if not commits:
+            return
+
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            commit_shas = [commit.commit_sha for commit in commits]
+
+            # 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 = []
+            new_commits_objects = []
+            for commit in commits:
+                if commit.commit_sha not in existing_commit_shas:
+                    new_commits_data.append({
+                        "commit_sha": commit.commit_sha,
+                        "repo_id": repo_id,
+                        "date": commit.date,
+                        "message": commit.message,
+                        "parent_commit_sha": commit.parent_commit_sha,
+                        "author": commit.author,
+                    })
+                    new_commits_objects.append(commit)
+
+            # Bulk insert new commits in chunks to avoid parameter limits
+            if new_commits_data:
+                chunk_size = 1000  # Conservative chunk size for parameter limits
+                for i in range(0, len(new_commits_data), chunk_size):
+                    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
+                await self._save_commits_files_bulk(session, new_commits_objects)
+
+    async def exists(self, commit_sha: str) -> bool:
+        """Check if a commit exists."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(db_entities.GitCommit.commit_sha).where(
+                db_entities.GitCommit.commit_sha == commit_sha
+            )
+            result = await session.scalar(stmt)
+            return result is not None
+
+    async def delete_by_repo_id(self, repo_id: int) -> None:
+        """Delete all commits for a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Get all commit SHAs for this repo
+            commit_shas_stmt = select(db_entities.GitCommit.commit_sha).where(
+                db_entities.GitCommit.repo_id == repo_id
+            )
+            commit_shas = (await session.scalars(commit_shas_stmt)).all()
+
+            # Delete snippet file associations first (they reference commit files)
+            for commit_sha in commit_shas:
+                del_snippet_files_stmt = delete(db_entities.SnippetV2File).where(
+                    db_entities.SnippetV2File.commit_sha == commit_sha
+                )
+                await session.execute(del_snippet_files_stmt)
+
+            # Delete commit files second (foreign key constraint)
+            for commit_sha in commit_shas:
+                del_files_stmt = delete(db_entities.GitCommitFile).where(
+                    db_entities.GitCommitFile.commit_sha == commit_sha
+                )
+                await session.execute(del_files_stmt)
+
+            # Delete commits
+            del_commits_stmt = delete(db_entities.GitCommit).where(
+                db_entities.GitCommit.repo_id == repo_id
+            )
+            await session.execute(del_commits_stmt)
+
+    async def count_by_repo_id(self, repo_id: int) -> int:
+        """Count the number of commits for a repository."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(func.count()).select_from(db_entities.GitCommit).where(
+                db_entities.GitCommit.repo_id == repo_id
+            )
+            result = await session.scalar(stmt)
+            return result or 0
+
+    async def _save_commit_files(
+        self, session: AsyncSession, commit: GitCommit
+    ) -> None:
+        """Save files for a single commit."""
+        if not commit.files:
+            return
+
+        # Check which files already exist
+        existing_files_stmt = select(
+            db_entities.GitCommitFile.commit_sha,
+            db_entities.GitCommitFile.path
+        ).where(
+            db_entities.GitCommitFile.commit_sha == commit.commit_sha
+        )
+        existing_file_keys = set(await session.execute(existing_files_stmt))
+
+        # Prepare new files for insert
+        new_files = []
+        for file in commit.files:
+            file_key = (commit.commit_sha, file.path)
+            if file_key not in existing_file_keys:
+                new_files.append({
+                    "commit_sha": commit.commit_sha,
+                    "path": file.path,
+                    "blob_sha": file.blob_sha,
+                    "extension": file.extension,
+                    "mime_type": file.mime_type,
+                    "size": file.size,
+                    "created_at": file.created_at,
+                })
+
+        # Bulk insert new files in chunks to avoid parameter limits
+        if new_files:
+            chunk_size = 1000  # Conservative chunk size for parameter limits
+            for i in range(0, len(new_files), chunk_size):
+                chunk = new_files[i : i + chunk_size]
+                stmt = insert(db_entities.GitCommitFile).values(chunk)
+                await session.execute(stmt)
+
+    async def _save_commits_files_bulk(
+        self, session: AsyncSession, commits: list[GitCommit]
+    ) -> None:
+        """Bulk save files for multiple commits."""
+        all_file_identifiers = [
+            (commit.commit_sha, file.path)
+            for commit in commits
+            for file in commit.files
+        ]
+
+        if not all_file_identifiers:
+            return
+
+        # Get existing files in chunks to avoid SQL parameter limits
+        existing_file_keys = await self._get_existing_file_keys_bulk(
+            session, all_file_identifiers
+        )
+
+        # Prepare new files for bulk insert
+        new_files = []
+        for commit in commits:
+            for file in commit.files:
+                file_key = (commit.commit_sha, file.path)
+                if file_key not in existing_file_keys:
+                    new_files.append({
+                        "commit_sha": commit.commit_sha,
+                        "path": file.path,
+                        "blob_sha": file.blob_sha,
+                        "extension": file.extension,
+                        "mime_type": file.mime_type,
+                        "size": file.size,
+                        "created_at": file.created_at,
+                    })
+
+        # Bulk insert new files in chunks
+        if new_files:
+            chunk_size = 1000
+            for i in range(0, len(new_files), chunk_size):
+                chunk = new_files[i : i + chunk_size]
+                stmt = insert(db_entities.GitCommitFile).values(chunk)
+                await session.execute(stmt)
+
+    async def _get_existing_file_keys_bulk(
+        self, session: AsyncSession, file_identifiers: list[tuple[str, str]]
+    ) -> set[tuple[str, str]]:
+        """Get existing file keys in chunks to avoid SQL parameter limits."""
+        chunk_size = 1000
+        existing_file_keys = set()
+
+        for i in range(0, len(file_identifiers), chunk_size):
+            chunk = file_identifiers[i : i + chunk_size]
+            commit_shas = [item[0] for item in chunk]
+            paths = [item[1] for item in chunk]
+
+            existing_files_stmt = select(
+                db_entities.GitCommitFile.commit_sha, db_entities.GitCommitFile.path
+            ).where(
+                db_entities.GitCommitFile.commit_sha.in_(commit_shas),
+                db_entities.GitCommitFile.path.in_(paths),
+            )
+
+            chunk_existing = await session.execute(existing_files_stmt)
+            for commit_sha, path in chunk_existing:
+                existing_file_keys.add((commit_sha, path))
+
+        return existing_file_keys