add data source field

Canner · Dec 12, 2024 · d9dca7f · d9dca7f
1 parent c989b8c
commit d9dca7f
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 3 deletions.
diff --git a/wren-core/core/src/mdl/builder.rs b/wren-core/core/src/mdl/builder.rs
@@ -1,7 +1,8 @@
 #![allow(dead_code)]
 
 use crate::mdl::manifest::{
-    Column, JoinType, Manifest, Metric, Model, Relationship, TimeGrain, TimeUnit, View,
+    Column, DataSource, JoinType, Manifest, Metric, Model, Relationship, TimeGrain,
+    TimeUnit, View,
 };
 use std::sync::Arc;
 
@@ -26,6 +27,7 @@ impl ManifestBuilder {
                 relationships: vec![],
                 metrics: vec![],
                 views: vec![],
+                data_source: None,
             },
         }
     }
@@ -60,6 +62,11 @@ impl ManifestBuilder {
         self
     }
 
+    pub fn data_source(mut self, data_source: DataSource) -> Self {
+        self.manifest.data_source = Some(data_source);
+        self
+    }
+
     pub fn build(self) -> Manifest {
         self.manifest
     }
@@ -328,7 +335,7 @@ mod test {
         TimeGrainBuilder, ViewBuilder,
     };
     use crate::mdl::manifest::{
-        Column, JoinType, Metric, Model, Relationship, TimeUnit, View,
+        Column, DataSource, JoinType, Metric, Model, Relationship, TimeUnit, View,
     };
     use std::sync::Arc;
 
@@ -546,6 +553,7 @@ mod test {
             .relationship(relationship)
             .metric(metric)
             .view(view)
+            .data_source(DataSource::Datafusion)
             .build();
 
         let json_str = serde_json::to_string(&expected).unwrap();

diff --git a/wren-core/core/src/mdl/manifest.rs b/wren-core/core/src/mdl/manifest.rs
@@ -18,6 +18,50 @@ pub struct Manifest {
     pub metrics: Vec<Arc<Metric>>,
     #[serde(default)]
     pub views: Vec<Arc<View>>,
+    #[serde(default)]
+    pub data_source: Option<DataSource>,
+}
+
+#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Hash, Clone, Copy)]
+#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
+pub enum DataSource {
+    #[serde(alias = "bigquery")]
+    BigQuery,
+    #[serde(alias = "clickhouse")]
+    Clickhouse,
+    #[serde(alias = "canner")]
+    Canner,
+    #[serde(alias = "trino")]
+    Trino,
+    #[serde(alias = "mssql")]
+    MSSQL,
+    #[serde(alias = "mysql")]
+    MySQL,
+    #[serde(alias = "postgres")]
+    Postgres,
+    #[serde(alias = "snowflake")]
+    Snowflake,
+    #[serde(alias = "datafusion")]
+    Datafusion,
+    #[serde(alias = "duckdb")]
+    DuckDB,
+}
+
+impl Display for DataSource {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DataSource::BigQuery => write!(f, "BIGQUERY"),
+            DataSource::Clickhouse => write!(f, "CLICKHOUSE"),
+            DataSource::Canner => write!(f, "CANNER"),
+            DataSource::Trino => write!(f, "TRINO"),
+            DataSource::MSSQL => write!(f, "MSSQL"),
+            DataSource::MySQL => write!(f, "MYSQL"),
+            DataSource::Postgres => write!(f, "POSTGRES"),
+            DataSource::Snowflake => write!(f, "SNOWFLAKE"),
+            DataSource::Datafusion => write!(f, "DATAFUSION"),
+            DataSource::DuckDB => write!(f, "DUCKDB"),
+        }
+    }
 }
 
 #[serde_as]

diff --git a/wren-core/core/tests/data/mdl.json b/wren-core/core/tests/data/mdl.json
@@ -138,5 +138,6 @@
       "name": "customer_view",
       "statement": "select * from test.test.customer"
     }
-  ]
+  ],
+  "dataSource": "datafusion"
 }