Merge pull request #56531 from kitaisreal/analyzer-log-used-row-policies

Analyzer log used row policies

src/Analyzer/QueryTreePassManager.cpp

@@ -166,7 +166,6 @@ private:
   * TODO: Support GROUP BY injective function elimination.
   * TODO: Support setting optimize_move_functions_out_of_any.
   * TODO: Support setting optimize_aggregators_of_group_by_keys.
-  * TODO: Support setting optimize_duplicate_order_by_and_distinct.
   * TODO: Support setting optimize_monotonous_functions_in_order_by.
   * TODO: Add optimizations based on function semantics. Example: SELECT * FROM test_table WHERE id != id. (id is not nullable column).
   */

src/Interpreters/InterpreterSelectQueryAnalyzer.cpp

@@ -257,4 +257,10 @@ void InterpreterSelectQueryAnalyzer::addStorageLimits(const StorageLimitsList &
     planner.addStorageLimits(storage_limits);
 }
 
+void InterpreterSelectQueryAnalyzer::extendQueryLogElemImpl(QueryLogElement & elem, const ASTPtr & /*ast*/, ContextPtr /*context*/) const
+{
+    for (const auto & used_row_policy : planner.getUsedRowPolicies())
+        elem.used_row_policies.emplace(used_row_policy);
+}
+
 }

src/Interpreters/InterpreterSelectQueryAnalyzer.h

@@ -57,14 +57,16 @@ public:
 
     void addStorageLimits(const StorageLimitsList & storage_limits);
 
+    void extendQueryLogElemImpl(QueryLogElement & elem, const ASTPtr & /*ast*/, ContextPtr /*context*/) const override;
+
     bool supportsTransactions() const override { return true; }
 
     bool ignoreLimits() const override { return select_query_options.ignore_limits; }
 
     bool ignoreQuota() const override { return select_query_options.ignore_quota; }
 
-
     const Planner & getPlanner() const { return planner; }
+
     Planner & getPlanner() { return planner; }
 
     const QueryTreeNodePtr & getQueryTree() const { return query_tree; }

src/Planner/Planner.cpp

@@ -1201,6 +1201,8 @@ void Planner::buildPlanForUnionNode()
     {
         Planner query_planner(query_node, select_query_options);
         query_planner.buildQueryPlanIfNeeded();
+        for (const auto & row_policy : query_planner.getUsedRowPolicies())
+            used_row_policies.insert(row_policy);
         auto query_node_plan = std::make_unique<QueryPlan>(std::move(query_planner).extractQueryPlan());
         query_plans_headers.push_back(query_node_plan->getCurrentDataStream().header);
         query_plans.push_back(std::move(query_node_plan));
@@ -1348,8 +1350,10 @@ void Planner::buildPlanForQueryNode()
         select_query_options,
         top_level_identifiers,
         planner_context);
+
     auto from_stage = join_tree_query_plan.from_stage;
     query_plan = std::move(join_tree_query_plan.query_plan);
+    used_row_policies = std::move(join_tree_query_plan.used_row_policies);
 
     LOG_TRACE(&Poco::Logger::get("Planner"), "Query {} from stage {} to stage {}{}",
         query_tree->formatConvertedASTForErrorMessage(),

src/Planner/Planner.h

@@ -44,6 +44,11 @@ public:
         return query_plan;
     }
 
+    const std::set<std::string> & getUsedRowPolicies() const
+    {
+        return used_row_policies;
+    }
+
     void buildQueryPlanIfNeeded();
 
     QueryPlan && extractQueryPlan() &&
@@ -70,6 +75,7 @@ private:
     PlannerContextPtr planner_context;
     QueryPlan query_plan;
     StorageLimitsList storage_limits;
+    std::set<std::string> used_row_policies;
 };
 
 }

src/Planner/PlannerJoinTree.cpp

@@ -126,8 +126,8 @@ bool shouldIgnoreQuotaAndLimits(const TableNode & table_node)
         return false;
     if (storage_id.database_name == DatabaseCatalog::SYSTEM_DATABASE)
     {
-        static const boost::container::flat_set<String> tables_ignoring_quota{"quotas", "quota_limits", "quota_usage", "quotas_usage", "one"};
-        if (tables_ignoring_quota.count(storage_id.table_name))
+        static const boost::container::flat_set<std::string_view> tables_ignoring_quota{"quotas", "quota_limits", "quota_usage", "quotas_usage", "one"};
+        if (tables_ignoring_quota.contains(storage_id.table_name))
             return true;
     }
     return false;
