fix: enable URL tables on the shared session

Preserve context and state identity, initialize the URL factory before publishing it, and avoid nested catalog wrappers under one write lock. Cover SQL configuration, registrations, aliases, and real FFI provider lifetimes; document the API change for #1708.

Generated-by: Codex (GPT-6)

Author: emecii

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

@@ -187,13 +187,9 @@ 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.
 
-`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
-one `session_id()` with divergent configuration. That is a bug rather than a
-design — tracked in
-[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708)
-— so do not cite it as precedent for deriving a replacement context.
+`SessionContext.enable_url_table` follows the same rule: it replaces only the
+catalog list through `state_ref()` and returns a handle sharing the original
+allocation. Its idempotence check and catalog replacement share one write lock.
 
 ## Rule 7 — installing a planner mutates the session, and says so
 

crates/core/src/context.rs

@@ -27,8 +27,11 @@ use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory};
+use datafusion::catalog::{
+    CatalogProvider, CatalogProviderList, DynamicFileCatalog, TableProviderFactory, UrlTableFactory,
+};
 use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_err};
+use datafusion::datasource::dynamic_file::DynamicListTableFactory;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::listing::{
@@ -423,15 +426,29 @@ impl PySessionContext {
     }
 
     pub fn enable_url_table(&self) -> PyResult<Self> {
-        // Pre-existing caveat, unrelated to query planners: this is the one
-        // method that mints a second `Arc<SessionContext>` for a session, and
-        // it also forks the session's state while keeping its id. Any weak
-        // `FFI_TaskContextProvider` handed out by the receiver stays bound to
-        // the receiver, so the returned context must not outlive it. See
-        // `set_session_query_planner` for why everything else mutates in place.
-        // Tracked as a bug in <https://github.com/apache/datafusion-python/issues/1708>.
+        let state_ref = self.ctx.state_ref();
+        {
+            // Check and replace under one lock so concurrent calls cannot nest
+            // wrappers or overwrite a newer catalog list.
+            let mut state = state_ref.write();
+            if !state.catalog_list().is::<DynamicFileCatalog>() {
+                let factory = Arc::new(DynamicListTableFactory::default());
+                // Bind before publishing the catalog: a reader must never see
+                // a factory whose session store has not been initialized.
+                factory
+                    .session_store()
+                    .with_state(self.ctx.state_weak_ref());
+                let catalog_list = Arc::new(DynamicFileCatalog::new(
+                    Arc::clone(state.catalog_list()),
+                    factory as Arc<dyn UrlTableFactory>,
+                ));
+                // Only the catalog changes. In particular, preserve the state
+                // and context allocations targeted by weak FFI providers.
+                state.register_catalog_list(catalog_list);
+            }
+        }
         Ok(PySessionContext {
-            ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()),
+            ctx: Arc::clone(&self.ctx),
             logical_codec: Arc::clone(&self.logical_codec),
             physical_codec: Arc::clone(&self.physical_codec),
         })
@@ -1434,14 +1451,13 @@ impl PySessionContext {
     /// time, and a payload written through one would resolve to the other on
     /// decode. See [`SESSION_CODEC_ID_PREFIX`].
     ///
-    /// Handles derived from one session — `with_python_udf_inlining`,
-    /// `with_logical_extension_codec`, [`Self::_install_extension_codecs`] —
-    /// report the same id even though their codec chains differ, so installing
-    /// two of them on one target is refused. That is the intended answer: they
-    /// share a `state_ref`, so their payloads would resolve against the same
-    /// session and are indistinguishable on decode. Every derivation shares the
-    /// session for exactly this reason; `enable_url_table` is the one that does
-    /// not, and it is tracked as a bug.
+    /// Handles derived from one session — `enable_url_table`,
+    /// `with_python_udf_inlining`, `with_logical_extension_codec`,
+    /// [`Self::_install_extension_codecs`] — report the same id even though
+    /// their codec chains may differ, so installing two of them on one target
+    /// is refused. That is the intended answer: they share a `state_ref`, so
+    /// their payloads would resolve against the same session and are
+    /// indistinguishable on decode.
     #[getter]
     pub fn __datafusion_codec_id__(&self) -> String {
         format!("{SESSION_CODEC_ID_PREFIX}{}", self.ctx.session_id())

docs/source/contributor-guide/ffi-internals.md

@@ -64,13 +64,10 @@ would close the cycle
 `SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and
 leak it.
 
