From 0bc09fdbefaf8b4e184f4c0803919103789acb0a Mon Sep 17 00:00:00 2001 From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com> Date: Mon, 12 Feb 2024 23:46:38 +0400 Subject: [PATCH] fix(cubesql): Support new Metabase meta queries --- .../cubesql/cubesql/src/compile/engine/udf.rs | 85 ++++++++++++++++--- rust/cubesql/cubesql/src/compile/mod.rs | 80 ++++++++++++++++- ...__compile__tests__has_table_privilege.snap | 12 +++ ...sts__has_table_privilege_default_user.snap | 12 +++ ...tests__metabase_table_privilege_query.snap | 12 +++ 5 files changed, 185 insertions(+), 16 deletions(-) create mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege.snap create mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege_default_user.snap create mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_privilege_query.snap diff --git a/rust/cubesql/cubesql/src/compile/engine/udf.rs b/rust/cubesql/cubesql/src/compile/engine/udf.rs index ce8729e663312..6302f0c496f3d 100644 --- a/rust/cubesql/cubesql/src/compile/engine/udf.rs +++ b/rust/cubesql/cubesql/src/compile/engine/udf.rs @@ -2824,6 +2824,77 @@ pub fn create_has_schema_privilege_udf(state: Arc) -> ScalarUDF { ) } +pub fn create_has_table_privilege_udf(state: Arc) -> ScalarUDF { + let fun = make_scalar_function(move |args: &[ArrayRef]| { + let (users, tables, privileges) = if args.len() == 3 { + ( + Some(downcast_string_arg!(args[0], "user", i32)), + downcast_string_arg!(args[1], "table", i32), + downcast_string_arg!(args[2], "privilege", i32), + ) + } else { + ( + None, + downcast_string_arg!(args[0], "table", i32), + downcast_string_arg!(args[1], "privilege", i32), + ) + }; + + let result = izip!(tables, privileges) + .enumerate() + .map(|(i, args)| { + Ok(match args { + (Some(_table), Some(privilege)) => { + match (users, state.user()) { + (Some(users), Some(session_user)) => { + let user = users.value(i); + if user != session_user { + return Err(DataFusionError::Execution(format!( + "role \"{}\" does not exist", + user + ))); + } + } + _ => (), + } + + // TODO: check if table exists + + match privilege { + "SELECT" => Some(true), + "UPDATE" | "INSERT" | "DELETE" => Some(false), + _ => { + return Err(DataFusionError::Execution(format!( + "unrecognized privilege type: \"{}\"", + privilege + ))) + } + } + } + _ => None, + }) + }) + .collect::>(); + + Ok(Arc::new(result?)) + }); + + let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Boolean))); + + ScalarUDF::new( + "has_table_privilege", + &Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8, DataType::Utf8]), + TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]), + ], + Volatility::Stable, + ), + &return_type, + &fun, + ) +} + pub fn create_pg_total_relation_size_udf() -> ScalarUDF { let fun = make_scalar_function(move |args: &[ArrayRef]| { assert!(args.len() == 1); @@ -3865,20 +3936,6 @@ pub fn register_fun_stubs(mut ctx: SessionContext) -> SessionContext { rettyp = Boolean, vol = Stable ); - register_fun_stub!( - udf, - "has_table_privilege", - tsigs = [ - [Utf8, Utf8], - [Oid, Utf8], - [Utf8, Utf8, Utf8], - [Utf8, Oid, Utf8], - [Oid, Utf8, Utf8], - [Oid, Oid, Utf8], - ], - rettyp = Boolean, - vol = Stable - ); register_fun_stub!( udf, "has_tablespace_privilege", diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index a94fc9cf1bbd1..391cb79370b45 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -51,8 +51,8 @@ use self::{ create_date_udf, create_dateadd_udf, create_datediff_udf, create_dayofmonth_udf, create_dayofweek_udf, create_dayofyear_udf, create_db_udf, create_ends_with_udf, create_format_type_udf, create_generate_series_udtf, create_generate_subscripts_udtf, - create_has_schema_privilege_udf, create_hour_udf, create_if_udf, - create_inet_server_addr_udf, create_instr_udf, create_interval_mul_udf, + create_has_schema_privilege_udf, create_has_table_privilege_udf, create_hour_udf, + create_if_udf, create_inet_server_addr_udf, create_instr_udf, create_interval_mul_udf, create_isnull_udf, create_json_build_object_udf, create_least_udf, create_locate_udf, create_makedate_udf, create_measure_udaf, create_minute_udf, create_pg_backend_pid_udf, create_pg_datetime_precision_udf, create_pg_encoding_to_char_udf, @@ -1325,6 +1325,7 @@ WHERE `TABLE_SCHEMA` = '{}'", ctx.register_udf(create_pg_my_temp_schema()); ctx.register_udf(create_pg_is_other_temp_schema()); ctx.register_udf(create_has_schema_privilege_udf(self.state.clone())); + ctx.register_udf(create_has_table_privilege_udf(self.state.clone())); ctx.register_udf(create_pg_total_relation_size_udf()); ctx.register_udf(create_cube_regclass_cast_udf()); ctx.register_udf(create_pg_get_serial_sequence_udf()); @@ -8952,6 +8953,43 @@ limit Ok(()) } + #[tokio::test] + async fn test_has_table_privilege_postgres() -> Result<(), CubeError> { + insta::assert_snapshot!( + "has_table_privilege", + execute_query( + "SELECT + relname, + has_table_privilege('ovr', relname, 'SELECT') \"select\", + has_table_privilege('ovr', relname, 'INSERT') \"insert\" + FROM pg_class + ORDER BY relname ASC + " + .to_string(), + DatabaseProtocol::PostgreSQL + ) + .await? + ); + + insta::assert_snapshot!( + "has_table_privilege_default_user", + execute_query( + "SELECT + relname, + has_table_privilege(relname, 'SELECT') \"select\", + has_table_privilege(relname, 'INSERT') \"insert\" + FROM pg_class + ORDER BY relname ASC + " + .to_string(), + DatabaseProtocol::PostgreSQL + ) + .await? + ); + + Ok(()) + } + #[tokio::test] async fn test_pg_total_relation_size() -> Result<(), CubeError> { insta::assert_snapshot!( @@ -21277,4 +21315,42 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), .sql; assert!(sql.contains("OFFSET 1\nLIMIT 2")); } + + #[tokio::test] + async fn test_metabase_table_privilege_query() -> Result<(), CubeError> { + insta::assert_snapshot!( + "metabase_table_privilege_query", + execute_query( + r#" + with table_privileges as ( + select + NULL as role, + t.schemaname as schema, + t.objectname as table, + pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'UPDATE') as update, + pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'SELECT') as select, + pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'INSERT') as insert, + pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'DELETE') as delete + from ( + select schemaname, tablename as objectname from pg_catalog.pg_tables + union + select schemaname, viewname as objectname from pg_catalog.pg_views + union + select schemaname, matviewname as objectname from pg_catalog.pg_matviews + ) t + where t.schemaname !~ '^pg_' + and t.schemaname <> 'information_schema' + and pg_catalog.has_schema_privilege(current_user, t.schemaname, 'USAGE') + ) + select t.* + from table_privileges t + order by t.schema, t.table + "#.to_string(), + DatabaseProtocol::PostgreSQL + ) + .await? + ); + + Ok(()) + } } diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege.snap new file mode 100644 index 0000000000000..8f5002c1479ef --- /dev/null +++ b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege.snap @@ -0,0 +1,12 @@ +--- +source: cubesql/src/compile/mod.rs +expression: "execute_query(\"SELECT\n relname,\n has_table_privilege('ovr', relname, 'SELECT') \\\"select\\\",\n has_table_privilege('ovr', relname, 'INSERT') \\\"insert\\\"\n FROM pg_class\n ORDER BY relname ASC\n \".to_string(),\n DatabaseProtocol::PostgreSQL).await?" +--- ++---------------------------+--------+--------+ +| relname | select | insert | ++---------------------------+--------+--------+ +| KibanaSampleDataEcommerce | true | false | +| Logs | true | false | +| NumberCube | true | false | +| WideCube | true | false | ++---------------------------+--------+--------+ diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege_default_user.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege_default_user.snap new file mode 100644 index 0000000000000..18f1b5d962232 --- /dev/null +++ b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege_default_user.snap @@ -0,0 +1,12 @@ +--- +source: cubesql/src/compile/mod.rs +expression: "execute_query(\"SELECT\n relname,\n has_table_privilege(relname, 'SELECT') \\\"select\\\",\n has_table_privilege(relname, 'INSERT') \\\"insert\\\"\n FROM pg_class\n ORDER BY relname ASC\n \".to_string(),\n DatabaseProtocol::PostgreSQL).await?" +--- ++---------------------------+--------+--------+ +| relname | select | insert | ++---------------------------+--------+--------+ +| KibanaSampleDataEcommerce | true | false | +| Logs | true | false | +| NumberCube | true | false | +| WideCube | true | false | ++---------------------------+--------+--------+ diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_privilege_query.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_privilege_query.snap new file mode 100644 index 0000000000000..158b259dc54d3 --- /dev/null +++ b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_privilege_query.snap @@ -0,0 +1,12 @@ +--- +source: cubesql/src/compile/mod.rs +expression: "execute_query(r#\"\n with table_privileges as (\n\t select\n\t NULL as role,\n\t t.schemaname as schema,\n\t t.objectname as table,\n\t pg_catalog.has_table_privilege(current_user, '\"' || t.schemaname || '\"' || '.' || '\"' || t.objectname || '\"', 'UPDATE') as update,\n\t pg_catalog.has_table_privilege(current_user, '\"' || t.schemaname || '\"' || '.' || '\"' || t.objectname || '\"', 'SELECT') as select,\n\t pg_catalog.has_table_privilege(current_user, '\"' || t.schemaname || '\"' || '.' || '\"' || t.objectname || '\"', 'INSERT') as insert,\n\t pg_catalog.has_table_privilege(current_user, '\"' || t.schemaname || '\"' || '.' || '\"' || t.objectname || '\"', 'DELETE') as delete\n\t from (\n\t select schemaname, tablename as objectname from pg_catalog.pg_tables\n\t union\n\t select schemaname, viewname as objectname from pg_catalog.pg_views\n\t union\n\t select schemaname, matviewname as objectname from pg_catalog.pg_matviews\n\t ) t\n\t where t.schemaname !~ '^pg_'\n\t and t.schemaname <> 'information_schema'\n\t and pg_catalog.has_schema_privilege(current_user, t.schemaname, 'USAGE')\n\t)\n\tselect t.*\n\tfrom table_privileges t\n order by t.schema, t.table\n \"#.to_string(),\n DatabaseProtocol::PostgreSQL).await?" +--- ++------+--------+---------------------------+--------+--------+--------+--------+ +| role | schema | table | update | select | insert | delete | ++------+--------+---------------------------+--------+--------+--------+--------+ +| NULL | public | KibanaSampleDataEcommerce | false | true | false | false | +| NULL | public | Logs | false | true | false | false | +| NULL | public | NumberCube | false | true | false | false | +| NULL | public | WideCube | false | true | false | false | ++------+--------+---------------------------+--------+--------+--------+--------+