Skip to content

Commit 547a41b

Browse files
ovrovr
authored
chore(cubesql): Cleanup code (#8592)
1 parent 733df2c commit 547a41b

File tree

4 files changed

+27
-174
lines changed

4 files changed

+27
-174
lines changed

‎rust/cubesql/cubesql/benches/benchmarks.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,19 @@
11
use criterion::{criterion_group, criterion_main, Criterion};
22
use cubesql::compile::test::{
33
get_test_tenant_ctx,
4-
rewrite_engine::{cube_context, query_to_logical_plan, rewrite_rules, rewrite_runner},
4+
rewrite_engine::{
5+
create_test_postgresql_cube_context, query_to_logical_plan, rewrite_rules, rewrite_runner,
6+
},
57
};
68
use itertools::Itertools;
79
use std::sync::Arc;
810

911
macro_rules! bench_func {
1012
($NAME:expr, $QUERY:expr, $CRITERION:expr) => {{
11-
let context = Arc::new(futures::executor::block_on(cube_context(
12-
get_test_tenant_ctx(),
13-
)));
13+
let context = Arc::new(
14+
futures::executor::block_on(create_test_postgresql_cube_context(get_test_tenant_ctx()))
15+
.unwrap(),
16+
);
1417
let plan = query_to_logical_plan($QUERY, &context);
1518
let rules = rewrite_rules(context.clone());
1619

‎rust/cubesql/cubesql/benches/large_model.rs‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@ use criterion::{criterion_group, criterion_main, Criterion};
55
use cubeclient::models::{V1CubeMeta, V1CubeMetaDimension, V1CubeMetaMeasure};
66
use cubesql::{
77
compile::test::{
8-
rewrite_engine::{cube_context, query_to_logical_plan, rewrite_rules, rewrite_runner},
8+
rewrite_engine::{
9+
create_test_postgresql_cube_context, query_to_logical_plan, rewrite_rules,
10+
rewrite_runner,
11+
},
912
sql_generator,
1013
},
1114
transport::MetaContext,
@@ -16,9 +19,13 @@ use uuid::Uuid;
1619

1720
macro_rules! bench_large_model {
1821
($DIMS:expr, $NAME:expr, $QUERY_FN:expr, $CRITERION:expr) => {{
19-
let context = Arc::new(futures::executor::block_on(cube_context(
20-
get_large_model_test_tenant_ctx($DIMS),
21-
)));
22+
let context = Arc::new(
23+
futures::executor::block_on(create_test_cube_context(
24+
create_test_postgresql_cube_context($DIMS),
25+
))
26+
.unwrap(),
27+
);
28+
2229
let plan = query_to_logical_plan($QUERY_FN($DIMS), &context);
2330
let rules = rewrite_rules(context.clone());
2431

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

Lines changed: 0 additions & 153 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc, time::Syste
77

88
use crate::{
99
compile::{
10-
engine::{df::planner::CubeQueryPlanner, udf::*, VariablesProvider},
1110
error::{CompilationError, CompilationResult},
1211
parser::parse_sql_to_statement,
1312
DatabaseVariable, DatabaseVariablesToUpdate,
@@ -23,16 +22,11 @@ use crate::{
2322
transport::{MetaContext, SpanId},
2423
};
2524
use datafusion::{
26-
execution::context::{
27-
default_session_builder, SessionConfig as DFSessionConfig,
28-
SessionContext as DFSessionContext,
29-
},
3025
logical_plan::{
3126
plan::{Analyze, Explain, ToStringifiedPlan},
3227
LogicalPlan, PlanType, ToDFSchema,
3328
},
3429
scalar::ScalarValue,
35-
variable::VarType,
3630
};
3731
use itertools::Itertools;
3832
use sqlparser::ast::{escape_single_quote_string, ObjectName};
@@ -671,153 +665,6 @@ impl QueryRouter {
671665
Ok(QueryPlan::MetaOk(flags, CommandCompletion::DropTable))
672666
}
673667

674-
pub fn create_execution_ctx(&self) -> DFSessionContext {
675-
let query_planner = Arc::new(CubeQueryPlanner::new(
676-
self.session_manager.server.transport.clone(),
677-
self.state.get_load_request_meta(),
678-
self.session_manager.server.config_obj.clone(),
679-
));
680-
let mut ctx = DFSessionContext::with_state(
681-
default_session_builder(
682-
DFSessionConfig::new()
683-
.create_default_catalog_and_schema(false)
684-
.with_information_schema(false)
685-
.with_default_catalog_and_schema("db", "public"),
686-
)
687-
.with_query_planner(query_planner),
688-
);
689-
690-
if self.state.protocol == DatabaseProtocol::MySQL {
691-
let system_variable_provider =
692-
VariablesProvider::new(self.state.clone(), self.session_manager.server.clone());
693-
let user_defined_variable_provider =
694-
VariablesProvider::new(self.state.clone(), self.session_manager.server.clone());
695-
696-
ctx.register_variable(VarType::System, Arc::new(system_variable_provider));
697-
ctx.register_variable(
698-
VarType::UserDefined,
699-
Arc::new(user_defined_variable_provider),
700-
);
701-
}
702-
703-
// udf
704-
if self.state.protocol == DatabaseProtocol::MySQL {
705-
ctx.register_udf(create_version_udf("8.0.25".to_string()));
706-
ctx.register_udf(create_db_udf("database".to_string(), self.state.clone()));
707-
ctx.register_udf(create_db_udf("schema".to_string(), self.state.clone()));
708-
ctx.register_udf(create_current_user_udf(
709-
self.state.clone(),
710-
"current_user",
711-
true,
712-
));
713-
ctx.register_udf(create_user_udf(self.state.clone()));
714-
} else if self.state.protocol == DatabaseProtocol::PostgreSQL {
715-
ctx.register_udf(create_version_udf(
716-
"PostgreSQL 14.1 on x86_64-cubesql".to_string(),
717-
));
718-
ctx.register_udf(create_db_udf(
719-
"current_database".to_string(),
720-
self.state.clone(),
721-
));
722-
ctx.register_udf(create_db_udf(
723-
"current_schema".to_string(),
724-
self.state.clone(),
725-
));
726-
ctx.register_udf(create_current_user_udf(
727-
self.state.clone(),
728-
"current_user",
729-
false,
730-
));
731-
ctx.register_udf(create_current_user_udf(self.state.clone(), "user", false));
732-
ctx.register_udf(create_session_user_udf(self.state.clone()));
733-
}
734-
735-
ctx.register_udf(create_connection_id_udf(self.state.clone()));
736-
ctx.register_udf(create_pg_backend_pid_udf(self.state.clone()));
737-
ctx.register_udf(create_instr_udf());
738-
ctx.register_udf(create_ucase_udf());
739-
ctx.register_udf(create_isnull_udf());
740-
ctx.register_udf(create_if_udf());
741-
ctx.register_udf(create_least_udf());
742-
ctx.register_udf(create_greatest_udf());
743-
ctx.register_udf(create_convert_tz_udf());
744-
ctx.register_udf(create_timediff_udf());
745-
ctx.register_udf(create_time_format_udf());
746-
ctx.register_udf(create_locate_udf());
747-
ctx.register_udf(create_date_udf());
748-
ctx.register_udf(create_makedate_udf());
749-
ctx.register_udf(create_year_udf());
750-
ctx.register_udf(create_quarter_udf());
751-
ctx.register_udf(create_hour_udf());
752-
ctx.register_udf(create_minute_udf());
753-
ctx.register_udf(create_second_udf());
754-
ctx.register_udf(create_dayofweek_udf());
755-
ctx.register_udf(create_dayofmonth_udf());
756-
ctx.register_udf(create_dayofyear_udf());
757-
ctx.register_udf(create_date_sub_udf());
758-
ctx.register_udf(create_date_add_udf());
759-
ctx.register_udf(create_str_to_date_udf());
760-
ctx.register_udf(create_current_timestamp_udf("current_timestamp"));
761-
ctx.register_udf(create_current_timestamp_udf("localtimestamp"));
762-
ctx.register_udf(create_current_schema_udf());
763-
ctx.register_udf(create_current_schemas_udf());
764-
ctx.register_udf(create_format_type_udf());
765-
ctx.register_udf(create_pg_datetime_precision_udf());
766-
ctx.register_udf(create_pg_numeric_precision_udf());
767-
ctx.register_udf(create_pg_numeric_scale_udf());
768-
ctx.register_udf(create_pg_get_userbyid_udf(self.state.clone()));
769-
ctx.register_udf(create_pg_get_expr_udf());
770-
ctx.register_udf(create_pg_table_is_visible_udf());
771-
ctx.register_udf(create_pg_type_is_visible_udf());
772-
ctx.register_udf(create_pg_get_constraintdef_udf());
773-
ctx.register_udf(create_pg_truetypid_udf());
774-
ctx.register_udf(create_pg_truetypmod_udf());
775-
ctx.register_udf(create_to_char_udf());
776-
ctx.register_udf(create_array_lower_udf());
777-
ctx.register_udf(create_array_upper_udf());
778-
ctx.register_udf(create_pg_my_temp_schema());
779-
ctx.register_udf(create_pg_is_other_temp_schema());
780-
ctx.register_udf(create_has_schema_privilege_udf(self.state.clone()));
781-
ctx.register_udf(create_has_table_privilege_udf(self.state.clone()));
782-
ctx.register_udf(create_has_any_column_privilege_udf(self.state.clone()));
783-
ctx.register_udf(create_pg_total_relation_size_udf());
784-
ctx.register_udf(create_cube_regclass_cast_udf());
785-
ctx.register_udf(create_pg_get_serial_sequence_udf());
786-
ctx.register_udf(create_json_build_object_udf());
787-
ctx.register_udf(create_regexp_substr_udf());
788-
ctx.register_udf(create_ends_with_udf());
789-
ctx.register_udf(create_position_udf());
790-
ctx.register_udf(create_date_to_timestamp_udf());
791-
ctx.register_udf(create_to_date_udf());
792-
ctx.register_udf(create_sha1_udf());
793-
ctx.register_udf(create_current_setting_udf());
794-
ctx.register_udf(create_quote_ident_udf());
795-
ctx.register_udf(create_pg_encoding_to_char_udf());
796-
ctx.register_udf(create_array_to_string_udf());
797-
ctx.register_udf(create_charindex_udf());
798-
ctx.register_udf(create_to_regtype_udf());
799-
ctx.register_udf(create_pg_get_indexdef_udf());
800-
ctx.register_udf(create_inet_server_addr_udf());
801-
802-
// udaf
803-
ctx.register_udaf(create_measure_udaf());
804-
805-
// udtf
806-
ctx.register_udtf(create_generate_series_udtf());
807-
ctx.register_udtf(create_unnest_udtf());
808-
ctx.register_udtf(create_generate_subscripts_udtf());
809-
ctx.register_udtf(create_pg_expandarray_udtf());
810-
811-
// redshift
812-
ctx.register_udf(create_datediff_udf());
813-
ctx.register_udf(create_dateadd_udf());
814-
815-
// fn stubs
816-
ctx = register_fun_stubs(ctx);
817-
818-
ctx
819-
}
820-
821668
async fn reauthenticate_if_needed(&self) -> CompilationResult<()> {
822669
if self.state.is_auth_context_expired() {
823670
let authenticate_response = self

‎rust/cubesql/cubesql/src/compile/test/rewrite_engine.rs‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,27 +16,23 @@ use crate::{
1616
rewriter::Rewriter,
1717
LogicalPlanLanguage,
1818
},
19-
rewrite_statement, CubeContext, DatabaseProtocol, QueryRouter,
19+
rewrite_statement, CompilationError, CubeContext, DatabaseProtocol, QueryEngine,
20+
SqlQueryEngine,
2021
},
2122
config::{ConfigObj, ConfigObjImpl},
2223
transport::MetaContext,
2324
};
2425

25-
pub async fn cube_context(meta: Arc<MetaContext>) -> CubeContext {
26+
pub async fn create_test_postgresql_cube_context(
27+
meta: Arc<MetaContext>,
28+
) -> Result<CubeContext, CompilationError> {
2629
let session = get_test_session(DatabaseProtocol::PostgreSQL, meta.clone()).await;
27-
let planner = QueryRouter::new(
28-
session.state.clone(),
29-
meta.clone(),
30-
session.session_manager.clone(),
31-
);
32-
let ctx = planner.create_execution_ctx();
33-
let df_state = Arc::new(ctx.state.write().clone());
30+
let query_engine = SqlQueryEngine::new(session.session_manager.clone());
3431

35-
CubeContext::new(
36-
df_state,
37-
meta.clone(),
38-
session.session_manager.clone(),
32+
query_engine.create_cube_ctx(
3933
session.state.clone(),
34+
meta,
35+
query_engine.create_session_ctx(session.state.clone())?,
4036
)
4137
}
4238

0 commit comments

Comments
 (0)