library-desk/src/clients/neo4j_client.py

"""
Neo4j async client for Library Desk.

Provides async Neo4j operations with:
- Connection pooling via AsyncGraphDatabase
- Session management with context managers
- Multi-tenancy support via user labels
- Automatic retry on transient failures
"""

from neo4j import AsyncGraphDatabase, AsyncDriver, AsyncSession, READ_ACCESS
from typing import Optional, List, Dict, Any
import logging

from src.core.multi_tenancy import get_neo4j_user_label

logger = logging.getLogger(__name__)


class Neo4jClient:
    """
    Async Neo4j client with connection pooling.

    Features:
    - Singleton driver instance (expensive to create)
    - Session-per-request pattern (lightweight)
    - Automatic transaction retry
    - Multi-tenancy via user-specific labels
    """

    def __init__(self, uri: str, user: str, password: str):
        """
        Initialize Neo4j client.

        Args:
            uri: Neo4j Bolt URI (e.g., "bolt://neo4j:7687")
            user: Neo4j username
            password: Neo4j password
        """
        self.uri = uri
        self._driver: Optional[AsyncDriver] = None
        self._auth = (user, password)

    async def connect(self):
        """
        Initialize connection pool.

        Should be called once at app startup.
        Driver handles connection pooling internally.
        """
        if not self._driver:
            self._driver = AsyncGraphDatabase.driver(
                self.uri,
                auth=self._auth,
                max_connection_pool_size=50,
                connection_timeout=30.0,
                max_transaction_retry_time=30.0
            )
            # Verify connectivity
            await self._driver.verify_connectivity()
            logger.info(f"Connected to Neo4j at {self.uri}")

    async def close(self):
        """
        Close connection pool.

        Should be called once at app shutdown.
        """
        if self._driver:
            await self._driver.close()
            self._driver = None
            logger.info("Closed Neo4j connection")

    async def execute_query(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Execute Cypher query and return results.

        Args:
            cypher: Cypher query string
            parameters: Query parameters

        Returns:
            List of result records as dictionaries

        Raises:
            Exception: If driver not initialized or query fails
        """
        if not self._driver:
            await self.connect()

        async with self._driver.session() as session:
            result = await session.run(cypher, parameters or {})
            records = await result.data()
            return records

    async def execute_read(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Execute Cypher in a READ-ONLY session.

        The session is opened with default_access_mode=READ_ACCESS, so the
        database rejects any write attempt (CREATE/MERGE/DELETE/SET/...)
        even if it slips past caller-side validation. Use this for any
        query built from untrusted input (e.g. the /query/graph endpoint).

        Args:
            cypher: Cypher query string
            parameters: Query parameters

        Returns:
            List of result records as dictionaries

        Raises:
            Exception: If driver not initialized, query fails, or the
                query attempts a write (rejected by the read session)
        """
        if not self._driver:
            await self.connect()

        async with self._driver.session(default_access_mode=READ_ACCESS) as session:
            result = await session.run(cypher, parameters or {})
            records = await result.data()
            return records

    async def execute_write(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Execute write transaction with automatic retry.

        Args:
            cypher: Cypher query string
            parameters: Query parameters

        Returns:
            List of result records as dictionaries
        """
        if not self._driver:
            await self.connect()

        async def write_tx(tx):
            result = await tx.run(cypher, parameters or {})
            return await result.data()

        async with self._driver.session() as session:
            return await session.execute_write(write_tx)

    # Multi-tenancy helpers

    def get_user_label(self, user: str) -> str:
        """
        Get Neo4j label for user's documents.

        Args:
            user: User identifier

        Returns:
            Neo4j label string (e.g., "User_Jpmschweitzer_Document")
        """
        return get_neo4j_user_label(user)

    # Document node operations

    async def create_document_node(
        self,
        user: str,
        doc_id: str,
        properties: Dict[str, Any]
    ) -> Optional[Dict[str, Any]]:
        """
        Create document node with user label.

        Node structure:
        (doc:Document:User_{user}_Document {
            id: "doc_123",
            source: "github",
            repository: "anthropic-cookbook",
            path: "skills/citation/guide.md",
            title: "Citation Guide",
            created_at: timestamp(),
            updated_at: timestamp(),
            content_hash: "sha256:..."
        })

        Args:
            user: User identifier
            doc_id: Unique document ID
            properties: Document properties

        Returns:
            Created node properties or None on failure
        """
        user_label = self.get_user_label(user)

        # Ensure required properties
        properties["id"] = doc_id
        if "created_at" not in properties:
            properties["created_at"] = "timestamp()"

        cypher = f"""
        CREATE (doc:Document:{user_label})
        SET doc = $properties
        SET doc.created_at = timestamp()
        SET doc.updated_at = timestamp()
        RETURN doc
        """

        try:
            result = await self.execute_write(cypher, {"properties": properties})
            return result[0]["doc"] if result else None
        except Exception as e:
            logger.error(f"Failed to create document node: {e}", exc_info=True)
            return None

    async def get_document_node(
        self,
        user: str,
        doc_id: str
    ) -> Optional[Dict[str, Any]]:
        """
        Get document node by ID.

        Args:
            user: User identifier
            doc_id: Document ID

        Returns:
            Document node properties or None if not found
        """
        user_label = self.get_user_label(user)

        cypher = f"""
        MATCH (doc:Document:{user_label} {{id: $doc_id}})
        RETURN doc
        """

        try:
            result = await self.execute_query(cypher, {"doc_id": doc_id})
            return result[0]["doc"] if result else None
        except Exception as e:
            logger.error(f"Failed to get document node: {e}", exc_info=True)
            return None

    async def delete_document_node(
        self,
        user: str,
        doc_id: str
    ) -> bool:
        """
        Delete document node and all its relationships.

        Args:
            user: User identifier
            doc_id: Document ID

        Returns:
            True if deleted, False otherwise
        """
        user_label = self.get_user_label(user)

        cypher = f"""
        MATCH (doc:Document:{user_label} {{id: $doc_id}})
        DETACH DELETE doc
        RETURN count(doc) as deleted
        """

        try:
            result = await self.execute_write(cypher, {"doc_id": doc_id})
            return result[0]["deleted"] > 0 if result else False
        except Exception as e:
            logger.error(f"Failed to delete document node: {e}", exc_info=True)
            return False

    async def find_similar_documents(
        self,
        user: str,
        doc_ids: List[str],
        max_depth: int = 2,
        limit: int = 20
    ) -> List[Dict[str, Any]]:
        """
        Find documents similar to given docs via graph traversal.

        Uses: Shared concepts, shared entities, citation links.

        Args:
            user: User identifier
            doc_ids: List of source document IDs
            max_depth: Maximum traversal depth
            limit: Maximum results to return

        Returns:
            List of similar documents with connection strength
        """
        user_label = self.get_user_label(user)

        cypher = f"""
        MATCH (source:Document:{user_label})
        WHERE source.id IN $doc_ids
        MATCH (source)-[*1..{max_depth}]-(related:Document:{user_label})
        WHERE related.id <> source.id AND NOT related.id IN $doc_ids
        WITH related, count(*) as connection_strength
        ORDER BY connection_strength DESC
        LIMIT $limit
        RETURN related, connection_strength
        """

        try:
            result = await self.execute_query(
                cypher,
                {"doc_ids": doc_ids, "limit": limit}
            )
            return result
        except Exception as e:
            logger.error(f"Failed to find similar documents: {e}", exc_info=True)
            return []

    async def list_user_documents(
        self,
        user: str,
        limit: int = 100,
        offset: int = 0
    ) -> List[Dict[str, Any]]:
        """
        List all documents for a user.

        Args:
            user: User identifier
            limit: Maximum results to return
            offset: Number of results to skip

        Returns:
            List of document nodes
        """
        user_label = self.get_user_label(user)

        cypher = f"""
        MATCH (doc:Document:{user_label})
        RETURN doc
        ORDER BY doc.created_at DESC
        SKIP $offset
        LIMIT $limit
        """

        try:
            result = await self.execute_query(
                cypher,
                {"offset": offset, "limit": limit}
            )
            return [r["doc"] for r in result]
        except Exception as e:
            logger.error(f"Failed to list documents: {e}", exc_info=True)
            return []

    # Concept/entity operations

    async def create_concept_node(
        self,
        concept_name: str,
        concept_type: str,
        properties: Optional[Dict[str, Any]] = None
    ) -> Optional[Dict[str, Any]]:
        """
        Create or update concept node.

        Args:
            concept_name: Concept name
            concept_type: Concept type (Technique, Tool, Pattern, etc.)
            properties: Additional properties

        Returns:
            Concept node properties
        """
        cypher = """
        MERGE (concept:Concept {name: $name})
        ON CREATE SET
            concept.type = $type,
            concept.first_seen = timestamp(),
            concept.mention_count = 1
        ON MATCH SET
            concept.mention_count = concept.mention_count + 1
        SET concept += $properties
        RETURN concept
        """

        try:
            result = await self.execute_write(
                cypher,
                {
                    "name": concept_name,
                    "type": concept_type,
                    "properties": properties or {}
                }
            )
            return result[0]["concept"] if result else None
        except Exception as e:
            logger.error(f"Failed to create concept node: {e}", exc_info=True)
            return None

    async def link_document_to_concept(
        self,
        user: str,
        doc_id: str,
        concept_name: str
    ) -> bool:
        """
        Create MENTIONS relationship between document and concept.

        Args:
            user: User identifier
            doc_id: Document ID
            concept_name: Concept name

        Returns:
            True if link created
        """
        user_label = self.get_user_label(user)

        cypher = f"""
        MATCH (doc:Document:{user_label} {{id: $doc_id}})
        MATCH (concept:Concept {{name: $concept_name}})
        MERGE (doc)-[r:MENTIONS]->(concept)
        ON CREATE SET r.count = 1
        ON MATCH SET r.count = r.count + 1
        RETURN r
        """

        try:
            result = await self.execute_write(
                cypher,
                {"doc_id": doc_id, "concept_name": concept_name}
            )
            return len(result) > 0
        except Exception as e:
            logger.error(f"Failed to link document to concept: {e}", exc_info=True)
            return False