markitect-main/markitect/prompts/dependencies/repository.py

"""
Repository interfaces and SQLite implementation for dependency persistence.

Implements FR-6.1: Directed dependency edges between artifacts.
Mirrors the SQLiteArtifactRepository pattern for consistency.
"""

import sqlite3
import json
from abc import ABC, abstractmethod
from pathlib import Path
from typing import List, Optional
from datetime import datetime

from markitect.prompts.dependencies.models import DependencyEdge, EdgeType


class DependencyRepositoryError(Exception):
    """Base exception for dependency repository errors."""
    pass


class DuplicateDependencyError(DependencyRepositoryError):
    """Raised when attempting to create a duplicate dependency edge."""
    pass


# SQL Schema for dependency tables
DEPENDENCY_TABLES_SQL = """
CREATE TABLE IF NOT EXISTS prompt_dependencies (
    id TEXT PRIMARY KEY,
    source_artifact_id TEXT NOT NULL,
    target_artifact_id TEXT NOT NULL,
    run_id TEXT NOT NULL,
    edge_type TEXT NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    UNIQUE(source_artifact_id, target_artifact_id, run_id)
);

CREATE INDEX IF NOT EXISTS idx_deps_source ON prompt_dependencies(source_artifact_id);
CREATE INDEX IF NOT EXISTS idx_deps_target ON prompt_dependencies(target_artifact_id);
CREATE INDEX IF NOT EXISTS idx_deps_run ON prompt_dependencies(run_id);
CREATE INDEX IF NOT EXISTS idx_deps_type ON prompt_dependencies(edge_type);
CREATE INDEX IF NOT EXISTS idx_deps_source_target ON prompt_dependencies(source_artifact_id, target_artifact_id);
"""


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

    Args:
        db_path: Path to the SQLite database file
    """
    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(DEPENDENCY_TABLES_SQL)
        conn.commit()
    finally:
        conn.close()


class IDependencyRepository(ABC):
    """
    Abstract interface for dependency edge persistence.

    Implements FR-6.1: Directed dependency edges between artifacts.
    """

    @abstractmethod
    def create(self, edge: DependencyEdge) -> DependencyEdge:
        """
        Persist a new dependency edge.

        Args:
            edge: Dependency edge to create

        Returns:
            Created edge

        Raises:
            DuplicateDependencyError: If edge with same source+target+run exists
            DependencyRepositoryError: On other persistence errors
        """
        pass

    @abstractmethod
    def get_by_id(self, edge_id: str) -> Optional[DependencyEdge]:
        """
        Retrieve edge by ID.

        Args:
            edge_id: Edge identifier

        Returns:
            Edge if found, None otherwise
        """
        pass

    @abstractmethod
    def get_by_source(self, source_artifact_id: str) -> List[DependencyEdge]:
        """
        Get all edges originating from a source artifact.

        Args:
            source_artifact_id: Source artifact ID

        Returns:
            List of edges from this source
        """
        pass

    @abstractmethod
    def get_by_target(self, target_artifact_id: str) -> List[DependencyEdge]:
        """
        Get all edges pointing to a target artifact.

        Args:
            target_artifact_id: Target artifact ID

        Returns:
            List of edges to this target
        """
        pass

    @abstractmethod
    def get_by_run(self, run_id: str) -> List[DependencyEdge]:
        """
        Get all edges created by a specific run.

        Args:
            run_id: Run identifier

        Returns:
            List of edges from this run
        """
        pass

    @abstractmethod
    def get_by_type(self, edge_type: EdgeType) -> List[DependencyEdge]:
        """
        Get all edges of a specific type.

        Args:
            edge_type: Edge type to filter by

        Returns:
            List of edges of this type
        """
        pass

    @abstractmethod
    def get_all(self) -> List[DependencyEdge]:
        """
        Get all dependency edges.

        Returns:
            List of all edges
        """
        pass

    @abstractmethod
    def delete(self, edge_id: str) -> bool:
        """
        Delete a dependency edge.

        Args:
            edge_id: Edge identifier

        Returns:
            True if deleted, False if not found
        """
        pass

    @abstractmethod
    def delete_by_run(self, run_id: str) -> int:
        """
        Delete all edges created by a specific run.

        Args:
            run_id: Run identifier

        Returns:
            Number of edges deleted
        """
        pass


class SQLiteDependencyRepository(IDependencyRepository):
    """
    SQLite implementation of dependency repository.

    Provides persistent storage for dependency edges with query support
    by source, target, run, and edge type.
    """

    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_dependency_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 _row_to_edge(self, row: sqlite3.Row) -> DependencyEdge:
        """Convert database row to DependencyEdge instance."""
        return DependencyEdge(
            id=row["id"],
            source_artifact_id=row["source_artifact_id"],
            target_artifact_id=row["target_artifact_id"],
            run_id=row["run_id"],
            edge_type=EdgeType(row["edge_type"]),
            created_at=datetime.fromisoformat(row["created_at"]),
        )

    def create(self, edge: DependencyEdge) -> DependencyEdge:
        """
        Persist a new dependency edge.

        Args:
            edge: Dependency edge to create

        Returns:
            Created edge

        Raises:
            DuplicateDependencyError: If edge with same source+target+run exists
            DependencyRepositoryError: On other persistence errors
        """
        conn = self._get_connection()
        try:
            conn.execute(
                """
                INSERT INTO prompt_dependencies (
                    id, source_artifact_id, target_artifact_id,
                    run_id, edge_type, created_at
                ) VALUES (?, ?, ?, ?, ?, ?)
                """,
                (
                    edge.id,
                    edge.source_artifact_id,
                    edge.target_artifact_id,
                    edge.run_id,
                    edge.edge_type.value,
                    edge.created_at.isoformat(),
                ),
            )
            conn.commit()
            return edge
        except sqlite3.IntegrityError as e:
            if "UNIQUE constraint" in str(e):
                raise DuplicateDependencyError(
                    f"Dependency edge from '{edge.source_artifact_id}' to "
                    f"'{edge.target_artifact_id}' already exists for run '{edge.run_id}'"
                )
            raise DependencyRepositoryError(f"Database integrity error: {e}")
        except Exception as e:
            raise DependencyRepositoryError(f"Failed to create dependency edge: {e}")
        finally:
            conn.close()

    def get_by_id(self, edge_id: str) -> Optional[DependencyEdge]:
        """
        Retrieve edge by ID.

        Args:
            edge_id: Edge identifier

        Returns:
            Edge if found, None otherwise
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_dependencies WHERE id = ?",
                (edge_id,),
            )
            row = cursor.fetchone()
            return self._row_to_edge(row) if row else None
        finally:
            conn.close()

    def get_by_source(self, source_artifact_id: str) -> List[DependencyEdge]:
        """
        Get all edges originating from a source artifact.

        Args:
            source_artifact_id: Source artifact ID

        Returns:
            List of edges from this source
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_dependencies WHERE source_artifact_id = ?",
                (source_artifact_id,),
            )
            return [self._row_to_edge(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def get_by_target(self, target_artifact_id: str) -> List[DependencyEdge]:
        """
        Get all edges pointing to a target artifact.

        Args:
            target_artifact_id: Target artifact ID

        Returns:
            List of edges to this target
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_dependencies WHERE target_artifact_id = ?",
                (target_artifact_id,),
            )
            return [self._row_to_edge(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def get_by_run(self, run_id: str) -> List[DependencyEdge]:
        """
        Get all edges created by a specific run.

        Args:
            run_id: Run identifier

        Returns:
            List of edges from this run
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_dependencies WHERE run_id = ?",
                (run_id,),
            )
            return [self._row_to_edge(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def get_by_type(self, edge_type: EdgeType) -> List[DependencyEdge]:
        """
        Get all edges of a specific type.

        Args:
            edge_type: Edge type to filter by

        Returns:
            List of edges of this type
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "SELECT * FROM prompt_dependencies WHERE edge_type = ?",
                (edge_type.value,),
            )
            return [self._row_to_edge(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def get_all(self) -> List[DependencyEdge]:
        """
        Get all dependency edges.

        Returns:
            List of all edges
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute("SELECT * FROM prompt_dependencies")
            return [self._row_to_edge(row) for row in cursor.fetchall()]
        finally:
            conn.close()

    def delete(self, edge_id: str) -> bool:
        """
        Delete a dependency edge.

        Args:
            edge_id: Edge identifier

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

    def delete_by_run(self, run_id: str) -> int:
        """
        Delete all edges created by a specific run.

        Args:
            run_id: Run identifier

        Returns:
            Number of edges deleted
        """
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "DELETE FROM prompt_dependencies WHERE run_id = ?",
                (run_id,),
            )
            conn.commit()
            return cursor.rowcount
        finally:
            conn.close()