consistent KeyErrors across backends

src/tracksdata/graph/_base_graph.py

@@ -163,6 +163,53 @@ def _validate_attributes(
                 f"{mode} attribute keys not found in attrs: '{missing_keys}'\nRequested keys: '{reference_keys}'"
             )
 
+    def _validate_attr_keys(
+        self,
+        attr_keys: Sequence[str] | str | None,
+        mode: Literal["node", "edge"],
+    ) -> None:
+        """
+        Validate that attribute keys being *read* exist on this graph.
+
+        Read-path counterpart of `_validate_attributes`. Raises `KeyError`, not `ValueError`,
+        because asking for a key that isn't there is a lookup miss (mirroring `dict[missing]`),
+        whereas `_validate_attributes` guards a *write* against a declared schema, which is a
+        bad-argument situation.
+
+        Without a central guard each backend leaks whatever its storage layer raises for an
+        unknown column -- `AttributeError` from SQLAlchemy's `getattr`, `KeyError` from a dict
+        lookup, or polars' `ColumnNotFoundError` (which is not a `KeyError` subclass) -- so the
+        same call raised three different types depending on the backend.
+
+        Parameters
+        ----------
+        attr_keys : Sequence[str] | str | None
+            The attribute keys to validate. `None` means "all keys" and is always valid.
+        mode : Literal["node", "edge"]
+            Whether to validate against node or edge attribute keys.
+
+        Raises
+        ------
+        KeyError
+            If any key is not a declared attribute key of this graph.
+        """
+        if attr_keys is None:
+            return
+
+        if isinstance(attr_keys, str):
+            attr_keys = [attr_keys]
+
+        # ``return_ids=True``: the id columns (node_id / edge_id / source_id / target_id) are
+        # legitimately requestable even though they are not user-declared attributes.
+        valid_keys = self.node_attr_keys(return_ids=True) if mode == "node" else self.edge_attr_keys(return_ids=True)
+
+        valid = set(valid_keys)
+        missing = [key for key in dict.fromkeys(attr_keys) if key not in valid]
+        if missing:
+            raise KeyError(
+                f"{mode} attribute key(s) {missing} not found. Available {mode} attribute keys: {sorted(valid)}"
+            )
+
     def add_node(
         self,
         attrs: dict[str, Any],

src/tracksdata/graph/_rustworkx_graph.py

@@ -7,7 +7,13 @@
 import polars as pl
 import rustworkx as rx
 
-from tracksdata.attrs import AttrComparison, AttrFilter, Filter, split_attr_comps
+from tracksdata.attrs import (
+    AttrComparison,
+    AttrFilter,
+    Filter,
+    attr_comps_to_strs,
+    split_attr_comps,
+)
 from tracksdata.constants import DEFAULT_ATTR_KEYS
 from tracksdata.graph._base_graph import BaseGraph
 from tracksdata.graph._mapped_graph_mixin import MappedGraphMixin
@@ -180,6 +186,9 @@ def __init__(
         self._include_targets = include_targets
         self._include_sources = include_sources
         self._node_attr_comps, self._edge_attr_comps = split_attr_comps(attr_comps)
+        # validate eagerly so a typo'd key points at the `filter()` call, not at the collect
+        graph._validate_attr_keys(attr_comps_to_strs(self._node_attr_comps), "node")
+        graph._validate_attr_keys(attr_comps_to_strs(self._edge_attr_comps), "edge")
 
     @cache_method
     def _current_node_ids(self) -> list[int]:
@@ -329,6 +338,8 @@ def edge_attrs(
         attr_keys: list[str] | None = None,
         unpack: bool = False,
     ) -> pl.DataFrame:
+        self._graph._validate_attr_keys(attr_keys, "edge")
+
         df = self._edge_attrs()
         if df.is_empty():
             return df
@@ -839,6 +850,9 @@ def _get_neighbors(
 
         if isinstance(attr_keys, str):
             attr_keys = [attr_keys]
+
+        self._validate_attr_keys(attr_keys, "node")
+
         valid_schema = None
         neighbors: dict[int, list[int]] | dict[int, pl.DataFrame] = {}
         for node_id in node_ids:
@@ -1163,6 +1177,8 @@ def _node_attrs_from_node_ids(
         if isinstance(attr_keys, str):
             attr_keys = [attr_keys]
 
+        self._validate_attr_keys(attr_keys, "node")
+
         node_attr_schemas = self._node_attr_schemas()
         pl_schema = {k: node_attr_schemas[k].dtype for k in attr_keys}
 
@@ -1232,6 +1248,8 @@ def edge_attrs(
         if attr_keys is None:
             attr_keys = self.edge_attr_keys()
 
+        self._validate_attr_keys(attr_keys, "edge")
+
         attr_keys = [DEFAULT_ATTR_KEYS.EDGE_ID, *attr_keys]
         attr_keys = list(dict.fromkeys(attr_keys))
 

src/tracksdata/graph/_sql_graph.py

@@ -6,7 +6,7 @@
 from collections.abc import Callable, Sequence
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 import cloudpickle
 import numpy as np
@@ -18,7 +18,13 @@
 from sqlalchemy.orm.query import Query
 from sqlalchemy.sql.type_api import TypeEngine
 
-from tracksdata.attrs import AttrComparison, AttrFilter, Filter, split_attr_comps
+from tracksdata.attrs import (
+    AttrComparison,
+    AttrFilter,
+    Filter,
+    attr_comps_to_strs,
+    split_attr_comps,
+)
 from tracksdata.constants import DEFAULT_ATTR_KEYS
 from tracksdata.graph._base_graph import BaseGraph
 from tracksdata.graph.filters._base_filter import BaseFilter
@@ -242,6 +248,9 @@ def __init__(
         super().__init__()
         self._graph = graph
         self._node_attr_comps, self._edge_attr_comps = split_attr_comps(attr_filters)
+        # validate eagerly so a typo'd key points at the `filter()` call, not at the collect
+        graph._validate_attr_keys(attr_comps_to_strs(self._node_attr_comps), "node")
+        graph._validate_attr_keys(attr_comps_to_strs(self._edge_attr_comps), "edge")
         self._include_targets = include_targets
         self._include_sources = include_sources
         self._id_set: _SQLIDSet | None = None
@@ -435,6 +444,8 @@ def _query_from_attr_keys(
         if attr_keys is not None:
             attr_keys = list(dict.fromkeys(attr_keys))
 
+            self._graph._validate_attr_keys(attr_keys, self._graph._mode_for_table(table))
+
             if extra_columns is not None:
                 attr_keys.extend(extra_columns)
 
@@ -1644,9 +1655,18 @@ def _physical_cols_for_query(
         logical_keys: Sequence[str],
         table_class: type[DeclarativeBase],
     ) -> list[Any]:
-        """Like :meth:`_physical_column_names`, but returning SQLAlchemy column objects."""
+        """Like :meth:`_physical_column_names`, but returning SQLAlchemy column objects.
+
+        Validates the keys first so an unknown one raises `KeyError` instead of the
+        `AttributeError` that ``getattr(table_class, ...)`` would leak from SQLAlchemy.
+        """
+        self._validate_attr_keys(logical_keys, self._mode_for_table(table_class))
         return [getattr(table_class, name) for name in self._physical_column_names(logical_keys, table_class)]
 
+    def _mode_for_table(self, table_class: type[DeclarativeBase]) -> Literal["node", "edge"]:
+        """Whether ``table_class`` is the node or the edge table."""
+        return "node" if table_class is self.Node else "edge"
+
     def node_attr_keys(self, return_ids: bool = False) -> list[str]:
         """
         Get the keys of the attributes of the nodes.

src/tracksdata/graph/_test/test_graph_backends.py

@@ -55,6 +55,122 @@ def test_edge_validation(graph_backend: BaseGraph) -> None:
         graph_backend.add_edge(0, 1, {"weight": 0.5})
 
 
+def _one_edge_graph(graph: BaseGraph) -> dict[str, int]:
+    """Declare one node attr and one edge attr, then build a 2-node / 1-edge graph."""
+    graph.add_node_attr_key("area", dtype=pl.Float64, default_value=0.0)
+    graph.add_edge_attr_key("w", dtype=pl.Float64, default_value=0.0)
+    a = graph.add_node({"t": 0, "area": 1.0})
+    b = graph.add_node({"t": 1, "area": 2.0})
+    e = graph.add_edge(a, b, {"w": 1.0})
+    return {"a": a, "b": b, "e": e}
+
+
+# Every public read path that accepts an attribute key. Each must report an unknown key the
+# same way on every backend; before this was centralized they raised three different types
+# (AttributeError from SQLAlchemy's getattr, KeyError from a dict lookup, and polars'
+# ColumnNotFoundError, which is not even a KeyError subclass).
+_MISSING_ATTR_ACCESSORS: dict[str, Callable[[BaseGraph, dict[str, int]], Any]] = {
+    "node_attrs": lambda g, i: g.node_attrs(attr_keys=["nope"]),
+    "edge_attrs": lambda g, i: g.edge_attrs(attr_keys=["nope"]),
+    "nodes[id][key]": lambda g, i: g.nodes[i["a"]]["nope"],
+    "edges[id][key]": lambda g, i: g.edges[i["e"]]["nope"],
+    "filter(NodeAttr)": lambda g, i: g.filter(NodeAttr("nope") == 1).node_ids(),
+    "filter(EdgeAttr)": lambda g, i: g.filter(EdgeAttr("nope") == 1).edge_ids(),
+    "successors": lambda g, i: g.successors(i["a"], attr_keys=["nope"], return_attrs=True),
+    "predecessors": lambda g, i: g.predecessors(i["b"], attr_keys=["nope"], return_attrs=True),
+    "filter().node_attrs": lambda g, i: g.filter(node_ids=[i["a"]]).node_attrs(attr_keys=["nope"]),
+    "filter().edge_attrs": lambda g, i: g.filter(node_ids=[i["a"], i["b"]]).edge_attrs(attr_keys=["nope"]),
+}
+
+
+@pytest.mark.parametrize(
+    "accessor",
+    list(_MISSING_ATTR_ACCESSORS.values()),
+    ids=list(_MISSING_ATTR_ACCESSORS.keys()),
+)
+def test_missing_attr_key_raises_key_error(
+    graph_backend: BaseGraph,
+    accessor: Callable[[BaseGraph, dict[str, int]], Any],
+) -> None:
+    """Reading an undeclared attribute key raises KeyError on every backend."""
+    ids = _one_edge_graph(graph_backend)
+
+    with pytest.raises(KeyError):
+        accessor(graph_backend, ids)
+
+
+def test_missing_attr_key_error_message(graph_backend: BaseGraph) -> None:
+    """The error names the unknown key and lists the valid ones.
+
+    Diagnosability is the point of the guard, not just the exception type.
+    """
+    _one_edge_graph(graph_backend)
+
+    with pytest.raises(KeyError) as exc_info:
+        graph_backend.node_attrs(attr_keys=["nope"])
+
+    message = str(exc_info.value)
+    assert "nope" in message
+    assert "area" in message, f"error should list the available keys, got: {message}"
+
+    with pytest.raises(KeyError) as exc_info:
+        graph_backend.edge_attrs(attr_keys=["nope"])
+
+    message = str(exc_info.value)
+    assert "nope" in message
+    assert "w" in message, f"error should list the available keys, got: {message}"
+
+
+def test_filter_missing_attr_key_raises_eagerly(graph_backend: BaseGraph) -> None:
+    """filter() validates at construction, not at collect time.
+
+    Without this the traceback points at the collect call rather than at the caller's typo.
+    """
+    _one_edge_graph(graph_backend)
+
+    with pytest.raises(KeyError):
+        graph_backend.filter(NodeAttr("nope") == 1)
+
+    with pytest.raises(KeyError):
+        graph_backend.filter(EdgeAttr("nope") == 1)
+
+
+def test_missing_attr_key_among_valid_ones_raises(graph_backend: BaseGraph) -> None:
+    """One bad key in an otherwise valid list still raises."""
+    _one_edge_graph(graph_backend)
+
+    with pytest.raises(KeyError):
+        graph_backend.node_attrs(attr_keys=["area", "nope"])
+
+
+def test_valid_attr_keys_are_not_rejected(graph_backend: BaseGraph) -> None:
+    """The guard must not reject legitimate keys, including the id columns."""
+    ids = _one_edge_graph(graph_backend)
+
+    assert graph_backend.node_attrs(attr_keys=["area"])["area"].to_list() == [1.0, 2.0]
+    assert graph_backend.node_attrs(attr_keys="area")["area"].to_list() == [1.0, 2.0]
+    assert graph_backend.edge_attrs(attr_keys=["w"])["w"].to_list() == [1.0]
+
+    # id columns are not user-declared attributes but are legitimately requestable
+    node_df = graph_backend.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, "area"])
+    assert DEFAULT_ATTR_KEYS.NODE_ID in node_df.columns
+
+    edge_df = graph_backend.edge_attrs(
+        attr_keys=[DEFAULT_ATTR_KEYS.EDGE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET, "w"]
+    )
+    assert DEFAULT_ATTR_KEYS.EDGE_ID in edge_df.columns
+
+    # attr_keys=None means "everything" and must stay valid
+    assert not graph_backend.node_attrs().is_empty()
+    assert not graph_backend.edge_attrs().is_empty()
+
+    # single-item accessors and filters with valid keys keep working
+    assert graph_backend.nodes[ids["a"]]["area"] == 1.0
+    assert graph_backend.edges[ids["e"]]["w"] == 1.0
+    assert graph_backend.filter(NodeAttr("area") == 1.0).node_ids() == [ids["a"]]
+    assert graph_backend.filter(EdgeAttr("w") == 1.0).edge_ids() == [ids["e"]]
+
+
 def test_add_node(graph_backend: BaseGraph) -> None:
     """Test adding nodes with various attributes."""
 
@@ -1109,14 +1225,10 @@ def test_sucessors_predecessors_edge_cases(graph_backend: BaseGraph) -> None:
     assert isinstance(predecessors_dict, dict)
     assert len(predecessors_dict) == 0
 
-    # Test with non-existent attribute keys (should work but return limited columns)
-    # This depends on implementation - some might raise errors, others might ignore
-    try:
-        successors_df = graph_backend.successors(node0, attr_keys=["nonexistent"], return_attrs=True)
-        assert isinstance(successors_df, pl.DataFrame)
-    except (KeyError, AttributeError):
-        # This is also acceptable behavior
-        pass
+    # A non-existent attribute key is a lookup miss on every backend.
+    # See test_missing_attr_key_raises_key_error for the full accessor matrix.
+    with pytest.raises(KeyError):
+        graph_backend.successors(node0, attr_keys=["nonexistent"], return_attrs=True)
 
 
 def test_match_method(graph_backend: BaseGraph) -> None: