Add ability to choose object storage cluster in select query

ianton-ru · ianton-ru · commit 3fafe6f27b76 · 2025-02-18T16:42:50.000+01:00
diff --git a/src/Databases/Iceberg/DatabaseIceberg.cpp b/src/Databases/Iceberg/DatabaseIceberg.cpp
@@ -43,7 +43,6 @@ namespace DatabaseIcebergSetting
 namespace Setting
 {
     extern const SettingsBool allow_experimental_database_iceberg;
-    extern const SettingsString object_storage_cluster;
 }
 
 namespace ErrorCodes
@@ -240,33 +239,18 @@ StoragePtr DatabaseIceberg::tryGetTable(const String & name, ContextPtr context_
 
     auto cluster_name = settings[DatabaseIcebergSetting::object_storage_cluster].value;
 
-    if (cluster_name.empty())
-    {
-        return std::make_shared<StorageObjectStorage>(
-            configuration,
-            configuration->createObjectStorage(context_, /* is_readonly */ false),
-            context_,
-            StorageID(getDatabaseName(), name),
-            /* columns */columns,
-            /* constraints */ConstraintsDescription{},
-            /* comment */"",
-            getFormatSettings(context_),
-            LoadingStrictnessLevel::CREATE,
-            /* distributed_processing */false,
-            /* partition_by */nullptr,
-            /* lazy_init */true);
-    }
-    else
-    {
-        return std::make_shared<StorageObjectStorageCluster>(
-            cluster_name,
-            configuration,
-            configuration->createObjectStorage(context_, /* is_readonly */ false),
-            StorageID(getDatabaseName(), name),
-            columns,
-            ConstraintsDescription{},
-            context_);
-    }
+    return std::make_shared<StorageObjectStorageCluster>(
+        cluster_name,
+        configuration,
+        configuration->createObjectStorage(context_, /* is_readonly */ false),
+        context_,
+        StorageID(getDatabaseName(), name),
+        /* columns */columns,
+        /* constraints */ConstraintsDescription{},
+        /* comment */"",
+        /* format_settings */ getFormatSettings(context_),
+        /* mode */ LoadingStrictnessLevel::CREATE,
+        /* partition_by */nullptr);
 }
 
 DatabaseTablesIteratorPtr DatabaseIceberg::getTablesIterator(
diff --git a/src/Storages/IStorageCluster.cpp b/src/Storages/IStorageCluster.cpp
@@ -36,6 +36,11 @@ namespace Setting
     extern const SettingsBool skip_unavailable_shards;
 }
 
+namespace ErrorCodes
+{
+    extern const int NOT_IMPLEMENTED;
+}
+
 IStorageCluster::IStorageCluster(
     const String & cluster_name_,
     const StorageID & table_id_,
@@ -65,6 +70,19 @@ void ReadFromCluster::createExtension(const ActionsDAG::Node * predicate)
     extension = storage->getTaskIteratorExtension(predicate, context);
 }
 
+void IStorageCluster::readFallBackToPure(
+    QueryPlan & /*query_plan*/,
+    const Names & /*column_names*/,
+    const StorageSnapshotPtr & /*storage_snapshot*/,
+    SelectQueryInfo & /*query_info*/,
+    ContextPtr /*context*/,
+    QueryProcessingStage::Enum /*processed_stage*/,
+    size_t /*max_block_size*/,
+    size_t /*num_streams*/)
+{
+    throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method readFallBackToPure is not supported by storage {}", getName());
+}
+
 /// The code executes on initiator
 void IStorageCluster::read(
     QueryPlan & query_plan,
@@ -73,13 +91,21 @@ void IStorageCluster::read(
     SelectQueryInfo & query_info,
     ContextPtr context,
     QueryProcessingStage::Enum processed_stage,
-    size_t /*max_block_size*/,
-    size_t /*num_streams*/)
+    size_t max_block_size,
+    size_t num_streams)
 {
+    auto cluster_name_ = getClusterName(context);
+
+    if (cluster_name_.empty())
+    {
+        readFallBackToPure(query_plan, column_names, storage_snapshot, query_info, context, processed_stage, max_block_size, num_streams);
+        return;
+    }
+
     storage_snapshot->check(column_names);
 
     updateBeforeRead(context);
-    auto cluster = getCluster(context);
+    auto cluster = getClusterImpl(context, cluster_name_);
 
     /// Calculate the header. This is significant, because some columns could be thrown away in some cases like query with count(*)
 
@@ -196,9 +222,9 @@ ContextPtr ReadFromCluster::updateSettings(const Settings & settings)
     return new_context;
 }
 
