Let extension bundles declare catalog providers

.ai/skills/ffi-capsule-protocol/SKILL.md

@@ -76,6 +76,25 @@ receiver had — the same session, and the same task-context provider, but not
 this call's codecs, not even your own. Read the host's codec chains in the
 planner hook, never in the extension hook.
 
+**That is why a bundle declares unresolved components, not wrapped ones.** The
+components it returns split by what their getter asks for:
+
+- Getters taking no argument — the three function kinds,
+  `__datafusion_physical_optimizer_rule__` — have nothing session-scoped to
+  bind, so a bundle may hand over either the raw exportable or an
+  already-wrapped object.
+- Getters taking the session or a codec — `__datafusion_table_function__`,
+  `__datafusion_table_provider__`, `__datafusion_catalog_provider__` — must be
+  handed over **unwrapped**, with a name. Wrapping one inside the components
+  hook would call its getter with the `ctx` that hook received, capturing a
+  chain missing every library in the call. The host wraps these itself, against
+  the handle carrying the final chains, which is the only place that chain
+  exists.
+
+`RecordingTableFunction` in `examples/datafusion-ffi-example/src/extension.rs`
+records the ids it was resolved against, so the difference is asserted rather
+than described.
+
 A *codec* must always be handed over as an object implementing its getter, never
 as the bare capsule the getter returns; `with_extensions` refuses a capsule.
 A codec's wire id — the string a payload names on decode, which has to mean the
@@ -187,6 +206,18 @@ an instruction to derive one first. It is not: the factories are handed the
 receiver, and the returned handle shares its allocation. There is nothing to
 keep alive separately and nothing to garbage-collect out from under a provider.
 
+It is also the one place where sharing an allocation has a cost, and the cost
+shapes how a component is added to it. Because the returned handle *is* the
+receiver's session, a failure part-way through has nothing to roll back to. So
+`with_extensions` does every fallible thing first — importing capsules,
+resolving names, running the planner hooks — and only then writes. **Adding a
+new kind of component means adding a resolve step, never a fallible commit
+step:** a `_resolve_extension_*` that returns an opaque carrier and an
+`_install_extension_*` that takes it and returns `()`. The one exception is
+table registration, whose insert goes through a `SchemaProvider` that a foreign
+library may implement; it is committed first so nothing else is written behind
+it. Do not add a second exception without the same justification.
+
 `SessionContext.enable_url_table` is the one method that mints a second
 allocation for a session. Its result must not outlive the receiver, and it also
 forks the session's `SessionState` while keeping its id, so two handles report

crates/core/src/context.rs

