From 4d8fc8db1b7a4d9ae8b15a4bc132dd3d1046fddf Mon Sep 17 00:00:00 2001 From: Teun Huijben Date: Thu, 30 Jul 2026 14:54:28 -0700 Subject: [PATCH 1/2] consistent KeyErrors acros backends --- src/tracksdata/graph/_base_graph.py | 47 +++++++ src/tracksdata/graph/_rustworkx_graph.py | 20 ++- src/tracksdata/graph/_sql_graph.py | 26 +++- .../graph/_test/test_graph_backends.py | 128 ++++++++++++++++-- 4 files changed, 209 insertions(+), 12 deletions(-) diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index 3a214e35..ec16e712 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/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], diff --git a/src/tracksdata/graph/_rustworkx_graph.py b/src/tracksdata/graph/_rustworkx_graph.py index cebcfd6f..d7466f08 100644 --- a/src/tracksdata/graph/_rustworkx_graph.py +++ b/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)) diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index b2f66b46..c1c6804e 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/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. diff --git a/src/tracksdata/graph/_test/test_graph_backends.py b/src/tracksdata/graph/_test/test_graph_backends.py index 3c8095c6..8b3306d8 100644 --- a/src/tracksdata/graph/_test/test_graph_backends.py +++ b/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: From dafe39636277407f4e6a9868ec1079730dd26227 Mon Sep 17 00:00:00 2001 From: Teun Huijben Date: Fri, 7 Aug 2026 17:13:42 -0700 Subject: [PATCH 2/2] use set arithmetic --- src/tracksdata/graph/_base_graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index ec16e712..c59ff35e 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/src/tracksdata/graph/_base_graph.py @@ -204,7 +204,7 @@ def _validate_attr_keys( 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] + missing = sorted(set(attr_keys) - valid) # sorted for consistent erorr message if missing: raise KeyError( f"{mode} attribute key(s) {missing} not found. Available {mode} attribute keys: {sorted(valid)}"