kodit 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

kodit/infrastructure/sqlalchemy/entities.py

@@ -2,12 +2,14 @@
 
 from datetime import UTC, datetime
 from enum import Enum
+from pathlib import Path
 from typing import Any
 
-from git import Actor
 from sqlalchemy import (
     DateTime,
+    Float,
     ForeignKey,
+    ForeignKeyConstraint,
     Integer,
     String,
     TypeDecorator,
@@ -43,6 +45,25 @@ class TZDateTime(TypeDecorator):
         return value
 
 
+class PathType(TypeDecorator):
+    """Path type that stores Path objects as strings."""
+
+    impl = String
+    cache_ok = True
+
+    def process_bind_param(self, value: Any, dialect: Any) -> Any:  # noqa: ARG002
+        """Process bind param - convert Path to string."""
+        if value is not None:
+            return str(value)
+        return value
+
+    def process_result_value(self, value: Any, dialect: Any) -> Any:  # noqa: ARG002
+        """Process result value - convert string to Path."""
+        if value is not None:
+            return Path(value)
+        return value
+
+
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models."""
 
@@ -62,203 +83,479 @@ class CommonMixin:
     )
 
 
-class SourceType(Enum):
-    """The type of source."""
+class EmbeddingType(Enum):
+    """Embedding type."""
+
+    CODE = 1
+    TEXT = 2
+
+
+class Embedding(Base, CommonMixin):
+    """Embedding model."""
+
+    __tablename__ = "embeddings"
+
+    snippet_id: Mapped[str] = mapped_column(String(64), index=True)
+    type: Mapped[EmbeddingType] = mapped_column(
+        SQLAlchemyEnum(EmbeddingType), index=True
+    )
+    embedding: Mapped[list[float]] = mapped_column(JSON)
+
+
+class Task(Base, CommonMixin):
+    """Queued tasks."""
+
+    __tablename__ = "tasks"
+
+    # dedup_key is used to deduplicate items in the queue
+    dedup_key: Mapped[str] = mapped_column(String(255), index=True)
+    # type represents what the task is meant to achieve
+    type: Mapped[str] = mapped_column(String(255), index=True)
+    # payload contains the task-specific payload data
+    payload: Mapped[dict] = mapped_column(JSON)
+    # priority is used to determine the order of the items in the queue
+    priority: Mapped[int] = mapped_column(Integer)
+
+    def __init__(
+        self,
+        dedup_key: str,
+        type: str,  # noqa: A002
+        payload: dict,
+        priority: int,
+    ) -> None:
+        """Initialize the queue item."""
+        super().__init__()
+        self.dedup_key = dedup_key
+        self.type = type
+        self.payload = payload
+        self.priority = priority
+
+
+class TaskStatus(Base):
+    """Task status model."""
 
-    UNKNOWN = 0
-    FOLDER = 1
-    GIT = 2
+    __tablename__ = "task_status"
+    id: Mapped[str] = mapped_column(
+        String(255), primary_key=True, index=True, nullable=False
+    )
+    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),
+    )
+    operation: Mapped[str] = mapped_column(String(255), index=True, nullable=False)
+    trackable_id: Mapped[int | None] = mapped_column(Integer, index=True, nullable=True)
+    trackable_type: Mapped[str | None] = mapped_column(
+        String(255), index=True, nullable=True
+    )
+    parent: Mapped[str | None] = mapped_column(
+        ForeignKey("task_status.id"), index=True, nullable=True
+    )
+    message: Mapped[str] = mapped_column(UnicodeText, default="")
+    state: Mapped[str] = mapped_column(String(255), default="")
+    error: Mapped[str] = mapped_column(UnicodeText, default="")
+    total: Mapped[int] = mapped_column(Integer, default=0)
+    current: Mapped[int] = mapped_column(Integer, default=0)
+
+    def __init__(  # noqa: PLR0913
+        self,
+        id: str,  # noqa: A002
+        operation: str,
+        created_at: datetime,
+        updated_at: datetime,
+        trackable_id: int | None,
+        trackable_type: str | None,
+        parent: str | None,
+        state: str,
+        error: str | None,
+        total: int,
+        current: int,
+        message: str,
+    ) -> None:
+        """Initialize the task status."""
+        super().__init__()
+        self.id = id
+        self.operation = operation
+        self.created_at = created_at
+        self.updated_at = updated_at
+        self.trackable_id = trackable_id
+        self.trackable_type = trackable_type
+        self.parent = parent
+        self.state = state
+        self.error = error or ""
+        self.total = total
+        self.current = current
+        self.message = message or ""
 
 
-class Source(Base, CommonMixin):
-    """Base model for tracking code sources.
+# Git-related entities for new GitRepo domain
 
-    This model serves as the parent table for different types of sources.
-    It provides common fields and relationships for all source types.
 
-    Attributes:
-        id: The unique identifier for the source.
-        created_at: Timestamp when the source was created.
-        updated_at: Timestamp when the source was last updated.
-        cloned_uri: A URI to a copy of the source on the local filesystem.
-        uri: The URI of the source.
+class GitRepo(Base, CommonMixin):
+    """Git repository model."""
 
-    """
+    __tablename__ = "git_repos"
 
-    __tablename__ = "sources"
-    uri: Mapped[str] = mapped_column(String(1024), index=True, unique=True)
-    cloned_path: Mapped[str] = mapped_column(String(1024), index=True)
-    type: Mapped[SourceType] = mapped_column(
-        SQLAlchemyEnum(SourceType), default=SourceType.UNKNOWN, index=True
+    sanitized_remote_uri: Mapped[str] = mapped_column(
+        String(1024), index=True, unique=True
     )
+    remote_uri: Mapped[str] = mapped_column(String(1024))
+    cloned_path: Mapped[Path | None] = mapped_column(PathType(1024), nullable=True)
+    last_scanned_at: Mapped[datetime | None] = mapped_column(TZDateTime, nullable=True)
+    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)
 
-    def __init__(self, uri: str, cloned_path: str, source_type: SourceType) -> None:
-        """Initialize a new Source instance for typing purposes."""
+    def __init__(  # noqa: PLR0913
+        self,
+        sanitized_remote_uri: str,
+        remote_uri: str,
+        cloned_path: Path | None,
+        last_scanned_at: datetime | None = None,
+        num_commits: int = 0,
+        num_branches: int = 0,
+        num_tags: int = 0,
+    ) -> None:
+        """Initialize Git repository."""
         super().__init__()
-        self.uri = uri
+        self.sanitized_remote_uri = sanitized_remote_uri
+        self.remote_uri = remote_uri
         self.cloned_path = cloned_path
-        self.type = source_type
+        self.last_scanned_at = last_scanned_at
+        self.num_commits = num_commits
+        self.num_branches = num_branches
+        self.num_tags = num_tags
 
 
-class Author(Base, CommonMixin):
-    """Author model."""
+class GitCommit(Base):
+    """Git commit model."""
 
-    __tablename__ = "authors"
+    __tablename__ = "git_commits"
 
-    __table_args__ = (UniqueConstraint("name", "email", name="uix_author"),)
+    commit_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),
+    )
+    repo_id: Mapped[int] = mapped_column(ForeignKey("git_repos.id"), index=True)
+    date: Mapped[datetime] = mapped_column(TZDateTime)
+    message: Mapped[str] = mapped_column(UnicodeText)
+    parent_commit_sha: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    author: Mapped[str] = mapped_column(String(255), index=True)
+
+    def __init__(  # noqa: PLR0913
+        self,
+        commit_sha: str,
+        repo_id: int,
+        date: datetime,
+        message: str,
+        parent_commit_sha: str | None,
+        author: str,
+    ) -> None:
+        """Initialize Git commit."""
+        super().__init__()
+        self.commit_sha = commit_sha
+        self.repo_id = repo_id
+        self.date = date
+        self.message = message
+        self.parent_commit_sha = parent_commit_sha
+        self.author = author
+
+
+class GitBranch(Base):
+    """Git branch model."""
 
-    name: Mapped[str] = mapped_column(String(255), index=True)
-    email: Mapped[str] = mapped_column(String(255), index=True)
+    __tablename__ = "git_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),
+    )
+    head_commit_sha: Mapped[str] = mapped_column(ForeignKey("git_commits.commit_sha"))
 
-    @staticmethod
-    def from_actor(actor: Actor) -> "Author":
-        """Create an Author from an Actor."""
-        return Author(name=actor.name, email=actor.email)
+    __table_args__ = (UniqueConstraint("repo_id", "name", name="uix_repo_branch"),)
 
+    def __init__(self, repo_id: int, name: str, head_commit_sha: str) -> None:
+        """Initialize Git branch."""
+        super().__init__()
+        self.repo_id = repo_id
+        self.name = name
+        self.head_commit_sha = head_commit_sha
 
-class AuthorFileMapping(Base, CommonMixin):
-    """Author file mapping model."""
 
-    __tablename__ = "author_file_mappings"
+class GitTrackingBranch(Base):
+    """Git tracking branch model."""
 
-    __table_args__ = (
-        UniqueConstraint("author_id", "file_id", name="uix_author_file_mapping"),
+    __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),
     )
 
-    author_id: Mapped[int] = mapped_column(ForeignKey("authors.id"), index=True)
-    file_id: Mapped[int] = mapped_column(ForeignKey("files.id"), index=True)
+    def __init__(self, repo_id: int, name: str) -> None:
+        """Initialize Git tracking branch."""
+        super().__init__()
+        self.repo_id = repo_id
+        self.name = name
 
 
-class File(Base, CommonMixin):
-    """File model."""
+class GitTag(Base):
+    """Git tag model."""
 
-    __tablename__ = "files"
+    __tablename__ = "git_tags"
+    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),
+    )
+    target_commit_sha: Mapped[str] = mapped_column(
+        ForeignKey("git_commits.commit_sha"), index=True
+    )
+
+    __table_args__ = (UniqueConstraint("repo_id", "name", name="uix_repo_tag"),)
+
+    def __init__(self, repo_id: int, name: str, target_commit_sha: str) -> None:
+        """Initialize Git tag."""
+        super().__init__()
+        self.repo_id = repo_id
+        self.name = name
+        self.target_commit_sha = target_commit_sha
+
+
+class GitCommitFile(Base):
+    """Files in a git commit (tree entries)."""
+
+    __tablename__ = "git_commit_files"
 
-    source_id: Mapped[int] = mapped_column(ForeignKey("sources.id"))
-    mime_type: Mapped[str] = mapped_column(String(255), default="", index=True)
-    uri: Mapped[str] = mapped_column(String(1024), default="", index=True)
-    cloned_path: Mapped[str] = mapped_column(String(1024), index=True)
-    sha256: Mapped[str] = mapped_column(String(64), default="", index=True)
-    size_bytes: Mapped[int] = mapped_column(Integer, default=0)
-    extension: Mapped[str] = mapped_column(String(255), default="", index=True)
-    file_processing_status: Mapped[int] = mapped_column(Integer, default=0)
+    commit_sha: Mapped[str] = mapped_column(
+        ForeignKey("git_commits.commit_sha"), primary_key=True
+    )
+    path: Mapped[str] = mapped_column(String(1024), primary_key=True)
+    blob_sha: Mapped[str] = mapped_column(String(64), index=True)
+    mime_type: Mapped[str] = mapped_column(String(255), index=True)
+    extension: Mapped[str] = mapped_column(String(255), index=True)
+    size: Mapped[int] = mapped_column(Integer)
+    created_at: Mapped[datetime] = mapped_column(TZDateTime, nullable=False)
+
+    __table_args__ = (UniqueConstraint("commit_sha", "path", name="uix_commit_file"),)
 
     def __init__(  # noqa: PLR0913
         self,
-        created_at: datetime,
-        updated_at: datetime,
-        source_id: int,
+        commit_sha: str,
+        path: str,
+        blob_sha: str,
         mime_type: str,
-        uri: str,
-        cloned_path: str,
-        sha256: str,
-        size_bytes: int,
         extension: str,
-        file_processing_status: int,
+        size: int,
+        created_at: datetime,
     ) -> None:
-        """Initialize a new File instance for typing purposes."""
+        """Initialize Git commit file."""
         super().__init__()
-        self.created_at = created_at
-        self.updated_at = updated_at
-        self.source_id = source_id
+        self.commit_sha = commit_sha
+        self.path = path
+        self.blob_sha = blob_sha
         self.mime_type = mime_type
-        self.uri = uri
-        self.cloned_path = cloned_path
-        self.sha256 = sha256
-        self.size_bytes = size_bytes
+        self.size = size
+        self.created_at = created_at
         self.extension = extension
-        self.file_processing_status = file_processing_status
 
 
-class EmbeddingType(Enum):
-    """Embedding type."""
+class SnippetV2(Base):
+    """SnippetV2 model for commit-based snippets."""
 
-    CODE = 1
-    TEXT = 2
+    __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)
 
-class Embedding(Base, CommonMixin):
-    """Embedding model."""
+    def __init__(
+        self,
+        sha: str,
+        content: str,
+        extension: str,
+    ) -> None:
+        """Initialize snippet."""
+        super().__init__()
+        self.sha = sha
+        self.content = content
+        self.extension = extension
 
-    __tablename__ = "embeddings"
 
-    snippet_id: Mapped[int] = mapped_column(ForeignKey("snippets.id"), index=True)
-    type: Mapped[EmbeddingType] = mapped_column(
-        SQLAlchemyEnum(EmbeddingType), index=True
+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",
+        ),
     )
-    embedding: Mapped[list[float]] = mapped_column(JSON)
+
+    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 Index(Base, CommonMixin):
-    """Index model."""
+class CommitSnippetV2(Base):
+    """Association table for commits and snippets v2."""
 
-    __tablename__ = "indexes"
+    __tablename__ = "commit_snippets_v2"
 
-    source_id: Mapped[int] = mapped_column(
-        ForeignKey("sources.id"), unique=True, index=True
+    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, source_id: int) -> None:
-        """Initialize the index."""
+    def __init__(self, commit_sha: str, snippet_sha: str) -> None:
+        """Initialize commit snippet association."""
         super().__init__()
-        self.source_id = source_id
+        self.commit_sha = commit_sha
+        self.snippet_sha = snippet_sha
 
 
-class Snippet(Base, CommonMixin):
-    """Snippet model."""
+# Enrichment model for SnippetV2
 
-    __tablename__ = "snippets"
 
-    file_id: Mapped[int] = mapped_column(ForeignKey("files.id"), index=True)
-    index_id: Mapped[int] = mapped_column(ForeignKey("indexes.id"), index=True)
-    content: Mapped[str] = mapped_column(UnicodeText, default="")
-    summary: Mapped[str] = mapped_column(UnicodeText, default="")
+class EnrichmentType(Enum):
+    """Enrichment type enum."""
+
+    UNKNOWN = "unknown"
+    SUMMARIZATION = "summarization"
+
+
+class Enrichment(Base, CommonMixin):
+    """Enrichment model for snippet enrichments."""
+
+    __tablename__ = "enrichments"
+
+    snippet_sha: Mapped[str] = mapped_column(ForeignKey("snippets_v2.sha"), index=True)
+    type: Mapped[EnrichmentType] = mapped_column(
+        SQLAlchemyEnum(EnrichmentType), index=True
+    )
+    content: Mapped[str] = mapped_column(UnicodeText)
+
+    __table_args__ = (
+        UniqueConstraint("snippet_sha", "type", name="uix_snippet_enrichment"),
+    )
 
     def __init__(
         self,
-        file_id: int,
-        index_id: int,
+        snippet_sha: str,
+        type: EnrichmentType,  # noqa: A002
         content: str,
-        summary: str = "",
     ) -> None:
-        """Initialize the snippet."""
+        """Initialize enrichment."""
         super().__init__()
-        self.file_id = file_id
-        self.index_id = index_id
+        self.snippet_sha = snippet_sha
+        self.type = type
         self.content = content
-        self.summary = summary
-
 
-class TaskType(Enum):
-    """Task type."""
 
-    INDEX_UPDATE = 1
-
-
-class Task(Base, CommonMixin):
-    """Queued tasks."""
+class CommitIndex(Base):
+    """Commit index model."""
 
-    __tablename__ = "tasks"
+    __tablename__ = "commit_indexes"
 
-    # dedup_key is used to deduplicate items in the queue
-    dedup_key: Mapped[str] = mapped_column(String(255), index=True)
-    # type represents what the task is meant to achieve
-    type: Mapped[TaskType] = mapped_column(SQLAlchemyEnum(TaskType), index=True)
-    # payload contains the task-specific payload data
-    payload: Mapped[dict] = mapped_column(JSON)
-    # priority is used to determine the order of the items in the queue
-    priority: Mapped[int] = mapped_column(Integer)
+    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),
+    )
+    commit_sha: Mapped[str] = mapped_column(String(64), primary_key=True)
+    status: Mapped[str] = mapped_column(String(255), index=True)
+    indexed_at: Mapped[datetime | None] = mapped_column(TZDateTime, nullable=True)
+    error_message: Mapped[str | None] = mapped_column(UnicodeText, nullable=True)
+    files_processed: Mapped[int] = mapped_column(Integer, default=0)
+    processing_time_seconds: Mapped[float] = mapped_column(Float, default=0.0)
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
-        dedup_key: str,
-        type: TaskType,  # noqa: A002
-        payload: dict,
-        priority: int,
+        commit_sha: str,
+        status: str,
+        indexed_at: datetime | None = None,
+        error_message: str | None = None,
+        files_processed: int = 0,
+        processing_time_seconds: float = 0.0,
     ) -> None:
-        """Initialize the queue item."""
+        """Initialize commit index."""
         super().__init__()
-        self.dedup_key = dedup_key
-        self.type = type
-        self.payload = payload
-        self.priority = priority
+        self.commit_sha = commit_sha
+        self.status = status
+        self.indexed_at = indexed_at
+        self.error_message = error_message
+        self.files_processed = files_processed
+        self.processing_time_seconds = processing_time_seconds