@@ -842,35 +842,9 @@ impl PySessionContext {
     pub fn register_catalog_provider(
         &self,
         name: &str,
-        mut provider: Bound<'_, PyAny>,
+        provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        if provider.hasattr("__datafusion_catalog_provider__")? {
-            let py = provider.py();
-            let ffi = self.ffi_logical_codec();
-            let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
-            provider = call_capsule_getter(
-                provider,
-                "__datafusion_catalog_provider__",
-                CapsuleGetterArg::LogicalCodec(&codec_capsule),
-            )?;
-        }
-
-        let provider = if let Ok(capsule) = provider.cast::<PyCapsule>() {
-            let data: NonNull<FFI_CatalogProvider> = capsule
-                .pointer_checked(Some(c"datafusion_catalog_provider"))?
-                .cast();
-            let provider = unsafe { data.as_ref() };
-            let provider: Arc<dyn CatalogProvider> = provider.into();
-            provider
-        } else {
-            match provider.extract::<PyCatalog>() {
-                Ok(py_catalog) => py_catalog.catalog,
-                Err(_) => Arc::new(RustWrappedPyCatalogProvider::new(
-                    provider.into(),
-                    self.ffi_logical_codec(),
-                )) as Arc<dyn CatalogProvider>,
-            }
-        };
+        let provider = self.resolve_catalog_provider(provider)?;
 
         let _ = self.ctx.register_catalog(name, provider);
 
@@ -1828,6 +1802,42 @@ impl PySessionContext {
         Ok(())
     }
 
+    /// Resolve the catalogs a `with_extensions` call declared.
+    ///
+    /// The fallible half. Each provider is imported against `self` — the handle
+    /// carrying the completed codec chains, since
+    /// `__datafusion_catalog_provider__` is handed the logical codec it will
+    /// serialize through.
+    ///
+    /// No name is refused here. `register_catalog` replaces rather than
+    /// rejects, and `datafusion` — the default catalog — always exists, so a
+    /// bundle replacing a catalog is ordinary rather than a mistake. Two
+    /// bundles claiming one name in the same call is refused on the Python
+    /// side, where both can be named.
+    ///
+    /// **Writes nothing.**
+    pub fn _resolve_extension_catalogs<'py>(
+        &self,
+        catalogs: Vec<(String, Bound<'py, PyAny>)>,
+    ) -> PyDataFusionResult<PyResolvedCatalogs> {
+        let catalogs = catalogs
+            .into_iter()
+            .map(|(name, provider)| Ok((name, self.resolve_catalog_provider(provider)?)))
+            .collect::<PyDataFusionResult<Vec<_>>>()?;
+        Ok(PyResolvedCatalogs { catalogs })
+    }
+
+    /// Commit the catalogs for a `with_extensions` call.
+    ///
+    /// Nothing here can fail: the providers were imported by
+    /// [`Self::_resolve_extension_catalogs`], and `register_catalog` returns
+    /// whichever provider it displaced rather than refusing.
+    pub fn _install_extension_catalogs(&self, resolved: PyRef<'_, PyResolvedCatalogs>) {
+        for (name, provider) in &resolved.catalogs {
+            let _ = self.ctx.register_catalog(name, Arc::clone(provider));
+        }
+    }
+
     /// Import the physical optimizer rules a `with_extensions` call declared.
     ///
     /// The fallible half of installing them, run while the call can still fail
@@ -1900,6 +1910,14 @@ struct ResolvedTable {
     provider: Arc<dyn TableProvider>,
 }
 
+/// Catalog providers imported for a `with_extensions` call.
+///
+/// Opaque to Python, like [`PyResolvedTables`] and [`PyPhysicalOptimizerRules`].
+#[pyclass(name = "ResolvedCatalogs", module = "datafusion._internal")]
+pub struct PyResolvedCatalogs {
+    catalogs: Vec<(String, Arc<dyn CatalogProvider>)>,
+}
+
 /// Physical optimizer rules imported for a `with_extensions` call.
 ///
 /// Opaque to Python, and deliberately not added to the module: it exists only
@@ -1912,6 +1930,50 @@ pub struct PyPhysicalOptimizerRules {
 }
 
 impl PySessionContext {
+    /// Turn whatever a caller offered as a catalog provider into one.
+    ///
+    /// The fallible half of registering a catalog, shared by
+    /// [`Self::register_catalog_provider`] and
+    /// [`Self::_resolve_extension_catalogs`] so both accept exactly the same
+    /// shapes: an object exposing `__datafusion_catalog_provider__`, a bare
+    /// capsule, a [`PyCatalog`], or a Python object implementing the provider
+    /// interface.
+    ///
+    /// The getter is handed **this context's** logical codec, so which handle
+    /// this is called on decides what the provider will serialize through.
+    fn resolve_catalog_provider(
+        &self,
+        mut provider: Bound<'_, PyAny>,
+    ) -> PyDataFusionResult<Arc<dyn CatalogProvider>> {
+        if provider.hasattr("__datafusion_catalog_provider__")? {
+            let py = provider.py();
+            let ffi = self.ffi_logical_codec();
+            let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
+            provider = call_capsule_getter(
+                provider,
+                "__datafusion_catalog_provider__",
+                CapsuleGetterArg::LogicalCodec(&codec_capsule),
+            )?;
+        }
+
+        Ok(if let Ok(capsule) = provider.cast::<PyCapsule>() {
+            let data: NonNull<FFI_CatalogProvider> = capsule
+                .pointer_checked(Some(c"datafusion_catalog_provider"))?
+                .cast();
+            let provider = unsafe { data.as_ref() };
+            let provider: Arc<dyn CatalogProvider> = provider.into();
+            provider
+        } else {
+            match provider.extract::<PyCatalog>() {
+                Ok(py_catalog) => py_catalog.catalog,
+                Err(_) => Arc::new(RustWrappedPyCatalogProvider::new(
+                    provider.into(),
+                    self.ffi_logical_codec(),
+                )) as Arc<dyn CatalogProvider>,
+            }
+        })
+    }
+
     /// Write the session's query planner, in place.
     ///
     /// Pass `Some(planner)` to install one, or `None` to rebuild whichever

docs/source/extension-guide/bundles.md

@@ -300,10 +300,10 @@ where the host can resolve them:
 - **Getters taking no argument** — the three function kinds and physical
   optimizer rules. Nothing is session-scoped, so a bundle may hand over either
   a wrapped object or the raw exportable.
-- **Getters taking the session or a codec** — table functions and table
-  providers. These are resolved by the host against the *finished* handle,
-  which is why you hand over the unwrapped value and a name rather than a
-  {py:class}`~datafusion.user_defined.TableFunction` you built yourself.
+- **Getters taking the session or a codec** — table functions, table providers,
+  and catalog providers. These are resolved by the host against the *finished*
+  handle, which is why you hand over the unwrapped value and a name rather than
+  a {py:class}`~datafusion.user_defined.TableFunction` you built yourself.
   Wrapping one inside your components hook binds it to the context that hook
   received, which has none of the call's codecs — so it would capture a chain
   missing every library in the call, including your own.
@@ -349,6 +349,8 @@ is nothing to refuse. See {doc}`other-components`.
 Tables go the other way. DataFusion refuses a duplicate table registration
 rather than replacing it, so a declared table name that is *already* on the
 session is an error too — a table cannot shadow one the way a function can.
+Catalogs are back on the function side of that line, and for a reason worth
+reading before you declare one: see {doc}`table-providers`.
 
 Your caller cannot rename your function, so stay out of the way: prefix the
 names with something tied to your library.

docs/source/extension-guide/table-providers.md

@@ -51,6 +51,20 @@ wrapped: `__datafusion_table_provider__` takes the session, and the one your
 components hook receives has none of the call's codecs yet. The host resolves
 it against the finished handle. See {ref}`extension_bundles_binding`.
 
+Catalogs work the same way, as `catalog_providers`:
+
+```python
+return SessionExtensionComponents(catalog_providers=(("engine", MyCatalog()),))
+```
+
+with one difference worth knowing. A declared **table** name that is already
+registered is an error, because DataFusion refuses a duplicate table rather
+than replacing it. A **catalog** name is not: `register_catalog` returns
+whichever provider it displaced, and the default `datafusion` catalog always
+exists — so replacing one is the usual way a library backs a session with its
+own metadata. Only two bundles claiming the same catalog name in one call is
+refused.
+
 Start with a table provider. Reach for the schema and catalog levels when your
 data source has its own namespace that should be browsable rather than
 registered table by table, and for the provider list only when your library is

examples/datafusion-ffi-example/python/tests/_test_session_extension.py

@@ -23,6 +23,7 @@
 import pytest
 from datafusion import SessionContext, SessionExtensionComponents
 from datafusion_ffi_example import (
+    MyCatalogExtension,
     MyDataExtension,
     MyFunctionExtension,
     MyLogicalExtensionCodec,
@@ -237,6 +238,42 @@ def test_a_table_name_already_registered_is_refused():
         ctx.udf("my_custom_is_null")
 
 
+def test_a_declared_catalog_is_queryable():
+    """A catalog declared by a bundle is reachable by its qualified name."""
+    ctx = SessionContext().with_extensions(MyCatalogExtension())
+
+    assert "declared_catalog" in ctx.catalog_names()
+    result = ctx.sql("SELECT * FROM declared_catalog.my_schema.my_table").collect()
+    assert result[0].num_rows > 0
+
+
+def test_four_libraries_install_in_one_call():
+    """The whole point, across a real FFI boundary.
+
+    Four independently declared bundles — functions, rules, a table and a table
+    function, a catalog — in one call, and a single query that touches three of
+    them while the fourth counts the planning it did.
+    """
+    rules = MyRuleExtension()
+    ctx = SessionContext().with_extensions(
+        MyFunctionExtension(),
+        rules,
+        MyDataExtension(),
+        MyCatalogExtension(),
+    )
+
+    result = ctx.sql(
+        'SELECT my_custom_is_null("A") AS n FROM declared_table '
+        "UNION ALL "
+        "SELECT my_custom_is_null(units) AS n "
+        "FROM declared_catalog.my_schema.my_table"
+    ).collect()
+
+    assert sum(batch.num_rows for batch in result) > 0
+    assert rules.first_calls() > 0
+    assert rules.second_calls() > 0
+
+
 def test_the_hook_returns_the_components_type():
     """The bundle builds a real dataclass, not a duck-typed stand-in.
 

examples/datafusion-ffi-example/src/extension.rs

@@ -21,6 +21,7 @@ use pyo3::types::{PyAnyMethods, PyCapsule, PyDict, PyDictMethods};
 use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods};
 
 use crate::aggregate_udf::MySumUDF;
+use crate::catalog_provider::MyCatalogProvider;
 use crate::physical_optimizer::MyPhysicalOptimizerRule;
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
@@ -230,3 +231,43 @@ impl MyDataExtension {
         components.call((), Some(&kwargs))
     }
 }
+
+/// A bundle contributing a catalog.
+///
+/// `__datafusion_catalog_provider__` takes the session and pulls the host's
+/// logical codec off it, so like a table provider it is handed over unresolved
+/// and the host binds it to the finished handle.
+#[pyclass(
+    from_py_object,
+    name = "MyCatalogExtension",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Debug, Clone, Default)]
+pub(crate) struct MyCatalogExtension {}
+
+#[pymethods]
+impl MyCatalogExtension {
+    #[new]
+    fn new() -> Self {
+        Self {}
+    }
+
+    fn __datafusion_session_components__<'py>(
+        &self,
+        py: Python<'py>,
+        ctx: Bound<'py, PyAny>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let _ = ctx;
+
+        let components = py
+            .import("datafusion")?
+            .getattr("SessionExtensionComponents")?;
+        let kwargs = PyDict::new(py);
+        kwargs.set_item(
+            "catalog_providers",
+            (("declared_catalog", Py::new(py, MyCatalogProvider::new()?)?),),
+        )?;
+        components.call((), Some(&kwargs))
+    }
+}

examples/datafusion-ffi-example/src/lib.rs

@@ -20,7 +20,7 @@ use pyo3::prelude::*;
 use crate::aggregate_udf::MySumUDF;
 use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList};
 use crate::config::MyConfig;
-use crate::extension::{MyDataExtension, MyFunctionExtension, MyRuleExtension};
+use crate::extension::{MyCatalogExtension, MyDataExtension, MyFunctionExtension, MyRuleExtension};
 use crate::logical_extension_codec::MyLogicalExtensionCodec;
 use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec};
 use crate::physical_extension_codec::MyPhysicalExtensionCodec;
@@ -68,5 +68,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyFunctionExtension>()?;
     m.add_class::<MyRuleExtension>()?;
     m.add_class::<MyDataExtension>()?;
+    m.add_class::<MyCatalogExtension>()?;
     Ok(())
 }