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

kodit/infrastructure/sqlalchemy/entities.py

@@ -2,12 +2,15 @@
 
 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,
+    Index,
     Integer,
     String,
     TypeDecorator,
@@ -43,6 +46,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,114 +84,6 @@ class CommonMixin:
     )
 
 
-class SourceType(Enum):
-    """The type of source."""
-
-    UNKNOWN = 0
-    FOLDER = 1
-    GIT = 2
-
-
-class Source(Base, CommonMixin):
-    """Base model for tracking code sources.
-
-    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.
-
-    """
-
-    __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
-    )
-
-    def __init__(self, uri: str, cloned_path: str, source_type: SourceType) -> None:
-        """Initialize a new Source instance for typing purposes."""
-        super().__init__()
-        self.uri = uri
-        self.cloned_path = cloned_path
-        self.type = source_type
-
-
-class Author(Base, CommonMixin):
-    """Author model."""
-
-    __tablename__ = "authors"
-
-    __table_args__ = (UniqueConstraint("name", "email", name="uix_author"),)
-
-    name: Mapped[str] = mapped_column(String(255), index=True)
-    email: Mapped[str] = mapped_column(String(255), index=True)
-
-    @staticmethod
-    def from_actor(actor: Actor) -> "Author":
-        """Create an Author from an Actor."""
-        return Author(name=actor.name, email=actor.email)
-
-
-class AuthorFileMapping(Base, CommonMixin):
-    """Author file mapping model."""
-
-    __tablename__ = "author_file_mappings"
-
-    __table_args__ = (
-        UniqueConstraint("author_id", "file_id", name="uix_author_file_mapping"),
-    )
-
-    author_id: Mapped[int] = mapped_column(ForeignKey("authors.id"), index=True)
-    file_id: Mapped[int] = mapped_column(ForeignKey("files.id"), index=True)
-
-
-class File(Base, CommonMixin):
-    """File model."""
-
-    __tablename__ = "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)
-
-    def __init__(  # noqa: PLR0913
-        self,
-        created_at: datetime,
-        updated_at: datetime,
-        source_id: int,
-        mime_type: str,
-        uri: str,
-        cloned_path: str,
-        sha256: str,
-        size_bytes: int,
-        extension: str,
-        file_processing_status: int,
-    ) -> None:
-        """Initialize a new File instance for typing purposes."""
-        super().__init__()
-        self.created_at = created_at
-        self.updated_at = updated_at
-        self.source_id = source_id
-        self.mime_type = mime_type
-        self.uri = uri
-        self.cloned_path = cloned_path
-        self.sha256 = sha256
-        self.size_bytes = size_bytes
-        self.extension = extension
-        self.file_processing_status = file_processing_status
-
-
 class EmbeddingType(Enum):
     """Embedding type."""
 
@@ -182,59 +96,13 @@ class Embedding(Base, CommonMixin):
 
     __tablename__ = "embeddings"
 
-    snippet_id: Mapped[int] = mapped_column(ForeignKey("snippets.id"), index=True)
+    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 Index(Base, CommonMixin):
-    """Index model."""
-
-    __tablename__ = "indexes"
-
-    source_id: Mapped[int] = mapped_column(
-        ForeignKey("sources.id"), unique=True, index=True
-    )
-
-    def __init__(self, source_id: int) -> None:
-        """Initialize the index."""
-        super().__init__()
-        self.source_id = source_id
-
-
-class Snippet(Base, CommonMixin):
-    """Snippet model."""
-
-    __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="")
-
-    def __init__(
-        self,
-        file_id: int,
-        index_id: int,
-        content: str,
-        summary: str = "",
-    ) -> None:
-        """Initialize the snippet."""
-        super().__init__()
-        self.file_id = file_id
-        self.index_id = index_id
-        self.content = content
-        self.summary = summary
-
-
-class TaskType(Enum):
-    """Task type."""
-
-    INDEX_UPDATE = 1
-
-
 class Task(Base, CommonMixin):
     """Queued tasks."""
 
@@ -243,7 +111,7 @@ class Task(Base, CommonMixin):
     # 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)
+    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
@@ -252,7 +120,7 @@ class Task(Base, CommonMixin):
     def __init__(
         self,
         dedup_key: str,
-        type: TaskType,  # noqa: A002
+        type: str,  # noqa: A002
         payload: dict,
         priority: int,
     ) -> None:
