diff --git a/pgdog/src/backend/replication/logical/publisher/table.rs b/pgdog/src/backend/replication/logical/publisher/table.rs index e4218df03..016afbaf6 100644 --- a/pgdog/src/backend/replication/logical/publisher/table.rs +++ b/pgdog/src/backend/replication/logical/publisher/table.rs @@ -457,7 +457,12 @@ impl Table { // Create new standalone connection for the copy. // let mut server = Server::connect(source, ServerOptions::new_replication()).await?; - let mut copy_sub = CopySubscriber::new(copy.statement(), dest, self.query_parser_engine)?; + let mut copy_sub = CopySubscriber::new( + copy.statement(), + dest, + #[cfg(not(feature = "new_parser"))] + self.query_parser_engine, + )?; copy_sub.connect().await?; // Create sync slot. diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 4d7f3e707..eeeccdb86 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -2,7 +2,11 @@ //! between N shards. use futures::future::join_all; +#[cfg(not(feature = "new_parser"))] use pg_query::{NodeEnum, parse_raw}; +#[cfg(feature = "new_parser")] +use pg_raw_parse::Node; +#[cfg(not(feature = "new_parser"))] use pgdog_config::QueryParserEngine; use tracing::debug; @@ -11,6 +15,8 @@ use crate::frontend::client::query_engine::two_pc::{ Manager, TwoPcTransaction, statement::phase_control, }; +#[cfg(feature = "new_parser")] +use crate::frontend::router::parser::Error as ParseError; use crate::{ backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection}, config::Role, @@ -44,29 +50,11 @@ impl CopySubscriber { /// 1. What kind of encoding we use. /// 2. Which column is used for sharding. /// - pub fn new( - copy_stmt: &CopyStatement, - cluster: &Cluster, - query_parser_engine: QueryParserEngine, - ) -> Result { - let stmt = match query_parser_engine { - QueryParserEngine::PgQueryProtobuf => { - pg_query::parse(copy_stmt.clone().copy_in().as_str()) - } - QueryParserEngine::PgQueryRaw => parse_raw(copy_stmt.clone().copy_in().as_str()), - }?; - let stmt = stmt - .protobuf - .stmts - .first() - .ok_or(Error::MissingData)? - .stmt - .as_ref() - .ok_or(Error::MissingData)? - .node - .as_ref() - .ok_or(Error::MissingData)?; - let copy = if let NodeEnum::CopyStmt(stmt) = stmt { + #[cfg(feature = "new_parser")] + pub fn new(copy_stmt: &CopyStatement, cluster: &Cluster) -> Result { + let ast = pg_raw_parse::parse(©_stmt.copy_in()).map_err(ParseError::from)?; + let stmt = ast.stmts().next().ok_or(ParseError::EmptyQuery)?; + let copy = if let Node::CopyStmt(stmt) = stmt { CopyParser::new(stmt, cluster).map_err(|_| Error::MissingData)? } else { return Err(Error::MissingData); @@ -82,6 +70,49 @@ impl CopySubscriber { }) } + cfg_select! { + not(feature = "new_parser") => { + pub fn new( + copy_stmt: &CopyStatement, + cluster: &Cluster, + query_parser_engine: QueryParserEngine, + ) -> Result { + let stmt = match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => { + pg_query::parse(copy_stmt.clone().copy_in().as_str()) + } + QueryParserEngine::PgQueryRaw => parse_raw(copy_stmt.clone().copy_in().as_str()), + }?; + let stmt = stmt + .protobuf + .stmts + .first() + .ok_or(Error::MissingData)? + .stmt + .as_ref() + .ok_or(Error::MissingData)? + .node + .as_ref() + .ok_or(Error::MissingData)?; + let copy = if let NodeEnum::CopyStmt(stmt) = stmt { + CopyParser::new(stmt, cluster).map_err(|_| Error::MissingData)? + } else { + return Err(Error::MissingData); + }; + + Ok(Self { + copy, + cluster: cluster.clone(), + buffer: vec![], + connections: vec![], + stmt: copy_stmt.clone(), + bytes_sharded: 0, + }) + } + } + _ => {} + } + /// Connect to all shards. One connection per primary. pub async fn connect(&mut self) -> Result<(), Error> { let mut servers = vec![]; @@ -398,9 +429,13 @@ mod test { .await .unwrap(); - let mut subscriber = - CopySubscriber::new(©, &cluster, config().config.general.query_parser_engine) - .unwrap(); + let mut subscriber = CopySubscriber::new( + ©, + &cluster, + #[cfg(not(feature = "new_parser"))] + config().config.general.query_parser_engine, + ) + .unwrap(); subscriber.start_copy().await.unwrap(); let header = CopyData::new(&Header::new().to_bytes()); diff --git a/pgdog/src/frontend/router/parser/column.rs b/pgdog/src/frontend/router/parser/column.rs index 00bd4a0f5..d9cd5dd76 100644 --- a/pgdog/src/frontend/router/parser/column.rs +++ b/pgdog/src/frontend/router/parser/column.rs @@ -1,9 +1,6 @@ //! Column name reference. -use pg_query::{ - Node, NodeEnum, - protobuf::{self, String as PgQueryString}, -}; +use pg_query::{Node, NodeEnum, protobuf::String as PgQueryString}; #[cfg(feature = "new_parser")] use pg_raw_parse::{list, nodes}; use std::fmt::{Display, Formatter, Result as FmtResult}; @@ -31,9 +28,10 @@ impl<'a> Column<'a> { }) } + #[cfg(not(feature = "new_parser"))] pub(crate) fn from_string(string: &'a Node) -> Result { match &string.node { - Some(NodeEnum::String(protobuf::String { sval })) => Ok(Self { + Some(NodeEnum::String(PgQueryString { sval })) => Ok(Self { name: sval.as_str(), ..Default::default() }), diff --git a/pgdog/src/frontend/router/parser/copy.rs b/pgdog/src/frontend/router/parser/copy.rs index 64ba9e8a8..54400e79a 100644 --- a/pgdog/src/frontend/router/parser/copy.rs +++ b/pgdog/src/frontend/router/parser/copy.rs @@ -1,6 +1,9 @@ //! Parse COPY statement. +#[cfg(not(feature = "new_parser"))] use pg_query::{NodeEnum, protobuf::CopyStmt}; +#[cfg(feature = "new_parser")] +use pg_raw_parse::nodes; use crate::{ backend::{Cluster, ShardingSchema}, @@ -94,7 +97,8 @@ impl Default for CopyParser { impl CopyParser { /// Create new copy parser from a COPY statement. - pub fn new(stmt: &CopyStmt, cluster: &Cluster) -> Result { + #[cfg(feature = "new_parser")] + pub fn new(stmt: &nodes::CopyStmt, cluster: &Cluster) -> Result { let mut parser = Self { is_from: stmt.is_from, ..Default::default() @@ -103,14 +107,14 @@ impl CopyParser { let mut format = CopyFormat::Text; let mut null_string = "\\N".to_owned(); - if let Some(ref rel) = stmt.relation { - let mut columns = vec![]; - - for column in &stmt.attlist { - if let Ok(column) = Column::from_string(column) { - columns.push(column); - } - } + if let Some(rel) = stmt.relation() { + let columns = stmt + .attlist() + .iter() + .map(|column| { + Column::from(column.as_str().expect("attlist is empty or all strings")) + }) + .collect::>(); let table = Table::from(rel); @@ -128,51 +132,39 @@ impl CopyParser { parser.columns = columns.len(); - for option in &stmt.options { - if let Some(NodeEnum::DefElem(ref elem)) = option.node { - match elem.defname.to_lowercase().as_str() { - "format" => { - if let Some(ref arg) = elem.arg - && let Some(NodeEnum::String(ref string)) = arg.node - { - match string.sval.to_lowercase().as_str() { - "binary" => { - parser.headers = true; - format = CopyFormat::Binary; - } - "csv" => { - if parser.delimiter.is_none() { - parser.delimiter = Some(','); - } - format = CopyFormat::Csv; - } - _ => (), - } - } + for elem in stmt.options() { + match elem.defname().unwrap_or_default().to_lowercase().as_str() { + "format" => match elem.arg().as_str().map(|s| s.to_lowercase()).as_deref() { + Some("binary") => { + parser.headers = true; + format = CopyFormat::Binary; } - - "delimiter" => { - if let Some(ref arg) = elem.arg - && let Some(NodeEnum::String(ref string)) = arg.node - { - parser.delimiter = Some(string.sval.chars().next().unwrap_or(',')); + Some("csv") => { + if parser.delimiter.is_none() { + parser.delimiter = Some(','); } + format = CopyFormat::Csv; } + _ => (), + }, - "header" => { - parser.headers = true; + "delimiter" => { + if let Some(string) = elem.arg().as_str() { + parser.delimiter = Some(string.chars().next().unwrap_or(',')); } + } - "null" => { - if let Some(ref arg) = elem.arg - && let Some(NodeEnum::String(ref string)) = arg.node - { - null_string = string.sval.clone(); - } - } + "header" => { + parser.headers = true; + } - _ => (), + "null" => { + if let Some(string) = elem.arg().as_str() { + null_string = string.to_owned(); + } } + + _ => (), } } } @@ -193,6 +185,110 @@ impl CopyParser { Ok(parser) } + cfg_select! { + not(feature = "new_parser") => { + pub fn new(stmt: &CopyStmt, cluster: &Cluster) -> Result { + let mut parser = Self { + is_from: stmt.is_from, + ..Default::default() + }; + + let mut format = CopyFormat::Text; + let mut null_string = "\\N".to_owned(); + + if let Some(ref rel) = stmt.relation { + let mut columns = vec![]; + + for column in &stmt.attlist { + if let Ok(column) = Column::from_string(column) { + columns.push(column); + } + } + + let table = Table::from(rel); + + // The CopyParser is used for replicating + // data during data-sync. This will ensure all rows + // are sent to the right schema-based shard. + if let Some(schema) = cluster.sharding_schema().schemas.get(table.schema()) { + parser.schema_shard = Some(schema.shard().into()); + } + + if let Some(key) = Tables::new(&cluster.sharding_schema()).key(table, &columns) { + parser.sharded_table = Some(key.table.clone()); + parser.sharded_column = key.position; + } + + parser.columns = columns.len(); + + for option in &stmt.options { + if let Some(NodeEnum::DefElem(ref elem)) = option.node { + match elem.defname.to_lowercase().as_str() { + "format" => { + if let Some(ref arg) = elem.arg + && let Some(NodeEnum::String(ref string)) = arg.node + { + match string.sval.to_lowercase().as_str() { + "binary" => { + parser.headers = true; + format = CopyFormat::Binary; + } + "csv" => { + if parser.delimiter.is_none() { + parser.delimiter = Some(','); + } + format = CopyFormat::Csv; + } + _ => (), + } + } + } + + "delimiter" => { + if let Some(ref arg) = elem.arg + && let Some(NodeEnum::String(ref string)) = arg.node + { + parser.delimiter = Some(string.sval.chars().next().unwrap_or(',')); + } + } + + "header" => { + parser.headers = true; + } + + "null" => { + if let Some(ref arg) = elem.arg + && let Some(NodeEnum::String(ref string)) = arg.node + { + null_string = string.sval.clone(); + } + } + + _ => (), + } + } + } + } + + parser.stream = if format == CopyFormat::Binary { + CopyStream::Binary(BinaryStream::default()) + } else { + CopyStream::Text(Box::new(CsvStream::new( + parser.delimiter(), + parser.headers, + format, + &null_string, + ))) + }; + parser.sharding_schema = cluster.sharding_schema(); + parser.null_string = null_string; + + Ok(parser) + } + } + _ => {} + } + #[inline] fn delimiter(&self) -> char { self.delimiter.unwrap_or('\t') @@ -309,22 +405,15 @@ impl CopyParser { #[cfg(test)] mod test { - use pg_query::parse; - use crate::config::config; + #[cfg(feature = "new_parser")] + use pg_raw_parse::{Node, Owned, make}; use super::*; #[test] fn test_copy_text() { - let copy = "COPY sharded (id, value) FROM STDIN"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN"); let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert_eq!(copy.delimiter(), '\t'); @@ -339,14 +428,7 @@ mod test { #[test] fn test_copy_csv() { - let copy = "COPY sharded (id, value) FROM STDIN CSV HEADER"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN CSV HEADER"); let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert!(copy.is_from); @@ -378,14 +460,7 @@ mod test { fn test_copy_csv_stream() { let copy_data = CopyData::new(b"id,value\n1,test\n6,test6\n"); - let copy = "COPY sharded (id, value) FROM STDIN CSV HEADER"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN CSV HEADER"); let mut copy = CopyParser::new(©, &Cluster::new_test(&config())).unwrap(); let rows = copy.shard(&[copy_data]).unwrap(); @@ -400,14 +475,7 @@ mod test { #[test] fn test_copy_csv_custom_null() { - let copy = "COPY sharded (id, value) FROM STDIN CSV NULL 'NULL'"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN CSV NULL 'NULL'"); let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert_eq!(copy.delimiter(), ','); @@ -426,14 +494,7 @@ mod test { fn test_copy_text_pg_dump_end_marker() { // pg_dump generates text format COPY with `\.` as end-of-copy marker. // This marker should be sent to all shards without extracting a sharding key. - let copy = "COPY sharded (id, value) FROM STDIN"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN"); let mut copy = CopyParser::new(©, &Cluster::new_test(&config())).unwrap(); let one = CopyData::new("1\tAlice\n".as_bytes()); @@ -455,14 +516,7 @@ mod test { // pg_dump text format uses `\N` to represent NULL values. // When the sharding key is NULL, route to all shards. // When a non-sharding column is NULL, route normally based on the key. - let copy = "COPY sharded (id, value) FROM STDIN"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN"); let mut copy = CopyParser::new(©, &Cluster::new_test(&config())).unwrap(); let one = CopyData::new("1\tAlice\n".as_bytes()); @@ -485,14 +539,7 @@ mod test { #[test] fn test_copy_text_composite_type_sharded() { // Test the same composite type but with sharding enabled (using the sharded table from config) - let copy = "COPY sharded (id, value) FROM STDIN"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN"); let mut copy = CopyParser::new(©, &Cluster::new_test(&config())).unwrap(); // Row where the value contains a composite type with commas and quotes @@ -512,14 +559,8 @@ mod test { #[test] fn test_copy_explicit_text_format() { // Test with explicit FORMAT text (like during resharding) - let copy = r#"COPY "public"."entity_values" ("id", "value_location") FROM STDIN WITH (FORMAT text)"#; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy_stmt = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let sql = r#"COPY "public"."entity_values" ("id", "value_location") FROM STDIN WITH (FORMAT text)"#; + let copy_stmt = parse(sql); let mut copy = CopyParser::new(©_stmt, &Cluster::default()).unwrap(); // Verify it's using tab delimiter (text format default) @@ -543,14 +584,7 @@ mod test { #[test] fn test_copy_binary() { - let copy = "COPY sharded (id, value) FROM STDIN (FORMAT 'binary')"; - let stmt = parse(copy).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - let copy = match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - }; - + let copy = parse("COPY sharded (id, value) FROM STDIN (FORMAT 'binary')"); let mut copy = CopyParser::new(©, &Cluster::new_test(&config())).unwrap(); assert!(copy.is_from); assert!(copy.headers); @@ -578,4 +612,27 @@ mod test { assert_eq!(sharded[2].message().data(), (-1_i16).to_be_bytes()); assert_eq!(sharded[2].shard(), &Shard::All) } + + #[cfg(feature = "new_parser")] + fn parse(sql: &str) -> Owned { + let stmt = pg_raw_parse::parse(sql).unwrap(); + match stmt.stmts().next() { + Some(Node::CopyStmt(copy)) => make::owned(|mem| mem.make_unique(copy)), + _ => panic!("not a copy"), + } + } + + cfg_select! { + not(feature = "new_parser") => { + fn parse(sql: &str) -> Box { + let stmt = pg_query::parse(sql).unwrap(); + let stmt = stmt.protobuf.stmts.first().unwrap(); + match stmt.stmt.clone().unwrap().node.unwrap() { + NodeEnum::CopyStmt(copy) => copy, + _ => panic!("not a copy"), + } + } + } + _ => {} + } } diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 81a9176ab..305dd1d94 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -38,16 +38,12 @@ mod update; #[cfg(feature = "new_parser")] use itertools::*; use multi_tenant::MultiTenantCheck; +use pg_query::NodeEnum; use pg_query::protobuf::Node as PgNode; #[cfg(feature = "new_parser")] -use pg_query::{NodeEnum, protobuf::CopyStmt}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; #[cfg(not(feature = "new_parser"))] -use pgdog_plugin::pg_query::{ - NodeEnum, - protobuf::{a_const::Val, *}, -}; +use pgdog_plugin::pg_query::protobuf::{a_const::Val, *}; use plugins::PluginOutput; use tracing::{debug, trace}; @@ -351,13 +347,7 @@ impl QueryParser { Node::SelectStmt(stmt) => self.select(&statement, stmt, context), - Node::CopyStmt(stmt) => Self::copy( - match &old_root { - Some(NodeEnum::CopyStmt(stmt)) => stmt, - _ => unreachable!(), - }, - context, - ), + Node::CopyStmt(stmt) => Self::copy(stmt, context), Node::InsertStmt(stmt) => self.insert(stmt.into(), context), Node::UpdateStmt(stmt) => self.update(stmt.into(), context), @@ -854,7 +844,8 @@ impl QueryParser { } /// Handle COPY command. - fn copy(stmt: &CopyStmt, context: &mut QueryParserContext) -> Result { + #[cfg(feature = "new_parser")] + fn copy(stmt: &nodes::CopyStmt, context: &mut QueryParserContext) -> Result { // Schema-based routing. // // We do this here as well because COPY TO STDOUT @@ -865,7 +856,7 @@ impl QueryParser { // but that's only used for logical replication during the first // phase of data-sync. // - let table = stmt.relation.as_ref().map(Table::from); + let table = stmt.relation().map(Table::from); if let Some(table) = table && let Some(schema) = context.sharding_schema.schemas.get(table.schema()) { @@ -897,6 +888,54 @@ impl QueryParser { } } + cfg_select! { + not(feature = "new_parser") => { + fn copy(stmt: &CopyStmt, context: &mut QueryParserContext) -> Result { + // Schema-based routing. + // + // We do this here as well because COPY
TO STDOUT + // doesn't use the CopyParser (doesn't need to, normally), + // so we need to handle this case here. + // + // The CopyParser itself has handling for schema-based sharding, + // but that's only used for logical replication during the first + // phase of data-sync. + // + let table = stmt.relation.as_ref().map(Table::from); + if let Some(table) = table + && let Some(schema) = context.sharding_schema.schemas.get(table.schema()) + { + let shard: Shard = schema.shard().into(); + context + .shards_calculator + .push(ShardWithPriority::new_table(shard)); + if !stmt.is_from { + return Ok(Command::Query(Route::read( + context.shards_calculator.shard(), + ))); + } else { + return Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))); + } + } + + let parser = CopyParser::new(stmt, context.router_context.cluster)?; + if !stmt.is_from { + context + .shards_calculator + .push(ShardWithPriority::new_table(Shard::All)); + Ok(Command::Query(Route::read( + context.shards_calculator.shard(), + ))) + } else { + Ok(Command::Copy(Box::new(parser))) + } + } + } + _ => {} + } + /// Handle INSERT statement. /// /// # Arguments