Antalya 26.1; Remote initiator improvements 2

src/Storages/IStorageCluster.cpp

@@ -60,6 +60,7 @@ namespace ErrorCodes
 {
     extern const int NOT_IMPLEMENTED;
     extern const int LOGICAL_ERROR;
+    extern const int BAD_ARGUMENTS;
 }
 
 namespace ErrorCodes
@@ -315,10 +316,40 @@ void IStorageCluster::read(
     size_t max_block_size,
     size_t num_streams)
 {
+    if (!isClusterSupported())
+    {
+        readFallBackToPure(query_plan, column_names, storage_snapshot, query_info, context, processed_stage, max_block_size, num_streams);
+        return;
+    }
+
     auto cluster_name_from_settings = getClusterName(context);
+    const auto & settings = context->getSettingsRef();
+    ASTPtr query_to_send = query_info.query;
 
-    if (!isClusterSupported() || cluster_name_from_settings.empty())
+    if (cluster_name_from_settings.empty())
     {
+        if (settings[Setting::object_storage_remote_initiator])
+        {
+            /// rewrite query to execute `remote('remote_host', s3(...))`
+            /// remote_host can execute query itself or make on-cluster query depends on own `object_storage_cluster` setting
+            updateConfigurationIfNeeded(context);
+            updateQueryWithJoinToSendIfNeeded(query_to_send, query_info.query_tree, context);
+            updateQueryToSendIfNeeded(query_to_send, storage_snapshot, context, /*make_cluster_function*/ false);
+
+            auto remote_initiator_cluster_name = settings[Setting::object_storage_remote_initiator_cluster].value;
+            if (remote_initiator_cluster_name.empty())
+                throw Exception(ErrorCodes::BAD_ARGUMENTS, "Setting 'object_storage_remote_initiator' can be used only with 'object_storage_remote_initiator_cluster' or 'object_storage_cluster'");
+
+            auto remote_initiator_cluster = getClusterImpl(context, remote_initiator_cluster_name);
+            auto storage_and_context = convertToRemote(remote_initiator_cluster, context, remote_initiator_cluster_name, query_to_send);
+            auto src_distributed = std::dynamic_pointer_cast<StorageDistributed>(storage_and_context.storage);
+            auto modified_query_info = query_info;
+            modified_query_info.cluster = src_distributed->getCluster();
+            auto new_storage_snapshot = storage_and_context.storage->getStorageSnapshot(storage_snapshot->metadata, storage_and_context.context);
+            storage_and_context.storage->read(query_plan, column_names, new_storage_snapshot, modified_query_info, storage_and_context.context, processed_stage, max_block_size, num_streams);
+            return;
+        }
+
         readFallBackToPure(query_plan, column_names, storage_snapshot, query_info, context, processed_stage, max_block_size, num_streams);
         return;
     }
@@ -327,12 +358,9 @@ void IStorageCluster::read(
 
     storage_snapshot->check(column_names);
 
-    const auto & settings = context->getSettingsRef();
-
     /// Calculate the header. This is significant, because some columns could be thrown away in some cases like query with count(*)
 
     SharedHeader sample_block;
-    ASTPtr query_to_send = query_info.query;
 
     updateQueryWithJoinToSendIfNeeded(query_to_send, query_info.query_tree, context);
 
@@ -347,7 +375,7 @@ void IStorageCluster::read(
         query_to_send = interpreter.getQueryInfo().query->clone();
     }
 
-    updateQueryToSendIfNeeded(query_to_send, storage_snapshot, context);
+    updateQueryToSendIfNeeded(query_to_send, storage_snapshot, context, /*make_cluster_function*/ true);
 
     /// In case the current node is not supposed to initiate the clustered query
     /// Sends this query to a remote initiator using the `remote` table function
@@ -431,8 +459,8 @@ IStorageCluster::RemoteCallVariables IStorageCluster::convertToRemote(
 
     /// Clean object_storage_remote_initiator setting to avoid infinite remote call
     auto new_context = Context::createCopy(context);
-    new_context->setSetting("object_storage_remote_initiator", false);
-    new_context->setSetting("object_storage_remote_initiator_cluster", String(""));
+    std::vector<std::string> settings_to_remove = {"object_storage_remote_initiator", "object_storage_remote_initiator_cluster"};
+    new_context->resetSettingsToDefaultValue(settings_to_remove);
 
     auto * select_query = query_to_send->as<ASTSelectQuery>();
     if (!select_query)
@@ -442,15 +470,18 @@ IStorageCluster::RemoteCallVariables IStorageCluster::convertToRemote(
     if (query_settings)
     {
         auto & settings_ast = query_settings->as<ASTSetQuery &>();
-        if (settings_ast.changes.removeSetting("object_storage_remote_initiator") && settings_ast.changes.empty())
-        {
+        bool settings_changed = false;
+        for (const auto & setting_to_remove : settings_to_remove)
+            settings_changed |= settings_ast.changes.removeSetting(setting_to_remove);
+        if (settings_changed && settings_ast.changes.empty())
             select_query->setExpression(ASTSelectQuery::Expression::SETTINGS, {});
-        }
     }
 
     ASTTableExpression * table_expression = extractTableExpressionASTPtrFromSelectQuery(query_to_send);
     if (!table_expression)
         throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't find table expression");
+    if (!table_expression->table_function)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't find table function in table expression");
 
     boost::intrusive_ptr<ASTFunction> remote_query;
 

src/Storages/IStorageCluster.h

@@ -62,7 +62,11 @@ class IStorageCluster : public IStorage
     virtual String getClusterName(ContextPtr /* context */) const { return getOriginalClusterName(); }
 
 protected:
-    virtual void updateQueryToSendIfNeeded(ASTPtr & /*query*/, const StorageSnapshotPtr & /*storage_snapshot*/, const ContextPtr & /*context*/) {}
+    virtual void updateQueryToSendIfNeeded(
+        ASTPtr & /*query*/,
+        const StorageSnapshotPtr & /*storage_snapshot*/,
+        const ContextPtr & /*context*/,
+        bool /*make_cluster_function*/) {}
     void updateQueryWithJoinToSendIfNeeded(ASTPtr & query_to_send, QueryTreeNodePtr query_tree, const ContextPtr & context);
 
     virtual void updateConfigurationIfNeeded(ContextPtr /* context */) {}

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -266,7 +266,7 @@ std::optional<UInt64> StorageObjectStorageCluster::totalBytes(ContextPtr query_c
     return configuration->totalBytes(query_context);
 }
 
-void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context)
+void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context, bool make_cluster_function)
 {
     // Change table engine on table function for distributed request
     // CREATE TABLE t (...) ENGINE=IcebergS3(...)
@@ -356,16 +356,6 @@ void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr
     auto function_ast = make_intrusive<ASTFunction>();
     function_ast->name = table_function_name;
 
-    auto cluster_name = getClusterName(context);
-
-    if (cluster_name.empty())
-    {
-        throw Exception(
-            ErrorCodes::LOGICAL_ERROR,
-            "Can't be here without cluster name, no cluster name in query {}",
-            query->formatForLogging());
-    }
-
     function_ast->arguments = configuration->createArgsWithAccessData();
     function_ast->children.push_back(function_ast->arguments);
     function_ast->setAlias(table_alias);
@@ -376,29 +366,43 @@ void StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr
     table_expression->table_function = function_ast_ptr;
     table_expression->children[0] = function_ast_ptr;
 
-    auto settings = select_query->settings();
-    if (settings)
+    if (make_cluster_function)
     {
-        auto & settings_ast = settings->as<ASTSetQuery &>();
-        settings_ast.changes.insertSetting("object_storage_cluster", cluster_name);
-    }
-    else
-    {
-        auto settings_ast_ptr = make_intrusive<ASTSetQuery>();
-        settings_ast_ptr->is_standalone = false;
-        settings_ast_ptr->changes.setSetting("object_storage_cluster", cluster_name);
-        select_query->setExpression(ASTSelectQuery::Expression::SETTINGS, std::move(settings_ast_ptr));
-    }
+        auto cluster_name = getClusterName(context);
 
-    cluster_name_in_settings = true;
+        if (cluster_name.empty())
+        {
+            throw Exception(
+                ErrorCodes::LOGICAL_ERROR,
+                "Can't be here without cluster name, no cluster name in query {}",
+                query->formatForLogging());
+        }
+
+        auto settings = select_query->settings();
+        if (settings)
+        {
+            auto & settings_ast = settings->as<ASTSetQuery &>();
+            settings_ast.changes.insertSetting("object_storage_cluster", cluster_name);
+        }
+        else
+        {
+            auto settings_ast_ptr = make_intrusive<ASTSetQuery>();
+            settings_ast_ptr->is_standalone = false;
+            settings_ast_ptr->changes.setSetting("object_storage_cluster", cluster_name);
+            select_query->setExpression(ASTSelectQuery::Expression::SETTINGS, std::move(settings_ast_ptr));
+        }
+
+        cluster_name_in_settings = true;
+    }
 }
 
 void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
     ASTPtr & query,
     const DB::StorageSnapshotPtr & storage_snapshot,
-    const ContextPtr & context)
+    const ContextPtr & context,
+    bool make_cluster_function)
 {
-    updateQueryForDistributedEngineIfNeeded(query, context);
+    updateQueryForDistributedEngineIfNeeded(query, context, make_cluster_function);
 
     auto * table_function = extractTableFunctionFromSelectQuery(query);
     if (!table_function)
@@ -441,60 +445,63 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
     {
         configuration->addStructureAndFormatToArgsIfNeeded(args, structure, configuration->getFormat(), context, /*with_structure=*/true);
 
-        /// Convert to old-stype *Cluster table function.
-        /// This allows to use old clickhouse versions in cluster.
-        static std::unordered_map<std::string, std::string> function_to_cluster_function = {
-            {"s3", "s3Cluster"},
-            {"azureBlobStorage", "azureBlobStorageCluster"},
-            {"hdfs", "hdfsCluster"},
-            {"iceberg", "icebergCluster"},
-            {"icebergS3", "icebergS3Cluster"},
-            {"icebergAzure", "icebergAzureCluster"},
-            {"icebergHDFS", "icebergHDFSCluster"},
-            {"icebergLocal", "icebergLocalCluster"},
-            {"deltaLake", "deltaLakeCluster"},
-            {"deltaLakeS3", "deltaLakeS3Cluster"},
-            {"deltaLakeAzure", "deltaLakeAzureCluster"},
-            {"hudi", "hudiCluster"},
-            {"paimonS3", "paimonS3Cluster"},
-            {"paimonAzure", "paimonAzureCluster"},
-        };
-
-        auto p = function_to_cluster_function.find(table_function->name);
-        if (p == function_to_cluster_function.end())
+        if (make_cluster_function)
         {
-            throw Exception(
-                ErrorCodes::LOGICAL_ERROR,
-                "Can't find cluster variant for table function {}",
-                table_function->name);
-        }
+            /// Convert to old-stype *Cluster table function.
+            /// This allows to use old clickhouse versions in cluster.
+            static std::unordered_map<std::string, std::string> function_to_cluster_function = {
+                {"s3", "s3Cluster"},
+                {"azureBlobStorage", "azureBlobStorageCluster"},
+                {"hdfs", "hdfsCluster"},
+                {"iceberg", "icebergCluster"},
+                {"icebergS3", "icebergS3Cluster"},
+                {"icebergAzure", "icebergAzureCluster"},
+                {"icebergHDFS", "icebergHDFSCluster"},
+                {"icebergLocal", "icebergLocalCluster"},
+                {"deltaLake", "deltaLakeCluster"},
+                {"deltaLakeS3", "deltaLakeS3Cluster"},
+                {"deltaLakeAzure", "deltaLakeAzureCluster"},
+                {"hudi", "hudiCluster"},
+                {"paimonS3", "paimonS3Cluster"},
+                {"paimonAzure", "paimonAzureCluster"},
+            };
+
+            auto p = function_to_cluster_function.find(table_function->name);
+            if (p == function_to_cluster_function.end())
+            {
+                throw Exception(
+                    ErrorCodes::LOGICAL_ERROR,
+                    "Can't find cluster variant for table function {}",
+                    table_function->name);
+            }
 
-        table_function->name = p->second;
+            table_function->name = p->second;
 
-        auto cluster_name = getClusterName(context);
-        auto cluster_name_arg = make_intrusive<ASTLiteral>(cluster_name);
-        args.insert(args.begin(), cluster_name_arg);
+            auto cluster_name = getClusterName(context);
+            auto cluster_name_arg = make_intrusive<ASTLiteral>(cluster_name);
+            args.insert(args.begin(), cluster_name_arg);
 
-        auto * select_query = query->as<ASTSelectQuery>();
-        if (!select_query)
-            throw Exception(
-                ErrorCodes::LOGICAL_ERROR,
-                "Expected SELECT query from table function {}",
-                configuration->getEngineName());
+            auto * select_query = query->as<ASTSelectQuery>();
+            if (!select_query)
+                throw Exception(
+                    ErrorCodes::LOGICAL_ERROR,
+                    "Expected SELECT query from table function {}",
+                    configuration->getEngineName());
 
-        auto settings = select_query->settings();
-        if (settings)
-        {
-            auto & settings_ast = settings->as<ASTSetQuery &>();
-            if (settings_ast.changes.removeSetting("object_storage_cluster") && settings_ast.changes.empty())
+            auto settings = select_query->settings();
+            if (settings)
             {
-                select_query->setExpression(ASTSelectQuery::Expression::SETTINGS, {});
+                auto & settings_ast = settings->as<ASTSetQuery &>();
+                if (settings_ast.changes.removeSetting("object_storage_cluster") && settings_ast.changes.empty())
+                {
+                    select_query->setExpression(ASTSelectQuery::Expression::SETTINGS, {});
+                }
+                /// No throw if not found - `object_storage_cluster` can be global setting.
             }
-            /// No throw if not found - `object_storage_cluster` can be global setting.
         }
     }
     else
-    {
+    { /// *Cluster function has cluster name as first argument. Temporary remove it before add structure and format
         ASTPtr cluster_name_arg = args.front();
         args.erase(args.begin());
         configuration->addStructureAndFormatToArgsIfNeeded(args, structure, configuration->getFormat(), context, /*with_structure=*/true);

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -171,7 +171,8 @@ class StorageObjectStorageCluster : public IStorageCluster
     void updateQueryToSendIfNeeded(
         ASTPtr & query,
         const StorageSnapshotPtr & storage_snapshot,
-        const ContextPtr & context) override;
+        const ContextPtr & context,
+        bool make_cluster_function) override;
 
     bool isClusterSupported() const override;
 
@@ -205,7 +206,7 @@ class StorageObjectStorageCluster : public IStorageCluster
     SELECT * FROM s3(...) SETTINGS object_storage_cluster='cluster'
     to make distributed request over cluster 'cluster'.
     */
-    void updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context);
+    void updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context, bool make_cluster_function);
 
     const String engine_name;
     StorageObjectStorageConfigurationPtr configuration;

