change enable_url_table as method

apache · Jan 7, 2025 · 99c2fd0 · 99c2fd0
1 parent 62ae451
commit 99c2fd0
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 26 deletions.
diff --git a/examples/create-context.py b/examples/create-context.py
@@ -38,6 +38,5 @@
 ctx = SessionContext(config, runtime)
 print(ctx)
 
-config = config.with_url_table(True)
-ctx = SessionContext(config, runtime)
+ctx = ctx.enable_url_table()
 print(ctx)
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
@@ -66,27 +66,13 @@ def __arrow_c_array__(  # noqa: D105
 class SessionConfig:
     """Session configuration options."""
 
-    def __init__(self, config_options: dict[str, str] | None = None, enable_url_table: bool = False) -> None:
+    def __init__(self, config_options: dict[str, str] | None = None) -> None:
         """Create a new :py:class:`SessionConfig` with the given configuration options.
 
         Args:
             config_options: Configuration options.
         """
         self.config_internal = SessionConfigInternal(config_options)
-        self.enable_url_table = enable_url_table
-
-    def with_url_table(self, enabled: bool = True) -> SessionConfig:
-
-        """Control if local files can be queried as tables.
-
-        Args:
-            enabled: Whether local files can be queried as tables.
-
-        Returns:
-            A new :py:class:`SessionConfig` object with the updated setting.
-        """
-        self.enable_url_table = enabled
-        return self
 
     def with_create_default_catalog_and_schema(
         self, enabled: bool = True
@@ -481,11 +467,22 @@ def __init__(
             ctx = SessionContext()
             df = ctx.read_csv("data.csv")
         """
-        enable_url_table = config.enable_url_table if config is not None else False
         config = config.config_internal if config is not None else None
         runtime = runtime.config_internal if runtime is not None else None
 
-        self.ctx = SessionContextInternal(config, runtime, enable_url_table)
+        self.ctx = SessionContextInternal(config, runtime)
+
+    def enable_url_table(self) -> "SessionContext":
+
+        """Control if local files can be queried as tables.
+
+        Returns:
+            A new :py:class:`SessionContext` object with url table enabled.
+        """
+        klass = self.__class__
+        obj = klass.__new__(klass)
+        obj.ctx = self.ctx.enable_url_table()
+        return obj
 
     def register_object_store(
         self, schema: str, store: Any, host: str | None = None

diff --git a/src/context.rs b/src/context.rs
@@ -272,12 +272,11 @@ pub struct PySessionContext {
 
 #[pymethods]
 impl PySessionContext {
-    #[pyo3(signature = (config=None, runtime=None, enable_url_table=false))]
+    #[pyo3(signature = (config=None, runtime=None))]
     #[new]
     pub fn new(
         config: Option<PySessionConfig>,
         runtime: Option<PyRuntimeConfig>,
-        enable_url_table: bool,
     ) -> PyResult<Self> {
         let config = if let Some(c) = config {
             c.config
@@ -295,11 +294,15 @@ impl PySessionContext {
             .with_runtime_env(runtime)
             .with_default_features()
             .build();
-        let mut ctx = SessionContext::new_with_state(session_state);
-        if enable_url_table {
-            ctx = ctx.enable_url_table();
-        }
-        Ok(PySessionContext { ctx })
+        Ok(PySessionContext {
+            ctx: SessionContext::new_with_state(session_state),
+        })
+    }
+
+    pub fn enable_url_table(&self) -> PyResult<Self> {
+        Ok(PySessionContext {
+            ctx: self.ctx.clone().enable_url_table(),
+        })
     }
 
     /// Register an object store with the given name