markitect-main/markitect/prompts/repositories/sqlite.py

"""
SQLite implementation of artifact repositories.

This module provides SQLite-backed implementations of the repository
interfaces for persistent storage of prompt artifacts.
"""

import sqlite3
import json
from pathlib import Path
from typing import List, Optional
from datetime import datetime

from markitect.prompts.repositories.interfaces import (
    IArtifactRepository,
    ArtifactNotFoundError,
    DuplicateArtifactError,
    RepositoryError,
)
from markitect.prompts.models import Artifact, ArtifactType, ArtifactMetadata


# SQL Schema for artifact tables
ARTIFACT_TABLES_SQL = """
-- Prompt artifacts table
CREATE TABLE IF NOT EXISTS prompt_artifacts (
    id TEXT PRIMARY KEY,
    space_id TEXT NOT NULL,
    name TEXT NOT NULL,
    artifact_type TEXT NOT NULL,
    content_digest TEXT NOT NULL,
    content_size INTEGER DEFAULT 0,
    metadata JSON,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    UNIQUE(space_id, name)
);

-- Indexes for performance
CREATE INDEX IF NOT EXISTS idx_artifacts_space ON prompt_artifacts(space_id);
CREATE INDEX IF NOT EXISTS idx_artifacts_digest ON prompt_artifacts(content_digest);
CREATE INDEX IF NOT EXISTS idx_artifacts_type ON prompt_artifacts(artifact_type);
CREATE INDEX IF NOT EXISTS idx_artifacts_name ON prompt_artifacts(space_id, name);
"""


def initialize_artifact_tables(db_path: str) -> None:
    """
    Initialize the artifact-related database tables.

    Args:
        db_path: Path to the SQLite database file
    """
    # Ensure directory exists
    db_dir = Path(db_path).parent
    if db_dir and not db_dir.exists():
        db_dir.mkdir(parents=True, exist_ok=True)

    conn = sqlite3.connect(db_path)
    try:
        conn.executescript(ARTIFACT_TABLES_SQL)
        conn.commit()
    finally:
        conn.close()