src/Storages/StorageFileCluster.cpp

@@ -81,7 +81,11 @@ StorageFileCluster::StorageFileCluster(
     setInMemoryMetadata(storage_metadata);
 }
 
-void StorageFileCluster::updateQueryToSendIfNeeded(DB::ASTPtr & query, const StorageSnapshotPtr & storage_snapshot, const DB::ContextPtr & context)
+void StorageFileCluster::updateQueryToSendIfNeeded(
+    DB::ASTPtr & query,
+    const StorageSnapshotPtr & storage_snapshot,
+    const DB::ContextPtr & context,
+    bool /*make_cluster_function*/)
 {
     auto * table_function = extractTableFunctionFromSelectQuery(query);
     if (!table_function)

src/Storages/StorageFileCluster.h

@@ -36,7 +36,11 @@ class StorageFileCluster : public IStorageCluster
         StorageMetadataPtr) const override;
 
 private:
-    void updateQueryToSendIfNeeded(ASTPtr & query, const StorageSnapshotPtr & storage_snapshot, const ContextPtr & context) override;
+    void updateQueryToSendIfNeeded(
+        ASTPtr & query,
+        const StorageSnapshotPtr & storage_snapshot,
+        const ContextPtr & context,
+        bool /*make_cluster_function*/) override;
 
     Strings paths;
     String filename;