@@ -441,7 +441,8 @@ void updatePrewhereOutputsIfNeeded(SelectQueryInfo & table_expression_query_info
 
 FilterDAGInfo buildRowPolicyFilterIfNeeded(const StoragePtr & storage,
     SelectQueryInfo & table_expression_query_info,
-    PlannerContextPtr & planner_context)
+    PlannerContextPtr & planner_context,
+    std::set<std::string> & used_row_policies)
 {
     auto storage_id = storage->getStorageID();
     const auto & query_context = planner_context->getQueryContext();
@@ -450,6 +451,12 @@ FilterDAGInfo buildRowPolicyFilterIfNeeded(const StoragePtr & storage,
     if (!row_policy_filter || row_policy_filter->empty())
         return {};
 
+    for (const auto & row_policy : row_policy_filter->policies)
+    {
+        auto name = row_policy->getFullName().toString();
+        used_row_policies.emplace(std::move(name));
+    }
+
     return buildFilterInfo(row_policy_filter->expression, table_expression_query_info.table_expression, planner_context);
 }
 
@@ -586,6 +593,7 @@ JoinTreeQueryPlan buildQueryPlanForTableExpression(QueryTreeNodePtr table_expres
     auto * union_node = table_expression->as<UnionNode>();
 
     QueryPlan query_plan;
+    std::set<std::string> used_row_policies;
 
     if (table_node || table_function_node)
     {
@@ -781,7 +789,7 @@ JoinTreeQueryPlan buildQueryPlanForTableExpression(QueryTreeNodePtr table_expres
                     }
                 };
 
-                auto row_policy_filter_info = buildRowPolicyFilterIfNeeded(storage, table_expression_query_info, planner_context);
+                auto row_policy_filter_info = buildRowPolicyFilterIfNeeded(storage, table_expression_query_info, planner_context, used_row_policies);
                 add_filter(row_policy_filter_info, "Row-level security filter");
                 if (row_policy_filter_info.actions)
                     table_expression_data.setRowLevelFilterActions(row_policy_filter_info.actions);
@@ -940,7 +948,7 @@ JoinTreeQueryPlan buildQueryPlanForTableExpression(QueryTreeNodePtr table_expres
         }
     }
 
-    return {std::move(query_plan), from_stage};
+    return {std::move(query_plan), from_stage, std::move(used_row_policies)};
 }
 
 JoinTreeQueryPlan buildQueryPlanForJoinNode(const QueryTreeNodePtr & join_table_expression,
@@ -1399,7 +1407,10 @@ JoinTreeQueryPlan buildQueryPlanForJoinNode(const QueryTreeNodePtr & join_table_
     drop_unused_columns_after_join_transform_step->setStepDescription("DROP unused columns after JOIN");
     result_plan.addStep(std::move(drop_unused_columns_after_join_transform_step));
 
-    return {std::move(result_plan), QueryProcessingStage::FetchColumns};
+    for (const auto & right_join_tree_query_plan_row_policy : right_join_tree_query_plan.used_row_policies)
+        left_join_tree_query_plan.used_row_policies.insert(right_join_tree_query_plan_row_policy);
+
+    return {std::move(result_plan), QueryProcessingStage::FetchColumns, std::move(left_join_tree_query_plan.used_row_policies)};
 }
 
 JoinTreeQueryPlan buildQueryPlanForArrayJoinNode(const QueryTreeNodePtr & array_join_table_expression,
@@ -1477,7 +1488,7 @@ JoinTreeQueryPlan buildQueryPlanForArrayJoinNode(const QueryTreeNodePtr & array_
     array_join_step->setStepDescription("ARRAY JOIN");
     plan.addStep(std::move(array_join_step));
 
-    return {std::move(plan), QueryProcessingStage::FetchColumns};
+    return {std::move(plan), QueryProcessingStage::FetchColumns, std::move(join_tree_query_plan.used_row_policies)};
 }
 
 }

src/Planner/PlannerJoinTree.h

@@ -15,6 +15,7 @@ struct JoinTreeQueryPlan
 {
     QueryPlan query_plan;
     QueryProcessingStage::Enum from_stage;
+    std::set<std::string> used_row_policies;
 };
 
 /// Build JOIN TREE query plan for query node

tests/analyzer_tech_debt.txt

@@ -28,7 +28,6 @@
 01925_test_storage_merge_aliases
 01947_mv_subquery
 01952_optimize_distributed_group_by_sharding_key
-02131_used_row_policies_in_query_log
 02139_MV_with_scalar_subquery
 02174_cte_scalar_cache_mv
 02302_s3_file_pruning