/cognee/infrastructure/databases/graph/postgres/adapter.py

"""Postgres graph adapter using two tables (graph_node, graph_edge) over SQLAlchemy + asyncpg."""

import asyncio
import json
from uuid import UUID
from datetime import datetime, timezone
from contextlib import asynccontextmanager
from typing import AsyncIterator, Dict, Any, List, Union, Optional, Tuple, Type

from sqlalchemy import NullPool, text, values, select, exists, func, String, case
from sqlalchemy import column as sa_column
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
from sqlalchemy.exc import DBAPIError
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
from asyncpg import DeadlockDetectedError

from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.engine import DataPoint
from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
from cognee.infrastructure.databases.relational import get_relational_config
from cognee.modules.storage.utils import JSONEncoder
from cognee.infrastructure.databases.provenance import (
    EdgeDeleteData,
    EdgeIdentity,
    NodeDeleteData,
)
from cognee.infrastructure.databases.provenance.source_refs import (
    get_dataset_id_from_source_ref_key,
    get_pipeline_run_id_from_source_run_ref,
    get_source_ref_key_from_source_run_ref,
)
from cognee.infrastructure.databases.provenance.source_ref_state import (
    ProvenanceAttachInputs,
    provenance_after_attach,
    provenance_after_remove,
    provenance_attach_inputs,
)

from .tables import _meta, _node_table, _edge_table, _metadata_table

logger = get_logger()

# Rows per INSERT statement for bulk node/edge writes. Bounds the size of the
# compiled SQLAlchemy statement and the asyncpg parameter buffer per execute, so
# a large single-batch write (e.g. the whole of War and Peace) streams to
# Postgres in fixed-size chunks instead of materializing one multi-thousand-row
# statement in memory. Does not change how many data points the pipeline batches.
_WRITE_CHUNK_SIZE = 1000


def _provenance_insert_values(inputs: ProvenanceAttachInputs) -> Dict[str, List[str]]:
    """Initial provenance arrays for a freshly INSERTed node/edge (single source ref)."""
    return {
        "source_ref_keys": inputs.add_keys,
        "source_dataset_ids": inputs.add_dataset_ids,
        "source_run_ids": inputs.add_run_ids,
        "source_run_refs": inputs.add_run_refs,
    }


def _provenance_conflict_set(table, inputs: ProvenanceAttachInputs) -> Dict[str, Any]:
    """``ON CONFLICT`` SET that set-merges one source ref into committed arrays.

    Postgres analogue of the Ladybug fold clause: a node/edge is created and
    stamped in one atomic upsert, so there is no read-then-write window (closes
    the write-then-attach gap and the concurrent lost update, COG-5522 #4/#8).
    The ``CASE`` guards read the *pre-update* ``source_ref_keys`` column, so the
    run ref/id are appended only when the key was not already present (Model A) —
    re-attaching an existing key adds no new run mapping. Dataset id is deduped
    independently against its own column.
    """
    sr_key = inputs.source_ref_key
    ds_id = inputs.add_dataset_ids[0]
    key_present = table.c.source_ref_keys.any(sr_key)
    ds_present = table.c.source_dataset_ids.any(ds_id)
    set_: Dict[str, Any] = {
        "source_ref_keys": case(
            (key_present, table.c.source_ref_keys),
            else_=func.array_append(table.c.source_ref_keys, sr_key),
        ),
        "source_dataset_ids": case(
            (ds_present, table.c.source_dataset_ids),
            else_=func.array_append(table.c.source_dataset_ids, ds_id),
        ),
    }
    # Run ref/id only exist when the write carried a pipeline_run_id; otherwise the
    # run columns are left untouched on conflict (the write is not rollbackable by run).
    if inputs.add_run_refs:
        set_["source_run_refs"] = case(
            (key_present, table.c.source_run_refs),
            else_=func.array_append(table.c.source_run_refs, inputs.add_run_refs[0]),
        )
        set_["source_run_ids"] = case(
            (key_present, table.c.source_run_ids),
            else_=func.array_append(table.c.source_run_ids, inputs.add_run_ids[0]),
        )
    return set_


