Merge pull request ClickHouse#79566 from jkartseva/part-operations-plain-disks

nikitamikhaylov · web-flow · commit 5eaa932b31e5 · 2025-05-08T23:06:18.000Z
Allow partition operations for tables on different plain_rewritable disks
diff --git a/src/Disks/IDisk.h b/src/Disks/IDisk.h
@@ -442,6 +442,9 @@ class IDisk : public Space
 
     virtual bool isReadOnly() const { return false; }
 
+    /// If the disk is plain object storage.
+    virtual bool isPlain() const { return false; }
+
     virtual bool isWriteOnce() const { return false; }
 
     virtual bool supportsHardLinks() const { return true; }
diff --git a/src/Disks/IStoragePolicy.h b/src/Disks/IStoragePolicy.h
@@ -59,6 +59,8 @@ class IStoragePolicy
     VolumePtr getVolumeByName(const String & volume_name) const;
     /// Checks if storage policy can be replaced by another one.
     virtual void checkCompatibleWith(const StoragePolicyPtr & new_storage_policy) const = 0;
+    /// If the policy allows table partition operations (move, replace) with the other storage policy.
+    virtual bool isCompatibleForPartitionOps(const StoragePolicyPtr & other) const = 0;
     /// Finds a volume index, which contains disk
     virtual std::optional<size_t> tryGetVolumeIndexByDiskName(const String & disk_name) const = 0;
     size_t getVolumeIndexByDiskName(const String & disk_name) const;
diff --git a/src/Disks/ObjectStorages/DiskObjectStorage.h b/src/Disks/ObjectStorages/DiskObjectStorage.h
@@ -205,7 +205,7 @@ friend class DiskObjectStorageRemoteMetadataRestoreHelper;
     /// with static files, so only read-only operations are allowed for this storage.
     bool isReadOnly() const override;
 
-    bool isPlain() const;
+    bool isPlain() const override;
 
     /// Is object write-once?
     /// For example: S3PlainObjectStorage is write once, this means that it
diff --git a/src/Disks/StoragePolicy.cpp b/src/Disks/StoragePolicy.cpp
@@ -12,6 +12,7 @@
 #include <Disks/VolumeJBOD.h>
 
 #include <algorithm>
+#include <ranges>
 #include <set>
 
 
@@ -392,6 +393,22 @@ void StoragePolicy::checkCompatibleWith(const StoragePolicyPtr & new_storage_pol
     }
 }
 
