Skip to content

Commit

Permalink
fix(cubesql): Support new Metabase meta queries
Browse files Browse the repository at this point in the history
  • Loading branch information
@MazterQyou
MazterQyou committed Feb 12, 2024
1 parent 00c2a6b commit 0bc09fd
Show file tree
Hide file tree
Showing 5 changed files with 185 additions and 16 deletions.
85 changes: 71 additions & 14 deletions rust/cubesql/cubesql/src/compile/engine/udf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2824,6 +2824,77 @@ pub fn create_has_schema_privilege_udf(state: Arc<SessionState>) -> ScalarUDF {
)
}

pub fn create_has_table_privilege_udf(state: Arc<SessionState>) -> 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::<Result<BooleanArray>>();

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);
Expand Down Expand Up @@ -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",
Expand Down
80 changes: 78 additions & 2 deletions rust/cubesql/cubesql/src/compile/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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!(
Expand Down Expand Up @@ -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(())
}
}
12 changes: 12 additions & 0 deletions rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege.snap
View file
Original file line number Diff line number Diff line change
@@ -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 |
+---------------------------+--------+--------+
12 changes: 12 additions & 0 deletions ...esql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege_default_user.snap
View file
Original file line number Diff line number Diff line change
@@ -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 |
+---------------------------+--------+--------+
12 changes: 12 additions & 0 deletions ...ubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_privilege_query.snap
View file
Original file line number Diff line number Diff line change
@@ -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 |
+------+--------+---------------------------+--------+--------+--------+--------+

0 comments on commit 0bc09fd

Please sign in to comment.