-`SessionContext.enable_url_table` is the one exception. It clones the
-underlying `SessionContext`, so the returned context has an allocation of its
-own and must not outlive the receiver. It also forks the session's state while
-keeping its id, so two handles report one `session_id()` with divergent
-configuration. That is a bug rather than a design, tracked in
-[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708);
-do not copy the pattern.
+`SessionContext.enable_url_table` follows the same rule. It replaces only the
+catalog list through `state_ref()` and returns a handle sharing the original
+allocation. The idempotence check and catalog replacement share one write lock,
+so concurrent calls cannot nest wrappers or overwrite a newer catalog list.
 
 (ffi_internals_rebinding)=
 

docs/source/extension-guide/sessions.md

@@ -52,6 +52,7 @@ divergence.
 
 ## What a derived context shares
 
+{py:meth}`~datafusion.SessionContext.enable_url_table`,
 {py:meth}`~datafusion.SessionContext.with_logical_extension_codec`,
 {py:meth}`~datafusion.SessionContext.with_physical_extension_codec`,
 {py:meth}`~datafusion.SessionContext.with_python_udf_inlining`, and
@@ -105,13 +106,6 @@ The same rule applies to a capsule you take off a context inside your own code:
 a codec capsule taken from a throwaway `SessionContext()` names a session that
 is already gone and fails on first use.
 
-:::{warning}
-{py:meth}`~datafusion.SessionContext.enable_url_table` is an exception to the
-one-session-one-allocation rule above: it clones the underlying
-`SessionContext`, so the returned context has an allocation of its own and must
-not outlive the receiver. It also forks the session's state while keeping its
-id, so two handles report one `session_id()` with divergent configuration. That
-is a bug rather than a design, tracked in
-[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708);
-do not build on the behaviour.
-:::
+Enabling URL tables takes effect on the shared session even if the returned
+handle is discarded. The returned handle keeps the receiver's codec settings
+unchanged, and repeated calls do not nest catalog wrappers.

docs/source/user-guide/upgrade-guides.md

@@ -21,6 +21,34 @@
 
 ## DataFusion 55.0.0
 
+### URL tables share the original session
+
+`SessionContext.enable_url_table()` now enables URL tables on the existing session
+and returns another handle on it. Previously it copied session state into a separate
+session while retaining the same session id. Configuration and function registrations
+could then diverge, and replacing the original handle could invalidate FFI providers.
+
+Before, callers had to use the returned context to query file paths:
+
+```python
+ctx = SessionContext()
+enabled = ctx.enable_url_table()
+# Only enabled could query local file paths as tables.
+```
+
+After, both handles share configuration, registrations, and URL table support:
+
+```python
+ctx = SessionContext()
+ctx.enable_url_table()  # Takes effect even when the returned handle is discarded.
+```
+
+Existing `ctx = ctx.enable_url_table()` calls continue to work and now retain FFI
+providers bound to the original session. Repeated calls have no effect. To keep a
+session without URL table support, create a separate `SessionContext` explicitly.
+
+### FFI codec hooks receive the session
+
 This release extends the change made in 52.0.0 to the remaining
 {ref}`extension_capsule_protocol` hook methods. Users who contribute their own
 `LogicalExtensionCodec` or

examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py

