Merge pull request #835 from Altinity/feature/antalya-25.3/object_storage_remote_initiator

25.3 Antalya port of #756 - object storage cluster function

Author: Enmk

src/Core/Settings.cpp

@@ -6135,8 +6135,12 @@ Possible values:
     /** Experimental tsToGrid aggregate function. */ \
     DECLARE(Bool, allow_experimental_ts_to_grid_aggregate_function, false, R"(
 Experimental tsToGrid aggregate function for Prometheus-like timeseries resampling. Cloud only
+)", EXPERIMENTAL) \
+    DECLARE(Bool, object_storage_remote_initiator, false, R"(
+Execute request to object storage as remote on one of object_storage_cluster nodes.
 )", EXPERIMENTAL) \
     \
+
     /* ####################################################### */ \
     /* ############ END OF EXPERIMENTAL FEATURES ############# */ \
     /* ####################################################### */ \

src/Core/SettingsChangesHistory.cpp

@@ -87,6 +87,7 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"allow_experimental_database_unity_catalog", false, false, "Allow experimental database engine DataLakeCatalog with catalog_type = 'unity'"},
             {"allow_experimental_database_glue_catalog", false, false, "Allow experimental database engine DataLakeCatalog with catalog_type = 'glue'"},
             {"use_page_cache_with_distributed_cache", false, false, "New setting"},
+            {"object_storage_remote_initiator", false, false, "New setting."},
             {"use_iceberg_metadata_files_cache", true, true, "New setting"},
             {"use_query_condition_cache", false, false, "New setting."},
             {"iceberg_timestamp_ms", 0, 0, "New setting."},

src/IO/ReadBufferFromS3.cpp

@@ -431,6 +431,12 @@ Aws::S3::Model::GetObjectResult ReadBufferFromS3::sendRequest(size_t attempt, si
             log, "Read S3 object. Bucket: {}, Key: {}, Version: {}, Offset: {}",
             bucket, key, version_id.empty() ? "Latest" : version_id, range_begin);
     }
+    else
+    {
+        LOG_TEST(
+            log, "Read S3 object. Bucket: {}, Key: {}, Version: {}",
+            bucket, key, version_id.empty() ? "Latest" : version_id);
+    }
 
     ProfileEvents::increment(ProfileEvents::S3GetObject);
     if (client_ptr->isClientForDisk())

src/Storages/IStorageCluster.cpp

@@ -1,5 +1,8 @@
 #include <Storages/IStorageCluster.h>
 
+#include <pcg_random.hpp>
+#include <Common/randomSeed.h>
+
 #include <Common/Exception.h>
 #include <Core/Settings.h>
 #include <Core/QueryProcessingStage.h>
@@ -12,6 +15,7 @@
 #include <Interpreters/AddDefaultDatabaseVisitor.h>
 #include <Interpreters/TranslateQualifiedNamesVisitor.h>
 #include <Interpreters/InterpreterSelectQueryAnalyzer.h>
+#include <Planner/Utils.h>
 #include <Processors/Sources/NullSource.h>
 #include <Processors/Sources/RemoteSource.h>
 #include <QueryPipeline/narrowPipe.h>
@@ -21,6 +25,9 @@
 #include <Storages/IStorage.h>
 #include <Storages/SelectQueryInfo.h>
 #include <Storages/StorageDictionary.h>
+#include <Storages/StorageDistributed.h>
+#include <TableFunctions/TableFunctionFactory.h>
+#include <Storages/extractTableFunctionFromSelectQuery.h>
 
 #include <algorithm>
 #include <memory>
@@ -39,6 +46,7 @@ namespace Setting
     extern const SettingsString cluster_for_parallel_replicas;
     extern const SettingsNonZeroUInt64 max_parallel_replicas;
     extern const SettingsUInt64 object_storage_max_nodes;
+    extern const SettingsBool object_storage_remote_initiator;
 }
 
 namespace ErrorCodes