class SQLiteArtifactRepository(IArtifactRepository):
    """
    SQLite implementation of artifact repository.

    Provides persistent storage for artifacts with content digest tracking.
    """

    def __init__(self, db_path: str):
        """
        Initialize repository with database path.

        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = db_path
        initialize_artifact_tables(db_path)

    def _get_connection(self) -> sqlite3.Connection:
        """Get a database connection."""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def create(self, artifact: Artifact) -> Artifact:
        """
        Persist a new artifact.

        Args:
            artifact: Artifact to create

        Returns:
            Created artifact

        Raises:
            DuplicateArtifactError: If artifact with same space_id+name exists
            RepositoryError: On other persistence errors
        """
        conn = self._get_connection()
        try:
            conn.execute(
                """
                INSERT INTO prompt_artifacts (
                    id, space_id, name, artifact_type, content_digest,
                    content_size, metadata, created_at, updated_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    artifact.id,
                    artifact.space_id,
                    artifact.name,
                    artifact.artifact_type.value,
                    artifact.content_digest,
                    artifact.content_size,
                    json.dumps(artifact.metadata.to_dict()),
                    artifact.created_at.isoformat(),
                    artifact.updated_at.isoformat(),
                ),
            )
            conn.commit()
            return artifact
        except sqlite3.IntegrityError as e:
            if "UNIQUE constraint" in str(e):
                raise DuplicateArtifactError(
                    f"Artifact '{artifact.name}' already exists in space '{artifact.space_id}'"
                )
            raise RepositoryError(f"Database integrity error: {e}")
        except Exception as e:
            raise RepositoryError(f"Failed to create artifact: {e}")
        finally:
            conn.close()

    def get_by_id(self, artifact_id: str) -> Optional[Artifact]:
        """
        Retrieve artifact by ID.

        Args:
            artifact_id: Artifact identifier

        Returns:
            Artifact if found, None otherwise
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_artifacts WHERE id = ?",
                (artifact_id,)
            )
            row = cursor.fetchone()
            return self._row_to_artifact(row) if row else None
        finally:
            conn.close()

    def get_by_name(self, space_id: str, name: str) -> Optional[Artifact]:
        """
        Retrieve artifact by space and name.

        Args:
            space_id: Space identifier
            name: Artifact name

        Returns:
            Artifact if found, None otherwise
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_artifacts WHERE space_id = ? AND name = ?",
                (space_id, name)
            )
            row = cursor.fetchone()
            return self._row_to_artifact(row) if row else None
        finally:
            conn.close()

    def get_by_digest(self, content_digest: str) -> List[Artifact]:
        """
        Find artifacts with matching content digest.

        Args:
            content_digest: SHA-256 digest to match

        Returns:
            List of artifacts with matching digest
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_artifacts WHERE content_digest = ?",
                (content_digest,)
            )
            return [self._row_to_artifact(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def list_by_space(
        self,
        space_id: str,
        artifact_type: Optional[ArtifactType] = None,
    ) -> List[Artifact]:
        """
        List all artifacts in a space.

        Args:
            space_id: Space identifier
            artifact_type: Optional type filter

        Returns:
            List of artifacts in space
        """
        conn = self._get_connection()
        try:
            if artifact_type:
                cursor = conn.execute(
                    "SELECT * FROM prompt_artifacts WHERE space_id = ? AND artifact_type = ? ORDER BY name",
                    (space_id, artifact_type.value)
                )
            else:
                cursor = conn.execute(
                    "SELECT * FROM prompt_artifacts WHERE space_id = ? ORDER BY name",
                    (space_id,)
                )
            return [self._row_to_artifact(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def update(self, artifact: Artifact) -> Artifact:
        """
        Update an existing artifact.

        Args:
            artifact: Artifact with updated data

        Returns:
            Updated artifact

        Raises:
            ArtifactNotFoundError: If artifact doesn't exist
            RepositoryError: On other persistence errors
        """
        conn = self._get_connection()
        try:
            # Update timestamp
            artifact.updated_at = datetime.utcnow()

            cursor = conn.execute(
                """
                UPDATE prompt_artifacts
                SET content_digest = ?, content_size = ?, metadata = ?, updated_at = ?
                WHERE id = ?
                """,
                (
                    artifact.content_digest,
                    artifact.content_size,
                    json.dumps(artifact.metadata.to_dict()),
                    artifact.updated_at.isoformat(),
                    artifact.id,
                ),
            )
            if cursor.rowcount == 0:
                raise ArtifactNotFoundError(f"Artifact with ID '{artifact.id}' not found")
            conn.commit()
            return artifact
        except ArtifactNotFoundError:
            raise
        except Exception as e:
            raise RepositoryError(f"Failed to update artifact: {e}")
        finally:
            conn.close()

    def delete(self, artifact_id: str) -> bool:
        """
        Delete an artifact.

        Args:
            artifact_id: Artifact identifier

        Returns:
            True if deleted, False if not found
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "DELETE FROM prompt_artifacts WHERE id = ?",
                (artifact_id,)
            )
            conn.commit()
            return cursor.rowcount > 0
        finally:
            conn.close()

    def exists(self, space_id: str, name: str) -> bool:
        """
        Check if artifact exists.

        Args:
            space_id: Space identifier
            name: Artifact name

        Returns:
            True if artifact exists
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT COUNT(*) FROM prompt_artifacts WHERE space_id = ? AND name = ?",
                (space_id, name)
            )
            count = cursor.fetchone()[0]
            return count > 0
        finally:
            conn.close()

    def _row_to_artifact(self, row: sqlite3.Row) -> Artifact:
        """Convert database row to Artifact instance."""
        metadata_dict = json.loads(row["metadata"]) if row["metadata"] else {}
        return Artifact(
            id=row["id"],
            space_id=row["space_id"],
            name=row["name"],
            artifact_type=ArtifactType(row["artifact_type"]),
            content_digest=row["content_digest"],
            content_size=row["content_size"],
            metadata=ArtifactMetadata.from_dict(metadata_dict),
            created_at=datetime.fromisoformat(row["created_at"]),
            updated_at=datetime.fromisoformat(row["updated_at"]),
        )