-ClusterPtr IStorageCluster::getCluster(ContextPtr context) const
+ClusterPtr IStorageCluster::getClusterImpl(ContextPtr context, const String & cluster_name_)
 {
-    return context->getCluster(cluster_name)->getClusterWithReplicasAsShards(context->getSettingsRef());
+    return context->getCluster(cluster_name_)->getClusterWithReplicasAsShards(context->getSettingsRef());
 }
 
 }
diff --git a/src/Storages/IStorageCluster.h b/src/Storages/IStorageCluster.h
@@ -29,10 +29,10 @@ class IStorageCluster : public IStorage
         SelectQueryInfo & query_info,
         ContextPtr context,
         QueryProcessingStage::Enum processed_stage,
-        size_t /*max_block_size*/,
-        size_t /*num_streams*/) override;
+        size_t max_block_size,
+        size_t num_streams) override;
 
-    ClusterPtr getCluster(ContextPtr context) const;
+    ClusterPtr getCluster(ContextPtr context) const { return getClusterImpl(context, cluster_name); }
     /// Query is needed for pruning by virtual columns (_file, _path)
     virtual RemoteQueryExecutor::Extension getTaskIteratorExtension(const ActionsDAG::Node * predicate, const ContextPtr & context) const = 0;
 
@@ -43,13 +43,26 @@ class IStorageCluster : public IStorage
     bool supportsOptimizationToSubcolumns() const override { return false; }
     bool supportsTrivialCountOptimization(const StorageSnapshotPtr &, ContextPtr) const override { return true; }
 
-    const String & getClusterName() const { return cluster_name; }
+    const String & getOriginalClusterName() const { return cluster_name; }
+    virtual String getClusterName(ContextPtr /* context */) const { return getOriginalClusterName(); }
 
 protected:
     virtual void updateBeforeRead(const ContextPtr &) {}
     virtual void updateQueryToSendIfNeeded(ASTPtr & /*query*/, const StorageSnapshotPtr & /*storage_snapshot*/, const ContextPtr & /*context*/) {}
 
+    virtual void readFallBackToPure(
+        QueryPlan & query_plan,
+        const Names & column_names,
+        const StorageSnapshotPtr & storage_snapshot,
+        SelectQueryInfo & query_info,
+        ContextPtr context,
+        QueryProcessingStage::Enum processed_stage,
+        size_t max_block_size,
+        size_t num_streams);
+
 private:
+    static ClusterPtr getClusterImpl(ContextPtr context, const String & cluster_name_);
+
     LoggerPtr log;
     String cluster_name;
 };
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp b/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp
@@ -23,6 +23,7 @@ namespace DB
 namespace Setting
 {
     extern const SettingsBool use_hive_partitioning;
+    extern const SettingsString object_storage_cluster;
 }
 
 namespace ErrorCodes
@@ -63,15 +64,24 @@ StorageObjectStorageCluster::StorageObjectStorageCluster(
     const String & cluster_name_,
     ConfigurationPtr configuration_,
     ObjectStoragePtr object_storage_,
+    ContextPtr context_,
     const StorageID & table_id_,
     const ColumnsDescription & columns_,
     const ConstraintsDescription & constraints_,
-    ContextPtr context_)
+    const String & comment_,
+    std::optional<FormatSettings> format_settings_,
+    LoadingStrictnessLevel mode_,
+    ASTPtr partition_by_
+)
     : IStorageCluster(
         cluster_name_, table_id_, getLogger(fmt::format("{}({})", configuration_->getEngineName(), table_id_.table_name)))
     , configuration{configuration_}
     , object_storage(object_storage_)
     , cluster_name_in_settings(false)