@@ -93,6 +93,43 @@ def probe_context(
     return ctx, logical_codec, physical_codec
 
 
+@pytest.mark.parametrize("keep_returned", [False, True])
+def test_enable_url_table_preserves_ffi_providers(keep_returned):
+    """A catalog's unreachable weak codec remains valid through URL enabling."""
+    ctx, logical_codec, physical_codec = probe_context(
+        logical_requires=HOST_ONLY_UDF, physical_requires=HOST_ONLY_UDF, max_rows=100
+    )
+    ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())
+    session_id = ctx.session_id()
+    alias = ctx.with_python_udf_inlining(enabled=True)
+    if keep_returned:
+        ctx = alias.enable_url_table()
+    else:
+        alias.enable_url_table()
+    del alias
+    gc.collect()
+
+    for _ in range(2):
+        # Force filter pushdown through a real foreign catalog before planning.
+        batches = ctx.sql(
+            "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5"
+        ).collect()
+        assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [
+            7,
+            10,
+            20,
+            30,
+        ]
+        assert logical_codec.table_provider_decode_calls() > 0
+        assert physical_codec.execution_plan_decode_calls() > 0
+        assert logical_codec.last_task_context_session_id() == session_id
+        assert physical_codec.last_task_context_session_id() == session_id
+        assert logical_codec.task_context_udf_resolutions() > 0
+        assert physical_codec.task_context_udf_resolutions() > 0
+        ctx = ctx.enable_url_table()
+        gc.collect()
+
+
 def test_logical_codec_resolves_a_host_registered_udf():
     """``try_decode_table_provider`` sees the host session's registry.
 

python/datafusion/context.py

@@ -551,7 +551,8 @@ class SessionContext:
     See :ref:`user_guide_concepts` in the online documentation for more information.
 
     **A context is a handle on a session, not the session itself.** The
-    ``with_*`` methods — :py:meth:`with_logical_extension_codec`,
+    derivation methods — :py:meth:`enable_url_table`,
+    :py:meth:`with_logical_extension_codec`,
     :py:meth:`with_physical_extension_codec`,
     :py:meth:`with_python_udf_inlining`, and :py:meth:`with_extensions` —
     return a new context wrapping the *same* underlying session. Only the
@@ -560,10 +561,11 @@ class SessionContext:
     either handle is visible to both.
 
     A few things therefore belong to the session rather than to a handle, and
-    take effect even if the handle that set them is discarded: the query
-    planner (see :py:meth:`set_query_planner`), and the rebuild of an installed
-    foreign planner that follows installing a codec. :ref:`extension_sessions`
-    in the online documentation works through when that matters.
+    take effect even if the handle that set them is discarded: URL table
+    support, the query planner (see :py:meth:`set_query_planner`), and the
+    rebuild of an installed foreign planner that follows installing a codec.
+    :ref:`extension_sessions` in the online documentation works through when
+    that matters.
 
     **Keep a context alive for as long as anything derived from it is in use.**
     A :py:class:`~datafusion.DataFrame`, logical plan, or exported capsule does
@@ -619,10 +621,26 @@ def global_ctx(cls) -> SessionContext:
         return wrapper
 
     def enable_url_table(self) -> SessionContext:
-        """Control if local files can be queried as tables.
+        """Enable querying local files as tables on the shared session.
+
+        The receiver and all handles sharing its session gain URL table support,
+        even if the returned handle is discarded. Repeated calls have no effect.
+        Registered catalogs, functions, configuration, and session identity are
+        preserved, as are FFI providers bound to the session.
 
         Returns:
-            A new :py:class:`SessionContext` object with url table enabled.
+            A new :py:class:`SessionContext` handle wrapping the same session.
+
+        Examples:
+            >>> ctx = SessionContext()
+            >>> enabled = ctx.enable_url_table()
+            >>> enabled.session_id() == ctx.session_id()
+            True
+            >>> ctx.sql("SET datafusion.execution.batch_size = 111").collect()
+            []
+            >>> batches = enabled.sql("SHOW datafusion.execution.batch_size").collect()
+            >>> batches[0].column(1).to_pylist()
+            ['111']
         """
         klass = self.__class__
         obj = klass.__new__(klass)

