Skip to content

Commit 0bc09fd

Browse files
MazterQyouMazterQyou
committed
fix(cubesql): Support new Metabase meta queries
1 parent 00c2a6b commit 0bc09fd

File tree

5 files changed

+185
-16
lines changed

5 files changed

+185
-16
lines changed

‎rust/cubesql/cubesql/src/compile/engine/udf.rs‎

Lines changed: 71 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2824,6 +2824,77 @@ pub fn create_has_schema_privilege_udf(state: Arc<SessionState>) -> ScalarUDF {
28242824
)
28252825
}
28262826

2827+
pub fn create_has_table_privilege_udf(state: Arc<SessionState>) -> ScalarUDF {
2828+
let fun = make_scalar_function(move |args: &[ArrayRef]| {
2829+
let (users, tables, privileges) = if args.len() == 3 {
2830+
(
2831+
Some(downcast_string_arg!(args[0], "user", i32)),
2832+
downcast_string_arg!(args[1], "table", i32),
2833+
downcast_string_arg!(args[2], "privilege", i32),
2834+
)
2835+
} else {
2836+
(
2837+
None,
2838+
downcast_string_arg!(args[0], "table", i32),
2839+
downcast_string_arg!(args[1], "privilege", i32),
2840+
)
2841+
};
2842+
2843+
let result = izip!(tables, privileges)
2844+
.enumerate()
2845+
.map(|(i, args)| {
2846+
Ok(match args {
2847+
(Some(_table), Some(privilege)) => {
2848+
match (users, state.user()) {
2849+
(Some(users), Some(session_user)) => {
2850+
let user = users.value(i);
2851+
if user != session_user {
2852+
return Err(DataFusionError::Execution(format!(
2853+
"role \"{}\" does not exist",
2854+
user
2855+
)));
2856+
}
2857+
}
2858+
_ => (),
2859+
}
2860+
2861+
// TODO: check if table exists
2862+
2863+
match privilege {
2864+
"SELECT" => Some(true),
2865+
"UPDATE" | "INSERT" | "DELETE" => Some(false),
2866+
_ => {
2867+
return Err(DataFusionError::Execution(format!(
2868+
"unrecognized privilege type: \"{}\"",
2869+
privilege
2870+
)))
2871+
}
2872+
}
2873+
}
2874+
_ => None,
2875+
})
2876+
})
2877+
.collect::<Result<BooleanArray>>();
2878+
2879+
Ok(Arc::new(result?))
2880+
});
2881+
2882+
let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Boolean)));
2883+
2884+
ScalarUDF::new(
2885+
"has_table_privilege",
2886+
&Signature::one_of(
2887+
vec![
2888+
TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8, DataType::Utf8]),
2889+
TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]),
2890+
],
2891+
Volatility::Stable,
2892+
),
2893+
&return_type,
2894+
&fun,
2895+
)
2896+
}
2897+
28272898
pub fn create_pg_total_relation_size_udf() -> ScalarUDF {
28282899
let fun = make_scalar_function(move |args: &[ArrayRef]| {
28292900
assert!(args.len() == 1);
@@ -3865,20 +3936,6 @@ pub fn register_fun_stubs(mut ctx: SessionContext) -> SessionContext {
38653936
rettyp = Boolean,
38663937
vol = Stable
38673938
);
3868-
register_fun_stub!(
3869-
udf,
3870-
"has_table_privilege",
3871-
tsigs = [
3872-
[Utf8, Utf8],
3873-
[Oid, Utf8],
3874-
[Utf8, Utf8, Utf8],
3875-
[Utf8, Oid, Utf8],
3876-
[Oid, Utf8, Utf8],
3877-
[Oid, Oid, Utf8],
3878-
],
3879-
rettyp = Boolean,
3880-
vol = Stable
3881-
);
38823939
register_fun_stub!(
38833940
udf,
38843941
"has_tablespace_privilege",

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

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,8 @@ use self::{
5151
create_date_udf, create_dateadd_udf, create_datediff_udf, create_dayofmonth_udf,
5252
create_dayofweek_udf, create_dayofyear_udf, create_db_udf, create_ends_with_udf,
5353
create_format_type_udf, create_generate_series_udtf, create_generate_subscripts_udtf,
54-
create_has_schema_privilege_udf, create_hour_udf, create_if_udf,
55-
create_inet_server_addr_udf, create_instr_udf, create_interval_mul_udf,
54+
create_has_schema_privilege_udf, create_has_table_privilege_udf, create_hour_udf,
55+
create_if_udf, create_inet_server_addr_udf, create_instr_udf, create_interval_mul_udf,
5656
create_isnull_udf, create_json_build_object_udf, create_least_udf, create_locate_udf,
5757
create_makedate_udf, create_measure_udaf, create_minute_udf, create_pg_backend_pid_udf,
5858
create_pg_datetime_precision_udf, create_pg_encoding_to_char_udf,
@@ -1325,6 +1325,7 @@ WHERE `TABLE_SCHEMA` = '{}'",
13251325
ctx.register_udf(create_pg_my_temp_schema());
13261326
ctx.register_udf(create_pg_is_other_temp_schema());
13271327
ctx.register_udf(create_has_schema_privilege_udf(self.state.clone()));
1328+
ctx.register_udf(create_has_table_privilege_udf(self.state.clone()));
13281329
ctx.register_udf(create_pg_total_relation_size_udf());
13291330
ctx.register_udf(create_cube_regclass_cast_udf());
13301331
ctx.register_udf(create_pg_get_serial_sequence_udf());
@@ -8952,6 +8953,43 @@ limit
89528953
Ok(())
89538954
}
89548955

8956+
#[tokio::test]
8957+
async fn test_has_table_privilege_postgres() -> Result<(), CubeError> {
8958+
insta::assert_snapshot!(
8959+
"has_table_privilege",
8960+
execute_query(
8961+
"SELECT
8962+
relname,
8963+
has_table_privilege('ovr', relname, 'SELECT') \"select\",
8964+
has_table_privilege('ovr', relname, 'INSERT') \"insert\"
8965+
FROM pg_class
8966+
ORDER BY relname ASC
8967+
"
8968+
.to_string(),
8969+
DatabaseProtocol::PostgreSQL
8970+
)
8971+
.await?
8972+
);
8973+
8974+
insta::assert_snapshot!(
8975+
"has_table_privilege_default_user",
8976+
execute_query(
8977+
"SELECT
8978+
relname,
8979+
has_table_privilege(relname, 'SELECT') \"select\",
8980+
has_table_privilege(relname, 'INSERT') \"insert\"
8981+
FROM pg_class
8982+
ORDER BY relname ASC
8983+
"
8984+
.to_string(),
8985+
DatabaseProtocol::PostgreSQL
8986+
)
8987+
.await?
8988+
);
8989+
8990+
Ok(())
8991+
}
8992+
89558993
#[tokio::test]
89568994
async fn test_pg_total_relation_size() -> Result<(), CubeError> {
89578995
insta::assert_snapshot!(
@@ -21277,4 +21315,42 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
2127721315
.sql;
2127821316
assert!(sql.contains("OFFSET 1\nLIMIT 2"));
2127921317
}
21318+
21319+
#[tokio::test]
21320+
async fn test_metabase_table_privilege_query() -> Result<(), CubeError> {
21321+
insta::assert_snapshot!(
21322+
"metabase_table_privilege_query",
21323+
execute_query(
21324+
r#"
21325+
with table_privileges as (
21326+
select
21327+
NULL as role,
21328+
t.schemaname as schema,
21329+
t.objectname as table,
21330+
pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'UPDATE') as update,
21331+
pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'SELECT') as select,
21332+
pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'INSERT') as insert,
21333+
pg_catalog.has_table_privilege(current_user, '"' || t.schemaname || '"' || '.' || '"' || t.objectname || '"', 'DELETE') as delete
21334+
from (
21335+
select schemaname, tablename as objectname from pg_catalog.pg_tables
21336+
union
21337+
select schemaname, viewname as objectname from pg_catalog.pg_views
21338+
union
21339+
select schemaname, matviewname as objectname from pg_catalog.pg_matviews
21340+
) t
21341+
where t.schemaname !~ '^pg_'
21342+
and t.schemaname <> 'information_schema'
21343+
and pg_catalog.has_schema_privilege(current_user, t.schemaname, 'USAGE')
21344+
)
21345+
select t.*
21346+
from table_privileges t
21347+
order by t.schema, t.table
21348+
"#.to_string(),
21349+
DatabaseProtocol::PostgreSQL
21350+
)
21351+
.await?
21352+
);
21353+
21354+
Ok(())
21355+
}
2128021356
}

‎rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege.snap‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
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?"
4+
---
5+
+---------------------------+--------+--------+
6+
| relname | select | insert |
7+
+---------------------------+--------+--------+
8+
| KibanaSampleDataEcommerce | true | false |
9+
| Logs | true | false |
10+
| NumberCube | true | false |
11+
| WideCube | true | false |
12+
+---------------------------+--------+--------+

‎rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__has_table_privilege_default_user.snap‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
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?"
4+
---
5+
+---------------------------+--------+--------+
6+
| relname | select | insert |
7+
+---------------------------+--------+--------+
8+
| KibanaSampleDataEcommerce | true | false |
9+
| Logs | true | false |
10+
| NumberCube | true | false |
11+
| WideCube | true | false |
12+
+---------------------------+--------+--------+

‎rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_privilege_query.snap‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
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?"
4+
---
5+
+------+--------+---------------------------+--------+--------+--------+--------+
6+
| role | schema | table | update | select | insert | delete |
7+
+------+--------+---------------------------+--------+--------+--------+--------+
8+
| NULL | public | KibanaSampleDataEcommerce | false | true | false | false |
9+
| NULL | public | Logs | false | true | false | false |
10+
| NULL | public | NumberCube | false | true | false | false |
11+
| NULL | public | WideCube | false | true | false | false |
12+
+------+--------+---------------------------+--------+--------+--------+--------+

0 commit comments

Comments
 (0)