@@ -96,15 +104,16 @@ void IStorageCluster::read(
 
     storage_snapshot->check(column_names);
 
-    updateBeforeRead(context);
-    auto cluster = getClusterImpl(context, cluster_name_from_settings, context->getSettingsRef()[Setting::object_storage_max_nodes]);
+    const auto & settings = context->getSettingsRef();
+
+    auto cluster = getClusterImpl(context, cluster_name_from_settings, settings[Setting::object_storage_max_nodes]);
 
     /// Calculate the header. This is significant, because some columns could be thrown away in some cases like query with count(*)
 
     Block sample_block;
     ASTPtr query_to_send = query_info.query;
 
-    if (context->getSettingsRef()[Setting::allow_experimental_analyzer])
+    if (settings[Setting::allow_experimental_analyzer])
     {
         sample_block = InterpreterSelectQueryAnalyzer::getSampleBlock(query_info.query, context, SelectQueryOptions(processed_stage));
     }
@@ -117,6 +126,17 @@ void IStorageCluster::read(
 
     updateQueryToSendIfNeeded(query_to_send, storage_snapshot, context);
 
+    if (settings[Setting::object_storage_remote_initiator])
+    {
+        auto storage_and_context = convertToRemote(cluster, context, cluster_name_from_settings, 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;
+    }
+
     RestoreQualifiedNamesVisitor::Data data;
     data.distributed_table = DatabaseAndTableWithAlias(*getTableExpression(query_info.query->as<ASTSelectQuery &>(), 0));
     data.remote_table.database = context->getCurrentDatabase();
@@ -144,6 +164,62 @@ void IStorageCluster::read(
     query_plan.addStep(std::move(reading));
 }
 
+IStorageCluster::RemoteCallVariables IStorageCluster::convertToRemote(
+    ClusterPtr cluster,
+    ContextPtr context,
+    const std::string & cluster_name_from_settings,
+    ASTPtr query_to_send)
+{
+    auto host_addresses = cluster->getShardsAddresses();
+    if (host_addresses.empty())
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Empty cluster {}", cluster_name_from_settings);
+
+    static pcg64 rng(randomSeed());
+    size_t shard_num = rng() % host_addresses.size();
+    auto shard_addresses = host_addresses[shard_num];
+    /// After getClusterImpl each shard must have exactly 1 replica
+    if (shard_addresses.size() != 1)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Size of shard {} in cluster {} is not equal 1", shard_num, cluster_name_from_settings);
+    auto host_name = shard_addresses[0].toString();
+
+    LOG_INFO(log, "Choose remote initiator '{}'", host_name);
+
+    bool secure = shard_addresses[0].secure == Protocol::Secure::Enable;
+    std::string remote_function_name = secure ? "remoteSecure" : "remote";
+
+    /// 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);
+
+    auto * select_query = query_to_send->as<ASTSelectQuery>();
+    if (!select_query)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Expected SELECT query");
+
+    auto query_settings = select_query->settings();
+    if (query_settings)
+    {
+        auto & settings_ast = query_settings->as<ASTSetQuery &>();
+        if (settings_ast.changes.removeSetting("object_storage_remote_initiator") && 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");
+
+    auto remote_query = makeASTFunction(remote_function_name, std::make_shared<ASTLiteral>(host_name), table_expression->table_function);
+
+    table_expression->table_function = remote_query;
+
+    auto remote_function = TableFunctionFactory::instance().get(remote_query, new_context);
+
+    auto storage = remote_function->execute(query_to_send, new_context, remote_function_name);
+
+    return RemoteCallVariables{storage, new_context};
+}
+
 SinkToStoragePtr IStorageCluster::write(
     const ASTPtr & query,
     const StorageMetadataPtr & metadata_snapshot,

src/Storages/IStorageCluster.h

@@ -58,9 +58,20 @@ class IStorageCluster : public IStorage
     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*/) {}
 
+    struct RemoteCallVariables
+    {
+        StoragePtr storage;
+        ContextPtr context;
+    };
+
+    RemoteCallVariables convertToRemote(
+        ClusterPtr cluster,
+        ContextPtr context,
+        const std::string & cluster_name_from_settings,
+        ASTPtr query_to_send);
+
     virtual void readFallBackToPure(
         QueryPlan & /* query_plan */,
         const Names & /* column_names */,

src/Storages/ObjectStorage/DataLakes/DataLakeConfiguration.h

@@ -85,6 +85,16 @@ class DataLakeConfiguration : public BaseStorageConfiguration, public std::enabl
         return std::nullopt;
     }
 
+    std::optional<String> tryGetSamplePathFromMetadata() const override
+    {
+        if (!current_metadata)
+            return std::nullopt;
+        auto data_files = current_metadata->getDataFiles();
+        if (!data_files.empty())
+            return data_files[0];
+        return std::nullopt;
+    }
+
     std::optional<size_t> totalRows() override
     {
         if (!current_metadata)
@@ -465,8 +475,14 @@ class StorageIcebergConfiguration : public StorageObjectStorage::Configuration,
         createDynamicStorage(type);
     }
 
+    std::optional<String> tryGetSamplePathFromMetadata() const override
+    {
+        return getImpl().tryGetSamplePathFromMetadata();
+    }
+
     virtual void assertInitialized() const override { return getImpl().assertInitialized(); }
 
+
 private:
     inline StorageObjectStorage::Configuration & getImpl() const
     {

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadata.h

@@ -41,6 +41,8 @@ class DeltaLakeMetadata final : public IDataLakeMetadata
 
     DeltaLakeMetadata(ObjectStoragePtr object_storage_, ConfigurationObserverPtr configuration_, ContextPtr context_);
 
+    Strings getDataFiles() const override { return data_files; }
+
     NamesAndTypesList getTableSchema() const override { return schema; }
 
     DeltaLakePartitionColumns getPartitionColumns() const { return partition_columns; }

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadataDeltaKernel.cpp

@@ -32,6 +32,11 @@ bool DeltaLakeMetadataDeltaKernel::update(const ContextPtr &)
     return table_snapshot->update();
 }
 
+Strings DeltaLakeMetadataDeltaKernel::getDataFiles() const
+{
+    throwNotImplemented("getDataFiles()");
+}
+
 ObjectIterator DeltaLakeMetadataDeltaKernel::iterate(
     const ActionsDAG * filter_dag,
     FileProgressCallback callback,

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadataDeltaKernel.h

@@ -40,6 +40,8 @@ class DeltaLakeMetadataDeltaKernel final : public IDataLakeMetadata
 
     bool update(const ContextPtr & context) override;
 
+    Strings getDataFiles() const override;
+
     NamesAndTypesList getTableSchema() const override;
 
     NamesAndTypesList getReadSchema() const override;

src/Storages/ObjectStorage/DataLakes/HudiMetadata.cpp

@@ -91,19 +91,19 @@ HudiMetadata::HudiMetadata(ObjectStoragePtr object_storage_, ConfigurationObserv
 {
 }
 
-Strings HudiMetadata::getDataFiles(const ActionsDAG *) const
+Strings HudiMetadata::getDataFiles() const
 {
     if (data_files.empty())
         data_files = getDataFilesImpl();
     return data_files;
 }
 
 ObjectIterator HudiMetadata::iterate(
-    const ActionsDAG * filter_dag,
+    const ActionsDAG * /* filter_dag */,
     FileProgressCallback callback,
     size_t /* list_batch_size */) const
 {
-    return createKeysIterator(getDataFiles(filter_dag), object_storage, callback);
+    return createKeysIterator(getDataFiles(), object_storage, callback);
 }
 
 }