chore: make bson copy to options more consistent (#2704)

GlareDB · Feb 27, 2024 · 0eabfad · 0eabfad
1 parent 8a38470
commit 0eabfad
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 6 deletions.
diff --git a/crates/protogen/src/metastore/types/options.rs b/crates/protogen/src/metastore/types/options.rs
@@ -1714,7 +1714,7 @@ pub enum CopyToFormatOptions {
     Parquet(CopyToFormatOptionsParquet),
     Lance(CopyToFormatOptionsLance),
     Json(CopyToFormatOptionsJson),
-    Bson,
+    Bson(CopyToFormatOptionsBson),
 }
 
 impl Default for CopyToFormatOptions {
@@ -1738,7 +1738,7 @@ impl CopyToFormatOptions {
             Self::Csv(_) => Self::CSV,
             Self::Parquet(_) => Self::PARQUET,
             Self::Json(_) => Self::JSON,
-            Self::Bson => Self::BSON,
+            Self::Bson(_) => Self::BSON,
             Self::Lance(_) => Self::LANCE,
         }
     }
@@ -1760,6 +1760,10 @@ pub struct CopyToFormatOptionsJson {
     pub array: bool,
 }
 
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+pub struct CopyToFormatOptionsBson {}
+
+
 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
 pub struct CopyToFormatOptionsLance {
     pub max_rows_per_file: Option<usize>,

diff --git a/crates/protogen/src/sqlexec/copy_to.rs b/crates/protogen/src/sqlexec/copy_to.rs
@@ -76,6 +76,8 @@ pub enum CopyToFormatOptionsEnum {
     Parquet(CopyToFormatOptionsParquet),
     #[prost(message, tag = "4")]
     Lance(CopyToFormatOptionsLance),
+    #[prost(message, tag = "5")]
+    Bson(CopyToFormatOptionsBson),
 }
 
 #[derive(Clone, PartialEq, Message)]
@@ -119,7 +121,7 @@ impl TryFrom<crate::metastore::types::options::CopyToFormatOptions> for CopyToFo
         value: crate::metastore::types::options::CopyToFormatOptions,
     ) -> Result<Self, Self::Error> {
         match value {
-            crate::metastore::types::options::CopyToFormatOptions::Bson => {
+            crate::metastore::types::options::CopyToFormatOptions::Bson(_) => {
                 Ok(CopyToFormatOptions::default())
             }
             crate::metastore::types::options::CopyToFormatOptions::Lance(opts) => {
@@ -198,6 +200,11 @@ impl TryFrom<CopyToFormatOptions> for crate::metastore::types::options::CopyToFo
                     crate::metastore::types::options::CopyToFormatOptionsJson { array: json.array },
                 ))
             }
+            CopyToFormatOptionsEnum::Bson(_) => {
+                Ok(crate::metastore::types::options::CopyToFormatOptions::Bson(
+                    crate::metastore::types::options::CopyToFormatOptionsBson {},
+                ))
+            }
             CopyToFormatOptionsEnum::Parquet(parquet) => Ok(
                 crate::metastore::types::options::CopyToFormatOptions::Parquet(
                     crate::metastore::types::options::CopyToFormatOptionsParquet {

diff --git a/crates/sqlexec/src/optimizer.rs b/crates/sqlexec/src/optimizer.rs
@@ -113,6 +113,7 @@ mod test {
         CopyToDestinationOptions,
         CopyToDestinationOptionsLocal,
         CopyToFormatOptions,
+        CopyToFormatOptionsBson,
     };
     use uuid::Uuid;
 
@@ -214,7 +215,7 @@ mod test {
                 dest: CopyToDestinationOptions::Local(CopyToDestinationOptionsLocal {
                     location: "/tmp".to_string(),
                 }),
-                format: CopyToFormatOptions::Bson,
+                format: CopyToFormatOptions::Bson(CopyToFormatOptionsBson {}),
             }
             .into_extension(),
         );

diff --git a/crates/sqlexec/src/planner/physical_plan/copy_to.rs b/crates/sqlexec/src/planner/physical_plan/copy_to.rs
@@ -244,7 +244,7 @@ fn get_sink_for_obj(
                 array: json_opts.array,
             },
         )),
-        CopyToFormatOptions::Bson => Box::new(BsonSink::from_obj_store(store, path)),
+        CopyToFormatOptions::Bson(_) => Box::new(BsonSink::from_obj_store(store, path)),
     };
     Ok(sink)
 }
diff --git a/crates/sqlexec/src/planner/session_planner.rs b/crates/sqlexec/src/planner/session_planner.rs
@@ -56,6 +56,7 @@ use protogen::metastore::types::options::{
     CopyToDestinationOptionsLocal,
     CopyToDestinationOptionsS3,
     CopyToFormatOptions,
+    CopyToFormatOptionsBson,
     CopyToFormatOptionsCsv,
     CopyToFormatOptionsJson,
     CopyToFormatOptionsLance,
@@ -1938,7 +1939,9 @@ impl<'a> SessionPlanner<'a> {
                 let array = m.remove_optional::<bool>("array")?.unwrap_or(false);
                 CopyToFormatOptions::Json(CopyToFormatOptionsJson { array })
             }
-            Some(CopyToFormatOptions::BSON) => CopyToFormatOptions::Bson {},
+            Some(CopyToFormatOptions::BSON) => {
+                CopyToFormatOptions::Bson(CopyToFormatOptionsBson {})
+            }
             Some(CopyToFormatOptions::LANCE) => {
                 CopyToFormatOptions::Lance(CopyToFormatOptionsLance {
                     max_rows_per_file: m.remove_optional("max_rows_per_file")?,