src/Storages/StorageURLCluster.cpp

@@ -112,7 +112,11 @@ StorageURLCluster::StorageURLCluster(
     setInMemoryMetadata(storage_metadata);
 }
 
-void StorageURLCluster::updateQueryToSendIfNeeded(ASTPtr & query, const StorageSnapshotPtr & storage_snapshot, const ContextPtr & context)
+void StorageURLCluster::updateQueryToSendIfNeeded(
+    ASTPtr & query,
+    const StorageSnapshotPtr & storage_snapshot,
+    const ContextPtr & context,
+    bool /*make_cluster_function*/)
 {
     auto * table_function = extractTableFunctionFromSelectQuery(query);
     if (!table_function)

src/Storages/StorageURLCluster.h

@@ -39,7 +39,11 @@ class StorageURLCluster : public IStorageCluster
         StorageMetadataPtr) const override;
 
 private:
-    void updateQueryToSendIfNeeded(ASTPtr & query, const StorageSnapshotPtr & storage_snapshot, const ContextPtr & context) override;
+    void updateQueryToSendIfNeeded(
+        ASTPtr & query,
+        const StorageSnapshotPtr & storage_snapshot,
+        const ContextPtr & context,
+        bool /*make_cluster_function*/) override;
 
     String uri;
     String format_name;