+bool StoragePolicy::isCompatibleForPartitionOps(const StoragePolicyPtr & other) const
+{
+    if (getName() == other->getName())
+        return true;
+
+    constexpr auto is_compatible = [](const auto & disk) -> bool
+    {
+        if (disk->isPlain())
+            return true;
+        if (auto delegate = disk->getDelegateDiskIfExists())
+            return delegate->isPlain();
+        return false;
+    };
+
+    return std::ranges::all_of(getDisks(), is_compatible) && std::ranges::all_of(other->getDisks(), is_compatible);
+}
 
 std::optional<size_t> StoragePolicy::tryGetVolumeIndexByDiskName(const String & disk_name) const
 {
diff --git a/src/Disks/StoragePolicy.h b/src/Disks/StoragePolicy.h
@@ -82,6 +82,9 @@ class StoragePolicy : public IStoragePolicy
     /// Checks if storage policy can be replaced by another one.
     void checkCompatibleWith(const StoragePolicyPtr & new_storage_policy) const override;
 
+    /// If the policy allows table partition operations (move, replace) with the other storage policy.
+    bool isCompatibleForPartitionOps(const StoragePolicyPtr & other) const override;
+
     /// Check if we have any volume with stopped merges
     bool hasAnyVolumeWithDisabledMerges() const override;
 
diff --git a/src/Storages/StorageMergeTree.cpp b/src/Storages/StorageMergeTree.cpp
@@ -2268,6 +2268,7 @@ void StorageMergeTree::replacePartitionFrom(const StoragePtr & source_table, con
 
     static const String TMP_PREFIX = "tmp_replace_from_";
 
+    bool are_policies_partition_op_compatible = getStoragePolicy()->isCompatibleForPartitionOps(source_table->getStoragePolicy());
     for (const DataPartPtr & src_part : src_parts)
     {
         if (is_all)
@@ -2294,7 +2295,7 @@ void StorageMergeTree::replacePartitionFrom(const StoragePtr & source_table, con
                 clone_params,
                 local_context->getReadSettings(),
                 local_context->getWriteSettings(),
-                true/*must_on_same_disk*/);
+                !are_policies_partition_op_compatible /*must_on_same_disk*/);
             dst_parts.emplace_back(std::move(dst_part));
             dst_parts_locks.emplace_back(std::move(part_lock));
         }
@@ -2371,13 +2372,19 @@ void StorageMergeTree::movePartitionToTable(const StoragePtr & dest_table, const
         throw Exception(ErrorCodes::NOT_IMPLEMENTED,
                         "Table {} supports movePartitionToTable only for MergeTree family of table engines. Got {}",
                         getStorageID().getNameForLogs(), dest_table->getName());
-    if (dest_table_storage->getStoragePolicy() != this->getStoragePolicy())
-        throw Exception(ErrorCodes::UNKNOWN_POLICY,
-                        "Destination table {} should have the same storage policy of source table {}. {}: {}, {}: {}",
-                        dest_table_storage->getStorageID().getNameForLogs(),
-                        getStorageID().getNameForLogs(), getStorageID().getNameForLogs(),
-                        this->getStoragePolicy()->getName(), dest_table_storage->getStorageID().getNameForLogs(),
-                        dest_table_storage->getStoragePolicy()->getName());
+    bool are_policies_partition_op_compatible = getStoragePolicy()->isCompatibleForPartitionOps(dest_table_storage->getStoragePolicy());
+
+    if (!are_policies_partition_op_compatible)
+        throw Exception(
+            ErrorCodes::UNKNOWN_POLICY,
+            "Destination table {} should have the same storage policy of source table, or the policies must be compatible for partition "
+            "operations {}. {}: {}, {}: {}",
+            dest_table_storage->getStorageID().getNameForLogs(),
+            getStorageID().getNameForLogs(),
+            getStorageID().getNameForLogs(),
+            this->getStoragePolicy()->getName(),
+            dest_table_storage->getStorageID().getNameForLogs(),
+            dest_table_storage->getStoragePolicy()->getName());
 
     // Use the same back-pressure (delay/throw) logic as for INSERTs to be consistent and avoid possibility of exceeding part limits using MOVE PARTITION queries
     dest_table_storage->delayInsertOrThrowIfNeeded(nullptr, local_context, true);
@@ -2437,7 +2444,7 @@ void StorageMergeTree::movePartitionToTable(const StoragePtr & dest_table, const
             clone_params,
             local_context->getReadSettings(),
             local_context->getWriteSettings(),
-            true/*must_on_same_disk*/
+            !are_policies_partition_op_compatible /*must_on_same_disk*/
         );
 
         dst_parts.emplace_back(std::move(dst_part));
diff --git a/tests/integration/test_s3_plain_rewritable_rotate_tables/configs/storage_conf.xml b/tests/integration/test_s3_plain_rewritable_rotate_tables/configs/storage_conf.xml
@@ -10,6 +10,19 @@
                 <access_key_id>minio</access_key_id>
                 <secret_access_key>ClickHouse_Minio_P@ssw0rd</secret_access_key>
             </disk_s3_plain_rewritable>
+            <disk_cache_s3_plain_rewritable>
+                <type>cache</type>
+                <disk>disk_s3_plain_rewritable</disk>
+                <path>disks/cache/</path>
+                <max_size>1000000000</max_size>
+                <cache_on_write_operations>1</cache_on_write_operations>
+            </disk_cache_s3_plain_rewritable>
+            <disk_encrypted>
+                <type>encrypted</type>
+                <disk>disk_cache_s3_plain_rewritable</disk>
+                <key>1234567812345678</key>
+                <path>disks/encrypted/</path>
+            </disk_encrypted>
         </disks>
         <policies>
             <s3_plain_rewritable>
@@ -19,6 +32,20 @@
                     </main>
                 </volumes>
             </s3_plain_rewritable>
+            <cache>
+                <volumes>
+                    <main>
+                        <disk>disk_cache_s3_plain_rewritable</disk>
+                    </main>
+                </volumes>
+            </cache>
+            <encrypted>
+                <volumes>
+                    <main>
+                        <disk>disk_encrypted</disk>
+                    </main>
+                </volumes>
+            </encrypted>
         </policies>
     </storage_configuration>
 </clickhouse>
diff --git a/tests/integration/test_s3_plain_rewritable_rotate_tables/test.py b/tests/integration/test_s3_plain_rewritable_rotate_tables/test.py
@@ -23,6 +23,37 @@ def randomize_name(table_name, random_suffix_length=8):
     return f"{table_name}_{''.join(random.choice(letters) for _ in range(random_suffix_length))}"
 
 
+s3 = """disk(
+    name={},
+    type='s3_plain_rewritable',
+    endpoint='http://minio1:9001/root/data/node1',
+    access_key_id='minio',
+    secret_access_key='ClickHouse_Minio_P@ssw0rd')"""
+
+cache = """disk(
+    type='cache',
+    disk={},
+    path='disks/cache/',
+    cache_on_write_operations=1,
+    max_size='1Gi')""".format(
+    s3
+)
+
+encrypted = """disk(
+    type='encrypted',
+    disk={},
+    path='disks/encrypted/',
+    key='1234567812345678')""".format(
+    cache
+)
+
+disk_defs = {
+    "s3_plain_rewritable": s3,
+    "cache": cache,
+    "encrypted": encrypted,
+}
+
+
 @pytest.fixture(scope="module", autouse=True)
 def start_cluster():
     cluster.add_instance(
@@ -51,7 +82,16 @@ def start_cluster():
         cluster.shutdown()
 
 
-def test_alter_partition_after_table_rotation():
+@pytest.mark.parametrize(
+    "storage_policy",
+    [
+        "s3_plain_rewritable",
+        "cache",
+        "encrypted",
+    ],
+)
+@pytest.mark.parametrize("partition_ops_on_same_disk", [1, 0])
+def test_alter_partition_after_table_rotation(start_cluster, storage_policy, partition_ops_on_same_disk):
     node1 = cluster.instances["node1"]
 
     def create_insert(node, table_name, insert_values):
@@ -63,7 +103,7 @@ def create_insert(node, table_name, insert_values):
             ) ENGINE=MergeTree()
             ORDER BY id
             PARTITION BY id%10
-            SETTINGS storage_policy='s3_plain_rewritable'
+            SETTINGS storage_policy='{storage_policy}'
             """
         )
 
@@ -80,44 +120,43 @@ def create_insert(node, table_name, insert_values):
 
     node1.query(f"DETACH TABLE {table1} PERMANENTLY SYNC")
 
-    disk_name = randomize_name("disk")
     node2 = cluster.instances["node2"]
     table2 = randomize_name("table2")
     node2.query(f"DROP TABLE IF EXISTS {table2} SYNC")
+
+    disk_name = randomize_name("disk")
+    disk_def = disk_defs[storage_policy].format(disk_name)
+    policy_def = (
+        f"disk={disk_def}"
+        if partition_ops_on_same_disk
+        else f"storage_policy='{storage_policy}'"
+    )
+
     node2.query(
         f"""CREATE TABLE {table2} (id Int64, data String)
         ENGINE=MergeTree()
         ORDER BY id
         PARTITION BY id%10
-        SETTINGS disk=disk(
-            name={disk_name},
-            type='s3_plain_rewritable',
-            endpoint='http://minio1:9001/root/data/node1',
-            access_key_id='minio',
-            secret_access_key='ClickHouse_Minio_P@ssw0rd')
+        SETTINGS {policy_def}
         """
     )
 
     rotated_table = f"{table1}_rotated"
     node2.query(f"""DROP TABLE IF EXISTS {rotated_table} SYNC""")
+
     node2.query(
         f"""ATTACH TABLE {rotated_table} UUID '{uuid1}' (id Int64, data String)
         ENGINE=MergeTree()
         ORDER BY id
         PARTITION BY id%10
-        SETTINGS disk=disk(
-            name={disk_name},
-            type='s3_plain_rewritable',
-            endpoint='http://minio1:9001/root/data/node1',
-            access_key_id='minio',
-            secret_access_key='ClickHouse_Minio_P@ssw0rd')
+        SETTINGS disk={disk_def}
         """
     )
 
     assert (
         int(
             node2.query(
-                f"SELECT count(*) FROM {rotated_table} WHERE _partition_id = '0'"
+                f"SELECT count(*) FROM {rotated_table} WHERE _partition_id= '0'"
             )
         )
         == 100
@@ -126,9 +165,10 @@ def create_insert(node, table_name, insert_values):
     assert int(node2.query(f"SELECT count(*) FROM {rotated_table}")) == 1000
 
     node2.query(f"""ALTER TABLE {rotated_table} MOVE PARTITION '0' TO TABLE {table2}""")
+    node2.query(f"""ALTER TABLE {table2} REPLACE PARTITION '1' FROM {rotated_table}""")
 
     assert int(node2.query(f"SELECT count(*) FROM {rotated_table}")) == 900
-    assert int(node2.query(f"SELECT count(*) FROM {table2}")) == 100
+    assert int(node2.query(f"SELECT count(*) FROM {table2}")) == 200
 
     node2.query(f"DROP TABLE {rotated_table} SYNC")
     node2.query(f"DROP TABLE {table2} SYNC")
