Merge pull request #7600 from Xuanwo/make-planner-depends-on-table-context

refactor: Make planner depends on TableContext trait

Author: mergify[bot]

src/query/catalog/src/table_context.rs

@@ -111,5 +111,7 @@ pub trait TableContext: Send + Sync {
     // Get user manager api.
     fn get_user_manager(&self) -> Arc<UserApiProvider>;
     fn get_cluster(&self) -> Arc<Cluster>;
+    async fn get_table(&self, catalog: &str, database: &str, table: &str)
+    -> Result<Arc<dyn Table>>;
     async fn get_processes_info(&self) -> Vec<ProcessInfo>;
 }

src/query/service/src/api/http/v1/logs.rs

@@ -25,6 +25,7 @@ use tokio_stream::StreamExt;
 use crate::sessions::QueryContext;
 use crate::sessions::SessionManager;
 use crate::sessions::SessionType;
+use crate::sessions::TableContext;
 use crate::storages::TableStreamReadWrap;
 use crate::storages::ToReadDataSourcePlan;
 

src/query/service/src/interpreters/interpreter_table_describe.rs

@@ -24,6 +24,7 @@ use common_streams::SendableDataBlockStream;
 
 use crate::interpreters::Interpreter;
 use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
 use crate::sql::executor::PhysicalScalar;
 use crate::sql::PlanParser;
 use crate::storages::view::view_table::QUERY;

src/query/service/src/interpreters/interpreter_table_exists.rs

@@ -24,6 +24,7 @@ use common_streams::SendableDataBlockStream;
 
 use crate::interpreters::Interpreter;
 use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
 
 pub struct ExistsTableInterpreter {
     ctx: Arc<QueryContext>,

src/query/service/src/interpreters/interpreter_table_truncate.rs

@@ -23,6 +23,7 @@ use common_streams::SendableDataBlockStream;
 
 use crate::interpreters::Interpreter;
 use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
 
 pub struct TruncateTableInterpreter {
     ctx: Arc<QueryContext>,

src/query/service/src/optimizers/optimizer_statistics_exact.rs

@@ -27,6 +27,7 @@ use common_legacy_planners::PlanRewriter;
 
 use crate::optimizers::Optimizer;
 use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
 use crate::storages::ToReadDataSourcePlan;
 
 struct StatisticsExactImpl<'a> {

src/query/service/src/sessions/query_ctx.rs

@@ -125,22 +125,6 @@ impl QueryContext {
         StageTable::try_create(table_info.clone())
     }
 
-    /// Fetch a Table by db and table name.
-    ///
-    /// It guaranteed to return a consistent result for multiple calls, in a same query.
-    /// E.g.:
-    /// ```sql
-    /// SELECT * FROM (SELECT * FROM db.table_name) as subquery_1, (SELECT * FROM db.table_name) AS subquery_2
-    /// ```
-    pub async fn get_table(
-        &self,
-        catalog: &str,
-        database: &str,
-        table: &str,
-    ) -> Result<Arc<dyn Table>> {
-        self.shared.get_table(catalog, database, table).await
-    }
-
     pub async fn set_current_database(&self, new_database_name: String) -> Result<()> {
         let tenant_id = self.get_tenant();
         let catalog = self.get_catalog(self.get_current_catalog().as_str())?;
@@ -418,6 +402,22 @@ impl TableContext for QueryContext {
         self.shared.get_cluster()
     }
 
+    /// Fetch a Table by db and table name.
+    ///
+    /// It guaranteed to return a consistent result for multiple calls, in a same query.
+    /// E.g.:
+    /// ```sql
+    /// SELECT * FROM (SELECT * FROM db.table_name) as subquery_1, (SELECT * FROM db.table_name) AS subquery_2
+    /// ```
+    async fn get_table(
+        &self,
+        catalog: &str,
+        database: &str,
+        table: &str,
+    ) -> Result<Arc<dyn Table>> {
+        self.shared.get_table(catalog, database, table).await
+    }
+
     // Get all the processes list info.
     async fn get_processes_info(&self) -> Vec<ProcessInfo> {
         SessionManager::instance().processes_info().await

src/query/service/src/sql/optimizer/heuristic/mod.rs

@@ -26,7 +26,7 @@ use once_cell::sync::Lazy;
 
 use super::rule::RuleID;
 use super::ColumnSet;
-use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
 use crate::sql::optimizer::heuristic::decorrelate::decorrelate_subquery;
 use crate::sql::optimizer::heuristic::implement::HeuristicImplementor;
 pub use crate::sql::optimizer::heuristic::rule_list::RuleList;
@@ -65,14 +65,14 @@ pub struct HeuristicOptimizer {
     rules: RuleList,
     implementor: HeuristicImplementor,
 
-    _ctx: Arc<QueryContext>,
+    _ctx: Arc<dyn TableContext>,
     bind_context: Box<BindContext>,
     metadata: MetadataRef,
 }
 
 impl HeuristicOptimizer {
     pub fn new(
-        ctx: Arc<QueryContext>,
+        ctx: Arc<dyn TableContext>,
         bind_context: Box<BindContext>,
         metadata: MetadataRef,
         rules: RuleList,

src/query/service/src/sql/optimizer/mod.rs

@@ -50,7 +50,7 @@ use self::util::contains_local_table_scan;
 use self::util::validate_distributed_query;
 use super::plans::Plan;
 use super::BindContext;
-use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
 pub use crate::sql::optimizer::heuristic::RuleList;
 pub use crate::sql::optimizer::rule::RuleID;
 use crate::sql::optimizer::rule::RuleSet;
@@ -74,7 +74,7 @@ impl OptimizerContext {
 }
 
 pub fn optimize(
-    ctx: Arc<QueryContext>,
+    ctx: Arc<dyn TableContext>,
     opt_ctx: Arc<OptimizerContext>,
     plan: Plan,
 ) -> Result<Plan> {
@@ -130,7 +130,7 @@ pub fn optimize(
 }
 
 pub fn optimize_query(
-    ctx: Arc<QueryContext>,
+    ctx: Arc<dyn TableContext>,
     opt_ctx: Arc<OptimizerContext>,
     metadata: MetadataRef,
     bind_context: Box<BindContext>,

src/query/service/src/sql/planner/binder/copy.rs

@@ -31,7 +31,6 @@ use common_meta_types::UserStageInfo;
 use common_storage::parse_uri_location;
 use common_storage::UriLocation;
 
-use crate::sessions::TableContext;
 use crate::sql::binder::Binder;
 use crate::sql::normalize_identifier;
 use crate::sql::plans::CopyPlanV2;