martinholovsky/claude-skills-generator/graph-database-expert

# tests/test_graph_queries.py
import pytest
from surrealdb import Surreal
@pytest.fixture
async def db():
    """Setup test database with graph schema."""
    db = Surreal("ws://localhost:8000/rpc")
    await db.connect()
    await db.signin({"user": "root", "pass": "root"})
    await db.use("test", "test")
    # Setup schema
    await db.query("""
        DEFINE TABLE person SCHEMAFULL;
        DEFINE FIELD name ON TABLE person TYPE string;
        DEFINE INDEX person_name ON TABLE person COLUMNS name;
        DEFINE TABLE follows SCHEMAFULL;
        DEFINE FIELD in ON TABLE follows TYPE record<person>;
        DEFINE FIELD out ON TABLE follows TYPE record<person>;
    """)
    yield db
    # Cleanup
    await db.query("REMOVE TABLE person; REMOVE TABLE follows;")
    await db.close()
@pytest.mark.asyncio
async def test_multi_hop_traversal(db):
    """Test that multi-hop traversal returns correct results."""
    # Arrange: Create test graph
    await db.query("""
        CREATE person:alice SET name = 'Alice';
        CREATE person:bob SET name = 'Bob';
        CREATE person:charlie SET name = 'Charlie';
        RELATE person:alice->follows->person:bob;
        RELATE person:bob->follows->person:charlie;
    """)
    # Act: Traverse 2 hops
    result = await db.query(
        "SELECT ->follows[..2]->person.name FROM person:alice"
    )
    # Assert: Should find Bob and Charlie
    names = result[0]['result'][0]['name']
    assert 'Bob' in names
    assert 'Charlie' in names
@pytest.mark.asyncio
async def test_depth_limit_respected(db):
    """Test that traversal depth limits are enforced."""
    # Arrange: Create chain of 5 nodes
    await db.query("""
        CREATE person:a SET name = 'A';
        CREATE person:b SET name = 'B';
        CREATE person:c SET name = 'C';
        CREATE person:d SET name = 'D';
        CREATE person:e SET name = 'E';
        RELATE person:a->follows->person:b;
        RELATE person:b->follows->person:c;
        RELATE person:c->follows->person:d;
        RELATE person:d->follows->person:e;
    """)
    # Act: Traverse only 2 hops
    result = await db.query(
        "SELECT ->follows[..2]->person.name FROM person:a"
    )
    # Assert: Should NOT include D or E
    names = result[0]['result'][0]['name']
    assert 'D' not in names
    assert 'E' not in names
@pytest.mark.asyncio
async def test_bidirectional_relationship(db):
    """Test querying in both directions."""
    # Arrange
    await db.query("""
        CREATE person:alice SET name = 'Alice';
        CREATE person:bob SET name = 'Bob';
        RELATE person:alice->follows->person:bob;
    """)
    # Act: Query both directions
    forward = await db.query(
        "SELECT ->follows->person.name FROM person:alice"
    )
    backward = await db.query(
        "SELECT <-follows<-person.name FROM person:bob"
    )
    # Assert
    assert 'Bob' in str(forward)
    assert 'Alice' in str(backward)
@pytest.mark.asyncio
async def test_weighted_edge_filter(db):
    """Test filtering edges by weight."""
    # Setup weighted edges
    await db.query("""
        DEFINE TABLE connected SCHEMAFULL;
        DEFINE FIELD in ON TABLE connected TYPE record<person>;
        DEFINE FIELD out ON TABLE connected TYPE record<person>;
        DEFINE FIELD weight ON TABLE connected TYPE float;
        CREATE person:alice SET name = 'Alice';
        CREATE person:bob SET name = 'Bob';
        CREATE person:charlie SET name = 'Charlie';
        RELATE person:alice->connected->person:bob SET weight = 0.9;
        RELATE person:alice->connected->person:charlie SET weight = 0.3;
    """)
    # Act: Filter by weight
    result = await db.query(
        "SELECT ->connected[WHERE weight > 0.5]->person.name FROM person:alice"
    )
    # Assert: Only Bob (high weight)
    assert 'Bob' in str(result)
    assert 'Charlie' not in str(result)

# src/graph/queries.py
from surrealdb import Surreal
class GraphQueryService:
    def __init__(self, db: Surreal):
        self.db = db
    async def get_connections(
        self,
        node_id: str,
        relationship: str,
        depth: int = 2,
        min_weight: float | None = None
    ) -> list[dict]:
        """Get connected nodes with depth limit."""
        if depth > 5:
            raise ValueError("Maximum depth is 5 to prevent runaway queries")
        # Build query with parameterization
        if min_weight is not None:
            query = f"""
                SELECT ->{relationship}[..{depth}][WHERE weight > $min_weight]->*.*
                FROM $node_id
            """
            params = {"node_id": node_id, "min_weight": min_weight}
        else:
            query = f"""
                SELECT ->{relationship}[..{depth}]->*.*
                FROM $node_id
            """
            params = {"node_id": node_id}
        result = await self.db.query(query, params)
        return result[0]['result']
    async def find_path(
        self,
        from_id: str,
        to_id: str,
        relationship: str,
        max_depth: int = 5
    ) -> list[str] | None:
        """Find shortest path between two nodes."""
        # BFS implementation with depth limit
        visited = set()
        queue = [(from_id, [from_id])]
        while queue and len(visited) < 1000:  # Safety limit
            current, path = queue.pop(0)
            if len(path) > max_depth:
                continue
            if current == to_id:
                return path
            if current in visited:
                continue
            visited.add(current)
            # Get neighbors
            result = await self.db.query(
                f"SELECT ->{relationship}->*.id FROM $node",
                {"node": current}
            )
            for neighbor in result[0]['result']:
                if neighbor not in visited:
                    queue.append((neighbor, path + [neighbor]))
        return None

# After tests pass, refactor for better performance
class GraphQueryService:
    def __init__(self, db: Surreal):
        self.db = db
        self._cache = {}  # Add caching
    async def get_connections_cached(
        self,
        node_id: str,
        relationship: str,
        depth: int = 2
    ) -> list[dict]:
        """Get connections with caching."""
        cache_key = f"{node_id}:{relationship}:{depth}"
        if cache_key in self._cache:
            return self._cache[cache_key]
        result = await self.get_connections(node_id, relationship, depth)
        self._cache[cache_key] = result
        return result