From 62ae45176f985c9238b972b7d0a6389f435cfe21 Mon Sep 17 00:00:00 2001
From: Chongchen Chen <chenkovsky@qq.com>
Date: Sun, 22 Dec 2024 15:09:54 +0800
Subject: [PATCH] feat: support enable_url_table config

---
 examples/create-context.py   |  4 ++++
 python/datafusion/context.py | 19 +++++++++++++++++--
 src/context.rs               | 11 +++++++----
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/examples/create-context.py b/examples/create-context.py
index 3184d4085..898eb7972 100644
--- a/examples/create-context.py
+++ b/examples/create-context.py
@@ -37,3 +37,7 @@
 )
 ctx = SessionContext(config, runtime)
 print(ctx)
+
+config = config.with_url_table(True)
+ctx = SessionContext(config, runtime)
+print(ctx)
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index a07b5d175..1e519a123 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -66,13 +66,27 @@ def __arrow_c_array__(  # noqa: D105
 class SessionConfig:
     """Session configuration options."""
 
-    def __init__(self, config_options: dict[str, str] | None = None) -> None:
+    def __init__(self, config_options: dict[str, str] | None = None, enable_url_table: bool = False) -> 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
@@ -467,10 +481,11 @@ 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)
+        self.ctx = SessionContextInternal(config, runtime, enable_url_table)
 
     def register_object_store(
         self, schema: str, store: Any, host: str | None = None
diff --git a/src/context.rs b/src/context.rs
index 8675e97df..a32c7d82a 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -272,11 +272,12 @@ pub struct PySessionContext {
 
 #[pymethods]
 impl PySessionContext {
-    #[pyo3(signature = (config=None, runtime=None))]
+    #[pyo3(signature = (config=None, runtime=None, enable_url_table=false))]
     #[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
@@ -294,9 +295,11 @@ impl PySessionContext {
             .with_runtime_env(runtime)
             .with_default_features()
             .build();
-        Ok(PySessionContext {
-            ctx: SessionContext::new_with_state(session_state),
-        })
+        let mut ctx = SessionContext::new_with_state(session_state);
+        if enable_url_table {
+            ctx = ctx.enable_url_table();
+        }
+        Ok(PySessionContext { ctx })
     }
 
     /// Register an object store with the given name