fix: Infer and scan BSON files from file path (#2410)

GlareDB · Jan 14, 2024 · c2131a3 · c2131a3
1 parent 4590005
commit c2131a3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/crates/datafusion_ext/src/planner/relation/mod.rs b/crates/datafusion_ext/src/planner/relation/mod.rs
@@ -255,6 +255,10 @@ fn infer_func_for_file(path: &str) -> Result<OwnedTableReference> {
             schema: "public".into(),
             table: "ndjson_scan".into(),
         },
+        "bson" => OwnedTableReference::Partial {
+            schema: "public".into(),
+            table: "read_bson".into(),
+        },
         ext => {
             return Err(DataFusionError::Plan(format!(
                 "unable to infer how to handle file extension: {ext}"

diff --git a/tests/tests/test_bson.py b/tests/tests/test_bson.py
@@ -74,7 +74,7 @@ def test_read_bson(
             f"create external table bson_beatles from bson options ( location='{data_path}', file_type='bson')"
         )
 
-    for from_clause in ["bson_beatles", f"read_bson('{data_path}')"]:
+    for from_clause in ["bson_beatles", f"read_bson('{data_path}')", f"'{data_path}'"]:
         with glaredb_connection.cursor() as curr:
             curr.execute(f"select count(*) from {from_clause}")
             r = curr.fetchone()