fix: route the last seven capsule getters through call_capsule_getter

`call_capsule_getter` claimed every capsule getter went through it, so the
mapping from a refused argument to a diagnosable error would live in one
place. Seven sites still called `getattr(...).call0()` or `.call1(...)`
directly, so the claim was false and four of them — catalog provider,
schema provider, catalog provider list, table provider factory — still
handed an out-of-date extension library a bare `TypeError`.

Those four take the host's logical extension codec rather than the session,
which is why they could not simply be passed through as they stood: the
diagnostic would have told a catalog author their method "must accept the
SessionContext", pointing them at the wrong parameter. Carry the argument
and its description together in a `CapsuleGetterArg` so one diagnostic can
serve getters that take a session, getters that take a codec, and getters
that take nothing. `Option<&Bound<PyAny>>` still converts into it, so the
documented `*_from_pycapsule` helper signatures are unchanged.

The three zero-argument getters (scalar, aggregate, window UDF) route
through as well. Nothing can be refused there, but the rule is easier to
follow with no exceptions to remember.

Also add `validate_pycapsule` to `table_provider_from_pycapsule` and
`ffi_logical_codec_from_pycapsule`, the two extraction sites that lacked
it. This is not redundant with `pointer_checked`, despite appearances:
`pointer_checked` bottoms out in CPython's `PyCapsule_GetPointer`, whose
error is the fixed string `PyCapsule_GetPointer called with incorrect
name` and names neither the expected capsule nor the one received. Say so
in a doc comment so it does not get "simplified" away later.

Drop the unused `datafusion-proto` dependency from the query planner
example while here.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

Cargo.lock