+    , comment(comment_)
+    , format_settings(format_settings_)
+    , mode(mode_)
+    , partition_by(partition_by_)
 {
     ColumnsDescription columns{columns_};
     std::string sample_path;
@@ -94,7 +104,7 @@ std::string StorageObjectStorageCluster::getName() const
     return configuration->getEngineName();
 }
 
-void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr & query)
+void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context)
 {
     // Change table engine on table function for distributed request
     // CREATE TABLE t (...) ENGINE=IcebergS3(...)
@@ -131,6 +141,7 @@ void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr
         {"S3", "s3"},
         {"Azure", "azureBlobStorage"},
         {"HDFS", "hdfs"},
+        {"Iceberg", "icebergS3"},
         {"IcebergS3", "icebergS3"},
         {"IcebergAzure", "icebergAzure"},
         {"IcebergHDFS", "icebergHDFS"},
@@ -153,7 +164,7 @@ void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr
     auto function_ast = std::make_shared<ASTFunction>();
     function_ast->name = table_function_name;
 
-    auto cluster_name = getClusterName();
+    auto cluster_name = getClusterName(context);
 
     if (cluster_name.empty())
     {
@@ -195,7 +206,7 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
     const DB::StorageSnapshotPtr & storage_snapshot,
     const ContextPtr & context)
 {
-    updateQueryForDistributedEngineIfNeeded(query);
+    updateQueryForDistributedEngineIfNeeded(query, context);
 
     ASTExpressionList * expression_list = extractTableFunctionArgumentsFromSelectQuery(query);
 
@@ -247,4 +258,39 @@ RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExten
     return RemoteQueryExecutor::Extension{ .task_iterator = std::move(callback) };
 }
 
+void StorageObjectStorageCluster::readFallBackToPure(
+    QueryPlan & query_plan,
+    const Names & column_names,
+    const StorageSnapshotPtr & storage_snapshot,
+    SelectQueryInfo & query_info,
+    ContextPtr context,
+    QueryProcessingStage::Enum processed_stage,
+    size_t max_block_size,
+    size_t num_streams)
+{
+    if (!pure_storage)
+        pure_storage = std::make_shared<StorageObjectStorage>(
+            configuration,
+            object_storage,
+            context,
+            getStorageID(),
+            getInMemoryMetadata().getColumns(),
+            getInMemoryMetadata().getConstraints(),
+            comment,
+            format_settings,
+            mode,
+            /* distributed_processing */false,
+            partition_by);
+
+    pure_storage->read(query_plan, column_names, storage_snapshot, query_info, context, processed_stage, max_block_size, num_streams);
+}
+
+String StorageObjectStorageCluster::getClusterName(ContextPtr context) const
+{
+    auto cluster_name_ = context->getSettingsRef()[Setting::object_storage_cluster].value;
+    if (cluster_name_.empty())
+        cluster_name_ = getOriginalClusterName();
+    return cluster_name_;
+}
+
 }
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageCluster.h b/src/Storages/ObjectStorage/StorageObjectStorageCluster.h
@@ -17,10 +17,15 @@ class StorageObjectStorageCluster : public IStorageCluster
         const String & cluster_name_,
         ConfigurationPtr configuration_,
         ObjectStoragePtr object_storage_,
+        ContextPtr context_,
         const StorageID & table_id_,
         const ColumnsDescription & columns_,
         const ConstraintsDescription & constraints_,
