Add atomic SessionContext.with_extensions API

crates/core/src/context.rs

@@ -1608,6 +1608,81 @@ impl PySessionContext {
         derived.set_session_query_planner(None);
         derived
     }
+
+    /// Create the destination context for a `with_extensions` transaction.
+    ///
+    /// Private support method for `SessionContext.with_extensions`. The
+    /// returned context is the single `Arc<SessionContext>` that every FFI
+    /// task-context provider created during the transaction must target;
+    /// `_install_extensions` later mutates its state in place rather than
+    /// deriving a new context.
+    pub fn _derive_for_extensions(&self) -> Self {
+        Self {
+            ctx: Arc::new(SessionContext::new_with_state(self.ctx.state())),
+            logical_codec: Arc::clone(&self.logical_codec),
+            physical_codec: Arc::clone(&self.physical_codec),
+        }
+    }
+
+    /// Commit a `with_extensions` transaction onto this context.
+    ///
+    /// Private support method for `SessionContext.with_extensions`; `self`
+    /// must be a context produced by `_derive_for_extensions`. Codec capsules
+    /// are imported and validated before any state change, so a failure
+    /// leaves the context untouched. The final state is written through this
+    /// context's own `state_ref()`, never a derived context, so FFI
+    /// task-context providers bound to it stay valid.
+    #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))]
+    pub fn _install_extensions<'py>(
+        slf: &Bound<'py, Self>,
+        logical_codecs: Vec<Bound<'py, PyAny>>,
+        physical_codecs: Vec<Bound<'py, PyAny>>,
+        planner: Option<Bound<'py, PyAny>>,
+    ) -> PyDataFusionResult<Self> {
+        // Chains are built as local values, so a codec that fails to import --
+        // or that collides with an id already installed -- leaves the session
+        // untouched. Nothing is borrowed across a call back into Python.
+        let (mut logical_codec, mut physical_codec) = {
+            let this = slf.borrow();
+            (
+                this.logical_codec.as_ref().clone(),
+                this.physical_codec.as_ref().clone(),
+            )
+        };
+
+        for codec in logical_codecs {
+            let id = resolve_codec_id(&codec, None, &logical_codec.codec_ids())?;
+            let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
+            let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
+            logical_codec = logical_codec.with_additional_codec(id, inner);
+        }
+        let logical_codec = Arc::new(logical_codec);
+
+        for codec in physical_codecs {
+            let id = resolve_codec_id(&codec, None, &physical_codec.codec_ids())?;
+            let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
+            let inner: Arc<dyn PhysicalExtensionCodec> = (&inner_ffi).into();
+            physical_codec = physical_codec.with_additional_codec(id, inner);
+        }
+        let physical_codec = Arc::new(physical_codec);
+
+        let planner = planner
+            .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any())))
+            .transpose()?;
+
+        let derived = Self {
+            ctx: Arc::clone(&slf.borrow().ctx),
+            logical_codec,
+            physical_codec,
+        };
+        // Bind the planner only once the codec chains are final, and through
+        // the derived handle so it carries them. Passing `None` still rebuilds
+        // whichever planner the session already holds against the new chains,
+        // exactly as `with_logical_extension_codec` does.
+        derived.set_session_query_planner(planner);
+
+        Ok(derived)
+    }
 }
 
 impl PySessionContext {

docs/source/contributor-guide/ffi.md

@@ -343,6 +343,61 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust
 `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and
 local build commands.
 
+### Extension bundles: `with_extensions`
+
+The chaining above works, but it makes the caller responsible for two things that are
+easy to get wrong: keeping every intermediate context alive, and installing the codecs
+before the planner. Every codec and planner capsule carries an
+`FFI_TaskContextProvider` holding a *weak* reference to the context it was built
+against, so a component bound to a `with_*` result that is then discarded fails at
+query time with `TaskContextProvider went out of scope over FFI boundary`.
+
+`SessionContext.with_extensions` removes both hazards. An extension library exposes a
+bundle object implementing `__datafusion_session_extension__`:
+
+```python
+class MyEngineExtension:
+    def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents:
+        # Create fresh components bound to `ctx` on every call. `ctx` is the
+        # exact context the host will return from with_extensions.
+        return SessionExtensionComponents(
+            logical_extension_codecs=(self._make_logical_codec(ctx),),
+            physical_extension_codecs=(self._make_physical_codec(ctx),),
+            query_planner=self._make_planner(ctx),
+        )
+```
+
+The host creates one destination context, passes it to every factory, installs all the
+codecs, binds the planner against the final codec chains, and returns that context in
+a single step:
+
+```python
+ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension())
+ctx.register_table("t", lib_a.TableProvider())
+ctx.register_udf(udf(lib_b.SomeUDF()))
+```
+
+Extensions are processed left to right and their codecs are appended to the chain in
+that order. As above, order affects only encoding — decoding routes by id. At most one
+extension per call may supply a query planner. If any factory raises, the source
+context is left exactly as it was.
+
+Bundle objects must be configuration-only: create fresh components on each call, never
+cache bound components, and do not retain the context passed in. Catalogs are shared
+with the source context, so registrations made during binding are not rolled back on
+failure.
+
+The returned context is the strong owner of every installed component's task-context
+provider, and dependent objects do not extend its lifetime. A `DataFrame`, logical
+plan, or capsule can outlive the context, but any operation that reaches an FFI codec
+after the context is collected fails with `TaskContextProvider went out of scope over
+FFI boundary`. Keep the context alive for as long as objects derived from it are in
+use.
+
+`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust
+implementation of the protocol, including taking the task-context provider off the
+supplied context and constructing a Python `SessionExtensionComponents`.
+
 ### Capsule getters receive the session they are installed on
 
 `__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and
@@ -516,6 +571,9 @@ the original handle rebinds the session's planner back to the original handle's
 instead, which is the trap
 `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins.
 
+`with_extensions` sidesteps the ordering question entirely: it installs every codec
+before it binds the planner, so there is no "afterwards" for a bundle's own planner.
+
 ## Alternative Approach
 
 Suppose you needed to expose some other features of DataFusion and you could not wait

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

@@ -31,6 +31,7 @@ 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/README.md

@@ -41,7 +41,22 @@ uv run pytest \
   examples/datafusion-ffi-query-planner-example/python/tests/_test*.py
 ```
 
-The integration test follows this setup:
+The preferred setup uses `SessionContext.with_extensions` with extension bundles:
+
+```python
+config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
+ctx = SessionContext(config).with_extensions(provider_bundle, MyPlannerExtension())
+ctx.register_table("numbers", provider)
+ctx.register_udf(provider_udf)
+```
+
+`MyPlannerExtension` implements the `__datafusion_session_extension__` protocol: it
+receives the destination context, binds fresh codec and planner components to that
+context's task-context provider, and returns them as `SessionExtensionComponents`.
+The host installs everything in one step, so no component can end up bound to an
+intermediate context that is later collected.
+
+The integration tests also cover the low-level chaining setup:
 
 ```python
 config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))