@@ -323,3 +191,379 @@ class TaskStatus(Base):
         self.total = total
         self.current = current
         self.message = message or ""
+
+
+# Git-related entities for new GitRepo domain
+
+
+class GitRepo(Base, CommonMixin):
+    """Git repository model."""
+
+    __tablename__ = "git_repos"
+
+    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__(  # 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.sanitized_remote_uri = sanitized_remote_uri
+        self.remote_uri = remote_uri
+        self.cloned_path = cloned_path
+        self.last_scanned_at = last_scanned_at
+        self.num_commits = num_commits
+        self.num_branches = num_branches
+        self.num_tags = num_tags
+
+
+class GitCommit(Base):
+    """Git commit model."""
+
+    __tablename__ = "git_commits"
+
+    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."""
+
+    __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"))
+
+    __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 GitTrackingBranch(Base):
+    """Git tracking branch model."""
+
+    __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),
+    )
+
+    def __init__(self, repo_id: int, name: str) -> None:
+        """Initialize Git tracking branch."""
+        super().__init__()
+        self.repo_id = repo_id
+        self.name = name
+
+
+class GitTag(Base):
+    """Git tag model."""
+
+    __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"
+
+    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,
+        commit_sha: str,
+        path: str,
+        blob_sha: str,
+        mime_type: str,
+        extension: str,
+        size: int,
+        created_at: datetime,
+    ) -> None:
+        """Initialize Git commit file."""
+        super().__init__()
+        self.commit_sha = commit_sha
+        self.path = path
+        self.blob_sha = blob_sha
+        self.mime_type = mime_type
+        self.size = size
+        self.created_at = created_at
+        self.extension = extension
+
+
+class SnippetV2(Base):
+    """SnippetV2 model for commit-based snippets."""
+
+    __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)
+
+    def __init__(
+        self,
+        sha: str,
+        content: str,
+        extension: str,
+    ) -> None:
+        """Initialize snippet."""
+        super().__init__()
+        self.sha = sha
+        self.content = content
+        self.extension = extension
+
+
+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",
+        ),
+    )
+
+    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 CommitSnippetV2(Base):
+    """Association table for commits and snippets v2."""
+
+    __tablename__ = "commit_snippets_v2"
+
+    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, commit_sha: str, snippet_sha: str) -> None:
+        """Initialize commit snippet association."""
+        super().__init__()
+        self.commit_sha = commit_sha
+        self.snippet_sha = snippet_sha
+
+
+class CommitIndex(Base):
+    """Commit index model."""
+
+    __tablename__ = "commit_indexes"
+
+    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__(  # noqa: PLR0913
+        self,
+        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 commit index."""
+        super().__init__()
+        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
+
+
+class EnrichmentV2(Base, CommonMixin):
+    """Generic enrichment entity."""
+
+    __tablename__ = "enrichments_v2"
+
+    type: Mapped[str] = mapped_column(String, nullable=False, index=True)
+    subtype: Mapped[str] = mapped_column(String, nullable=False, index=True)
+    content: Mapped[str] = mapped_column(UnicodeText, nullable=False)
+
+    __table_args__ = (Index("idx_type_subtype", "type", "subtype"),)
+
+
+class EnrichmentAssociation(Base, CommonMixin):
+    """Polymorphic association between enrichments and entities."""
+
+    __tablename__ = "enrichment_associations"
+
+    enrichment_id: Mapped[int] = mapped_column(
+        ForeignKey("enrichments_v2.id", ondelete="CASCADE"),
+        nullable=False,
+        index=True,
+    )
+    entity_type: Mapped[str] = mapped_column(
+        String(50),
+        nullable=False,
+        index=True,
+    )
+    entity_id: Mapped[str] = mapped_column(
+        String(255),
+        nullable=False,
+        index=True,
+    )
+
+    __table_args__ = (
+        UniqueConstraint(
+            "entity_type",
+            "entity_id",
+            "enrichment_id",
+            name="uix_entity_enrichment",
+        ),
+        Index("idx_entity_lookup", "entity_type", "entity_id"),
+        {"sqlite_autoincrement": True},
+    )