-        ContextPtr context_);
+        const String & comment_,
+        std::optional<FormatSettings> format_settings_,
+        LoadingStrictnessLevel mode_,
+        ASTPtr partition_by_ = nullptr
+    );
 
     std::string getName() const override;
 
@@ -31,12 +36,24 @@ class StorageObjectStorageCluster : public IStorageCluster
 
     void setClusterNameInSettings(bool cluster_name_in_settings_) { cluster_name_in_settings = cluster_name_in_settings_; }
 
+    String getClusterName(ContextPtr context) const override;
+
 private:
     void updateQueryToSendIfNeeded(
         ASTPtr & query,
         const StorageSnapshotPtr & storage_snapshot,
         const ContextPtr & context) override;
 
+    void readFallBackToPure(
+        QueryPlan & query_plan,
+        const Names & column_names,
+        const StorageSnapshotPtr & storage_snapshot,
+        SelectQueryInfo & query_info,
+        ContextPtr context,
+        QueryProcessingStage::Enum processed_stage,
+        size_t max_block_size,
+        size_t num_streams) override;
+
     /*
     In case the table was created with `object_storage_cluster` setting,
     modify the AST query object so that it uses the table function implementation
@@ -49,12 +66,18 @@ class StorageObjectStorageCluster : public IStorageCluster
     SELECT * FROM s3(...) SETTINGS object_storage_cluster='cluster'
     to make distributed request over cluster 'cluster'.
     */
-    void updateQueryForDistributedEngineIfNeeded(ASTPtr & query);
+    void updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context);
 
     const String engine_name;
     const StorageObjectStorage::ConfigurationPtr configuration;
     const ObjectStoragePtr object_storage;
     bool cluster_name_in_settings;
+
+    std::shared_ptr<StorageObjectStorage> pure_storage;
+    String comment;
+    std::optional<FormatSettings> format_settings;
+    LoadingStrictnessLevel mode;
+    ASTPtr partition_by;
 };
 
 }
diff --git a/src/Storages/ObjectStorage/registerStorageObjectStorage.cpp b/src/Storages/ObjectStorage/registerStorageObjectStorage.cpp
@@ -20,11 +20,6 @@ namespace ErrorCodes
     extern const int BAD_ARGUMENTS;
 }
 
-namespace Setting
-{
-    extern const SettingsString object_storage_cluster;
-}
-
 namespace StorageObjectStorageSetting
 {
     extern const StorageObjectStorageSettingsString object_storage_cluster;
@@ -74,32 +69,18 @@ createStorageObjectStorage(const StorageFactory::Arguments & args, StorageObject
     if (args.storage_def->partition_by)
         partition_by = args.storage_def->partition_by->clone();
 
-    if (cluster_name.empty())
-    {
-        return std::make_shared<StorageObjectStorage>(
-            configuration,
-            configuration->createObjectStorage(context, /* is_readonly */ false),
-            args.getContext(),
-            args.table_id,
-            args.columns,
-            args.constraints,
-            args.comment,
-            format_settings,
-            args.mode,
-            /* distributed_processing */ false,
-            partition_by);
-    }
-    else
-    {
-        return std::make_shared<StorageObjectStorageCluster>(
-            cluster_name,
-            configuration,
-            configuration->createObjectStorage(context, /* is_readonly */ false),
-            args.table_id,
-            args.columns,
-            args.constraints,
-            args.getContext());
-    }
+    return std::make_shared<StorageObjectStorageCluster>(
+        cluster_name,
+        configuration,
+        configuration->createObjectStorage(context, /* is_readonly */ false),
+        args.getContext(),
+        args.table_id,
+        args.columns,
+        args.constraints,
+        args.comment,
+        format_settings,
+        args.mode,
+        partition_by);
 }
 
 #endif
diff --git a/src/TableFunctions/TableFunctionObjectStorageCluster.cpp b/src/TableFunctions/TableFunctionObjectStorageCluster.cpp
diff --git a/tests/integration/test_storage_iceberg/test.py b/tests/integration/test_storage_iceberg/test.py
