From 882ba2e1a1e63f3069b870f7426ceb4d84b4d82e Mon Sep 17 00:00:00 2001
From: AlexYue <yj976240184@gmail.com>
Date: Wed, 31 Jan 2024 22:52:48 +0800
Subject: [PATCH] [fix](Cooldown) enhance the policy existence check logic when
 drop storage policy (#30404)

---
 .../org/apache/doris/policy/PolicyMgr.java    | 10 +++-
 .../cold_heat_separation/policy/drop.groovy   | 55 ++++++++++++++++++-
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/policy/PolicyMgr.java b/fe/fe-core/src/main/java/org/apache/doris/policy/PolicyMgr.java
index 42083ee6cd..3527b77c1d 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/policy/PolicyMgr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/policy/PolicyMgr.java
@@ -26,6 +26,7 @@ import org.apache.doris.analysis.UserIdentity;
 import org.apache.doris.catalog.Database;
 import org.apache.doris.catalog.Env;
 import org.apache.doris.catalog.OlapTable;
+import org.apache.doris.catalog.PartitionInfo;
 import org.apache.doris.catalog.Table;
 import org.apache.doris.cluster.ClusterNamespace;
 import org.apache.doris.common.AnalysisException;
@@ -153,9 +154,14 @@ public class PolicyMgr implements Writable {
                 List<Table> tables = db.getTables();
                 for (Table table : tables) {
                     if (table instanceof OlapTable) {
-                        if (((OlapTable) table).getStoragePolicy().equals(dropPolicyLog.getPolicyName())) {
-                            throw new DdlException("the policy " + dropPolicyLog.getPolicyName() + " is used by table: "
+                        OlapTable olapTable = (OlapTable) table;
+                        PartitionInfo partitionInfo = olapTable.getPartitionInfo();
+                        for (Long partitionId : olapTable.getPartitionIds()) {
+                            String policyName = partitionInfo.getDataProperty(partitionId).getStoragePolicy();
+                            if (policyName.equals(dropPolicyLog.getPolicyName())) {
+                                throw new DdlException("the policy " + policyName + " is used by table: "
                                     + table.getName());
+                            }
                         }
                     }
                 }
diff --git a/regression-test/suites/cold_heat_separation/policy/drop.groovy b/regression-test/suites/cold_heat_separation/policy/drop.groovy
index 20b53d42e0..6879e4891a 100644
--- a/regression-test/suites/cold_heat_separation/policy/drop.groovy
+++ b/regression-test/suites/cold_heat_separation/policy/drop.groovy
@@ -110,15 +110,29 @@ suite("drop_policy") {
         """
         assertEquals(storage_exist.call("drop_policy_test_has_table_binded"), true)
 
+        def create_succ_3 = try_sql """
+            CREATE STORAGE POLICY IF NOT EXISTS drop_policy_test_has_table_bind_1
+            PROPERTIES(
+            "storage_resource" = "${resource_table_use}",
+            "cooldown_datetime" = "2025-06-08 00:00:00"
+            );
+        """
+        assertEquals(storage_exist.call("drop_policy_test_has_table_bind_1"), true)
+
         // success
         def create_table_use_created_policy = try_sql """
             CREATE TABLE IF NOT EXISTS create_table_binding_created_policy
             (
                 k1 BIGINT,
-                k2 LARGEINT,
+                k2 date,
                 v1 VARCHAR(2048)
             )
-            UNIQUE KEY(k1)
+            UNIQUE KEY(k1, k2)
+            PARTITION BY RANGE(k2)(
+                partition p1 VALUES LESS THAN ("2014-01-01") ("storage_policy" = "drop_policy_test_has_table_bind_1"),
+                partition p2 VALUES LESS THAN ("2015-01-01"),
+                partition p3 VALUES LESS THAN ("2016-01-01")
+            )
             DISTRIBUTED BY HASH (k1) BUCKETS 3
             PROPERTIES(
                 "storage_policy" = "drop_policy_test_has_table_binded",
@@ -136,14 +150,51 @@ suite("drop_policy") {
         // fail to drop, there are tables using this policy
         assertEquals(drop_policy_fail_ret, null)
 
+        create_table_use_created_policy = try_sql """
+            CREATE TABLE IF NOT EXISTS create_table_binding_created_policy_1
+            (
+                k1 BIGINT,
+                k2 date,
+                v1 VARCHAR(2048)
+            )
+            UNIQUE KEY(k1, k2)
+            PARTITION BY RANGE(k2)(
+                partition p1 VALUES LESS THAN ("2014-01-01") ("storage_policy" = "drop_policy_test_has_table_bind_1"),
+                partition p2 VALUES LESS THAN ("2015-01-01"),
+                partition p3 VALUES LESS THAN ("2016-01-01")
+            )
+            DISTRIBUTED BY HASH (k1) BUCKETS 3
+            PROPERTIES(
+                "replication_num" = "1",
+                "enable_unique_key_merge_on_write" = "false"
+            );
+        """
+        // storage policy is disabled on mow table
+
+        assertEquals(create_table_use_created_policy.size(), 1);
+
+        drop_policy_fail_ret = try_sql """
+            DROP STORAGE POLICY drop_policy_test_has_table_bind_1
+        """
+        // fail to drop, there are partitions using this policy
+        assertEquals(drop_policy_fail_ret, null)
+
         sql """
         DROP TABLE create_table_binding_created_policy;
         """
 
+        sql """
+        DROP TABLE create_table_binding_created_policy_1;
+        """
+
         sql """
         DROP STORAGE POLICY drop_policy_test_has_table_binded;
         """
 
+        sql """
+        DROP STORAGE POLICY drop_policy_test_has_table_bind_1;
+        """
+
         sql """
         DROP RESOURCE ${resource_table_use};
         """