fix: standardize usage of TABLE_PROVIDER_CAPSULE_NAME to table_provider_capsule_name function

Author: kosiew

examples/datafusion-ffi-example/Cargo.toml

@@ -23,7 +23,12 @@ edition = "2021"
 [dependencies]
 datafusion = { version = "49.0.2" }
 datafusion-ffi = { version = "49.0.2" }
-pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py39"] }
+datafusion-python = { path = "../../" }
+pyo3 = { version = "0.25", features = [
+    "extension-module",
+    "abi3",
+    "abi3-py39",
+] }
 arrow = { version = "55.0.0" }
 arrow-array = { version = "55.0.0" }
 arrow-schema = { version = "55.0.0" }

examples/datafusion-ffi-example/src/table_provider.rs

@@ -20,6 +20,7 @@ use arrow_schema::{DataType, Field, Schema};
 use datafusion::catalog::MemTable;
 use datafusion::error::{DataFusionError, Result as DataFusionResult};
 use datafusion_ffi::table_provider::FFI_TableProvider;
+use datafusion_python::utils::table_provider_capsule_name;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::types::PyCapsule;
 use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
@@ -91,13 +92,11 @@ impl MyTableProvider {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = cr"datafusion_table_provider".into();
-
         let provider = self
             .create_table()
             .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
         let provider = FFI_TableProvider::new(Arc::new(provider), false, None);
 
-        PyCapsule::new(py, provider, Some(name))
+        PyCapsule::new(py, provider, Some(table_provider_capsule_name().to_owned()))
     }
 }

src/catalog.rs

@@ -18,7 +18,8 @@
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::utils::{
-    get_tokio_runtime, try_table_provider_from_object, validate_pycapsule, wait_for_future,
+    get_tokio_runtime, table_provider_capsule_name, try_table_provider_from_object,
+    validate_pycapsule, wait_for_future,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -36,7 +37,6 @@ use pyo3::types::PyCapsule;
 use pyo3::IntoPyObjectExt;
 use std::any::Any;
 use std::collections::HashSet;
-use std::ffi::CString;
 use std::sync::Arc;
 
 #[pyclass(name = "RawCatalog", module = "datafusion.catalog", subclass)]
@@ -257,14 +257,13 @@ impl PyTable {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = CString::new("datafusion_table_provider").unwrap();
         let runtime = get_tokio_runtime().0.handle().clone();
 
         let provider = Arc::clone(&self.table);
         let provider: Arc<dyn TableProvider + Send> = provider;
         let provider = FFI_TableProvider::new(provider, false, Some(runtime));
 
-        PyCapsule::new(py, provider, Some(name.clone()))
+        PyCapsule::new(py, provider, Some(table_provider_capsule_name().to_owned()))
     }
 
     fn __repr__(&self) -> PyResult<String> {

src/dataframe.rs

@@ -51,7 +51,8 @@ use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{
-    get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, validate_pycapsule, wait_for_future,
+    get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, table_provider_capsule_name,
+    validate_pycapsule, wait_for_future,
 };
 use crate::{
     errors::PyDataFusionResult,
@@ -83,12 +84,10 @@ impl PyTableProvider {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = CString::new("datafusion_table_provider").unwrap();
-
         let runtime = get_tokio_runtime().0.handle().clone();
         let provider = FFI_TableProvider::new(Arc::clone(&self.provider), false, Some(runtime));
 
-        PyCapsule::new(py, provider, Some(name.clone()))
+        PyCapsule::new(py, provider, Some(table_provider_capsule_name().to_owned()))
     }
 }
 

src/utils.rs

@@ -27,12 +27,26 @@ use datafusion::{
 use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::prelude::*;
 use pyo3::{exceptions::PyValueError, types::PyCapsule};
+use std::ffi::CString;
 use std::{
+    ffi::CStr,
     future::Future,
     sync::{Arc, OnceLock},
     time::Duration,
 };
 use tokio::{runtime::Runtime, time::sleep};
+
+/// Return a static CStr for the PyCapsule name.
+///
+/// We create this lazily from a `CString` to avoid unsafe const
+/// initialization from a byte literal and to satisfy compiler lints.
+pub fn table_provider_capsule_name() -> &'static CStr {
+    static NAME: OnceLock<CString> = OnceLock::new();
+    NAME.get_or_init(|| CString::new("datafusion_table_provider").unwrap())
+        .as_c_str()
+}
+
+pub const TABLE_PROVIDER_CAPSULE_NAME_STR: &str = "datafusion_table_provider";
 /// Utility to get the Tokio Runtime from Python
 #[inline]
 pub(crate) fn get_tokio_runtime() -> &'static TokioRuntime {
@@ -125,7 +139,7 @@ pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyRe
 pub(crate) fn foreign_table_provider_from_capsule(
     capsule: &Bound<PyCapsule>,
 ) -> PyResult<ForeignTableProvider> {
-    validate_pycapsule(capsule, "datafusion_table_provider")?;
+    validate_pycapsule(capsule, TABLE_PROVIDER_CAPSULE_NAME_STR)?;
     Ok(unsafe { capsule.reference::<FFI_TableProvider>() }.into())
 }