class PostgresAdapter(GraphDBInterface):
    """Graph-as-tables adapter backed by Postgres, accessed via SQLAlchemy async sessions."""

    _ALLOWED_FILTER_ATTRS = {"id", "name", "type"}

    def __init__(self, connection_string: str) -> None:
        """Create engine and sessionmaker from a Postgres connection string."""
        self.db_uri = connection_string

        relational_config = get_relational_config()
        pool_args: dict = dict(relational_config.pool_args) if relational_config.pool_args else {}
        if pool_args.get("poolclass", "").lower() == "nullpool":
            pool_args["poolclass"] = NullPool
        else:
            # QueuePool defaults, mirroring SQLAlchemyAdapter: pre-ping detects
            # connections killed when another process drops/recreates a per-dataset
            # database (in-process cache eviction cannot reach other workers' pools);
            # recycle refreshes idle connections before NAT/load balancers cut them.
            # Pool sizing is deliberately leaner than the relational adapter's:
            # per-dataset graph engines multiply with datasets, so retain almost no
            # idle connections and serve bursts from overflow connections, which
            # close on release instead of idling.
            pool_args.setdefault("pool_size", 2)
            pool_args.setdefault("max_overflow", 20)
            pool_args.setdefault("pool_pre_ping", True)
            pool_args.setdefault("pool_recycle", 280)
            pool_args.setdefault("pool_timeout", 280)
        # Managed Postgres (e.g. Neon, RDS) requires SSL;
        # reuse the relational DATABASE_CONNECT_ARGS (asyncpg `ssl`) for the graph
        # engine too. Empty dict is a no-op for in-cluster Postgres.
        connect_args: dict = (
            dict(relational_config.database_connect_args)
            if relational_config.database_connect_args
            else {}
        )

        # Serialize JSONB columns once, at execute time, with the UUID/datetime-aware
        # encoder. This lets add_nodes/add_edges pass raw property dicts straight through
        # instead of doing a per-row json.loads(json.dumps(...)) round-trip, which on a
        # large single-batch write (e.g. War and Peace) generated millions of transient
        # dict/string allocations -> pymalloc arena fragmentation and cyclic-GC thrash.
        self.engine = create_async_engine(
            self.db_uri,
            json_serializer=lambda obj: json.dumps(obj, cls=JSONEncoder),
            connect_args=connect_args,
            **pool_args,
        )
        self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)
        self._write_lock = asyncio.Lock()

    async def close(self) -> None:
        """Dispose the connection pool. Called by ``closing_lru_cache`` on eviction."""
        await self.engine.dispose(close=True)

    async def initialize(self) -> None:
        """Create tables and indexes if they do not exist.

        This creates a fresh schema (including the graph-provenance columns defined
        in tables.py). Adding those columns to a graph_node/graph_edge left over from
        a pre-provenance release is handled by the ``postgres_graph_provenance_columns``
        data migration (create_all cannot ALTER an existing table).
        """
        async with self.engine.begin() as conn:
            await conn.run_sync(_meta.create_all, checkfirst=True)

    @asynccontextmanager
    async def _session(self) -> AsyncIterator[Any]:
        """Yield an async session from the underlying engine."""
        async with self.sessionmaker() as session:
            yield session

    def _serialize_properties(self, props: Dict[str, Any]) -> str:
        """Serialize a dict to a JSON string, handling datetimes and UUIDs."""
        return json.dumps(props, cls=JSONEncoder)

    def _parse_node_row(self, row) -> Dict[str, Any]:
        """Convert a (id, name, type, properties) row to a merged dict."""
        data = {"id": row.id, "name": row.name, "type": row.type}
        if row.properties is not None:
            props = (
                row.properties if isinstance(row.properties, dict) else json.loads(row.properties)
            )
            data.update(props)
        return data

    async def query(self, query_str: str, params: Optional[dict] = None) -> List[Any]:
        """Not supported. Use typed adapter methods or a Cypher-capable graph backend.

        Raises:
        -------
            NotImplementedError
        """
        raise NotImplementedError(
            "The Postgres graph backend does not support raw Cypher queries. "
            "Use a Cypher-capable graph backend (Neo4j, Ladybug) for raw query support, "
            "or use the typed adapter methods (add_nodes, get_neighbors, etc.)."
        )

    async def is_empty(self) -> bool:
        """Check whether the graph contains any nodes.

        Returns:
        --------
            bool: True if the graph has no nodes.
        """
        await self.initialize()
        async with self._session() as session:
            result = await session.execute(text("SELECT EXISTS(SELECT 1 FROM graph_node LIMIT 1)"))
            return not result.scalar()

    async def add_node(
        self, node: Union[DataPoint, str], properties: Optional[Dict[str, Any]] = None
    ) -> None:
        """Add a single node. Delegates to add_nodes.

        Parameters:
        -----------
            node: A DataPoint instance or a string node ID.
            properties: Optional property dict when node is a string ID.
        """
        if isinstance(node, str):
            props = properties or {}
            props.setdefault("id", node)
            await self.add_nodes([(node, props)])
        else:
            await self.add_nodes([node])

    async def add_nodes(
        self,
        nodes: Union[List[Tuple[str, Dict]], List[DataPoint]],
        source_ref_key: Optional[str] = None,
        pipeline_run_id: Optional[str] = None,
    ) -> None:
        """Add multiple nodes via batch upsert.

        Parameters:
        -----------
            nodes: A list of (id, properties) tuples or DataPoint instances.
        """
        if not nodes:
            return

        now = datetime.now(timezone.utc)
        core_keys = {"id", "name", "type"}

        rows = []
        for node in nodes:
            if isinstance(node, tuple):
                props = {**(node[1] or {}), "id": node[0]}
            elif hasattr(node, "model_dump"):
                props = node.model_dump()
            else:
                props = vars(node)

            extra = {k: v for k, v in props.items() if k not in core_keys}
            rows.append(
                {
                    "id": str(props.get("id", "")),
                    "name": str(props.get("name", "")),
                    "type": str(props.get("type", "")),
                    "properties": extra,
                    "created_at": now,
                    "updated_at": now,
                }
            )

        # Deduplicate by id (last wins) to avoid ON CONFLICT errors within one batch
        rows = list({r["id"]: r for r in rows}.values())

        # Fold graph provenance into the same upsert when a source ref is supplied
        # (Model A, atomic). Without one (plain add_node / non-provenance write) the
        # provenance columns are left to their '{}' default on insert and untouched
        # on conflict, so a non-provenance re-write never clobbers existing refs.
        provenance_set: Dict[str, Any] = {}
        if source_ref_key is not None:
            inputs = provenance_attach_inputs(source_ref_key, pipeline_run_id)
            insert_prov = _provenance_insert_values(inputs)
            for r in rows:
                r.update(insert_prov)
            provenance_set = _provenance_conflict_set(_node_table, inputs)

        async with self._write_lock:
            async with self._session() as session:
                for i in range(0, len(rows), _WRITE_CHUNK_SIZE):
                    chunk = rows[i : i + _WRITE_CHUNK_SIZE]
                    stmt = pg_insert(_node_table).values(chunk)
                    stmt = stmt.on_conflict_do_update(
                        index_elements=["id"],
                        set_={
                            "name": stmt.excluded.name,
                            "type": stmt.excluded.type,
                            "properties": stmt.excluded.properties,
                            "updated_at": func.now(),
                            **provenance_set,
                        },
                    )
                    await session.execute(stmt)
                await session.commit()

    async def delete_node(self, node_id: str) -> None:
        """Delete a single node. Delegates to delete_nodes.

        Parameters:
        -----------
            node_id: The ID of the node to delete.
        """
        await self.delete_nodes([node_id])

    async def delete_nodes(self, node_ids: List[str]) -> None:
        """Delete multiple nodes by ID. Cascade-deletes connected edges.

        Parameters:
        -----------
            node_ids: List of node IDs to delete.
        """
        if not node_ids:
            return
        async with self._write_lock:
            async with self._session() as session:
                await session.execute(
                    text("DELETE FROM graph_node WHERE id = ANY(:ids)"), {"ids": node_ids}
                )
                await session.commit()

    async def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve a single node by ID.

        Parameters:
        -----------
            node_id: The ID of the node to retrieve.

        Returns:
        --------
            A property dict for the node, or None if not found.
        """
        results = await self.get_nodes([node_id])
        return results[0] if results else None

    async def has_node(self, node_id: str) -> bool:
        """Return True when a node with the given id exists."""
        async with self._session() as session:
            result = await session.execute(
                text("SELECT EXISTS(SELECT 1 FROM graph_node WHERE id = :id)"),
                {"id": node_id},
            )
            return bool(result.scalar())

    async def get_nodes(self, node_ids: List[str]) -> List[Dict[str, Any]]:
        """Retrieve multiple nodes by ID.

        Parameters:
        -----------
            node_ids: List of node IDs to retrieve.

        Returns:
        --------
            A list of property dicts, one per found node.
        """
        if not node_ids:
            return []
        async with self._session() as session:
            result = await session.execute(
                text("SELECT id, name, type, properties FROM graph_node WHERE id = ANY(:ids)"),
                {"ids": node_ids},
            )
            return [self._parse_node_row(row) for row in result.fetchall()]

    async def add_edge(
        self,
        source_id: str,
        target_id: str,
        relationship_name: str,
        properties: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Add a single edge. Delegates to add_edges.

        Parameters:
        -----------
            source_id: Source node ID.
            target_id: Target node ID.
            relationship_name: The edge label.
            properties: Optional property dict for the edge.
        """
        await self.add_edges(
            [(str(source_id), str(target_id), relationship_name, properties or {})]
        )

    async def add_edges(
        self,
        edges: Union[List[Tuple[str, str, str, Optional[Dict[str, Any]]]], List],
        source_ref_key: Optional[str] = None,
        pipeline_run_id: Optional[str] = None,
    ) -> None:
        """Add multiple edges via batch upsert.

        Parameters:
        -----------
            edges: A list of (source_id, target_id, relationship_name, properties) tuples.
        """
        if not edges:
            return

        now = datetime.now(timezone.utc)

        rows = []
        for edge in edges:
            raw_props = edge[3] if len(edge) > 3 and edge[3] else {}
            rows.append(
                {
                    "source_id": str(edge[0]),
                    "target_id": str(edge[1]),
                    "relationship_name": edge[2],
                    "properties": raw_props,
                    "created_at": now,
                    "updated_at": now,
                }
            )

        # Deduplicate by composite key (last wins) to avoid ON CONFLICT errors within one batch
        rows = list(
            {(r["source_id"], r["target_id"], r["relationship_name"]): r for r in rows}.values()
        )

        # Fold graph provenance into the same upsert (see add_nodes for the rationale).
        provenance_set: Dict[str, Any] = {}
        if source_ref_key is not None:
            inputs = provenance_attach_inputs(source_ref_key, pipeline_run_id)
            insert_prov = _provenance_insert_values(inputs)
            for r in rows:
                r.update(insert_prov)
            provenance_set = _provenance_conflict_set(_edge_table, inputs)

        async with self._write_lock:
            async with self._session() as session:
                for i in range(0, len(rows), _WRITE_CHUNK_SIZE):
                    chunk = rows[i : i + _WRITE_CHUNK_SIZE]
                    stmt = pg_insert(_edge_table).values(chunk)
                    stmt = stmt.on_conflict_do_update(
                        index_elements=["source_id", "target_id", "relationship_name"],
                        set_={
                            "properties": stmt.excluded.properties,
                            "updated_at": func.now(),
                            **provenance_set,
                        },
                    )
                    await session.execute(stmt)
                await session.commit()

    async def has_edge(self, source_id: str, target_id: str, relationship_name: str) -> bool:
        """Check whether a single edge exists.

        Parameters:
        -----------
            source_id: Source node ID.
            target_id: Target node ID.
            relationship_name: The edge label.

        Returns:
        --------
            True if the edge exists.
        """
        result = await self.has_edges([(str(source_id), str(target_id), relationship_name)])
        return len(result) > 0

    async def has_edges(self, edges: List[Tuple[str, str, str]]) -> List[Tuple[str, str, str]]:
        """Check which of the given edges exist.

        Parameters:
        -----------
            edges: A list of (source_id, target_id, relationship_name) tuples to check.

        Returns:
        --------
            The subset of input tuples that exist in the database.
        """
        if not edges:
            return []

        # asyncpg caps bind parameters at 32767; each edge uses 3 params.
        CHUNK_SIZE = 10_000
        found: List[Tuple[str, str, str]] = []

        async with self._session() as session:
            for i in range(0, len(edges), CHUNK_SIZE):
                chunk = edges[i : i + CHUNK_SIZE]
                candidates = values(
                    sa_column("src", String),
                    sa_column("tgt", String),
                    sa_column("rel", String),
                    name="q",
                ).data([(str(s), str(t), str(r)) for s, t, r in chunk])

                stmt = select(candidates.c.src, candidates.c.tgt, candidates.c.rel).where(
                    exists(
                        select(text("1"))
                        .select_from(_edge_table)
                        .where(_edge_table.c.source_id == candidates.c.src)
                        .where(_edge_table.c.target_id == candidates.c.tgt)
                        .where(_edge_table.c.relationship_name == candidates.c.rel)
                    )
                )

                result = await session.execute(stmt)
                found.extend((row[0], row[1], row[2]) for row in result.fetchall())

        return found

    async def get_edges(self, node_id: str) -> List[Tuple[Dict[str, Any], str, Dict[str, Any]]]:
        """Retrieve all edges connected to a node.

        Parameters:
        -----------
            node_id: The ID of the node.

        Returns:
        --------
            A list of (source_dict, relationship_name, target_dict) tuples.
        """
        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT
                        n.id, n.name, n.type, n.properties,
                        e.relationship_name,
                        m.id, m.name, m.type, m.properties
                    FROM graph_edge e
                    JOIN graph_node n ON n.id = e.source_id
                    JOIN graph_node m ON m.id = e.target_id
                    WHERE e.source_id = :nid OR e.target_id = :nid
                """),
                {"nid": node_id},
            )
            edges = []
            for row in result.fetchall():
                src = {"id": row[0], "name": row[1], "type": row[2]}
                if row[3]:
                    src.update(row[3] if isinstance(row[3], dict) else json.loads(row[3]))
                tgt = {"id": row[5], "name": row[6], "type": row[7]}
                if row[8]:
                    tgt.update(row[8] if isinstance(row[8], dict) else json.loads(row[8]))
                edges.append((src, row[4], tgt))
            return edges

    async def get_neighbors(self, node_id: str) -> List[Dict[str, Any]]:
        """Retrieve all nodes directly connected to a given node.

        Parameters:
        -----------
            node_id: The ID of the node.

        Returns:
        --------
            A list of property dicts for neighboring nodes.
        """
        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT DISTINCT m.id, m.name, m.type, m.properties
                    FROM graph_edge e
                    JOIN graph_node m ON m.id = CASE
                        WHEN e.source_id = :nid THEN e.target_id
                        ELSE e.source_id
                    END
                    WHERE e.source_id = :nid OR e.target_id = :nid
                """),
                {"nid": node_id},
            )
            return [self._parse_node_row(row) for row in result.fetchall()]

    async def get_connections(
        self, node_id: Union[str, UUID]
    ) -> List[Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]]:
        """Retrieve all connections (source, edge, target) for a node.

        Parameters:
        -----------
            node_id: The ID of the node.

        Returns:
        --------
            A list of (source_dict, edge_dict, target_dict) tuples.
        """
        nid = str(node_id)

        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT
                        n.id, n.name, n.type, n.properties,
                        e.relationship_name, e.properties AS edge_props,
                        m.id, m.name, m.type, m.properties
                    FROM graph_edge e
                    JOIN graph_node n ON n.id = e.source_id
                    JOIN graph_node m ON m.id = e.target_id
                    WHERE e.source_id = :nid OR e.target_id = :nid
                """),
                {"nid": nid},
            )

            connections = []
            for row in result.fetchall():
                src = {"id": row[0], "name": row[1], "type": row[2]}
                if row[3]:
                    src.update(row[3] if isinstance(row[3], dict) else json.loads(row[3]))

                edge = {"relationship_name": row[4]}
                if row[5]:
                    edge_props = row[5] if isinstance(row[5], dict) else json.loads(row[5])
                    edge.update(edge_props)

                tgt = {"id": row[6], "name": row[7], "type": row[8]}
                if row[9]:
                    tgt.update(row[9] if isinstance(row[9], dict) else json.loads(row[9]))

                connections.append((src, edge, tgt))
            return connections

    async def get_graph_data(
        self,
    ) -> Tuple[List[Tuple[str, Dict[str, Any]]], List[Tuple[str, str, str, Dict[str, Any]]]]:
        """Retrieve all nodes and edges in the graph.

        Returns:
        --------
            A tuple of (nodes, edges) where nodes are (id, props) and
            edges are (source_id, target_id, relationship_name, props).
        """
        async with self._session() as session:
            node_result = await session.execute(
                text("SELECT id, name, type, properties FROM graph_node")
            )
            nodes = []
            for row in node_result.fetchall():
                data = {"name": row[1], "type": row[2]}
                if row[3]:
                    data.update(row[3] if isinstance(row[3], dict) else json.loads(row[3]))
                nodes.append((row[0], data))

            if not nodes:
                return [], []

            edge_result = await session.execute(
                text("""
                    SELECT source_id, target_id, relationship_name, properties
                    FROM graph_edge
                """)
            )
            edges = []
            for row in edge_result.fetchall():
                props = {}
                if row[3]:
                    props = row[3] if isinstance(row[3], dict) else json.loads(row[3])
                edges.append((row[0], row[1], row[2], props))

            return nodes, edges

    async def get_id_filtered_graph_data(
        self, target_ids: List[str]
    ) -> Tuple[List[Tuple[str, Dict[str, Any]]], List[Tuple[str, str, str, Dict[str, Any]]]]:
        """Retrieve the subgraph touching target_ids: edges with either endpoint
        in the set, plus all endpoint nodes of those edges (edge-driven,
        matching the Ladybug/Neo4j contract). Lets CogneeGraph project only the
        vector-search neighborhood instead of the full graph.
        """
        if not target_ids:
            return [], []
        ids = [str(i) for i in target_ids]

        async with self._session() as session:
            edge_result = await session.execute(
                text("""
                    SELECT source_id, target_id, relationship_name, properties
                    FROM graph_edge
                    WHERE source_id = ANY(:ids) OR target_id = ANY(:ids)
                """),
                {"ids": ids},
            )
            edges = []
            endpoint_ids = set()
            for row in edge_result.fetchall():
                props = {}
                if row[3]:
                    props = row[3] if isinstance(row[3], dict) else json.loads(row[3])
                endpoint_ids.update((row[0], row[1]))
                edges.append((row[0], row[1], row[2], props))

            if not endpoint_ids:
                return [], []

            node_result = await session.execute(
                text("SELECT id, name, type, properties FROM graph_node WHERE id = ANY(:ids)"),
                {"ids": list(endpoint_ids)},
            )
            nodes = []
            for row in node_result.fetchall():
                data = {"name": row[1], "type": row[2]}
                if row[3]:
                    data.update(row[3] if isinstance(row[3], dict) else json.loads(row[3]))
                nodes.append((row[0], data))

            return nodes, edges

    async def get_filtered_graph_data(
        self, attribute_filters: List[Dict[str, List[Union[str, int]]]]
    ) -> Tuple[List[Tuple[str, Dict]], List[Tuple[str, str, str, Dict]]]:
        """Retrieve nodes matching attribute filters, plus edges between them.

        Parameters:
        -----------
            attribute_filters: A list of {attr: [values]} dicts. Only 'id',
                'name', and 'type' are valid filter attributes.

        Returns:
        --------
            A tuple of (nodes, edges) matching the filters.
        """
        if not attribute_filters:
            return await self.get_graph_data()

        # Validate attribute names against whitelist to prevent SQL injection
        where_parts = []
        params = {}
        for i, filter_dict in enumerate(attribute_filters):
            for attr, filter_values in filter_dict.items():
                if attr not in self._ALLOWED_FILTER_ATTRS:
                    raise ValueError(f"Invalid filter attribute: {attr!r}")
                param = f"filt_{i}_{attr}"
                where_parts.append(f"n.{attr} = ANY(:{param})")
                params[param] = filter_values

        if not where_parts:
            return await self.get_graph_data()

        where_clause = " AND ".join(where_parts)

        async with self._session() as session:
            result = await session.execute(
                text(f"""
                    WITH filtered_nodes AS (
                        SELECT id, name, type, properties
                        FROM graph_node n
                        WHERE {where_clause}
                    )
                    SELECT 'node' AS kind, fn.id, fn.name, fn.type, fn.properties,
                           NULL AS source_id, NULL AS target_id,
                           NULL AS relationship_name, NULL AS edge_props
                    FROM filtered_nodes fn
                    UNION ALL
                    SELECT 'edge', NULL, NULL, NULL, NULL,
                           e.source_id, e.target_id,
                           e.relationship_name, e.properties
                    FROM graph_edge e
                    WHERE e.source_id IN (SELECT id FROM filtered_nodes)
                      AND e.target_id IN (SELECT id FROM filtered_nodes)
                """),
                params,
            )

            nodes = []
            edges = []
            for row in result.fetchall():
                if row[0] == "node":
                    data = {"name": row[2], "type": row[3]}
                    if row[4]:
                        data.update(row[4] if isinstance(row[4], dict) else json.loads(row[4]))
                    nodes.append((row[1], data))
                else:
                    props = {}
                    if row[8]:
                        props = row[8] if isinstance(row[8], dict) else json.loads(row[8])
                    edges.append((row[5], row[6], row[7], props))

            return nodes, edges

    async def get_nodeset_subgraph(
        self, node_type: Type[Any], node_name: List[str], node_name_filter_operator: str = "OR"
    ) -> Tuple[List[Tuple[str, dict]], List[Tuple[str, str, str, dict]]]:
        """Retrieve a subgraph containing matching nodes, their neighbors, and interconnecting edges.

        Parameters:
        -----------
            node_type: The DataPoint subclass whose __name__ is the type label.
            node_name: List of node names to match.

        Returns:
        --------
            A tuple of (nodes, edges) for the subgraph.
        """
        label = node_type.__name__

        # OR: neighbor of any primary node qualifies
        # AND: neighbor must be connected to every primary node
        if node_name_filter_operator == "OR":
            neighbor_cte = """
                    neighbor_ids AS (
                        SELECT DISTINCT CASE
                            WHEN e.source_id IN (SELECT id FROM primary_nodes)
                            THEN e.target_id ELSE e.source_id
                        END AS id
                        FROM graph_edge e
                        WHERE e.source_id IN (SELECT id FROM primary_nodes)
                           OR e.target_id IN (SELECT id FROM primary_nodes)
                    )"""
        else:
            neighbor_cte = """
                    neighbor_ids AS (
                        SELECT nbr_id AS id FROM (
                            SELECT CASE
                                WHEN e.source_id IN (SELECT id FROM primary_nodes)
                                THEN e.target_id ELSE e.source_id
                            END AS nbr_id,
                            CASE
                                WHEN e.source_id IN (SELECT id FROM primary_nodes)
                                THEN e.source_id ELSE e.target_id
                            END AS primary_id
                            FROM graph_edge e
                            WHERE e.source_id IN (SELECT id FROM primary_nodes)
                               OR e.target_id IN (SELECT id FROM primary_nodes)
                        ) sub
                        GROUP BY nbr_id
                        HAVING COUNT(DISTINCT primary_id) = :primary_count
                    )"""

        query_str = f"""
                    WITH primary_nodes AS (
                        SELECT DISTINCT id
                        FROM graph_node
                        WHERE type = :label AND name = ANY(:names)
                    ),
                    {neighbor_cte},
                    all_ids AS (
                        SELECT id FROM primary_nodes
                        UNION
                        SELECT id FROM neighbor_ids
                    )
                    SELECT 'node' AS kind,
                           n.id, n.name, n.type, n.properties,
                           NULL AS source_id, NULL AS target_id,
                           NULL AS relationship_name, NULL AS edge_props
                    FROM graph_node n
                    WHERE n.id IN (SELECT id FROM all_ids)
                    UNION ALL
                    SELECT 'edge', NULL, NULL, NULL, NULL,
                           e.source_id, e.target_id,
                           e.relationship_name, e.properties
                    FROM graph_edge e
                    WHERE e.source_id IN (SELECT id FROM all_ids)
                      AND e.target_id IN (SELECT id FROM all_ids)
                """

        params = {"label": label, "names": node_name}
        if node_name_filter_operator != "OR":
            params["primary_count"] = len(node_name)

        async with self._session() as session:
            result = await session.execute(text(query_str), params)

            nodes = []
            edges = []
            for row in result.fetchall():
                if row[0] == "node":
                    data = {"name": row[2], "type": row[3]}
                    if row[4]:
                        data.update(row[4] if isinstance(row[4], dict) else json.loads(row[4]))
                    nodes.append((row[1], data))
                else:
                    props = {}
                    if row[8]:
                        props = row[8] if isinstance(row[8], dict) else json.loads(row[8])
                    edges.append((row[5], row[6], row[7], props))

            return nodes, edges

    async def get_graph_metrics(self, include_optional: bool = False) -> Dict[str, Any]:
        """Compute graph metrics (node/edge counts, degree, density, components).

        Parameters:
        -----------
            include_optional: If True, also compute self-loop count.

        Returns:
        --------
            A dict of metric names to values. Diameter, avg shortest path,
            and clustering return -1 (not computed).
        """
        async with self._session() as session:
            n_result = await session.execute(text("SELECT count(*) FROM graph_node"))
            num_nodes = n_result.scalar()
            e_result = await session.execute(text("SELECT count(*) FROM graph_edge"))
            num_edges = e_result.scalar()

            mean_degree = (2 * num_edges) / num_nodes if num_nodes else None
            edge_density = num_edges / (num_nodes * (num_nodes - 1)) if num_nodes > 1 else 0

            # Connected components via recursive CTE
            comp_result = await session.execute(
                text("""
                WITH RECURSIVE component AS (
                    SELECT id AS node_id, id AS comp_root
                    FROM graph_node
                    UNION
                    SELECT
                        CASE WHEN e.source_id = c.node_id THEN e.target_id ELSE e.source_id END,
                        c.comp_root
                    FROM component c
                    JOIN graph_edge e ON e.source_id = c.node_id OR e.target_id = c.node_id
                ),
                node_comp AS (
                    SELECT node_id, MIN(comp_root) AS comp_id
                    FROM component
                    GROUP BY node_id
                )
                SELECT comp_id, count(*) AS sz
                FROM node_comp
                GROUP BY comp_id
                ORDER BY sz DESC
            """)
            )
            comp_rows = comp_result.fetchall()
            num_components = len(comp_rows)
            component_sizes = [row[1] for row in comp_rows]

            metrics = {
                "num_nodes": num_nodes,
                "num_edges": num_edges,
                "mean_degree": mean_degree,
                "edge_density": edge_density,
                "num_connected_components": num_components,
                "sizes_of_connected_components": component_sizes,
            }

            if include_optional:
                sl_result = await session.execute(
                    text("SELECT count(*) FROM graph_edge WHERE source_id = target_id")
                )
                metrics["num_selfloops"] = sl_result.scalar()
                metrics["diameter"] = -1
                metrics["avg_shortest_path_length"] = -1
                metrics["avg_clustering"] = -1
            else:
                metrics["num_selfloops"] = -1
                metrics["diameter"] = -1
                metrics["avg_shortest_path_length"] = -1
                metrics["avg_clustering"] = -1

            return metrics

    async def get_neighborhood(
        self,
        node_ids: List[str],
        depth: int = 1,
        edge_types: Optional[List[str]] = None,
    ) -> Tuple[List[Tuple[str, Dict[str, Any]]], List[Tuple[str, str, str, Dict[str, Any]]]]:
        """Get the k-hop neighborhood subgraph around seed nodes.

        Uses a single recursive CTE query to collect all node IDs within
        `depth` hops, then returns nodes and edges for that subgraph.
        """
        if not node_ids:
            return [], []

        # Optional edge type filter for the CTE traversal
        edge_filter = ""
        if edge_types:
            placeholders = ", ".join(f":et_{i}" for i in range(len(edge_types)))
            edge_filter = f"AND e.relationship_name IN ({placeholders})"

        # Single query: recursive CTE finds reachable IDs, then joins
        # nodes and edges in two unioned result sets distinguished by 'kind'
        query_str = f"""
            WITH RECURSIVE neighborhood(id, hops) AS (
                SELECT unnest(CAST(:seeds AS text[])), 0
              UNION
                SELECT CASE WHEN e.source_id = n.id THEN e.target_id
                            ELSE e.source_id END,
                       n.hops + 1
                FROM neighborhood n
                JOIN graph_edge e ON (e.source_id = n.id OR e.target_id = n.id)
                    {edge_filter}
                WHERE n.hops < :depth
            ),
            ids AS (SELECT DISTINCT id FROM neighborhood)

            SELECT 'node' AS kind,
                   gn.id, gn.name, gn.type, gn.properties,
                   NULL AS source_id, NULL AS target_id,
                   NULL AS relationship_name, NULL AS edge_properties
            FROM graph_node gn
            JOIN ids ON gn.id = ids.id

            UNION ALL

            SELECT 'edge' AS kind,
                   NULL, NULL, NULL, NULL,
                   ge.source_id, ge.target_id,
                   ge.relationship_name, ge.properties
            FROM graph_edge ge
            WHERE ge.source_id IN (SELECT id FROM ids)
              AND ge.target_id IN (SELECT id FROM ids)
        """

        params: Dict[str, Any] = {"seeds": list(node_ids), "depth": depth}
        if edge_types:
            for i, et in enumerate(edge_types):
                params[f"et_{i}"] = et

        async with self._session() as session:
            result = await session.execute(text(query_str), params)

            nodes = []
            edges = []
            for row in result.fetchall():
                if row.kind == "node":
                    data = self._parse_node_row(row)
                    data.pop("id", None)
                    nodes.append((row.id, data))
                else:
                    props = {}
                    if row.edge_properties is not None:
                        props = (
                            row.edge_properties
                            if isinstance(row.edge_properties, dict)
                            else json.loads(row.edge_properties)
                        )
                    edges.append((row.source_id, row.target_id, row.relationship_name, props))

            return nodes, edges

    async def delete_graph(self) -> None:
        """Delete all nodes and edges from the graph."""
        await self.initialize()
        async with self._write_lock:
            async with self._session() as session:
                await session.execute(text("TRUNCATE graph_edge, graph_node CASCADE"))
                await session.commit()

    # ------------------------------------------------------------------ #
    # Graph provenance (COG-5522 Part 1).                                  #
    #                                                                      #
    # The four provenance fields live in declared ``text[]`` columns on    #
    # both graph_node and graph_edge — never inside the JSON ``properties`` #
    # blob — so delete/rollback can filter by source ref, dataset id, or    #
    # pipeline run id with an array-membership scan. The pure set-merge /   #
    # derive logic is shared with every other adapter via                  #
    # ``provenance.source_ref_state``; only the storage I/O is per-backend. #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _props_dict(raw: Any) -> Dict[str, Any]:
        """Decode a JSONB ``properties`` value, tolerating dict, str, or empty."""
        if not raw:
            return {}
        return raw if isinstance(raw, dict) else json.loads(raw)

    @staticmethod
    def _node_identity_row(node_id: str) -> dict:
        return {"id": node_id}

    @staticmethod
    def _edge_identity_row(edge: EdgeIdentity) -> dict:
        return {"s": edge.source_id, "t": edge.target_id, "rel": edge.relationship_name}

    async def _read_node_provenance(
        self, node_ids: List[str]
    ) -> Dict[str, Tuple[List[str], List[str]]]:
        """Return ``{node_id: (source_ref_keys, source_run_refs)}`` for existing nodes."""
        async with self._session() as session:
            result = await session.execute(
                text(
                    "SELECT id, source_ref_keys, source_run_refs "
                    "FROM graph_node WHERE id = ANY(:ids)"
                ),
                {"ids": list(node_ids)},
            )
            return {row[0]: (list(row[1] or []), list(row[2] or [])) for row in result.fetchall()}

    async def _write_node_provenance(self, batch: List[dict]) -> None:
        """Overwrite the four provenance columns for each ``{id, refs, ...}`` row.

        Casts are explicit so asyncpg sends Python lists as ``text[]`` instead of
        failing to infer the parameter type.
        """
        if not batch:
            return
        async with self._session() as session:
            for row in batch:
                await session.execute(
                    text(
                        "UPDATE graph_node SET "
                        "source_ref_keys = CAST(:refs AS text[]), "
                        "source_dataset_ids = CAST(:datasets AS text[]), "
                        "source_run_ids = CAST(:runs AS text[]), "
                        "source_run_refs = CAST(:run_refs AS text[]), "
                        "updated_at = now() WHERE id = :id"
                    ),
                    {
                        "id": row["id"],
                        "refs": row["refs"],
                        "datasets": row["datasets"],
                        "runs": row["runs"],
                        "run_refs": row["run_refs"],
                    },
                )
            await session.commit()

    async def _read_edge_provenance(
        self, edges: List[EdgeIdentity]
    ) -> Dict[EdgeIdentity, Tuple[List[str], List[str]]]:
        """Return ``{edge: (source_ref_keys, source_run_refs)}`` for existing edges."""
        src = [e.source_id for e in edges]
        tgt = [e.target_id for e in edges]
        rel = [e.relationship_name for e in edges]
        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT e.source_id, e.target_id, e.relationship_name,
                           e.source_ref_keys, e.source_run_refs
                    FROM graph_edge e
                    JOIN unnest(CAST(:src AS text[]), CAST(:tgt AS text[]), CAST(:rel AS text[]))
                         AS q(s, t, r)
                      ON e.source_id = q.s AND e.target_id = q.t AND e.relationship_name = q.r
                """),
                {"src": src, "tgt": tgt, "rel": rel},
            )
            out: Dict[EdgeIdentity, Tuple[List[str], List[str]]] = {}
            for row in result.fetchall():
                edge = EdgeIdentity(source_id=row[0], target_id=row[1], relationship_name=row[2])
                out[edge] = (list(row[3] or []), list(row[4] or []))
            return out

    async def _write_edge_provenance(self, batch: List[dict]) -> None:
        """Overwrite the four provenance columns for each ``{s, t, rel, refs, ...}`` row."""
        if not batch:
            return
        async with self._session() as session:
            for row in batch:
                await session.execute(
                    text(
                        "UPDATE graph_edge SET "
                        "source_ref_keys = CAST(:refs AS text[]), "
                        "source_dataset_ids = CAST(:datasets AS text[]), "
                        "source_run_ids = CAST(:runs AS text[]), "
                        "source_run_refs = CAST(:run_refs AS text[]), "
                        "updated_at = now() "
                        "WHERE source_id = :s AND target_id = :t AND relationship_name = :rel"
                    ),
                    {
                        "s": row["s"],
                        "t": row["t"],
                        "rel": row["rel"],
                        "refs": row["refs"],
                        "datasets": row["datasets"],
                        "runs": row["runs"],
                        "run_refs": row["run_refs"],
                    },
                )
            await session.commit()

    async def _apply_source_ref_change(
        self,
        artifacts,
        read_provenance,
        write_provenance,
        identity_row,
        transition,
    ) -> None:
        """Read each artifact's provenance, apply a pure transition, write it back.

        Shared by attach/remove for both nodes and edges. The write lock serializes
        this read-modify-write within one adapter instance so concurrent explicit
        attach/remove calls do not overwrite each other's provenance updates.
        """
        if not artifacts:
            return
        async with self._write_lock:
            current = await read_provenance(artifacts)
            batch = []
            for identity, (keys, run_refs) in current.items():
                cols = transition(keys, run_refs)
                batch.append(
                    {
                        **identity_row(identity),
                        "refs": cols.source_ref_keys,
                        "datasets": cols.source_dataset_ids,
                        "runs": cols.source_run_ids,
                        "run_refs": cols.source_run_refs,
                    }
                )
            await write_provenance(batch)

    async def attach_node_source_refs(
        self,
        node_ids: list[str],
        source_ref_keys: list[str],
        pipeline_run_id: str | None = None,
    ) -> None:
        if not source_ref_keys:
            return
        add_keys = list(source_ref_keys)
        await self._apply_source_ref_change(
            node_ids,
            self._read_node_provenance,
            self._write_node_provenance,
            self._node_identity_row,
            lambda keys, run_refs: provenance_after_attach(
                keys, run_refs, add_keys, pipeline_run_id
            ),
        )

    async def attach_edge_source_refs(
        self,
        edges: list[EdgeIdentity],
        source_ref_keys: list[str],
        pipeline_run_id: str | None = None,
    ) -> None:
        if not source_ref_keys:
            return
        add_keys = list(source_ref_keys)
        await self._apply_source_ref_change(
            edges,
            self._read_edge_provenance,
            self._write_edge_provenance,
            self._edge_identity_row,
            lambda keys, run_refs: provenance_after_attach(
                keys, run_refs, add_keys, pipeline_run_id
            ),
        )

    async def remove_node_source_refs(
        self,
        node_ids: list[str],
        source_ref_keys: list[str],
    ) -> None:
        if not source_ref_keys:
            return
        remove_keys = list(source_ref_keys)
        await self._apply_source_ref_change(
            node_ids,
            self._read_node_provenance,
            self._write_node_provenance,
            self._node_identity_row,
            lambda keys, run_refs: provenance_after_remove(keys, run_refs, remove_keys),
        )

    async def remove_edge_source_refs(
        self,
        edges: list[EdgeIdentity],
        source_ref_keys: list[str],
    ) -> None:
        if not source_ref_keys:
            return
        remove_keys = list(source_ref_keys)
        await self._apply_source_ref_change(
            edges,
            self._read_edge_provenance,
            self._write_edge_provenance,
            self._edge_identity_row,
            lambda keys, run_refs: provenance_after_remove(keys, run_refs, remove_keys),
        )

    async def delete_edge_triples(self, edges: list[EdgeIdentity]) -> None:
        """Delete edges by (source, target, relationship); keep the endpoint nodes."""
        if not edges:
            return
        src = [e.source_id for e in edges]
        tgt = [e.target_id for e in edges]
        rel = [e.relationship_name for e in edges]
        async with self._write_lock:
            async with self._session() as session:
                await session.execute(
                    text("""
                        DELETE FROM graph_edge e
                        USING unnest(CAST(:src AS text[]), CAST(:tgt AS text[]),
                                     CAST(:rel AS text[])) AS q(s, t, r)
                        WHERE e.source_id = q.s AND e.target_id = q.t
                          AND e.relationship_name = q.r
                    """),
                    {"src": src, "tgt": tgt, "rel": rel},
                )
                await session.commit()

    async def get_node_delete_data(self, node_ids: list[str]) -> dict[str, NodeDeleteData]:
        if not node_ids:
            return {}
        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT id, name, type, properties,
                           source_ref_keys, source_dataset_ids, source_run_ids, source_run_refs
                    FROM graph_node WHERE id = ANY(:ids)
                """),
                {"ids": list(node_ids)},
            )
            out: dict[str, NodeDeleteData] = {}
            for row in result.fetchall():
                properties = self._props_dict(row[3])
                # Reconstruct the flat payload the way get_node does: core columns
                # merged over the JSON blob.
                properties["id"] = row[0]
                properties["name"] = row[1]
                properties["type"] = row[2]
                metadata = properties.get("metadata") or {}
                indexed_fields = (
                    list(metadata.get("index_fields") or []) if isinstance(metadata, dict) else []
                )
                out[row[0]] = NodeDeleteData(
                    node_id=row[0],
                    node_type=row[2] or "",
                    indexed_fields=indexed_fields,
                    node_properties=properties,
                    source_ref_keys=list(row[4] or []),
                    source_dataset_ids=list(row[5] or []),
                    source_run_ids=list(row[6] or []),
                    source_run_refs=list(row[7] or []),
                )
            return out

    async def get_edge_delete_data(
        self, edges: list[EdgeIdentity]
    ) -> dict[EdgeIdentity, EdgeDeleteData]:
        if not edges:
            return {}
        src = [e.source_id for e in edges]
        tgt = [e.target_id for e in edges]
        rel = [e.relationship_name for e in edges]
        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT e.source_id, e.target_id, e.relationship_name, e.properties,
                           e.source_ref_keys, e.source_dataset_ids,
                           e.source_run_ids, e.source_run_refs
                    FROM graph_edge e
                    JOIN unnest(CAST(:src AS text[]), CAST(:tgt AS text[]), CAST(:rel AS text[]))
                         AS q(s, t, r)
                      ON e.source_id = q.s AND e.target_id = q.t AND e.relationship_name = q.r
                """),
                {"src": src, "tgt": tgt, "rel": rel},
            )
            rows = result.fetchall()

        # Lazy import: prepare_edges_for_storage lives in the modules layer, whose
        # package __init__ imports get_graph_engine -> this adapter. Importing it at
        # module load would create a cycle; at delete-time it is safe.
        from cognee.modules.graph.utils.prepare_edges_for_storage import get_edge_retrieval_text

        out: dict[EdgeIdentity, EdgeDeleteData] = {}
        for row in rows:
            edge = EdgeIdentity(source_id=row[0], target_id=row[1], relationship_name=row[2])
            properties = self._props_dict(row[3])
            # Stored edge_text wins; fall back to relationship_name when absent.
            edge_text = get_edge_retrieval_text(properties.get("edge_text"), edge.relationship_name)
            out[edge] = EdgeDeleteData(
                edge=edge,
                edge_text=edge_text,
                edge_properties=properties,
                source_ref_keys=list(row[4] or []),
                source_dataset_ids=list(row[5] or []),
                source_run_ids=list(row[6] or []),
                source_run_refs=list(row[7] or []),
            )
        return out

    async def find_nodes_by_source_ref(self, source_ref_key: str) -> list[str]:
        async with self._session() as session:
            result = await session.execute(
                text("SELECT id FROM graph_node WHERE :token = ANY(source_ref_keys)"),
                {"token": source_ref_key},
            )
            return [row[0] for row in result.fetchall()]

    async def find_edges_by_source_ref(self, source_ref_key: str) -> list[EdgeIdentity]:
        async with self._session() as session:
            result = await session.execute(
                text(
                    "SELECT source_id, target_id, relationship_name "
                    "FROM graph_edge WHERE :token = ANY(source_ref_keys)"
                ),
                {"token": source_ref_key},
            )
            return [
                EdgeIdentity(source_id=row[0], target_id=row[1], relationship_name=row[2])
                for row in result.fetchall()
            ]

    async def find_node_source_refs_by_dataset(self, dataset_id: str) -> dict[str, list[str]]:
        async with self._session() as session:
            result = await session.execute(
                text(
                    "SELECT id, source_ref_keys "
                    "FROM graph_node WHERE :token = ANY(source_dataset_ids)"
                ),
                {"token": dataset_id},
            )
            out: dict[str, list[str]] = {}
            for row in result.fetchall():
                owned = [
                    key
                    for key in (row[1] or [])
                    if str(get_dataset_id_from_source_ref_key(key)) == dataset_id
                ]
                if owned:
                    out[row[0]] = owned
            return out

    async def find_edge_source_refs_by_dataset(
        self, dataset_id: str
    ) -> dict[EdgeIdentity, list[str]]:
        async with self._session() as session:
            result = await session.execute(
                text(
                    "SELECT source_id, target_id, relationship_name, source_ref_keys "
                    "FROM graph_edge WHERE :token = ANY(source_dataset_ids)"
                ),
                {"token": dataset_id},
            )
            out: dict[EdgeIdentity, list[str]] = {}
            for row in result.fetchall():
                owned = [
                    key
                    for key in (row[3] or [])
                    if str(get_dataset_id_from_source_ref_key(key)) == dataset_id
                ]
                if owned:
                    edge = EdgeIdentity(
                        source_id=row[0], target_id=row[1], relationship_name=row[2]
                    )
                    out[edge] = owned
            return out

    async def find_node_source_refs_by_pipeline_run(
        self, pipeline_run_id: str
    ) -> dict[str, list[str]]:
        async with self._session() as session:
            result = await session.execute(
                text(
                    "SELECT id, source_run_refs FROM graph_node WHERE :token = ANY(source_run_ids)"
                ),
                {"token": pipeline_run_id},
            )
            out: dict[str, list[str]] = {}
            for row in result.fetchall():
                contributed = [
                    get_source_ref_key_from_source_run_ref(ref)
                    for ref in (row[1] or [])
                    if str(get_pipeline_run_id_from_source_run_ref(ref)) == pipeline_run_id
                ]
                if contributed:
                    out[row[0]] = contributed
            return out

    async def find_edge_source_refs_by_pipeline_run(
        self, pipeline_run_id: str
    ) -> dict[EdgeIdentity, list[str]]:
        async with self._session() as session:
            result = await session.execute(
                text(
                    "SELECT source_id, target_id, relationship_name, source_run_refs "
                    "FROM graph_edge WHERE :token = ANY(source_run_ids)"
                ),
                {"token": pipeline_run_id},
            )
            out: dict[EdgeIdentity, list[str]] = {}
            for row in result.fetchall():
                contributed = [
                    get_source_ref_key_from_source_run_ref(ref)
                    for ref in (row[3] or [])
                    if str(get_pipeline_run_id_from_source_run_ref(ref)) == pipeline_run_id
                ]
                if contributed:
                    edge = EdgeIdentity(
                        source_id=row[0], target_id=row[1], relationship_name=row[2]
                    )
                    out[edge] = contributed
            return out

    async def set_graph_metadata(self, metadata: dict[str, str]) -> None:
        if not metadata:
            return
        await self.initialize()
        async with self._write_lock:
            async with self._session() as session:
                for key, value in metadata.items():
                    stmt = pg_insert(_metadata_table).values(key=str(key), value=str(value))
                    stmt = stmt.on_conflict_do_update(
                        index_elements=["key"],
                        set_={"value": stmt.excluded.value},
                    )
                    await session.execute(stmt)
                await session.commit()

    async def get_graph_metadata(self) -> dict[str, str]:
        await self.initialize()
        async with self._session() as session:
            result = await session.execute(text("SELECT key, value FROM graph_metadata"))
            return {row[0]: row[1] for row in result.fetchall()}

    async def remove_belongs_to_set_tags(
        self,
        tags: List[str],
        node_ids: Optional[List[str]] = None,
    ) -> None:
        """Strip ``tags`` from each node's ``belongs_to_set`` property array.

        Keeps the graph node's denormalized membership list consistent with the
        additive belongs_to_set edges after a NodeSet (or its dataset) is deleted.
        ``belongs_to_set`` lives inside the JSONB ``properties`` blob (it is not a
        core column), so this is a read-filter-write over that array. When
        ``node_ids`` is given, only those nodes are reconciled.
        """
        if not tags:
            return None
        if node_ids is not None and not node_ids:
            return None

        tag_set = set(tags)
        async with self._session() as session:
            if node_ids is not None:
                result = await session.execute(
                    text("SELECT id, properties FROM graph_node WHERE id = ANY(:ids)"),
                    {"ids": [str(nid) for nid in node_ids]},
                )
            else:
                result = await session.execute(text("SELECT id, properties FROM graph_node"))
            rows = result.fetchall()

        updates = []
        for row in rows:
            properties = self._props_dict(row[1])
            current = properties.get("belongs_to_set")
            if not isinstance(current, list) or not any(tag in tag_set for tag in current):
                continue
            properties["belongs_to_set"] = [tag for tag in current if tag not in tag_set]
            updates.append({"id": row[0], "properties": json.dumps(properties, cls=JSONEncoder)})

        if updates:
            async with self._write_lock:
                async with self._session() as session:
                    for update in updates:
                        await session.execute(
                            text(
                                "UPDATE graph_node SET properties = CAST(:p AS jsonb), "
                                "updated_at = now() WHERE id = :id"
                            ),
                            {"id": update["id"], "p": update["properties"]},
                        )
                    await session.commit()
        return None

    async def get_triplets_batch(self, offset: int, limit: int) -> List[Dict[str, Any]]:
        """Retrieve a batch of (source, relationship, target) triplets.

        Parameters:
        -----------
            offset: Number of triplets to skip.
            limit: Maximum number of triplets to return.

        Returns:
        --------
            A list of dicts with 'start_node', 'relationship_properties',
            and 'end_node' keys.
        """
        if offset < 0:
            raise ValueError(f"Offset must be non-negative, got {offset}")
        if limit < 0:
            raise ValueError(f"Limit must be non-negative, got {limit}")

        async with self._session() as session:
            result = await session.execute(
                text("""
                    SELECT
                        s.id, s.name, s.type, s.properties,
                        e.relationship_name, e.properties AS edge_props,
                        t.id, t.name, t.type, t.properties
                    FROM graph_edge e
                    JOIN graph_node s ON s.id = e.source_id
                    JOIN graph_node t ON t.id = e.target_id
                    ORDER BY e.source_id, e.target_id, e.relationship_name
                    OFFSET :off LIMIT :lim
                """),
                {"off": offset, "lim": limit},
            )

            triplets = []
            for row in result.fetchall():
                start_node = {"id": row[0], "name": row[1], "type": row[2]}
                if row[3]:
                    start_node.update(row[3] if isinstance(row[3], dict) else json.loads(row[3]))

                rel = {"relationship_name": row[4]}
                if row[5]:
                    rel_props = row[5] if isinstance(row[5], dict) else json.loads(row[5])
                    rel.update(rel_props)

                end_node = {"id": row[6], "name": row[7], "type": row[8]}
                if row[9]:
                    end_node.update(row[9] if isinstance(row[9], dict) else json.loads(row[9]))

                triplets.append(
                    {
                        "start_node": start_node,
                        "relationship_properties": rel,
                        "end_node": end_node,
                    }
                )
            return triplets