feat: support get/put stmt (datafuselabs#198)

youngsofun · Aug 29, 2023 · 4fa9b95 · 4fa9b95
1 parent 5c39cc9
commit 4fa9b95
Show file tree

Hide file tree

Showing 20 changed files with 273 additions and 16 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,7 +11,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "Apache-2.0"
 authors = ["Databend Authors <opensource@datafuselabs.com>"]
@@ -20,7 +20,7 @@ keywords = ["databend", "database"]
 repository = "https://github.com/datafuselabs/bendsql"
 
 [workspace.dependencies]
-databend-client = { path = "core", version = "0.6.2" }
-databend-driver = { path = "driver", version = "0.6.2" }
-databend-driver-macros = { path = "macros", version = "0.6.2" }
-databend-sql = { path = "sql", version = "0.6.2" }
+databend-client = { path = "core", version = "0.6.3" }
+databend-driver = { path = "driver", version = "0.6.3" }
+databend-driver-macros = { path = "macros", version = "0.6.3" }
+databend-sql = { path = "sql", version = "0.6.3" }
diff --git a/bindings/nodejs/npm/darwin-arm64/package.json b/bindings/nodejs/npm/darwin-arm64/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@databend-driver/lib-darwin-arm64",
   "repository": "https://github.com/datafuselabs/bendsql.git",
-  "version": "0.6.2",
+  "version": "0.6.3",
   "os": [
     "darwin"
   ],

diff --git a/bindings/nodejs/npm/darwin-x64/package.json b/bindings/nodejs/npm/darwin-x64/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@databend-driver/lib-darwin-x64",
   "repository": "https://github.com/datafuselabs/bendsql.git",
-  "version": "0.6.2",
+  "version": "0.6.3",
   "os": [
     "darwin"
   ],

diff --git a/bindings/nodejs/npm/linux-x64-gnu/package.json b/bindings/nodejs/npm/linux-x64-gnu/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@databend-driver/lib-linux-x64-gnu",
   "repository": "https://github.com/datafuselabs/bendsql.git",
-  "version": "0.6.2",
+  "version": "0.6.3",
   "os": [
     "linux"
   ],

diff --git a/bindings/nodejs/npm/win32-x64-msvc/package.json b/bindings/nodejs/npm/win32-x64-msvc/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@databend-driver/lib-win32-x64-msvc",
   "repository": "https://github.com/datafuselabs/bendsql.git",
-  "version": "0.6.2",
+  "version": "0.6.3",
   "os": [
     "win32"
   ],

diff --git a/bindings/nodejs/package.json b/bindings/nodejs/package.json
@@ -1,7 +1,7 @@
 {
   "name": "databend-driver",
   "author": "Databend Authors <opensource@datafuselabs.com>",
-  "version": "0.6.2",
+  "version": "0.6.3",
   "license": "Apache-2.0",
   "main": "index.js",
   "types": "index.d.ts",

diff --git a/cli/README.md b/cli/README.md
@@ -90,6 +90,13 @@ Bye
 2
 ```
 
+### Put local files into stage
+
+```
+create stage s_temp;
+put fs:///tmp/a*.txt @s_temp/abc;
+```
+
 ## Features
 
 - basic keywords highlight

diff --git a/cli/src/ast/mod.rs b/cli/src/ast/mod.rs
@@ -17,7 +17,13 @@ mod tokenizer;
 use sqlformat::{Indent, QueryParams};
 pub use tokenizer::*;
 
+use crate::session::QueryKind;
+
 pub fn format_query(query: &str) -> String {
+    let kind = QueryKind::from(query);
+    if kind == QueryKind::Put || kind == QueryKind::Get {
+        return query.to_owned();
+    }
     let options = sqlformat::FormatOptions {
         indent: Indent::Spaces(2),
         uppercase: true,

diff --git a/cli/src/ast/tokenizer.rs b/cli/src/ast/tokenizer.rs
@@ -476,6 +476,8 @@ pub enum TokenKind {
     TABLE_FUNCTIONS,
     #[token("FUSE", ignore(ascii_case))]
     FUSE,
+    #[token("GET", ignore(ascii_case))]
+    GET,
     #[token("GLOBAL", ignore(ascii_case))]
     GLOBAL,
     #[token("GRAPH", ignore(ascii_case))]
@@ -634,6 +636,8 @@ pub enum TokenKind {
     PROCESSLIST,
     #[token("PURGE", ignore(ascii_case))]
     PURGE,
+    #[token("PUT", ignore(ascii_case))]
+    PUT,
     #[token("QUARTER", ignore(ascii_case))]
     QUARTER,
     #[token("QUERY", ignore(ascii_case))]

diff --git a/cli/src/session.rs b/cli/src/session.rs
@@ -205,7 +205,12 @@ impl Session {
             return vec![];
         }
 
-        if self.query.is_empty() && (line.starts_with('.') || line == "exit" || line == "quit") {
+        if self.query.is_empty()
+            && (line.starts_with('.')
+                || line == "exit"
+                || line == "quit"
+                || line.to_uppercase().starts_with("PUT"))
+        {
             return vec![line.to_owned()];
         }
 
@@ -316,13 +321,33 @@ impl Session {
                 }
                 Ok(false)
             }
-            _ => {
+            other => {
                 let replace_newline = !if self.settings.replace_newline {
                     false
                 } else {
                     replace_newline_in_box_display(query)
                 };
-                let (schema, data) = self.conn.query_iter_ext(query).await?;
+
+                let (schema, data) = match other.0 {
+                    QueryKind::Put => {
+                        let args: Vec<String> = get_put_get_args(query);
+                        if args.len() != 3 {
+                            eprintln!("Put args are invalid, must be 2 argruments");
+                            return Ok(false);
+                        }
+                        self.conn.put_files(&args[1], &args[2]).await?
+                    }
+                    QueryKind::Get => {
+                        let args: Vec<String> = get_put_get_args(query);
+                        if args.len() != 3 {
+                            eprintln!("Put args are invalid, must be 2 argruments");
+                            return Ok(false);
+                        }
+                        self.conn.get_files(&args[1], &args[2]).await?
+                    }
+                    _ => self.conn.query_iter_ext(query).await?,
+                };
+
                 let mut displayer = FormatDisplay::new(
                     &self.settings,
                     query,
@@ -413,6 +438,8 @@ pub enum QueryKind {
     Query,
     Update,
     Explain,
+    Put,
+    Get,
 }
 
 impl From<&str> for QueryKind {
@@ -421,6 +448,8 @@ impl From<&str> for QueryKind {
         match tz.next() {
             Some(Ok(t)) => match t.kind {
                 TokenKind::EXPLAIN => QueryKind::Explain,
+                TokenKind::PUT => QueryKind::Put,
+                TokenKind::GET => QueryKind::Get,
                 TokenKind::ALTER
                 | TokenKind::DELETE
                 | TokenKind::UPDATE
@@ -436,6 +465,13 @@ impl From<&str> for QueryKind {
     }
 }
 
+fn get_put_get_args(query: &str) -> Vec<String> {
+    query
+        .split_ascii_whitespace()
+        .map(|x| x.to_owned())
+        .collect()
+}
+
 fn replace_newline_in_box_display(query: &str) -> bool {
     let mut tz = Tokenizer::new(query);
     match tz.next() {

diff --git a/cli/tests/http/04-put.result b/cli/tests/http/04-put.result
@@ -0,0 +1,4 @@
+/tmp/a1.txt	SUCCESS
+/tmp/a2.txt	SUCCESS
+/tmp/edf/a1.txt	SUCCESS
+/tmp/edf/a2.txt	SUCCESS
diff --git a/cli/tests/http/04-put.sh b/cli/tests/http/04-put.sh
@@ -0,0 +1,9 @@
+#!/bin/bash
+
+echo "CREATE STAGE ss_temp" | ${BENDSQL}
+
+echo "ABCD" > /tmp/a1.txt
+echo "ABCD" > /tmp/a2.txt
+
+echo 'put fs:///tmp/a*.txt @ss_temp/abc' | ${BENDSQL}
+echo 'get @ss_temp/abc fs:///tmp/edf' | ${BENDSQL}
diff --git a/core/src/client.rs b/core/src/client.rs
@@ -80,7 +80,7 @@ impl TryFrom<&str> for StageLocation {
 
 #[derive(Clone)]
 pub struct APIClient {
-    cli: HttpClient,
+    pub cli: HttpClient,
     endpoint: Url,
     pub host: String,
     pub port: u16,
@@ -437,7 +437,7 @@ impl APIClient {
         }
     }
 
-    async fn upload_to_stage_with_stream(
+    pub async fn upload_to_stage_with_stream(
         &self,
         stage_location: &str,
         data: impl AsyncRead + Send + Sync + 'static,

diff --git a/driver/Cargo.toml b/driver/Cargo.toml
@@ -33,6 +33,8 @@ databend-sql = { workspace = true }
 async-trait = "0.1.68"
 chrono = { version = "0.4.26", default-features = false, features = ["clock"] }
 dyn-clone = "1.0.11"
+futures = "0.3.28"
+glob = "0.3.1"
 http = "0.2.9"
 percent-encoding = "2.3.0"
 serde = { version = "1.0.164", default-features = false, features = ["derive"] }

diff --git a/driver/src/conn.rs b/driver/src/conn.rs
@@ -97,5 +97,25 @@ pub trait Connection: DynClone + Send + Sync {
         file_format_options: Option<BTreeMap<&str, &str>>,
         copy_options: Option<BTreeMap<&str, &str>>,
     ) -> Result<QueryProgress>;
+
+    async fn put_files(
+        &self,
+        _local_file: &str,
+        _stage_path: &str,
+    ) -> Result<(Schema, RowProgressIterator)> {
+        Err(Error::IO(
+            "PUT statement only available in HTTP API".to_owned(),
+        ))
+    }
+
+    async fn get_files(
+        &self,
+        _stage_path: &str,
+        _local_file: &str,
+    ) -> Result<(Schema, RowProgressIterator)> {
+        Err(Error::IO(
+            "GET statement only available in HTTP API".to_owned(),
+        ))
+    }
 }
 dyn_clone::clone_trait_object!(Connection);