python/tests/test_context.py

@@ -20,6 +20,7 @@
 import gzip
 import pathlib
 import shutil
+from concurrent.futures import ThreadPoolExecutor
 
 import pyarrow as pa
 import pyarrow.dataset as ds
@@ -43,6 +44,91 @@ def test_create_context_no_args():
     SessionContext()
 
 
+@pytest.mark.parametrize("discard_returned", [False, True])
+def test_enable_url_table_shares_session(tmp_path, discard_returned):
+    """URL tables, configuration, and registrations belong to all aliases."""
+    ctx = SessionContext()
+    original_config = (
+        ctx.sql("SHOW datafusion.catalog.create_default_catalog_and_schema")
+        .collect()[0]
+        .column(1)
+        .to_pylist()
+    )
+    alias = ctx.with_python_udf_inlining(enabled=False)
+    session_id = ctx.session_id()
+    ctx.sql("CREATE SCHEMA existing").collect()
+    ctx.register_record_batches("existing.numbers", [[pa.record_batch({"n": [7]})]])
+    ctx.register_udf(
+        udf(lambda x: x, [pa.int64()], pa.int64(), "immutable", "identity")
+    )
+    path = tmp_path / "numbers.csv"
+    path.write_text("n\n7\n")
+
+    if discard_returned:
+        alias.enable_url_table()
+        returned = ctx
+    else:
+        returned = alias.enable_url_table()
+
+    # Both directions must see subsequent SETs, not only the initial snapshot.
+    for writer, reader, value in [(ctx, returned, 111), (returned, alias, 222)]:
+        writer.sql(f"SET datafusion.execution.batch_size = {value}").collect()
+        assert reader.sql("SHOW datafusion.execution.batch_size").collect()[0].column(
+            1
+        ).to_pylist() == [str(value)]
+
+    returned.register_udf(
+        udf(lambda x: x, [pa.int64()], pa.int64(), "immutable", "later_identity")
+    )
+    for handle in (ctx, alias, returned):
+        assert handle.session_id() == session_id
+        assert (
+            handle.sql("SHOW datafusion.catalog.create_default_catalog_and_schema")
+            .collect()[0]
+            .column(1)
+            .to_pylist()
+            == original_config
+        )
+        assert handle.sql("SELECT identity(n) FROM existing.numbers").collect()[
+            0
+        ].column(0).to_pylist() == [7]
+        assert handle.sql("SELECT later_identity(9)").collect()[0].column(
+            0
+        ).to_pylist() == [9]
+        assert handle.sql(f'SELECT n FROM "{path}"').collect()[0].column(
+            0
+        ).to_pylist() == [7]
+        handle.enable_url_table()
+        handle.enable_url_table()
+        assert handle.sql(f'SELECT n FROM "{path}"').collect()[0].column(
+            0
+        ).to_pylist() == [7]
+
+
+def test_enable_url_table_from_multiple_aliases(tmp_path):
+    """Enabling through multiple handles preserves concurrent registrations."""
+    ctx = SessionContext()
+    path = tmp_path / "numbers.csv"
+    path.write_text("n\n7\n")
+    aliases = [ctx.with_python_udf_inlining(enabled=False) for _ in range(4)]
+
+    def enable_and_register(index):
+        alias = aliases[index]
+        for _ in range(4):
+            alias.enable_url_table()
+            alias.sql(f'SELECT n FROM "{path}"').collect()
+        alias.register_udf(
+            udf(lambda x: x, [pa.int64()], pa.int64(), "immutable", f"identity_{index}")
+        )
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        list(executor.map(enable_and_register, range(4)))
+    for index in range(4):
+        assert ctx.sql(f"SELECT identity_{index}(7)").collect()[0].column(
+            0
+        ).to_pylist() == [7]
+
+
 def test_create_context_session_config_only():
     SessionContext(config=SessionConfig())