Skip to content

Commit 76dca65

Browse files
MazterQyouMazterQyou
committed
chore(cubesql): Support Microstrategy and Grafana introspection
1 parent 15c1548 commit 76dca65

File tree

7 files changed

+192
-15
lines changed

7 files changed

+192
-15
lines changed

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2969,3 +2969,72 @@ pub fn create_sha1_udf() -> ScalarUDF {
29692969
&fun,
29702970
)
29712971
}
2972+
2973+
pub fn create_current_setting_udf() -> ScalarUDF {
2974+
let fun = make_scalar_function(move |args: &[ArrayRef]| {
2975+
assert!(args.len() == 1);
2976+
2977+
let setting_names = downcast_string_arg!(args[0], "str", i32);
2978+
2979+
let result = setting_names
2980+
.iter()
2981+
.map(|setting_name| {
2982+
if let Some(setting_name) = setting_name {
2983+
Ok(Some(match setting_name.to_ascii_lowercase().as_str() {
2984+
"max_index_keys" => "32".to_string(), // Taken from PostgreSQL
2985+
"search_path" => "\"$user\", public".to_string(), // Taken from PostgreSQL
2986+
setting_name => Err(DataFusionError::Execution(format!(
2987+
"unrecognized configuration parameter \"{}\"",
2988+
setting_name
2989+
)))?,
2990+
}))
2991+
} else {
2992+
Ok(None)
2993+
}
2994+
})
2995+
.collect::<Result<StringArray>>()?;
2996+
2997+
Ok(Arc::new(result))
2998+
});
2999+
3000+
let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Utf8)));
3001+
3002+
ScalarUDF::new(
3003+
"current_setting",
3004+
&Signature::exact(vec![DataType::Utf8], Volatility::Stable),
3005+
&return_type,
3006+
&fun,
3007+
)
3008+
}
3009+
3010+
pub fn create_quote_ident_udf() -> ScalarUDF {
3011+
let fun = make_scalar_function(move |args: &[ArrayRef]| {
3012+
assert!(args.len() == 1);
3013+
3014+
let idents = downcast_string_arg!(args[0], "str", i32);
3015+
3016+
let re = Regex::new(r"^[a-z_][a-z0-9_]*$").unwrap();
3017+
let result = idents
3018+
.iter()
3019+
.map(|ident| {
3020+
ident.map(|ident| {
3021+
if re.is_match(ident) {
3022+
return ident.to_string();
3023+
}
3024+
format!("\"{}\"", ident.replace("\"", "\"\""))
3025+
})
3026+
})
3027+
.collect::<StringArray>();
3028+
3029+
Ok(Arc::new(result))
3030+
});
3031+
3032+
let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Utf8)));
3033+
3034+
ScalarUDF::new(
3035+
"quote_ident",
3036+
&Signature::exact(vec![DataType::Utf8], Volatility::Immutable),
3037+
&return_type,
3038+
&fun,
3039+
)
3040+
}

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

