diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 37be0bc9d..0cfbd3c9b 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.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,19 @@ 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 a new +parameter on `_commit_extensions` whose commit cannot fail. 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 diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 9a92c6055..0268b523e 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -844,35 +844,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::() { - let data: NonNull = capsule - .pointer_checked(Some(c"datafusion_catalog_provider"))? - .cast(); - let provider = unsafe { data.as_ref() }; - let provider: Arc = provider.into(); - provider - } else { - match provider.extract::() { - Ok(py_catalog) => py_catalog.catalog, - Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( - provider.into(), - self.ffi_logical_codec(), - )) as Arc, - } - }; + let provider = self.resolve_catalog_provider(provider)?; let _ = self.ctx.register_catalog(name, provider); @@ -1775,6 +1749,7 @@ impl PySessionContext { udafs: Vec, udwfs: Vec, udtfs: Vec, + catalogs: PyRef<'_, PyResolvedCatalogs>, rules: PyRef<'_, PyPhysicalOptimizerRules>, ) -> PyDataFusionResult<()> { let py = slf.py(); @@ -1824,6 +1799,12 @@ impl PySessionContext { for udtf in udtfs { this.register_udtf(udtf); } + // Nothing here can fail: the providers were imported by + // [`Self::_resolve_extension_catalogs`], and `register_catalog` + // returns whichever provider it displaced rather than refusing. + for (name, provider) in &catalogs.catalogs { + let _ = this.ctx.register_catalog(name, Arc::clone(provider)); + } // Rules accumulate rather than replace, so unlike a planner there is // no composition order to get right and no collision to refuse. All // of them go on in **one** `SessionState` rebuild. @@ -1923,6 +1904,39 @@ impl PySessionContext { Ok(PyResolvedTables { tables: resolved }) } + /// 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 { + let catalogs = catalogs + .into_iter() + .map(|(name, provider)| { + // The name identifies the culprit declaration, as for tables: + // a getter that returns junk fails in the importer, which + // knows neither the catalog nor the bundle. + let provider = self.resolve_catalog_provider(provider).map_err(|err| { + exec_datafusion_err!("Resolving the declared catalog {name}: {err}") + })?; + Ok((name, provider)) + }) + .collect::>>()?; + Ok(PyResolvedCatalogs { catalogs }) + } + /// Import the physical optimizer rules a `with_extensions` call declared. /// /// The fallible half of installing them, run while the call can still fail @@ -1964,6 +1978,15 @@ struct ResolvedTable { provider: Arc, } +/// Catalog providers imported for a `with_extensions` call. +/// +/// Opaque to Python, like [`PyResolvedTables`] and [`PyPhysicalOptimizerRules`], +/// and `frozen` for the same reason as both: the commit only reads. +#[pyclass(frozen, name = "ResolvedCatalogs", module = "datafusion._internal")] +pub struct PyResolvedCatalogs { + catalogs: Vec<(String, Arc)>, +} + /// Physical optimizer rules imported for a `with_extensions` call. /// /// Opaque to Python, and deliberately not added to the module: it exists only @@ -1984,6 +2007,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> { + 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::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_catalog_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider + } else { + match provider.extract::() { + Ok(py_catalog) => py_catalog.catalog, + Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( + provider.into(), + self.ffi_logical_codec(), + )) as Arc, + } + }) + } + /// Write the session's query planner, in place. /// /// Pass `Some(planner)` to install one, or `None` to rebuild whichever diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index a225f53f6..da003facc 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -129,12 +129,13 @@ A call therefore splits into a part that may fail and a part that may not: even though it can fail on a bad capsule or a duplicate id. 3. **Resolve.** Every declared function is wrapped and every name is checked, every declared table has its provider imported and its destination schema - resolved, every declared physical optimizer rule has its capsule imported, - and every `__datafusion_session_planner__` runs against the completed - chains. + resolved, every declared catalog has its provider imported, every declared + physical optimizer rule has its capsule imported, and every + `__datafusion_session_planner__` runs against the completed chains. 4. **Commit.** The tables are inserted, the planner is bound, the functions - are registered, and the optimizer rules are installed in a single - `SessionState` rebuild. + are registered, the catalogs are registered — `register_catalog` replaces + rather than refuses, so it cannot fail — and the optimizer rules are + installed in a single `SessionState` rebuild. Only step 4 touches the session, and every step that can fail happens before it — with one honest exception: a table insert goes through a diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 11dc278c7..33d3c685f 100644 --- a/docs/source/extension-guide/bundles.md +++ b/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. @@ -348,6 +348,8 @@ is nothing to refuse. See {doc}`other-components`. Tables go the other way: a declared table name that is *already* on the session is an error too, so 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`. Strictly, whether a duplicate registration replaces or refuses is the `SchemaProvider`'s own call, and the in-memory one a session starts with diff --git a/docs/source/extension-guide/table-providers.md b/docs/source/extension-guide/table-providers.md index 700ca1954..9e09e6a85 100644 --- a/docs/source/extension-guide/table-providers.md +++ b/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 diff --git a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py index c2e3a1dfa..6cd7cf27e 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py +++ b/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, @@ -256,6 +257,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. diff --git a/examples/datafusion-ffi-example/src/extension.rs b/examples/datafusion-ffi-example/src/extension.rs index 7b5d668fb..148439738 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/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; @@ -244,3 +245,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> { + 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)) + } +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 25ebb2f3e..a7dceb3e9 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/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::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index f46deb11e..e8e792eee 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -308,10 +308,16 @@ def _collect_contributions( declared: dict[str, list[tuple[int, object, Any]]] = { kind.field: [] for kind in _FUNCTION_KINDS } - # The rest are not function kinds: tables and table functions are named by - # the bundle rather than by the value, and rules carry no name at all, so - # each has its own collision rule -- or none -- and its own installer. - for field_name in ("udtfs", "table_providers", "physical_optimizer_rules"): + # The rest are not function kinds: tables, table functions, and catalogs are + # named by the bundle rather than by the value, and rules carry no name at + # all, so each has its own collision rule -- or none -- and its own + # installer. + for field_name in ( + "udtfs", + "table_providers", + "catalog_providers", + "physical_optimizer_rules", + ): declared[field_name] = [] for position, extension in enumerate(extensions): if not isinstance(extension, SessionComponentsExportable): @@ -2229,17 +2235,17 @@ def with_extensions( a declared function or optimizer rule does not expose its capsule getter and is not already a wrapper. ValueError: If two codecs claim the same id, if two extensions - declare a function, table, or table function of one kind under - the same name, or if a getter returns a capsule of the wrong - kind. See :py:meth:`with_logical_extension_codec` for how ids - are assigned. + declare a function, table, table function, or catalog of one + kind under the same name, or if a getter returns a capsule of + the wrong kind. See :py:meth:`with_logical_extension_codec` + for how ids are assigned. RuntimeError: If a getter is present but returns something that is not a ``PyCapsule`` at all. The message comes from the importer and does not name the bundle, because by then the declaration has already been accepted as the right shape. - Exception: If a declared table cannot be resolved — the name is - already registered, two declarations resolve to one table, the - schema is unknown, or the value is not a table. + Exception: If a declared table or catalog cannot be resolved — + name taken, two declarations resolving to one table, unknown + schema, value not a table, or a capsule of the wrong kind. Examples: The returned handle is a different object sharing one session, and @@ -2327,6 +2333,18 @@ def with_extensions( ) ] ) + # Bound to `new` for the same reason: the catalog getter is handed the + # logical codec its provider will serialize through. Unlike a table a + # catalog may replace one the session holds, so only names claimed + # twice within the call are refused. + resolved_catalogs = new.ctx._resolve_extension_catalogs( + [ + pair + for _, _, pair in _reject_repeated_names( + declared["catalog_providers"], "catalog" + ) + ] + ) # Rules accumulate, so there is no name to check and nothing to refuse # -- only the capsules to import while failing is still free. resolved_rules = _resolve_declared_rules( @@ -2353,6 +2371,7 @@ def with_extensions( [function._udaf for function in resolved["udafs"]], [function._udwf for function in resolved["udwfs"]], [table_function._udtf for table_function in resolved_udtfs], + resolved_catalogs, resolved_rules, ) return new diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 8e89c5c43..020d36eef 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -381,6 +381,19 @@ class SessionExtensionComponents: replaced: see :ref:`extension_bundles_collisions`. """ + catalog_providers: tuple[tuple[str, Any], ...] = _components("catalog") + """Catalogs to register, as ``(name, provider)`` pairs. + + Anything + :py:meth:`~datafusion.context.SessionContext.register_catalog_provider` + accepts. Bound to the finished context like :py:attr:`table_providers`. + + Two extensions claiming one name in the same call is refused. Replacing a + catalog the session already has is not — the default ``datafusion`` catalog + always exists, and swapping it is the usual way a library backs a session + with its own metadata. + """ + physical_optimizer_rules: tuple[PhysicalOptimizerRuleExportable, ...] = _components( "optimizer rule" ) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index b160b5cae..ef397fbf1 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -22,6 +22,7 @@ import shutil from dataclasses import fields +import datafusion.catalog import pyarrow as pa import pyarrow.compute as pc import pyarrow.dataset as ds @@ -1843,6 +1844,149 @@ def test_with_extensions_rejects_a_table_that_is_not_a_table_by_name(ctx): ctx.udf("double") +class _Schema(datafusion.catalog.SchemaProvider): + """One table, enough to prove a declared catalog is reachable from SQL.""" + + def __init__(self, table): + self.tables = {"t": table} + + def table_names(self) -> set[str]: + return set(self.tables) + + def register_table(self, name, table): + self.tables[name] = table + + def deregister_table(self, name, cascade: bool = True): + del self.tables[name] + + def table(self, name): + return self.tables.get(name) + + def table_exist(self, name) -> bool: + return name in self.tables + + +class _Catalog(datafusion.catalog.CatalogProvider): + def __init__(self, table): + self.schemas = {"s": _Schema(table)} + + def schema_names(self) -> set[str]: + return set(self.schemas) + + def schema(self, name): + return self.schemas.get(name) + + def register_schema(self, name, schema): + self.schemas[name] = schema + + def deregister_schema(self, name, cascade: bool): + del self.schemas[name] + + +class _CatalogExtension: + """Contributes catalogs as ``(name, provider)`` pairs.""" + + def __init__(self, catalog_providers=()): + self._catalog_providers = catalog_providers + + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents(catalog_providers=self._catalog_providers) + + +def test_with_extensions_registers_a_declared_catalog(ctx): + """A declared catalog is queryable through its qualified name.""" + table = ctx.from_pydict({"a": [1, 2, 3]}).into_view() + + result = ctx.with_extensions( + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)) + ) + + assert "engine" in result.catalog_names() + assert ( + result.sql("SELECT sum(a) FROM engine.s.t").collect()[0].column(0)[0].as_py() + == 6 + ) + + +def test_with_extensions_rejects_a_catalog_two_extensions_claim(ctx): + """One name, two bundles: refused with both named.""" + table = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(ValueError, match=r"catalog named 'engine'"): + ctx.with_extensions( + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)), + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)), + ) + + assert "engine" not in ctx.catalog_names() + + +def test_with_extensions_allows_replacing_an_existing_catalog(ctx): + """Replacing a catalog the session holds is ordinary, unlike a table. + + ``register_catalog`` returns whichever provider it displaced rather than + refusing, and the default ``datafusion`` catalog always exists — so a + library backing a session with its own metadata has to be able to do this. + """ + table = ctx.from_pydict({"a": [1, 2, 3]}).into_view() + + result = ctx.with_extensions( + _CatalogExtension(catalog_providers=(("datafusion", _Catalog(table)),)) + ) + + assert ( + result.sql("SELECT sum(a) FROM datafusion.s.t") + .collect()[0] + .column(0)[0] + .as_py() + == 6 + ) + + +def test_with_extensions_registers_no_catalog_when_a_later_hook_raises(ctx): + """Catalogs share the transaction, even though they write to a shared list.""" + + class BoomPlanner: + def __datafusion_session_planner__(self, ctx, fallback): + msg = "boom" + raise RuntimeError(msg) + + table = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions( + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)), + BoomPlanner(), + ) + + assert "engine" not in ctx.catalog_names() + + +def test_with_extensions_rejects_a_catalog_capsule_of_the_wrong_kind(ctx): + """A catalog whose getter returns the wrong capsule is refused by name. + + A non-capsule return is not refusable — any object can be a catalog + provider, so it falls through to the duck-typed wrap, like a table getter's + falls through to the dataset one. A capsule of the wrong *kind* is the + resolve failure that remains, and the declared name — unique within the + call — is what points at the culprit declaration, as for tables. + """ + + class WrongCapsuleCatalog: + def __datafusion_catalog_provider__(self, codec): + return ctx.__datafusion_query_planner__() + + with pytest.raises(Exception, match=r"declared catalog engine"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + _CatalogExtension(catalog_providers=(("engine", WrongCapsuleCatalog()),)), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + assert "engine" not in ctx.catalog_names() + + def test_session_extension_components_rejects_a_single_optimizer_rule(): """The same for rules, naming what that field holds.""" with pytest.raises( @@ -1925,7 +2069,7 @@ def test_every_component_field_has_an_installer(): Reaching into private names on purpose: the two sides answer different questions. The metadata says which fields are collections to normalize; - ``_FUNCTION_KINDS`` and the four fields named here say which of them + ``_FUNCTION_KINDS`` and the fields named here say which of them ``with_extensions`` knows how to install. Nothing observable from outside can tell you they have drifted, because the symptom is silence. """ @@ -1942,6 +2086,7 @@ def test_every_component_field_has_an_installer(): "function": {kind.field for kind in _FUNCTION_KINDS}, "table function": {"udtfs"}, "table": {"table_providers"}, + "catalog": {"catalog_providers"}, "optimizer rule": {"physical_optimizer_rules"}, } diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 3a3eb3ca8..f0c86a401 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -46,6 +46,8 @@ # Tables, split the same way: import the provider and resolve the # destination schema during resolution, insert at commit. "_resolve_extension_tables", + # Catalogs, sharing their import half with register_catalog_provider. + "_resolve_extension_catalogs", } )