Skip to content

Commit

Permalink
feat(cubesql): Workaround for Metabase introspection query
Browse files Browse the repository at this point in the history
  • Loading branch information
@ovr
ovr committed Jun 21, 2022
1 parent 904171e commit ee7b3cf
Show file tree
Hide file tree
Showing 5 changed files with 60 additions and 11 deletions.
42 changes: 41 additions & 1 deletion rust/cubesql/cubesql/src/compile/mod.rs
View file
Expand Up @@ -1445,6 +1445,13 @@ struct QueryPlanner {
session_manager: Arc<SessionManager>,
}

lazy_static! {
static ref METABASE_WORKAROUND: regex::Regex = regex::Regex::new(
r#"SELECT true AS "_" FROM "public"\."(?P<tblname>.*)" WHERE 1 <> 1 LIMIT 0"#
)
.unwrap();
}

impl QueryPlanner {
pub fn new(
state: Arc<SessionState>,
Expand All @@ -1466,6 +1473,24 @@ impl QueryPlanner {
stmt: &ast::Statement,
q: &Box<ast::Query>,
) -> CompilationResult<QueryPlan> {
// Metabase
if let Some(c) = METABASE_WORKAROUND.captures(&stmt.to_string()) {
let tblname = c.name("tblname").unwrap().as_str();
if self.meta.find_cube_with_name(tblname).is_some() {
return Ok(QueryPlan::MetaTabular(
StatusFlags::empty(),
Box::new(dataframe::DataFrame::new(
vec![dataframe::Column::new(
"_".to_string(),
ColumnType::Int8,
ColumnFlags::empty(),
)],
vec![],
)),
));
}
}

// TODO move CUBESQL_REWRITE_ENGINE env to config
let rewrite_engine = env::var("CUBESQL_REWRITE_ENGINE")
.ok()
Expand Down Expand Up @@ -1619,7 +1644,7 @@ impl QueryPlanner {
));
};

if let Some(cube) = self.meta.find_cube_with_name(table_name.clone()) {
if let Some(cube) = self.meta.find_cube_with_name(&table_name) {
let mut ctx = QueryContext::new(&cube);
let mut builder = compile_select(select, &mut ctx)?;

Expand Down Expand Up @@ -8696,6 +8721,21 @@ ORDER BY \"COUNT(count)\" DESC"
Ok(())
}

#[tokio::test]
async fn test_metabase_table_exists() -> Result<(), CubeError> {
insta::assert_snapshot!(
"metabase_table_exists",
execute_query(
r#"SELECT TRUE AS "_" FROM "public"."KibanaSampleDataEcommerce" WHERE 1 <> 1 LIMIT 0;"#
.to_string(),
DatabaseProtocol::PostgreSQL,
)
.await?
);

Ok(())
}

#[tokio::test]
async fn test_subquery_with_same_name_excel() -> Result<(), CubeError> {
insta::assert_snapshot!(
Expand Down
2 changes: 1 addition & 1 deletion rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs
View file
Expand Up @@ -1279,7 +1279,7 @@ impl FilterRules {
for cube in var_iter!(egraph[subst[cube_var]], FilterReplacerCube) {
if let Some(cube) = cube
.as_ref()
.and_then(|cube| meta_context.find_cube_with_name(cube.to_string()))
.and_then(|cube| meta_context.find_cube_with_name(cube))
{
for column in var_iter!(egraph[subst[column_var]], ColumnExprColumn).cloned() {
for table_name in
Expand Down
12 changes: 6 additions & 6 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/split.rs
View file
Expand Up @@ -570,7 +570,7 @@ impl SplitRules {
for fun in var_iter!(egraph[subst[fun_expr_var]], AggregateFunctionExprFun).cloned()
{
if fun == AggregateFunction::Min || fun == AggregateFunction::Max {
if let Some(cube) = meta.find_cube_with_name(cube.to_string()) {
if let Some(cube) = meta.find_cube_with_name(&cube) {
for column in var_iter!(egraph[subst[column_var]], ColumnExprColumn) {
if let Some(dimension) = cube.lookup_dimension(&column.name) {
if is_time_dimension && dimension._type == "time"
Expand Down Expand Up @@ -614,7 +614,7 @@ impl SplitRules {
)
.cloned()
{
if let Some(cube) = meta.find_cube_with_name(cube) {
if let Some(cube) = meta.find_cube_with_name(&cube) {
for column in column_var
.map(|column_var| {
var_iter!(egraph[subst[column_var]], ColumnExprColumn)
Expand Down Expand Up @@ -646,7 +646,7 @@ impl SplitRules {
)
.cloned()
{
if let Some(cube) = meta.find_cube_with_name(cube) {
if let Some(cube) = meta.find_cube_with_name(&cube) {
if cube
.lookup_measure(&MemberRules::default_count_measure_name())
.is_none()
Expand Down Expand Up @@ -769,7 +769,7 @@ impl SplitRules {
var_iter!(egraph[subst[cube_var]], OuterAggregateSplitReplacerCube).cloned()
{
if let Some(name) = original_expr_name(egraph, subst[original_expr_var]) {
if let Some(cube) = meta.find_cube_with_name(cube) {
if let Some(cube) = meta.find_cube_with_name(&cube) {
for column in
var_iter!(egraph[subst[column_var]], ColumnExprColumn).cloned()
{
Expand Down Expand Up @@ -830,7 +830,7 @@ impl SplitRules {
for cube in
var_iter!(egraph[subst[cube_var]], OuterAggregateSplitReplacerCube).cloned()
{
if let Some(cube) = meta.find_cube_with_name(cube) {
if let Some(cube) = meta.find_cube_with_name(&cube) {
if cube
.lookup_measure(&MemberRules::default_count_measure_name())
.is_none()
Expand Down Expand Up @@ -861,7 +861,7 @@ impl SplitRules {
for cube in var_iter!(egraph[subst[cube_var]], OuterAggregateSplitReplacerCube).cloned()
{
if let Some(name) = original_expr_name(egraph, subst[original_expr_var]) {
if let Some(cube) = meta.find_cube_with_name(cube) {
if let Some(cube) = meta.find_cube_with_name(&cube) {
for column in
var_iter!(egraph[subst[column_var]], ColumnExprColumn).cloned()
{
Expand Down
9 changes: 9 additions & 0 deletions ...cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__metabase_table_exists.snap
View file
@@ -0,0 +1,9 @@
---
source: cubesql/src/compile/mod.rs
assertion_line: 8726
expression: "execute_query(r#\"SELECT TRUE AS \"_\" FROM \"public\".\"KibanaSampleDataEcommerce\" WHERE 1 <> 1 LIMIT 0;\"#.to_string(),\n DatabaseProtocol::PostgreSQL).await?"
---
+---+
| _ |
+---+
+---+
6 changes: 3 additions & 3 deletions rust/cubesql/cubesql/src/transport/ctx.rs
View file
Expand Up @@ -57,7 +57,7 @@ impl MetaContext {
Self { cubes, tables }
}

pub fn find_cube_with_name(&self, name: String) -> Option<V1CubeMeta> {
pub fn find_cube_with_name(&self, name: &str) -> Option<V1CubeMeta> {
for cube in self.cubes.iter() {
if cube.name.eq(&name) {
return Some(cube.clone());
Expand All @@ -69,15 +69,15 @@ impl MetaContext {

pub fn find_measure_with_name(&self, name: String) -> Option<V1CubeMetaMeasure> {
let cube_and_member_name = name.split(".").collect::<Vec<_>>();
if let Some(cube) = self.find_cube_with_name(cube_and_member_name[0].to_string()) {
if let Some(cube) = self.find_cube_with_name(cube_and_member_name[0]) {
cube.lookup_measure(cube_and_member_name[1]).cloned()
} else {
None
}
}

pub fn find_df_data_type(&self, member_name: String) -> Option<DataType> {
self.find_cube_with_name(member_name.split(".").next()?.to_string())?
self.find_cube_with_name(member_name.split(".").next()?)?
.df_data_type(member_name.as_str())
}

Expand Down

0 comments on commit ee7b3cf

Please sign in to comment.