Lines changed: 75 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -43,26 +43,27 @@ use self::{
4343
udf::{
4444
create_array_lower_udf, create_array_upper_udf, create_connection_id_udf,
4545
create_convert_tz_udf, create_cube_regclass_cast_udf, create_current_schema_udf,
46-
create_current_schemas_udf, create_current_timestamp_udf, create_current_user_udf,
47-
create_date_add_udf, create_date_sub_udf, create_date_to_timestamp_udf,
48-
create_date_udf, create_dateadd_udf, create_datediff_udf, create_dayofmonth_udf,
49-
create_dayofweek_udf, create_dayofyear_udf, create_db_udf, create_ends_with_udf,
50-
create_format_type_udf, create_generate_series_udtf, create_generate_subscripts_udtf,
51-
create_has_schema_privilege_udf, create_hour_udf, create_if_udf, create_instr_udf,
52-
create_interval_mul_udf, create_isnull_udf, create_json_build_object_udf,
53-
create_least_udf, create_locate_udf, create_makedate_udf, create_measure_udaf,
54-
create_minute_udf, create_pg_backend_pid_udf, create_pg_datetime_precision_udf,
55-
create_pg_expandarray_udtf, create_pg_get_constraintdef_udf, create_pg_get_expr_udf,
46+
create_current_schemas_udf, create_current_setting_udf, create_current_timestamp_udf,
47+
create_current_user_udf, create_date_add_udf, create_date_sub_udf,
48+
create_date_to_timestamp_udf, create_date_udf, create_dateadd_udf, create_datediff_udf,
49+
create_dayofmonth_udf, create_dayofweek_udf, create_dayofyear_udf, create_db_udf,
50+
create_ends_with_udf, create_format_type_udf, create_generate_series_udtf,
51+
create_generate_subscripts_udtf, create_has_schema_privilege_udf, create_hour_udf,
52+
create_if_udf, create_instr_udf, create_interval_mul_udf, create_isnull_udf,
53+
create_json_build_object_udf, create_least_udf, create_locate_udf, create_makedate_udf,
54+
create_measure_udaf, create_minute_udf, create_pg_backend_pid_udf,
55+
create_pg_datetime_precision_udf, create_pg_expandarray_udtf,
56+
create_pg_get_constraintdef_udf, create_pg_get_expr_udf,
5657
create_pg_get_serial_sequence_udf, create_pg_get_userbyid_udf,
5758
create_pg_is_other_temp_schema, create_pg_my_temp_schema,
5859
create_pg_numeric_precision_udf, create_pg_numeric_scale_udf,
5960
create_pg_table_is_visible_udf, create_pg_total_relation_size_udf,
6061
create_pg_truetypid_udf, create_pg_truetypmod_udf, create_pg_type_is_visible_udf,
61-
create_position_udf, create_quarter_udf, create_regexp_substr_udf, create_second_udf,
62-
create_session_user_udf, create_sha1_udf, create_str_to_date_udf,
63-
create_time_format_udf, create_timediff_udf, create_to_char_udf, create_to_date_udf,
64-
create_ucase_udf, create_unnest_udtf, create_user_udf, create_version_udf,
65-
create_year_udf,
62+
create_position_udf, create_quarter_udf, create_quote_ident_udf,
63+
create_regexp_substr_udf, create_second_udf, create_session_user_udf, create_sha1_udf,
64+
create_str_to_date_udf, create_time_format_udf, create_timediff_udf,
65+
create_to_char_udf, create_to_date_udf, create_ucase_udf, create_unnest_udtf,
66+
create_user_udf, create_version_udf, create_year_udf,
6667
},
6768
},
6869
parser::parse_sql_to_statement,
@@ -1174,6 +1175,8 @@ WHERE `TABLE_SCHEMA` = '{}'",
11741175
ctx.register_udf(create_date_to_timestamp_udf());
11751176
ctx.register_udf(create_to_date_udf());
11761177
ctx.register_udf(create_sha1_udf());
1178+
ctx.register_udf(create_current_setting_udf());
1179+
ctx.register_udf(create_quote_ident_udf());
11771180

11781181
// udaf
11791182
ctx.register_udaf(create_measure_udaf());
@@ -9308,6 +9311,63 @@ ORDER BY \"COUNT(count)\" DESC"
93089311
Ok(())
93099312
}
93109313

