Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(cubesql): Support new Metabase meta queries #7749

Merged
merged 1 commit into from
Feb 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 71 additions & 14 deletions rust/cubesql/cubesql/src/compile/engine/udf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2824,6 +2824,77 @@
)
}

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!(

Check warning on line 2852 in rust/cubesql/cubesql/src/compile/engine/udf.rs

View check run for this annotation

Codecov / codecov/patch

rust/cubesql/cubesql/src/compile/engine/udf.rs#L2852

Added line #L2852 was not covered by tests
"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!(

Check warning on line 2867 in rust/cubesql/cubesql/src/compile/engine/udf.rs

View check run for this annotation

Codecov / codecov/patch

rust/cubesql/cubesql/src/compile/engine/udf.rs#L2867

Added line #L2867 was not covered by tests
"unrecognized privilege type: \"{}\"",
privilege
)))
}
}
}
_ => None,

Check warning on line 2874 in rust/cubesql/cubesql/src/compile/engine/udf.rs

View check run for this annotation

Codecov / codecov/patch

rust/cubesql/cubesql/src/compile/engine/udf.rs#L2874

Added line #L2874 was not covered by tests
})
})
.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 @@
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
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(())
}
}
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 |
+---------------------------+--------+--------+
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 |
+---------------------------+--------+--------+
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 |
+------+--------+---------------------------+--------+--------+--------+--------+
Loading