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

@@ -154,8 +154,9 @@ guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the
 weak handle during logical optimization, before plan serialization could fail
 first for an unrelated reason.
 
-`SessionContext.enable_url_table` is the one method that mints a second
-allocation for a session. Its result must not outlive the receiver.
+`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
 
@@ -185,7 +186,7 @@ pins that; changing it should be deliberate.
 
 ## Where the truth is
 
-- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat.
+- `docs/source/contributor-guide/ffi.md` — the protocol and session ownership.
 - `docs/source/user-guide/upgrade-guides.md` — every past migration.
 - `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch,
   and the two unframed cases from Rule 8.

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,13 +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. 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.
+        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),
         })

docs/source/contributor-guide/ffi.md

@@ -417,18 +417,21 @@ codec simply retain the session that built it: a codec handed to a provider is r
 registered straight back into that same session, which 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.
+`SessionContext.enable_url_table` also preserves that allocation. It wraps the
+existing catalog list in place, so previously exported weak providers remain valid
+while any handle on the session is alive. Repeated calls do not nest catalog wrappers.
 
 ### What a derived context shares
 
-`with_logical_extension_codec`, `with_physical_extension_codec`, and
+`enable_url_table`, `with_logical_extension_codec`, `with_physical_extension_codec`, and
 `with_python_udf_inlining` return a new `SessionContext` wrapping the *same* underlying
 session. Only the Python-side codec settings differ; catalogs, tables, registered
 functions, and configuration are the one shared session, so a registration on either
 side is visible to both.
 
+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.
+
 `set_query_planner` does not return anything. The query planner lives in `SessionState`,
 so it is a property of the session rather than of a handle on it, and installing one is
 visible to every context sharing that session — including ones a `with_*` call returned

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}`ffi` hook
 methods. Users who contribute their own `LogicalExtensionCodec` or
 `PhysicalExtensionCodec` via FFI must update

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

@@ -76,6 +76,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

@@ -594,10 +594,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

@@ -19,6 +19,7 @@
 import gzip
 import pathlib
 import shutil
+from concurrent.futures import ThreadPoolExecutor
 
 import pyarrow as pa
 import pyarrow.dataset as ds
@@ -41,6 +42,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())