feat: impl insert data from query (#1025)

* feat: refactor insertion in datanode

* feat: supports inserting data by select query

* feat: impl cast operation for vector

* feat: streaming insert from select query results

* chore: minor changes

* fix: remove unwrap

* test: insert_to_requsts

* test: test_execute_insert_by_select

* fix: cast operation for vectors

* fix: test

* fix: typo

* chore: by CR comments

* fix: test_statement_to_request
This commit is contained in:
dennis zhuang
2023-02-17 17:56:12 +08:00
committed by GitHub
parent 7787cfdd42
commit a9c8584c98
16 changed files with 769 additions and 119 deletions

View File

@@ -11,7 +11,7 @@
// 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.
#![feature(box_patterns)]
#![feature(assert_matches)]
pub mod ast;

View File

@@ -11,12 +11,12 @@
// 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 sqlparser::ast::{ObjectName, SetExpr, Statement, UnaryOperator, Values};
use sqlparser::ast::{ObjectName, Query, SetExpr, Statement, UnaryOperator, Values};
use sqlparser::parser::ParserError;
use crate::ast::{Expr, Value};
use crate::error::{self, Result};
use crate::statements::query::Query as GtQuery;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Insert {
@@ -39,16 +39,43 @@ impl Insert {
}
}
pub fn values(&self) -> Result<Vec<Vec<Value>>> {
pub fn values_body(&self) -> Result<Option<Vec<Vec<Value>>>> {
let values = match &self.inner {
Statement::Insert { source, .. } => match &*source.body {
SetExpr::Values(Values { rows, .. }) => sql_exprs_to_values(rows)?,
_ => unreachable!(),
},
_ => unreachable!(),
Statement::Insert {
source:
box Query {
body: box SetExpr::Values(Values { rows, .. }),
..
},
..
} => Some(sql_exprs_to_values(rows)?),
_ => None,
};
Ok(values)
}
pub fn query_body(&self) -> Result<Option<GtQuery>> {
Ok(match &self.inner {
Statement::Insert {
source: box query, ..
} => Some(query.clone().try_into()?),
_ => None,
})
}
pub fn is_insert_select(&self) -> bool {
matches!(
self.inner,
Statement::Insert {
source: box Query {
body: box SetExpr::Select { .. },
..
},
..
}
)
}
}
fn sql_exprs_to_values(exprs: &Vec<Vec<Expr>>) -> Result<Vec<Vec<Value>>> {
@@ -113,11 +140,10 @@ mod tests {
use super::*;
use crate::parser::ParserContext;
use crate::statements::statement::Statement;
#[test]
fn test_insert_value_with_unary_op() {
use crate::statements::statement::Statement;
// insert "-1"
let sql = "INSERT INTO my_table VALUES(-1)";
let stmt = ParserContext::create_with_dialect(sql, &GenericDialect {})
@@ -125,7 +151,7 @@ mod tests {
.remove(0);
match stmt {
Statement::Insert(insert) => {
let values = insert.values().unwrap();
let values = insert.values_body().unwrap().unwrap();
assert_eq!(values, vec![vec![Value::Number("-1".to_string(), false)]]);
}
_ => unreachable!(),
@@ -138,7 +164,7 @@ mod tests {
.remove(0);
match stmt {
Statement::Insert(insert) => {
let values = insert.values().unwrap();
let values = insert.values_body().unwrap().unwrap();
assert_eq!(values, vec![vec![Value::Number("1".to_string(), false)]]);
}
_ => unreachable!(),
@@ -147,8 +173,6 @@ mod tests {
#[test]
fn test_insert_value_with_default() {
use crate::statements::statement::Statement;
// insert "default"
let sql = "INSERT INTO my_table VALUES(default)";
let stmt = ParserContext::create_with_dialect(sql, &GenericDialect {})
@@ -156,7 +180,7 @@ mod tests {
.remove(0);
match stmt {
Statement::Insert(insert) => {
let values = insert.values().unwrap();
let values = insert.values_body().unwrap().unwrap();
assert_eq!(values, vec![vec![Value::Placeholder("default".to_owned())]]);
}
_ => unreachable!(),
@@ -165,8 +189,6 @@ mod tests {
#[test]
fn test_insert_value_with_default_uppercase() {
use crate::statements::statement::Statement;
// insert "DEFAULT"
let sql = "INSERT INTO my_table VALUES(DEFAULT)";
let stmt = ParserContext::create_with_dialect(sql, &GenericDialect {})
@@ -174,7 +196,7 @@ mod tests {
.remove(0);
match stmt {
Statement::Insert(insert) => {
let values = insert.values().unwrap();
let values = insert.values_body().unwrap().unwrap();
assert_eq!(values, vec![vec![Value::Placeholder("DEFAULT".to_owned())]]);
}
_ => unreachable!(),
@@ -183,8 +205,6 @@ mod tests {
#[test]
fn test_insert_value_with_quoted_string() {
use crate::statements::statement::Statement;
// insert "'default'"
let sql = "INSERT INTO my_table VALUES('default')";
let stmt = ParserContext::create_with_dialect(sql, &GenericDialect {})
@@ -192,7 +212,7 @@ mod tests {
.remove(0);
match stmt {
Statement::Insert(insert) => {
let values = insert.values().unwrap();
let values = insert.values_body().unwrap().unwrap();
assert_eq!(
values,
vec![vec![Value::SingleQuotedString("default".to_owned())]]
@@ -201,4 +221,26 @@ mod tests {
_ => unreachable!(),
}
}
#[test]
fn test_insert_select() {
let sql = "INSERT INTO my_table select * from other_table";
let stmt = ParserContext::create_with_dialect(sql, &GenericDialect {})
.unwrap()
.remove(0);
match stmt {
Statement::Insert(insert) => {
assert!(insert.is_insert_select());
let q = insert.query_body().unwrap().unwrap();
assert!(matches!(
q.inner,
Query {
body: box SetExpr::Select { .. },
..
}
));
}
_ => unreachable!(),
}
}
}