Add atomic SessionContext.with_extensions API

crates/core/src/context.rs

@@ -1534,6 +1534,90 @@ impl PySessionContext {
             physical_codec,
         }
     }
+
+    /// 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>(
+        &self,
+        logical_codecs: Vec<Bound<'py, PyAny>>,
+        physical_codecs: Vec<Bound<'py, PyAny>>,
+        planner: Option<Bound<'py, PyAny>>,
+    ) -> PyDataFusionResult<Self> {
+        let mut logical_codec = self.logical_codec.as_ref().clone();
+        for codec in logical_codecs {
+            let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?;
+            let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
+            logical_codec = logical_codec.with_additional_codec(inner);
+        }
+        let logical_codec = Arc::new(logical_codec);
+
+        let mut physical_codec = self.physical_codec.as_ref().clone();
+        for codec in physical_codecs {
+            let inner = physical_codec_from_pycapsule(&codec)?;
+            physical_codec = physical_codec.with_additional_codec(inner);
+        }
+        let physical_codec = Arc::new(physical_codec);
+
+        // Bind the planner only after the codec chains are final. Both FFI
+        // codec wrappers target this exact context so their weak task-context
+        // providers stay valid for as long as the returned context lives.
+        let ffi_logical = Self::ffi_logical_codec_for(&self.ctx, &logical_codec);
+        let ffi_physical = Self::ffi_physical_codec_for(&self.ctx, &physical_codec);
+        let query_planner: Option<Arc<dyn QueryPlanner + Send + Sync>> = match planner {
+            Some(planner) => {
+                let planner = ffi_query_planner_from_pycapsule(&planner)?;
+                let planner: Arc<dyn QueryPlanner + Send + Sync> = (&planner).into();
+                let planner =
+                    FFI_QueryPlanner::new_with_ffi_codecs(planner, ffi_logical, ffi_physical);
+                Some(Arc::new(RuntimeAwareQueryPlanner { planner }))
+            }
+            None => {
+                let state = self.ctx.state();
+                let planner_any: &dyn std::any::Any = state.query_planner().as_ref();
+                planner_any
+                    .downcast_ref::<RuntimeAwareQueryPlanner>()
+                    .map(|p| {
+                        Arc::new(p.with_ffi_codecs(ffi_logical, ffi_physical))
+                            as Arc<dyn QueryPlanner + Send + Sync>
+                    })
+            }
+        };
+
+        if let Some(query_planner) = query_planner {
+            let state = SessionStateBuilder::new_from_existing(self.ctx.state())
+                .with_query_planner(query_planner)
+                .build();
+            *self.ctx.state_ref().write() = state;
+        }
+
+        Ok(Self {
+            ctx: Arc::clone(&self.ctx),
+            logical_codec,
+            physical_codec,
+        })
+    }
 }
 
 impl PySessionContext {

docs/source/contributor-guide/ffi.md

@@ -283,8 +283,67 @@ planner are rebound automatically, but a planner wrapped inside another planner
 fallback is opaque and keeps the codecs it was exported with. **Install all extension
 codecs before exporting or chaining planners.**
 
-Putting it together for a session using two extension libraries that each provide
-tables, functions, and a query planner:
+### Extension bundles: `with_extensions`
+
+Codec and planner capsules carry an `FFI_TaskContextProvider` holding a *weak*
+reference to the `SessionContext` they were created against. A capsule does not keep
+that context alive, and a component bound to one context cannot be rebound to another.
+Chaining the low-level `with_*` methods by hand therefore risks binding components to
+an intermediate context that is later garbage collected, which fails at query time
+with `TaskContextProvider went out of scope over FFI boundary` — or worse, silently
+reads stale session state.
+
+`SessionContext.with_extensions` avoids this by construction. An extension library
+exposes a bundle object implementing the `__datafusion_session_extension__` protocol:
+
+```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
+codecs, binds the planner against the final codec chains, and returns the 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 prepend to the codec chain, so codecs from
+later extensions are consulted first. At most one extension may supply a query
+planner. If any factory fails, the source context's state is unchanged.
+
+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 this protocol, including extracting the task-context provider from
+the supplied context and constructing a Python `SessionExtensionComponents`.
+
+### Advanced: chaining the low-level methods
+
+The `with_logical_extension_codec`, `with_physical_extension_codec`, and
+`with_query_planner` methods remain available for advanced use. Putting them together
+for a session using two extension libraries that each provide tables, functions, and
+a query planner:
 
 ```python
 ctx = SessionContext(config)
@@ -307,6 +366,12 @@ ctx.register_table("t", lib_a.TableProvider())
 ctx.register_udf(udf(lib_b.SomeUDF()))
 ```
 
+When chaining by hand, keep the final context assigned to `ctx` as the single owner:
+components created against earlier intermediate contexts (for example a codec
+constructed with a context that is later discarded) hold weak references that break
+once that intermediate context is collected. Prefer `with_extensions` whenever the
+extension library provides a bundle.
+
 ## Alternative Approach
 
 Suppose you needed to expose some other features of DataFusion and you could not wait

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(PlannerConfig(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(PlannerConfig(max_rows=3))