Merge pull request #65 from Embucket/snowflake

osipovartem · web-flow · commit db27e6c9ee0e · 2024-12-24T14:32:09.000+03:00
Convert timestamp
diff --git a/crates/control_plane/src/service.rs b/crates/control_plane/src/service.rs
@@ -5,6 +5,8 @@ use crate::repository::{StorageProfileRepository, WarehouseRepository};
 use crate::sql::functions::common::convert_record_batches;
 use crate::sql::sql::SqlExecutor;
 use arrow::record_batch::RecordBatch;
+use arrow_json::writer::JsonArray;
+use arrow_json::WriterBuilder;
 use async_trait::async_trait;
 use base64::Engine;
 use bytes::Bytes;
@@ -240,10 +242,11 @@ impl ControlService for ControlServiceImpl {
         let (records, _) = self
             .query(warehouse_id, database_name, _table_name, query)
             .await?;
-        // println!("{records:?}");
 
         let buf = Vec::new();
-        let mut writer = arrow_json::ArrayWriter::new(buf);
+        let write_builder = WriterBuilder::new().with_explicit_nulls(true);
+        let mut writer = write_builder.build::<_, JsonArray>(buf);
+
         let record_refs: Vec<&RecordBatch> = records.iter().collect();
         writer.write_batches(&record_refs).unwrap();
         writer.finish().unwrap();
@@ -301,7 +304,9 @@ impl ControlService for ControlServiceImpl {
 
         // We use json format since there is a bug between arrow and nanoarrow
         let buf = Vec::new();
-        let mut writer = arrow_json::ArrayWriter::new(buf);
+        let write_builder = WriterBuilder::new().with_explicit_nulls(true);
+        let mut writer = write_builder.build::<_, JsonArray>(buf);
+
         let record_refs: Vec<&RecordBatch> = records.iter().collect();
         writer.write_batches(&record_refs).unwrap();
         writer.finish().unwrap();
diff --git a/crates/control_plane/src/sql/functions/common.rs b/crates/control_plane/src/sql/functions/common.rs
@@ -1,6 +1,9 @@
 use crate::models::ColumnInfo;
-use arrow::array::{Array, UnionArray};
-use arrow::datatypes::{Field, Schema};
+use arrow::array::{
+    Array, Int32Array, Int64Array, StructArray, TimestampMicrosecondArray,
+    TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UnionArray,
+};
+use arrow::datatypes::{Field, Schema, TimeUnit};
 use arrow::record_batch::RecordBatch;
 use datafusion::arrow::array::ArrayRef;
 use datafusion::arrow::datatypes::DataType;
@@ -30,26 +33,108 @@ pub fn convert_record_batches(
         for (i, column) in batch.columns().iter().enumerate() {
             let metadata = column_infos[i].to_metadata();
             let field = batch.schema().field(i).clone();
-            let converted_column = if let Some(union_array) =
-                column.as_any().downcast_ref::<UnionArray>()
-            {
-                if let Some((data_type, array)) = first_non_empty_type(union_array) {
-                    fields.push(Field::new(field.name(), data_type, field.is_nullable()).with_metadata(metadata));
-                    array
-                } else {
+            let converted_column = match field.data_type() {
+                DataType::Union(..) => {
+                    if let Some(union_array) = column.as_any().downcast_ref::<UnionArray>() {
+                        if let Some((data_type, array)) = first_non_empty_type(union_array) {
+                            fields.push(
+                                Field::new(field.name(), data_type, field.is_nullable())
+                                    .with_metadata(metadata),
+                            );
+                            array
+                        } else {
+                            fields.push(field.clone().with_metadata(metadata));
+                            Arc::clone(column)
+                        }
+                    } else {
+                        fields.push(field.clone().with_metadata(metadata));
+                        Arc::clone(column)
+                    }
+                }
+                DataType::Timestamp(unit, _) => {
+                    let converted_column = convert_timestamp_to_struct(column, unit);
+                    fields.push(
+                        Field::new(
+                            field.name(),
+                            converted_column.data_type().clone(),
+                            field.is_nullable(),
+                        )
+                            .with_metadata(metadata),
+                    );
+                    Arc::clone(&converted_column)
+                }
+                _ => {
                     fields.push(field.clone().with_metadata(metadata));
                     Arc::clone(column)
                 }
-            } else {
-                fields.push(field.clone().with_metadata(metadata));
-                Arc::clone(column)
             };
             columns.push(converted_column);
         }
         let new_schema = Arc::new(Schema::new(fields));
+        println!("new schema: {:?}", new_schema);
+        println!("columns: {:?}", columns);
         let converted_batch = RecordBatch::try_new(new_schema, columns)?;
         converted_batches.push(converted_batch);
     }
 
     Ok((converted_batches.clone(), column_infos))
 }
+
+fn convert_timestamp_to_struct(column: &ArrayRef, unit: &TimeUnit) -> ArrayRef {
+    let (epoch, fraction) = match unit {
+        TimeUnit::Second => {
+            let array = column
+                .as_any()
+                .downcast_ref::<TimestampSecondArray>()
+                .unwrap();
+            let epoch: Int64Array = array.clone().unary(|x| x);
+            let fraction: Int32Array = Int32Array::from(vec![0; column.len()]);
+            (epoch, fraction)
+        }
+        TimeUnit::Millisecond => {
+            let array = column
+                .as_any()
+                .downcast_ref::<TimestampMillisecondArray>()
+                .unwrap();
+            let epoch: Int64Array = array.clone().unary(|x| x / 1_000);
+            let fraction: Int32Array = array.clone().unary(|x| (x % 1_000 * 1_000_000) as i32);
+            (epoch, fraction)
+        }
+        TimeUnit::Microsecond => {
+            let array = column
+                .as_any()
+                .downcast_ref::<TimestampMicrosecondArray>()
+                .unwrap();
+            let epoch: Int64Array = array.clone().unary(|x| x / 1_000_000);
+            let fraction: Int32Array = array.clone().unary(|x| (x % 1_000_000 * 1_000) as i32);
+            (epoch, fraction)
+        }
+        TimeUnit::Nanosecond => {
+            let array = column
+                .as_any()
+                .downcast_ref::<TimestampNanosecondArray>()
+                .unwrap();
+            let epoch: Int64Array = array.clone().unary(|x| x / 1_000_000_000);
+            let fraction: Int32Array = array.clone().unary(|x| (x % 1_000_000_000) as i32);
+            (epoch, fraction)
+        }
+    };
+
+    // let timezone = Int32Array::from(vec![1440; column.len()]); // Assuming UTC timezone
+    let struct_array = StructArray::from(vec![
+        (
+            Arc::new(Field::new("epoch", DataType::Int64, false)),
+            Arc::new(epoch) as ArrayRef,
+        ),
+        (
+            Arc::new(Field::new("fraction", DataType::Int32, false)),
+            Arc::new(fraction) as ArrayRef,
+        ),
+        // (
+        //     Arc::new(Field::new("timezone", DataType::Int32, false)),
+        //     Arc::new(timezone) as ArrayRef,
+        // ),
+    ]);
+
+    Arc::new(struct_array) as ArrayRef
+}
diff --git a/crates/control_plane/src/sql/planner.rs b/crates/control_plane/src/sql/planner.rs
@@ -70,11 +70,7 @@ where
         match self.handle_custom_statement(statement.clone()) {
             Ok(plan) => return Ok(plan),
             Err(e) => {
-                eprintln!(
-                    "Custom statement parsing skipped: {} with err {}",
-                    statement.to_string(),
-                    e
-                );
+                eprintln!("Custom statement parsing skipped: {statement} ");
             }
         }
 
@@ -105,7 +101,8 @@ where
                 all_constraints.extend(inline_constraints);
                 // Build column default values
                 let column_defaults = self.build_column_defaults(&columns, planner_context)?;
-
+                println!("column_defaults: {:?}", column_defaults);
+                println!("statement 11: {:?}", statement);
                 let has_columns = !columns.is_empty();
                 let schema = self.build_schema(columns.clone())?.to_dfschema_ref()?;
                 if has_columns {
@@ -240,7 +237,7 @@ where
             SQLDataType::Char(_)
             | SQLDataType::Text
             | SQLDataType::String(_) => Ok(DataType::Utf8),
-            SQLDataType::Timestamp(None, tz_info) => {
+            SQLDataType::Timestamp(precision, tz_info) => {
                 let tz = if matches!(tz_info, TimezoneInfo::Tz)
                     || matches!(tz_info, TimezoneInfo::WithTimeZone)
                 {
@@ -252,7 +249,14 @@ where
                     // Timestamp Without Time zone
                     None
                 };
-                Ok(DataType::Timestamp(TimeUnit::Nanosecond, tz.map(Into::into)))
+                let precision = match precision {
+                    Some(0) => TimeUnit::Second,
+                    Some(3) => TimeUnit::Millisecond,
+                    Some(6) => TimeUnit::Microsecond,
+                    None | Some(9) => TimeUnit::Nanosecond,
+                    _ => unreachable!(),
+                };
+                Ok(DataType::Timestamp(precision, tz.map(Into::into)))
             }
             SQLDataType::Date => Ok(DataType::Date32),
             SQLDataType::Time(None, tz_info) => {
diff --git a/crates/control_plane/src/sql/sql.rs b/crates/control_plane/src/sql/sql.rs
@@ -44,7 +44,6 @@ impl SqlExecutor {
         let dialect = state.config().options().sql_parser.dialect.as_str();
         // Update query to use custom JSON functions
         let query = self.preprocess_query(query);
-        println!("Query: {}", query);
         let mut statement = state.sql_to_statement(&query, dialect)?;
         // statement = self.update_statement_references(statement, warehouse_name);
         // query = statement.to_string();
@@ -168,19 +167,21 @@ impl SqlExecutor {
                 .unwrap();
 
             // we don't need physical table for transient tables
-            if !transient {
-                // Copy data from InMemory table to created table
-                let insert_query =
-                    format!("INSERT INTO {new_table_full_name} SELECT * FROM {new_table_name}");
-                let result = self.execute_with_custom_plan(&insert_query, warehouse_name).await?;
-                // self.ctx.sql(&insert_query).await?.collect().await?;
+            // if !transient {
+            // Copy data from InMemory table to created table
+            let insert_query =
+                format!("INSERT INTO {new_table_full_name} SELECT * FROM {new_table_name}");
+            let result = self
+                .execute_with_custom_plan(&insert_query, warehouse_name)
+                .await?;
+            // self.ctx.sql(&insert_query).await?.collect().await?;
 
-                // Drop InMemory table
-                let drop_query = format!("DROP TABLE {new_table_name}");
-                self.ctx.sql(&drop_query).await?.collect().await?;
-                return Ok(result);
-            }
-            Ok(created_entity_response())
+            // Drop InMemory table
+            let drop_query = format!("DROP TABLE {new_table_name}");
+            self.ctx.sql(&drop_query).await?.collect().await?;
+            return Ok(result);
+            // }
+            // Ok(created_entity_response())
         } else {
             Err(datafusion::error::DataFusionError::NotImplemented(
                 "Only CREATE TABLE statements are supported".to_string(),
@@ -225,19 +226,23 @@ impl SqlExecutor {
         Ok(created_entity_response())
     }
 
-    pub async fn get_custom_logical_plan(&self, query: &String, warehouse_name: &str) -> Result<LogicalPlan> {
+    pub async fn get_custom_logical_plan(
+        &self,
+        query: &String,
+        warehouse_name: &str,
+    ) -> Result<LogicalPlan> {
         let state = self.ctx.state();
         let dialect = state.config().options().sql_parser.dialect.as_str();
         let mut statement = state.sql_to_statement(query, dialect)?;
+        println!("raw query: {:?}", statement.to_string());
         statement = self.update_statement_references(statement, warehouse_name);
-        println!("modified query: {:?}", statement);
+        println!("modified query: {:?}", statement.to_string());
 
         if let DFStatement::Statement(s) = statement.clone() {
             let mut ctx_provider = CustomContextProvider {
                 state: &state,
                 tables: HashMap::new(),
             };
-
             let references = state.resolve_table_references(&statement)?;
             println!("References: {:?}", references);
             for reference in references {
@@ -311,9 +316,15 @@ impl SqlExecutor {
             })
     }
 
-    pub async fn execute_with_custom_plan(&self, query: &String, warehouse_name: &str) -> Result<Vec<RecordBatch>> {
+    pub async fn execute_with_custom_plan(
+        &self,
+        query: &String,
+        warehouse_name: &str,
+    ) -> Result<Vec<RecordBatch>> {
         let plan = self.get_custom_logical_plan(query, warehouse_name).await?;
-        self.ctx.execute_logical_plan(plan).await?.collect().await
+        let res = self.ctx.execute_logical_plan(plan).await?.collect().await;
+        println!("Result: {:?}", res);
+        res
     }
 
     pub fn update_statement_references(
@@ -355,7 +366,6 @@ impl SqlExecutor {
                 }
                 Statement::Query(mut query) => {
                     self.update_tables_in_query(query.as_mut(), warehouse_name);
-                    println!("Query: {:?}", query);
                     DFStatement::Statement(Box::new(Statement::Query(query)))
                 }
                 _ => statement,
diff --git a/crates/nexus/src/http/dbt/handlers.rs b/crates/nexus/src/http/dbt/handlers.rs
@@ -13,6 +13,8 @@ use regex::Regex;
 use serde_json::json;
 use std::io::Read;
 use std::result::Result;
+use tokio::fs::OpenOptions;
+use tokio::io::AsyncWriteExt;
 use uuid::Uuid;
 
 pub async fn login(
@@ -76,6 +78,10 @@ pub async fn query(
     let (params, sql_query) = body_json.get_sql_text();
     println!("Query raw: {:?}", body_json.sql_text);
 
+    // if let Err(e) = log_query(body_json.sql_text.as_str()).await {
+    //     eprintln!("Failed to log query: {}", e);
+    // }
+
     let token = match extract_token(&headers) {
         Some(token) => token,
         None => {
@@ -148,3 +154,14 @@ pub fn extract_token(headers: &HeaderMap) -> Option<String> {
         })
     })
 }
+
+async fn log_query(query: &str) -> Result<(), std::io::Error> {
+    let mut file = OpenOptions::new()
+        .create(true)
+        .append(true)
+        .open("queries.log")
+        .await.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+    file.write_all(query.as_bytes()).await?;
+    file.write_all(b"\n").await?;
+    Ok(())
+}
