feat(cubesql): Support SET ROLE name

Author: MazterQyou

rust/cubesql/Cargo.lock

@@ -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 = "00ed6a6d469a69f57dd1f08c1fda3f7c2cf12d80", default-features = false, features = ["regex_expressions", "unicode_expressions"] }
+datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "be967a63967f23bdd946a45925d07af0a11dcc39", default-features = false, features = ["regex_expressions", "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 = "ac5fc7cbf4368cd4c9e6c2af0e0e35d2eff7cc16" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "037562f7975fcc92b725efbfbaaf050272668593" }
 lazy_static = "1.4.0"
 base64 = "0.13.0"
 tokio = { version = "1.0", features = ["full", "rt", "tracing"] }

rust/cubesql/cubesql/Cargo.toml

@@ -1083,6 +1083,16 @@ impl AsyncTestSuite for PostgresIntegrationTestSuite {
         )
         .await?;
 
+        self.test_simple_query(r#"SET ROLE "cube""#.to_string(), |messages| {
+            assert_eq!(messages.len(), 1);
+
+            // SET
+            if let SimpleQueryMessage::Row(_) = messages[0] {
+                panic!("Must be CommandComplete command, (SET is used)")
+            }
+        })
+        .await?;
+
         // Tableau Desktop
         self.test_simple_query(
             r#"SET DateStyle = 'ISO';SET extra_float_digits = 2;show transaction_isolation"#

rust/cubesql/cubesql/e2e/tests/postgres.rs

@@ -401,6 +401,7 @@ impl QueryPlanner {
                 StatusFlags::empty(),
                 CommandCompletion::Select(0),
             )),
+            (ast::Statement::SetRole { role_name, .. }, _) => self.set_role_to_plan(role_name),
             (ast::Statement::SetVariable { key_values }, _) => {
                 self.set_variable_to_plan(&key_values)
             }
@@ -938,6 +939,28 @@ WHERE `TABLE_SCHEMA` = '{}'",
         ))
     }
 
+    fn set_role_to_plan(
+        &self,
+        role_name: &Option<ast::Ident>,
+    ) -> Result<QueryPlan, CompilationError> {
+        let flags = StatusFlags::SERVER_STATE_CHANGED;
+        let role_name = role_name
+            .as_ref()
+            .map(|role_name| role_name.value.clone())
+            .unwrap_or("none".to_string());
+        let variable =
+            DatabaseVariable::system("role".to_string(), ScalarValue::Utf8(Some(role_name)), None);
+        self.state.set_variables(vec![variable]);
+        match self.state.protocol {
+            DatabaseProtocol::PostgreSQL => Ok(QueryPlan::MetaOk(flags, CommandCompletion::Set)),
+            // TODO: Verify that it's possible to use MetaOk too...
+            DatabaseProtocol::MySQL => Ok(QueryPlan::MetaTabular(
+                flags,
+                Box::new(dataframe::DataFrame::new(vec![], vec![])),
+            )),
+        }
+    }
+
     fn set_variable_to_plan(
         &self,
         key_values: &Vec<ast::SetVariableKeyValue>,
@@ -7085,6 +7108,16 @@ ORDER BY \"COUNT(count)\" DESC"
             .0
         );
 
+        insta::assert_snapshot!(
+            "pg_set_role_show",
+            execute_queries_with_flags(
+                vec!["SET ROLE NONE".to_string(), "SHOW ROLE".to_string()],
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+            .0
+        );
+
         Ok(())
     }
 

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

@@ -0,0 +1,9 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_queries_with_flags(vec![\"SET ROLE NONE\".to_string(),\n                    \"SHOW ROLE\".to_string()],\n                DatabaseProtocol::PostgreSQL).await?.0"
+---
++---------+
+| setting |
++---------+
+| none    |
++---------+

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

@@ -96,5 +96,14 @@ pub fn defaults() -> DatabaseVariables {
         ),
     );
 
+    variables.insert(
+        "role".to_string(),
+        DatabaseVariable::system(
+            "role".to_string(),
+            ScalarValue::Utf8(Some("none".to_string())),
+            None,
+        ),
+    );
+
     variables
 }