feat(cubesql): Support DISCARD [ALL | PLANS | SEQUENCES | TEMPORARY |… (

#4560)

rust/cubesql/cubesql/Cargo.toml

@@ -9,12 +9,12 @@ documentation = "https://cube.dev/docs"
 homepage = "https://cube.dev"
 
 [dependencies]
-datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "0944add79f290612bbf21bad428770ad277e0a00", default-features = false, features = ["unicode_expressions"] }
+datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "c152dcb817abcaac8c959a0e515aec9f8ddfd83a", default-features = false, features = ["unicode_expressions"] }
 anyhow = "1.0"
 thiserror = "1.0"
 cubeclient = { path = "../cubeclient" }
 pg-srv = { path = "../pg-srv" }
-sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "90bbbcf041160cdefd7ddc0f95e6b2b358dfcef2" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "47a7b66c64417662c6c30f0decdb0a900fc97a56" }
 lazy_static = "1.4.0"
 base64 = "0.13.0"
 tokio = { version = "1.0", features = ["full", "rt"] }

rust/cubesql/cubesql/src/compile/mod.rs

@@ -1741,6 +1741,12 @@ impl QueryPlanner {
                     CommandCompletion::Rollback,
                 ))
             }
+            (ast::Statement::Discard { object_type }, DatabaseProtocol::PostgreSQL) => {
+                Ok(QueryPlan::MetaOk(
+                    StatusFlags::empty(),
+                    CommandCompletion::Discard(object_type.to_string()),
+                ))
+            }
             _ => Err(CompilationError::Unsupported(format!(
                 "Unsupported query type: {}",
                 stmt.to_string()
@@ -6772,6 +6778,40 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_discard_postgres() -> Result<(), CubeError> {
+        insta::assert_snapshot!(
+            "discard_postgres_all",
+            execute_query("DISCARD ALL;".to_string(), DatabaseProtocol::PostgreSQL).await?
+        );
+        insta::assert_snapshot!(
+            "discard_postgres_plans",
+            execute_query("DISCARD PLANS;".to_string(), DatabaseProtocol::PostgreSQL).await?
+        );
+        insta::assert_snapshot!(
+            "discard_postgres_sequences",
+            execute_query(
+                "DISCARD SEQUENCES;".to_string(),
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+        );
+        insta::assert_snapshot!(
+            "discard_postgres_temporary",
+            execute_query(
+                "DISCARD TEMPORARY;".to_string(),
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+        );
+        insta::assert_snapshot!(
+            "discard_postgres_temp",
+            execute_query("DISCARD TEMP;".to_string(), DatabaseProtocol::PostgreSQL).await?
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn superset_meta_queries() -> Result<(), CubeError> {
         init_logger();

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__discard_postgres_all.snap

@@ -0,0 +1,5 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"DISCARD ALL;\".to_string(), DatabaseProtocol::PostgreSQL).await?"
+---
+

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__discard_postgres_plans.snap

@@ -0,0 +1,5 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"DISCARD PLANS;\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
+

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__discard_postgres_sequences.snap

@@ -0,0 +1,5 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"DISCARD SEQUENCES;\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
+

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__discard_postgres_temp.snap

@@ -0,0 +1,5 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"DISCARD TEMP;\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
+

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__discard_postgres_temporary.snap

@@ -0,0 +1,5 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"DISCARD TEMPORARY;\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
+

rust/cubesql/cubesql/src/sql/postgres/shim.rs

@@ -11,6 +11,7 @@ use crate::{
         extended::Portal,
         session::DatabaseProtocol,
         statement::{StatementParamsFinder, StatementPlaceholderReplacer},
+        types::CommandCompletion,
         writer::BatchWriter,
         AuthContext, Session,
     },
@@ -19,7 +20,7 @@ use crate::{
 use log::{debug, error, trace};
 use pg_srv::{
     buffer, protocol,
-    protocol::{ErrorCode, ErrorResponse, Format},
+    protocol::{CommandComplete, ErrorCode, ErrorResponse, Format},
     PgType, PgTypeId, ProtocolError,
 };
 use tokio::{io::AsyncWriteExt, net::TcpStream};
@@ -602,6 +603,8 @@ impl AsyncPostgresShim {
         let mut writer = BatchWriter::new(portal.get_format());
         let completion = portal.execute(&mut writer, 0).await?;
 
+        self.handle_command_complete(&completion);
+
         if writer.has_data() {
             buffer::write_direct(&mut self.socket, writer).await?;
         };
@@ -633,6 +636,15 @@ impl AsyncPostgresShim {
             Err(CubeError::internal("must be auth".to_string()))
         }
     }
+
+    fn handle_command_complete(&mut self, completion: &CommandComplete) {
+        if completion == &CommandCompletion::Discard("ALL".to_string()).to_pg_command()
+            || completion == &CommandCompletion::Discard("PLANS".to_string()).to_pg_command()
+        {
+            self.statements = HashMap::new();
+            self.portals = HashMap::new();
+        }
+    }
 }
 
 impl Drop for AsyncPostgresShim {

rust/cubesql/cubesql/src/sql/types.rs

@@ -92,6 +92,7 @@ pub enum CommandCompletion {
     Rollback,
     Set,
     Select(u32),
+    Discard(String),
 }
 
 impl CommandCompletion {
@@ -103,6 +104,7 @@ impl CommandCompletion {
             CommandCompletion::Set => CommandComplete::Plain("SET".to_string()),
             CommandCompletion::Use => CommandComplete::Plain("USE".to_string()),
             CommandCompletion::Select(rows) => CommandComplete::Select(rows),
+            CommandCompletion::Discard(tp) => CommandComplete::Plain(format!("DISCARD {}", tp)),
         }
     }
 }

rust/cubesql/pg-srv/src/protocol.rs

@@ -253,6 +253,7 @@ impl Serialize for ParseComplete {
     }
 }
 
+#[derive(PartialEq)]
 pub enum CommandComplete {
     Select(u32),
     Plain(String),