@@ -30,7 +30,8 @@ use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::schema_provider::FFI_SchemaProvider;
 use datafusion_python_util::{
-    create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, wait_for_future,
+    CapsuleGetterArg, call_capsule_getter, create_logical_extension_capsule,
+    ffi_logical_codec_from_pycapsule, wait_for_future,
 };
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyKeyError;
@@ -626,9 +627,11 @@ fn extract_catalog_provider_from_pyobj(
     if catalog_provider.hasattr("__datafusion_catalog_provider__")? {
         let py = catalog_provider.py();
         let codec_capsule = create_logical_extension_capsule(py, codec)?;
-        catalog_provider = catalog_provider
-            .getattr("__datafusion_catalog_provider__")?
-            .call1((codec_capsule,))?;
+        catalog_provider = call_capsule_getter(
+            catalog_provider,
+            "__datafusion_catalog_provider__",
+            CapsuleGetterArg::LogicalCodec(&codec_capsule),
+        )?;
     }
 
     let provider = if let Ok(capsule) = catalog_provider.cast::<PyCapsule>() {
@@ -658,9 +661,11 @@ fn extract_schema_provider_from_pyobj(
     if schema_provider.hasattr("__datafusion_schema_provider__")? {
         let py = schema_provider.py();
         let codec_capsule = create_logical_extension_capsule(py, codec)?;
-        schema_provider = schema_provider
-            .getattr("__datafusion_schema_provider__")?
-            .call1((codec_capsule,))?;
+        schema_provider = call_capsule_getter(
+            schema_provider,
+            "__datafusion_schema_provider__",
+            CapsuleGetterArg::LogicalCodec(&codec_capsule),
+        )?;
     }
 
     let provider = if let Ok(capsule) = schema_provider.cast::<PyCapsule>() {

crates/core/src/catalog.rs

@@ -58,10 +58,11 @@ use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_python_util::{
-    create_logical_extension_capsule, create_physical_extension_capsule,
-    create_query_planner_capsule, ffi_logical_codec_from_pycapsule,
-    ffi_physical_codec_from_pycapsule, ffi_query_planner_from_pycapsule, get_global_ctx,
-    get_tokio_runtime, physical_optimizer_rule_from_pycapsule, spawn_future, wait_for_future,
+    CapsuleGetterArg, call_capsule_getter, create_logical_extension_capsule,
+    create_physical_extension_capsule, create_query_planner_capsule,
+    ffi_logical_codec_from_pycapsule, ffi_physical_codec_from_pycapsule,
+    ffi_query_planner_from_pycapsule, get_global_ctx, get_tokio_runtime,
+    physical_optimizer_rule_from_pycapsule, spawn_future, wait_for_future,
 };
 use object_store::ObjectStore;
 use pyo3::IntoPyObjectExt;
@@ -725,9 +726,11 @@ impl PySessionContext {
             let py = factory.py();
             let ffi = self.ffi_logical_codec();
             let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
-            factory = factory
-                .getattr("__datafusion_table_provider_factory__")?
-                .call1((codec_capsule,))?;
+            factory = call_capsule_getter(
+                factory,
+                "__datafusion_table_provider_factory__",
+                CapsuleGetterArg::LogicalCodec(&codec_capsule),
+            )?;
         }
 
         let factory: Arc<dyn TableProviderFactory> =
@@ -760,9 +763,11 @@ impl PySessionContext {
             let py = provider.py();
             let ffi = self.ffi_logical_codec();
             let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
-            provider = provider
-                .getattr("__datafusion_catalog_provider_list__")?
-                .call1((codec_capsule,))?;
+            provider = call_capsule_getter(
+                provider,
+                "__datafusion_catalog_provider_list__",
+                CapsuleGetterArg::LogicalCodec(&codec_capsule),
+            )?;
         }
 
         let provider = if let Ok(capsule) = provider.cast::<PyCapsule>() {
@@ -796,9 +801,11 @@ impl PySessionContext {
             let py = provider.py();
             let ffi = self.ffi_logical_codec();
             let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
-            provider = provider
-                .getattr("__datafusion_catalog_provider__")?
-                .call1((codec_capsule,))?;
+            provider = call_capsule_getter(
+                provider,
+                "__datafusion_catalog_provider__",
+                CapsuleGetterArg::LogicalCodec(&codec_capsule),
+            )?;
         }
 
         let provider = if let Ok(capsule) = provider.cast::<PyCapsule>() {

crates/core/src/context.rs

@@ -28,7 +28,7 @@ use datafusion::logical_expr::{
     Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, Signature, Volatility,
 };
 use datafusion_ffi::udaf::FFI_AggregateUDF;
-use datafusion_python_util::parse_volatility;
+use datafusion_python_util::{CapsuleGetterArg, call_capsule_getter, parse_volatility};
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple};
 
@@ -365,7 +365,11 @@ impl PyAggregateUDF {
         }
 
         if func.hasattr("__datafusion_aggregate_udf__")? {
-            let capsule = func.getattr("__datafusion_aggregate_udf__")?.call0()?;
+            let capsule = call_capsule_getter(
+                func.clone(),
+                "__datafusion_aggregate_udf__",
+                CapsuleGetterArg::None,
+            )?;
             let capsule = capsule.cast::<PyCapsule>().map_err(py_datafusion_err)?;
             let function = aggregate_udf_from_capsule(capsule)?;
             return Ok(Self { function });

crates/core/src/udaf.rs

@@ -31,7 +31,7 @@ use datafusion::logical_expr::{
     Volatility,
 };
 use datafusion_ffi::udf::FFI_ScalarUDF;
-use datafusion_python_util::parse_volatility;
+use datafusion_python_util::{CapsuleGetterArg, call_capsule_getter, parse_volatility};
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple};
 
@@ -248,7 +248,11 @@ impl PyScalarUDF {
     #[staticmethod]
     pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
         if func.hasattr("__datafusion_scalar_udf__")? {
-            let capsule = func.getattr("__datafusion_scalar_udf__")?.call0()?;
+            let capsule = call_capsule_getter(
+                func.clone(),
+                "__datafusion_scalar_udf__",
+                CapsuleGetterArg::None,
+            )?;
             let capsule = capsule.cast::<PyCapsule>().map_err(to_datafusion_err)?;
             let data: NonNull<FFI_ScalarUDF> = capsule
                 .pointer_checked(Some(c"datafusion_scalar_udf"))?

crates/core/src/udf.rs

@@ -30,7 +30,7 @@ use datafusion::logical_expr::{
 };
 use datafusion::scalar::ScalarValue;
 use datafusion_ffi::udwf::FFI_WindowUDF;
-use datafusion_python_util::parse_volatility;
+use datafusion_python_util::{CapsuleGetterArg, call_capsule_getter, parse_volatility};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyList, PyTuple};
@@ -262,11 +262,8 @@ impl PyWindowUDF {
 
     #[staticmethod]
     pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
-        let capsule = if func.hasattr("__datafusion_window_udf__")? {
-            func.getattr("__datafusion_window_udf__")?.call0()?
-        } else {
-            func
-        };
+        let capsule =
+            call_capsule_getter(func, "__datafusion_window_udf__", CapsuleGetterArg::None)?;
 
         let capsule = capsule.cast::<PyCapsule>().map_err(to_datafusion_err)?;
         let data: NonNull<FFI_WindowUDF> = capsule

crates/core/src/udwf.rs

@@ -156,6 +156,17 @@ pub fn parse_volatility(value: &str) -> PyDataFusionResult<Volatility> {
     })
 }
 
+/// Check that `capsule` carries the name `name`, and say so usefully if not.
+///
+/// This looks redundant next to `capsule.pointer_checked(Some(name))`, which
+/// also rejects a mismatched name, and it is not. `pointer_checked` bottoms out
+/// in CPython's `PyCapsule_GetPointer`, whose error is the fixed string
+/// `PyCapsule_GetPointer called with incorrect name` — it names neither what
+/// was expected nor what was found. Handing an extension author that message
+/// tells them nothing about which capsule they got wrong.
+///
+/// So call this first at every extraction site. `pointer_checked` still gets
+/// the name, because it is what actually guards the cast.
 pub fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyResult<()> {
     let capsule_name = capsule.name()?;
     if capsule_name.is_none() {
@@ -228,6 +239,7 @@ pub fn table_provider_from_pycapsule<'py>(
     obj = call_capsule_getter(obj, "__datafusion_table_provider__", Some(&session))?;
 
     if let Ok(capsule) = obj.cast::<PyCapsule>() {
+        validate_pycapsule(capsule, "datafusion_table_provider")?;
         let data: NonNull<FFI_TableProvider> = capsule
             .pointer_checked(Some(c"datafusion_table_provider"))?
             .cast();
@@ -250,35 +262,75 @@ pub fn create_logical_extension_capsule<'py>(
     PyCapsule::new_with_value(py, codec, cr"datafusion_logical_extension_codec")
 }
 
-/// Calls `obj.__<attr_name>__(session)`, or `obj.__<attr_name>__()` when no
-/// session is supplied.
+/// The single positional argument datafusion-python hands a capsule getter.
+///
+/// Carrying the argument and its description together is what lets one
+/// diagnostic serve the whole family: the getters do not all take a session,
+/// and telling the author of a catalog provider that their method "must accept
+/// the SessionContext" would send them to fix the wrong signature.
+#[derive(Clone, Copy)]
+pub enum CapsuleGetterArg<'a, 'py> {
+    /// The getter takes no arguments, so it cannot refuse one. A `TypeError`
+    /// from one of these is always the getter's own and is never rewritten.
+    None,
+    /// The `SessionContext` the object is being installed on.
+    Session(&'a Bound<'py, PyAny>),
+    /// The host's logical extension codec, as a
+    /// `datafusion_logical_extension_codec` capsule.
+    LogicalCodec(&'a Bound<'py, PyAny>),
+}
+
+impl<'a, 'py> CapsuleGetterArg<'a, 'py> {
+    fn value(self) -> Option<&'a Bound<'py, PyAny>> {
+        match self {
+            Self::None => None,
+            Self::Session(value) | Self::LogicalCodec(value) => Some(value),
+        }
+    }
+
+    fn expected(self) -> &'static str {
+        match self {
+            Self::None => "",
+            Self::Session(_) => "the SessionContext it is being installed on",
+            Self::LogicalCodec(_) => "the host's logical extension codec",
+        }
+    }
+}
+
+impl<'a, 'py> From<Option<&'a Bound<'py, PyAny>>> for CapsuleGetterArg<'a, 'py> {
+    fn from(session: Option<&'a Bound<'py, PyAny>>) -> Self {
+        session.map_or(Self::None, Self::Session)
+    }
+}
+
+/// Calls `obj.__<attr_name>__(arg)`, or `obj.__<attr_name>__()` for
+/// [`CapsuleGetterArg::None`]. Returns `obj` untouched when it has no such
+/// attribute, so a raw capsule passes straight through.
 ///
-/// The session is how an exporting library obtains the codecs and task context
-/// of the session it is being installed on, instead of inventing one of its
-/// own. `None` is for the reverse direction, where `obj` *is* a session and is
-/// being asked for what it holds.
-/// Every capsule getter must go through here rather than calling `getattr`
+/// **Every** capsule getter goes through here rather than calling `getattr`
 /// directly, so that the mapping from a refused argument to a useful error
 /// lives in one place. Three importers previously each had their own copy and
-/// each missed later corrections to it.
-pub fn call_capsule_getter<'py>(
+/// each missed later corrections to it; the getters that take no argument route
+/// through as well, so the rule has no exceptions to remember.
+pub fn call_capsule_getter<'a, 'py: 'a>(
     obj: Bound<'py, PyAny>,
     attr_name: &str,
-    session: Option<&Bound<'py, PyAny>>,
+    arg: impl Into<CapsuleGetterArg<'a, 'py>>,
 ) -> PyResult<Bound<'py, PyAny>> {
     if !obj.hasattr(attr_name)? {
         return Ok(obj);
     }
 
+    let arg = arg.into();
     let getter = obj.getattr(attr_name)?;
-    let result = match session {
-        Some(session) => getter.call1((session,)),
+    let result = match arg.value() {
+        Some(value) => getter.call1((value,)),
         None => getter.call0(),
     };
 
     result.map_err(|err| {
         let py = obj.py();
-        if session.is_none() || !err.get_type(py).is(PyType::new::<PyTypeError>(py)) {
+        if arg.value().is_none() || !err.get_type(py).is(PyType::new::<PyTypeError>(py)) {
             return err;
         }
 
@@ -296,8 +348,9 @@ pub fn call_capsule_getter<'py>(
         }
 
         let import_err = PyImportError::new_err(format!(
-            "Incompatible libraries. `{attr_name}` must accept the SessionContext it \
-             is being installed on. Upgrade the library providing this object."
+            "Incompatible libraries. `{attr_name}` must accept {}. \
+             Upgrade the library providing this object.",
+            arg.expected()
         ));
         // Keep the original reachable as `__cause__` rather than discarding it.
         import_err.set_cause(py, Some(err));
@@ -312,6 +365,7 @@ pub fn ffi_logical_codec_from_pycapsule<'py>(
     let capsule = call_capsule_getter(obj, "__datafusion_logical_extension_codec__", session)?;
 
     let capsule = capsule.cast::<PyCapsule>()?;
+    validate_pycapsule(capsule, "datafusion_logical_extension_codec")?;
     let data: NonNull<FFI_LogicalExtensionCodec> = capsule
         .pointer_checked(Some(c"datafusion_logical_extension_codec"))?
         .cast();
@@ -350,7 +404,7 @@ pub fn ffi_task_context_provider_from_pycapsule(
     let capsule = call_capsule_getter(
         session.clone(),
         "__datafusion_task_context_provider__",
-        None,
+        CapsuleGetterArg::None,
     )?;
 
     let capsule = capsule.cast::<PyCapsule>()?;

crates/util/src/lib.rs

@@ -31,7 +31,6 @@ datafusion = { workspace = true }
 datafusion-catalog = { workspace = true, default-features = false }
 datafusion-common = { workspace = true, default-features = false }
 datafusion-ffi = { workspace = true }
-datafusion-proto = { workspace = true }
 datafusion-session = { workspace = true }
 async-trait = { workspace = true }
 datafusion-python-util.workspace = true

examples/datafusion-ffi-query-planner-example/Cargo.toml

@@ -1301,6 +1301,42 @@ def __datafusion_table_provider__(self):
     assert isinstance(excinfo.value.__cause__, TypeError)
 
 
+def test_catalog_provider_getter_arity_error_names_the_codec(ctx):
+    """A getter refusing the codec capsule is diagnosed, not left as a TypeError.
+
+    `__datafusion_catalog_provider__` is handed the host's logical extension
+    codec, not the session. It used to call `getattr(...).call1(...)` directly
+    and so produced a bare `TypeError` for an out-of-date library; it now routes
+    through `call_capsule_getter` like the rest of the family.
+
+    The message has to name the codec rather than the SessionContext, or it
+    would send the author to change the wrong parameter.
+    """
+
+    class PreCodecCatalogProvider:
+        def __datafusion_catalog_provider__(self):
+            msg = "should never be called"
+            raise AssertionError(msg)
+
+    with pytest.raises(ImportError, match="__datafusion_catalog_provider__") as excinfo:
+        ctx.register_catalog_provider("old_sig", PreCodecCatalogProvider())
+    assert "logical extension codec" in str(excinfo.value)
+    assert "SessionContext" not in str(excinfo.value)
+    assert isinstance(excinfo.value.__cause__, TypeError)
+
+
+def test_type_error_inside_a_catalog_provider_getter_propagates(ctx):
+    """A correctly-signed catalog getter's own TypeError survives unchanged."""
+
+    class RaisesTypeError:
+        def __datafusion_catalog_provider__(self, codec):
+            msg = "bad cast inside the getter"
+            raise TypeError(msg)
+
+    with pytest.raises(TypeError, match="bad cast inside the getter"):
+        ctx.register_catalog_provider("raises", RaisesTypeError())
+
+
 def test_type_error_inside_a_table_provider_getter_propagates(ctx):
     """A correctly-signed provider getter's own TypeError survives unchanged."""