Merge branch 'main' into check-view-dep

Author: mergify[bot]

Cargo.lock

@@ -106,6 +106,12 @@ pub enum StageFileCompression {
     None,
 }
 
+impl StageFileFormatType {
+    pub fn has_inner_schema(&self) -> bool {
+        matches!(self, StageFileFormatType::Parquet)
+    }
+}
+
 impl Default for StageFileCompression {
     fn default() -> Self {
         Self::None

src/meta/app/src/principal/user_stage.rs

@@ -98,8 +98,7 @@ storages-common-table-meta = { path = "../storages/common/table-meta" }
 # Crates.io dependencies
 aho-corasick = { version = "0.7.20" }
 async-channel = "1.7.1"
-# Wait for https://github.com/64bit/async-openai/pull/60
-async-openai = { version = "0.9.5", git = "https://github.com/Xuanwo/async-openai", rev = "22fd9ca0c9c86a2c57462c5fd32aaff407d6b000" }
+async-openai = "0.10.0"
 async-stream = "0.3.3"
 async-trait = { version = "0.1.57", package = "async-trait-fn" }
 base64 = "0.21.0"
@@ -135,7 +134,7 @@ scopeguard = "1.1.0"
 serde = { workspace = true }
 serde_json = { workspace = true }
 serde_urlencoded = "0.7.1"
-socket2 = "0.5.1"
+socket2 = "0.4.7"
 strength_reduce = "0.2.4"
 tempfile = { version = "3.4.0", optional = true }
 time = "0.3.14"

src/query/service/Cargo.toml

@@ -24,6 +24,7 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::infer_table_schema;
 use common_expression::DataField;
+use common_expression::DataSchema;
 use common_expression::DataSchemaRef;
 use common_expression::DataSchemaRefExt;
 use common_meta_app::principal::StageInfo;
@@ -41,6 +42,7 @@ use tracing::info;
 use crate::interpreters::common::append2table;
 use crate::interpreters::Interpreter;
 use crate::interpreters::SelectInterpreter;
+use crate::pipelines::processors::transforms::TransformRuntimeCastSchema;
 use crate::pipelines::processors::TransformCastSchema;
 use crate::pipelines::processors::TransformLimit;
 use crate::pipelines::PipelineBuildResult;
@@ -334,6 +336,26 @@ impl CopyInterpreter {
             )?;
         }
 
+        if stage_table_info
+            .stage_info
+            .file_format_options
+            .format
+            .has_inner_schema()
+        {
+            let dst_schema: Arc<DataSchema> = Arc::new(to_table.schema().into());
+            let func_ctx = self.ctx.get_function_context()?;
+            build_res.main_pipeline.add_transform(
+                |transform_input_port, transform_output_port| {
+                    TransformRuntimeCastSchema::try_create(
+                        transform_input_port,
+                        transform_output_port,
+                        dst_schema.clone(),
+                        func_ctx,
+                    )
+                },
+            )?;
+        }
+
         // Build append data pipeline.
         to_table.append_data(
             ctx.clone(),

src/query/service/src/interpreters/interpreter_copy.rs

@@ -378,20 +378,23 @@ impl Interpreter for InsertInterpreter {
                     .format
                     .exec_stream(input_context.clone(), &mut build_res.main_pipeline)?;
 
-                if Ok(StageFileFormatType::Parquet) == StageFileFormatType::from_str(format) {
-                    let dest_schema = plan.schema();
-                    let func_ctx = self.ctx.get_function_context()?;
-
-                    build_res.main_pipeline.add_transform(
-                        |transform_input_port, transform_output_port| {
-                            TransformRuntimeCastSchema::try_create(
-                                transform_input_port,
-                                transform_output_port,
-                                dest_schema.clone(),
-                                func_ctx,
-                            )
-                        },
-                    )?;
+                match StageFileFormatType::from_str(format) {
+                    Ok(f) if f.has_inner_schema() => {
+                        let dest_schema = plan.schema();
+                        let func_ctx = self.ctx.get_function_context()?;
+
+                        build_res.main_pipeline.add_transform(
+                            |transform_input_port, transform_output_port| {
+                                TransformRuntimeCastSchema::try_create(
+                                    transform_input_port,
+                                    transform_output_port,
+                                    dest_schema.clone(),
+                                    func_ctx,
+                                )
+                            },
+                        )?;
+                    }
+                    _ => {}
                 }
             }
             InsertInputSource::StreamingWithFileFormat(format_options, _, input_context) => {
@@ -400,7 +403,7 @@ impl Interpreter for InsertInterpreter {
                     .format
                     .exec_stream(input_context.clone(), &mut build_res.main_pipeline)?;
 
-                if StageFileFormatType::Parquet == format_options.format {
+                if format_options.format.has_inner_schema() {
                     let dest_schema = plan.schema();
                     let func_ctx = self.ctx.get_function_context()?;
 

src/query/service/src/interpreters/interpreter_insert.rs

@@ -28,6 +28,7 @@ use futures::future::AbortRegistration;
 use futures::future::Abortable;
 use futures::StreamExt;
 use opensrv_mysql::*;
+use socket2::SockRef;
 use socket2::TcpKeepalive;
 use tokio_stream::wrappers::TcpListenerStream;
 use tracing::error;
@@ -92,7 +93,7 @@ impl MySQLHandler {
         sessions: Arc<SessionManager>,
         executor: Arc<Runtime>,
         socket: TcpStream,
-        _keepalive: TcpKeepalive,
+        keepalive: TcpKeepalive,
     ) {
         executor.spawn(async move {
             match sessions.create_session(SessionType::MySQL).await {
@@ -103,11 +104,10 @@ impl MySQLHandler {
                 Ok(session) => {
                     info!("MySQL connection coming: {:?}", socket.peer_addr());
 
-                    // FIXME: tokio TcpStream doesn't implement `AsFd` anymore, this call should be refactored.
-                    // if let Err(e) = SockRef::from(&socket).set_tcp_keepalive(&keepalive)
-                    // {
-                    //     warn!("failed to set socket option keepalive {}", e);
-                    // }
+                    // TcpStream must implement AsFd for socket2 0.5, wait https://github.com/tokio-rs/tokio/pull/5514
+                    if let Err(e) = SockRef::from(&socket).set_tcp_keepalive(&keepalive) {
+                        warn!("failed to set socket option keepalive {}", e);
+                    }
 
                     if let Err(error) = MySQLConnection::run_on_stream(session, socket) {
                         error!("Unexpected error occurred during query: {:?}", error);

src/query/service/src/servers/mysql/mysql_handler.rs

@@ -14,10 +14,7 @@
 
 use std::any::Any;
 use std::sync::Arc;
-use std::time::Duration;
 
-use async_openai::types::CreateCompletionRequestArgs;
-use async_openai::Client;
 use chrono::NaiveDateTime;
 use chrono::TimeZone;
 use chrono::Utc;
@@ -50,6 +47,9 @@ use common_storages_fuse::TableContext;
 use common_storages_view::view_table::VIEW_ENGINE;
 use tracing::info;
 
+use crate::table_functions::openai::AIModel;
+use crate::table_functions::openai::OpenAI;
+
 pub struct GPT2SQLTable {
     prompt: String,
     api_key: String,
@@ -233,37 +233,16 @@ impl AsyncSource for GPT2SQLSource {
         template.push("#".to_string());
         template.push("SELECT".to_string());
 
-        let model = "code-davinci-002";
-        let timeout = Duration::from_secs(30);
         let prompt = template.join("");
-        let api_key = self.api_key.clone();
         info!("openai request prompt: {}", prompt);
 
-        // Client
-        let http_client = reqwest::ClientBuilder::new()
-            .user_agent("databend")
-            .timeout(timeout)
-            .build()
-            .map_err(|e| ErrorCode::Internal(format!("openai http error: {:?}", e)))?;
-        let client = Client::new()
-            .with_api_key(api_key)
-            .with_http_client(http_client);
-
-        // Request
-        let request = CreateCompletionRequestArgs::default()
-            .model(model)
-            .prompt(prompt)
-            .temperature(0.0)
-            .max_tokens(150_u16)
-            .top_p(1.0)
-            .frequency_penalty(0.0)
-            .presence_penalty(0.0)
-            .stop(["#", ";"])
-            .build()
-            .map_err(|e| ErrorCode::Internal(format!("openai request error: {:?}", e)))?;
-
         // Response.
-        let response = client
+        let api_key = self.api_key.clone();
+        let openai = OpenAI::create(api_key, AIModel::CodeDavinci002);
+        let request = openai.completion_request(prompt)?;
+
+        let response = openai
+            .client()?
             .completions()
             .create(request)
             .await

src/query/service/src/table_functions/openai/gpt_to_sql.rs

@@ -13,5 +13,9 @@
 //  limitations under the License.
 
 mod gpt_to_sql;
+#[allow(clippy::module_inception)]
+mod openai;
 
 pub use gpt_to_sql::GPT2SQLTable;
+pub use openai::AIModel;
+pub use openai::OpenAI;

src/query/service/src/table_functions/openai/mod.rs

@@ -0,0 +1,73 @@
+//  Copyright 2023 Datafuse Labs.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+use std::time::Duration;
+
+use async_openai::types::CreateCompletionRequest;
+use async_openai::types::CreateCompletionRequestArgs;
+use async_openai::Client;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+pub enum AIModel {
+    CodeDavinci002,
+}
+
+// https://platform.openai.com/examples
+impl ToString for AIModel {
+    fn to_string(&self) -> String {
+        match self {
+            AIModel::CodeDavinci002 => "code-davinci-002".to_string(),
+        }
+    }
+}
+
+pub struct OpenAI {
+    api_key: String,
+    model: AIModel,
+}
+
+impl OpenAI {
+    pub fn create(api_key: String, model: AIModel) -> Self {
+        OpenAI { api_key, model }
+    }
+
+    pub fn client(&self) -> Result<Client> {
+        let timeout = Duration::from_secs(30);
+        // Client
+        let http_client = reqwest::ClientBuilder::new()
+            .user_agent("databend")
+            .timeout(timeout)
+            .build()
+            .map_err(|e| ErrorCode::Internal(format!("openai http error: {:?}", e)))?;
+
+        Ok(Client::new()
+            .with_api_key(self.api_key.clone())
+            .with_http_client(http_client))
+    }
+
+    pub fn completion_request(&self, prompt: String) -> Result<CreateCompletionRequest> {
+        CreateCompletionRequestArgs::default()
+            .model(self.model.to_string())
+            .prompt(prompt)
+            .temperature(0.0)
+            .max_tokens(150_u16)
+            .top_p(1.0)
+            .frequency_penalty(0.0)
+            .presence_penalty(0.0)
+            .stop(["#", ";"])
+            .build()
+            .map_err(|e| ErrorCode::Internal(format!("openai completion request error: {:?}", e)))
+    }
+}

src/query/service/src/table_functions/openai/openai.rs

@@ -6,3 +6,4 @@
 199	2020.0	769
 199	2020.0	769
 199	2020.0	769
+199	2020.0000	769.00