9314+
#[tokio::test]
9315+
async fn test_current_setting() -> Result<(), CubeError> {
9316+
insta::assert_snapshot!(
9317+
"current_setting",
9318+
execute_query(
9319+
"SELECT current_setting('max_index_keys'), current_setting('search_path')"
9320+
.to_string(),
9321+
DatabaseProtocol::PostgreSQL
9322+
)
9323+
.await?
9324+
);
9325+
9326+
Ok(())
9327+
}
9328+
9329+
#[tokio::test]
9330+
async fn test_quote_ident() -> Result<(), CubeError> {
9331+
insta::assert_snapshot!(
9332+
"quote_ident",
9333+
execute_query(
9334+
"SELECT quote_ident('pg_catalog') i1, quote_ident('Foo bar') i2".to_string(),
9335+
DatabaseProtocol::PostgreSQL
9336+
)
9337+
.await?
9338+
);
9339+
9340+
Ok(())
9341+
}
9342+
9343+
#[tokio::test]
9344+
async fn test_subquery_current_schema() -> Result<(), CubeError> {
9345+
insta::assert_snapshot!(
9346+
"microstrategy_subquery_current_schema",
9347+
execute_query(
9348+
"SELECT t.oid FROM pg_catalog.pg_type AS t JOIN pg_catalog.pg_namespace AS n ON t.typnamespace = n.oid WHERE t.typname = 'citext' AND (n.nspname = (SELECT current_schema()) OR n.nspname = 'public')".to_string(),
9349+
DatabaseProtocol::PostgreSQL
9350+
)
9351+
.await?
9352+
);
9353+
9354+
Ok(())
9355+
}
9356+
9357+
#[tokio::test]
9358+
async fn test_insubquery_where_tables() -> Result<(), CubeError> {
9359+
insta::assert_snapshot!(
9360+
"grafana_insubquery_where_tables",
9361+
execute_query(
9362+
r#"SELECT quote_ident(table_name) AS "table" FROM information_schema.tables WHERE quote_ident(table_schema) NOT IN ('information_schema', 'pg_catalog', '_timescaledb_cache', '_timescaledb_catalog', '_timescaledb_internal', '_timescaledb_config', 'timescaledb_information', 'timescaledb_experimental') AND table_type = 'BASE TABLE' AND quote_ident(table_schema) IN (SELECT CASE WHEN TRIM(s[i]) = '"$user"' THEN user ELSE TRIM(s[i]) END FROM generate_series(array_lower(string_to_array(current_setting('search_path'), ','), 1), array_upper(string_to_array(current_setting('search_path'), ','), 1)) AS i, string_to_array(current_setting('search_path'), ',') AS s)"#.to_string(),
9363+
DatabaseProtocol::PostgreSQL
9364+
)
9365+
.await?
9366+
);
9367+
9368+
Ok(())
9369+
}
9370+
93119371
#[tokio::test]
93129372
async fn test_metabase_substring() -> Result<(), CubeError> {
93139373
init_logger();

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,17 @@ pub fn parse_sql_to_statements(
244244
"",
245245
);
246246

247+
// Microstrategy
248+
// TODO: Support Subquery Node
249+
let query = query.replace("= (SELECT current_schema())", "= current_schema()");
250+
251+
// Grafana
252+
// TODO: Support InSubquery Node
253+
let query = query.replace(
254+
"WHERE quote_ident(table_schema) NOT IN ('information_schema', 'pg_catalog', '_timescaledb_cache', '_timescaledb_catalog', '_timescaledb_internal', '_timescaledb_config', 'timescaledb_information', 'timescaledb_experimental') AND table_type = 'BASE TABLE' AND quote_ident(table_schema) IN (SELECT CASE WHEN TRIM(s[i]) = '\"$user\"' THEN user ELSE TRIM(s[i]) END FROM generate_series(array_lower(string_to_array(current_setting('search_path'), ','), 1), array_upper(string_to_array(current_setting('search_path'), ','), 1)) AS i, string_to_array(current_setting('search_path'), ',') AS s)",
255+
"WHERE quote_ident(table_schema) IN (current_user, current_schema()) AND table_type = 'BASE TABLE'"
256+
);
257+
247258
if let Some(qtrace) = qtrace {
248259
qtrace.set_replaced_query(&query)
249260
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
expression: "execute_query(\"SELECT current_setting('max_index_keys'), current_setting('search_path')\".to_string(),\n DatabaseProtocol::PostgreSQL).await?"
4+
---
5+
+-----------------------------------------+--------------------------------------+
6+
| current_setting(Utf8("max_index_keys")) | current_setting(Utf8("search_path")) |
7+
+-----------------------------------------+--------------------------------------+
8+
| 32 | "$user", public |
9+
+-----------------------------------------+--------------------------------------+

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
expression: "execute_query(r#\"SELECT quote_ident(table_name) AS \"table\" FROM information_schema.tables WHERE quote_ident(table_schema) NOT IN ('information_schema', 'pg_catalog', '_timescaledb_cache', '_timescaledb_catalog', '_timescaledb_internal', '_timescaledb_config', 'timescaledb_information', 'timescaledb_experimental') AND table_type = 'BASE TABLE' AND quote_ident(table_schema) IN (SELECT CASE WHEN TRIM(s[i]) = '\"$user\"' THEN user ELSE TRIM(s[i]) END FROM generate_series(array_lower(string_to_array(current_setting('search_path'), ','), 1), array_upper(string_to_array(current_setting('search_path'), ','), 1)) AS i, string_to_array(current_setting('search_path'), ',') AS s)\"#.to_string(),\n DatabaseProtocol::PostgreSQL).await?"
4+
---
5+
+-----------------------------+
6+
| table |
7+
+-----------------------------+
8+
| "KibanaSampleDataEcommerce" |
9+
| "Logs" |
10+
| "NumberCube" |
11+
+-----------------------------+

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
expression: "execute_query(\"SELECT t.oid FROM pg_catalog.pg_type AS t JOIN pg_catalog.pg_namespace AS n ON t.typnamespace = n.oid WHERE t.typname = 'citext' AND (n.nspname = (SELECT current_schema()) OR n.nspname = 'public')\".to_string(),\n DatabaseProtocol::PostgreSQL).await?"
4+
---
5+
+-----+
6+
| oid |
7+
+-----+
8+
+-----+

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
expression: "execute_query(\"SELECT quote_ident('pg_catalog') i1, quote_ident('Foo bar') i2\".to_string(),\n DatabaseProtocol::PostgreSQL).await?"
4+
---
5+
+------------+-----------+
6+
| i1 | i2 |
7+
+------------+-----------+
8+
| pg_catalog | "Foo bar" |
9+
+------------+-----------+

0 commit comments

Comments
 (0)