Distributed request to tables with Object Storage Engines

src/Databases/Iceberg/DatabaseIceberg.cpp

@@ -12,6 +12,7 @@
 #include <Storages/ConstraintsDescription.h>
 #include <Storages/StorageNull.h>
 #include <Storages/ObjectStorage/DataLakes/DataLakeConfiguration.h>
+#include <Storages/ObjectStorage/StorageObjectStorageCluster.h>
 
 #include <Interpreters/evaluateConstantExpression.h>
 #include <Interpreters/Context.h>
@@ -37,6 +38,7 @@ namespace DatabaseIcebergSetting
     extern const DatabaseIcebergSettingsString storage_endpoint;
     extern const DatabaseIcebergSettingsString oauth_server_uri;
     extern const DatabaseIcebergSettingsBool vended_credentials;
+    extern const DatabaseIcebergSettingsString object_storage_cluster;
 }
 namespace Setting
 {
@@ -235,19 +237,20 @@ StoragePtr DatabaseIceberg::tryGetTable(const String & name, ContextPtr context_
     /// no table structure in table definition AST.
     StorageObjectStorage::Configuration::initialize(*configuration, args, context_, /* with_table_structure */false, std::move(storage_settings));
 
-    return std::make_shared<StorageObjectStorage>(
+    auto cluster_name = settings[DatabaseIcebergSetting::object_storage_cluster].value;
+
+    return std::make_shared<StorageObjectStorageCluster>(
+        cluster_name,
         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);
+        /* format_settings */ getFormatSettings(context_),
+        /* mode */ LoadingStrictnessLevel::CREATE,
+        /* partition_by */nullptr);
 }
 
 DatabaseTablesIteratorPtr DatabaseIceberg::getTablesIterator(

src/Databases/Iceberg/DatabaseIcebergSettings.cpp

@@ -23,6 +23,7 @@ namespace ErrorCodes
     DECLARE(String, warehouse, "", "Warehouse name inside the catalog", 0)             \
     DECLARE(String, auth_header, "", "Authorization header of format 'Authorization: <scheme> <auth_info>'", 0)             \
     DECLARE(String, storage_endpoint, "", "Object storage endpoint", 0) \
+    DECLARE(String, object_storage_cluster, "", "Cluster for distributed requests", 0) \
 
 #define LIST_OF_DATABASE_ICEBERG_SETTINGS(M, ALIAS) \
     DATABASE_ICEBERG_RELATED_SETTINGS(M, ALIAS)

src/Parsers/FunctionSecretArgumentsFinder.h

@@ -102,7 +102,7 @@ class FunctionSecretArgumentsFinder
         }
         else if ((function->name() == "s3") || (function->name() == "cosn") || (function->name() == "oss") ||
                  (function->name() == "deltaLake") || (function->name() == "hudi") || (function->name() == "iceberg") ||
-                 (function->name() == "gcs"))
+                 (function->name() == "icebergS3") || (function->name() == "gcs"))
         {
             /// s3('url', 'aws_access_key_id', 'aws_secret_access_key', ...)
             findS3FunctionSecretArguments(/* is_cluster_function= */ false);

src/Storages/IStorage.h

@@ -71,6 +71,9 @@ class RestorerFromBackup;
 
 class ConditionSelectivityEstimator;
 
+class IObjectStorage;
+using ObjectStoragePtr = std::shared_ptr<IObjectStorage>;
+
 struct ColumnSize
 {
     size_t marks = 0;

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_,
@@ -73,13 +78,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_from_settings = getClusterName(context);
+
+    if (cluster_name_from_settings.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_from_settings);
 
     /// Calculate the header. This is significant, because some columns could be thrown away in some cases like query with count(*)
 
@@ -126,6 +139,20 @@ void IStorageCluster::read(
     query_plan.addStep(std::move(reading));
 }
 
+SinkToStoragePtr IStorageCluster::write(
+    const ASTPtr & query,
+    const StorageMetadataPtr & metadata_snapshot,
+    ContextPtr context,
+    bool async_insert)
+{
+    auto cluster_name_from_settings = getClusterName(context);
+
+    if (cluster_name_from_settings.empty())
+        return writeFallBackToPure(query, metadata_snapshot, context, async_insert);
+
+    throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method write is not supported by storage {}", getName());
+}
+
 void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &)
 {
     createExtension(nullptr);
@@ -196,9 +223,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());
 }
 
 }

src/Storages/IStorageCluster.h

@@ -29,10 +29,16 @@ 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;
+    SinkToStoragePtr write(
+        const ASTPtr & query,
+        const StorageMetadataPtr & metadata_snapshot,
+        ContextPtr context,
+        bool async_insert) override;
+
+    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,11 +49,38 @@ class IStorageCluster : public IStorage
     bool supportsOptimizationToSubcolumns() const override { return false; }
     bool supportsTrivialCountOptimization(const StorageSnapshotPtr &, ContextPtr) const override { return true; }
 
+    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 */)
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method readFallBackToPure is not supported by storage {}", getName());
+    }
+
+    virtual SinkToStoragePtr writeFallBackToPure(
+        const ASTPtr & /*query*/,
+        const StorageMetadataPtr & /*metadata_snapshot*/,
+        ContextPtr /*context*/,
+        bool /*async_insert*/)
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method writeFallBackToPure is not supported by storage {}", getName());
+    }
+
 private:
+    static ClusterPtr getClusterImpl(ContextPtr context, const String & cluster_name_);
+
     LoggerPtr log;
     String cluster_name;
 };

src/Storages/ObjectStorage/Azure/Configuration.cpp

@@ -132,8 +132,6 @@ void StorageAzureConfiguration::fromNamedCollection(const NamedCollection & coll
 
     String connection_url;
     String container_name;
-    std::optional<String> account_name;
-    std::optional<String> account_key;
 
     if (collection.has("connection_string"))
         connection_url = collection.get<String>("connection_string");
@@ -173,14 +171,10 @@ void StorageAzureConfiguration::fromAST(ASTs & engine_args, ContextPtr context,
 
     std::unordered_map<std::string_view, size_t> engine_args_to_idx;
 
-
     String connection_url = checkAndGetLiteralArgument<String>(engine_args[0], "connection_string/storage_account_url");
     String container_name = checkAndGetLiteralArgument<String>(engine_args[1], "container");
     blob_path = checkAndGetLiteralArgument<String>(engine_args[2], "blobpath");
 
-    std::optional<String> account_name;
-    std::optional<String> account_key;
-
     auto is_format_arg = [] (const std::string & s) -> bool
     {
         return s == "auto" || FormatFactory::instance().getAllFormats().contains(Poco::toLower(s));
@@ -444,6 +438,22 @@ void StorageAzureConfiguration::addStructureAndFormatToArgsIfNeeded(
     }
 }
 
+ASTPtr StorageAzureConfiguration::createArgsWithAccessData() const
+{
+    auto arguments = std::make_shared<ASTExpressionList>();
+
+    arguments->children.push_back(std::make_shared<ASTLiteral>(connection_params.endpoint.storage_account_url));
+    arguments->children.push_back(std::make_shared<ASTIdentifier>(connection_params.endpoint.container_name));
+    arguments->children.push_back(std::make_shared<ASTLiteral>(blob_path));
+    if (account_name && account_key)
+    {
+        arguments->children.push_back(std::make_shared<ASTLiteral>(*account_name));
+        arguments->children.push_back(std::make_shared<ASTLiteral>(*account_key));
+    }
+
+    return arguments;
+}
+
 }
 
 #endif

src/Storages/ObjectStorage/Azure/Configuration.h

@@ -79,13 +79,17 @@ class StorageAzureConfiguration : public StorageObjectStorage::Configuration
         ContextPtr context,
         bool with_structure) override;
 
+    ASTPtr createArgsWithAccessData() const override;
+
 protected:
     void fromNamedCollection(const NamedCollection & collection, ContextPtr context) override;
     void fromAST(ASTs & args, ContextPtr context, bool with_structure) override;
 
     std::string blob_path;
     std::vector<String> blobs_paths;
     AzureBlobStorage::ConnectionParams connection_params;
+    std::optional<std::string> account_name;
+    std::optional<std::string> account_key;
 };
 
 }

src/Storages/ObjectStorage/HDFS/Configuration.cpp

@@ -235,6 +235,13 @@ void StorageHDFSConfiguration::addStructureAndFormatToArgsIfNeeded(
     }
 }
 
+ASTPtr StorageHDFSConfiguration::createArgsWithAccessData() const
+{
+    auto arguments = std::make_shared<ASTExpressionList>();
+    arguments->children.push_back(std::make_shared<ASTLiteral>(url + path));
+    return arguments;
+}
+
 }
 
 #endif

src/Storages/ObjectStorage/HDFS/Configuration.h

@@ -65,6 +65,8 @@ class StorageHDFSConfiguration : public StorageObjectStorage::Configuration
         ContextPtr context,
         bool with_structure) override;
 
+    ASTPtr createArgsWithAccessData() const override;
+
 private:
     void fromNamedCollection(const NamedCollection &, ContextPtr context) override;
     void fromAST(ASTs & args, ContextPtr, bool /* with_structure */) override;

src/Storages/ObjectStorage/S3/Configuration.cpp

@@ -363,11 +363,11 @@ void StorageS3Configuration::fromAST(ASTs & args, ContextPtr context, bool with_
 
     if (engine_args_to_idx.contains("format"))
     {
-        format = checkAndGetLiteralArgument<String>(args[engine_args_to_idx["format"]], "format");
+        auto format_ = checkAndGetLiteralArgument<String>(args[engine_args_to_idx["format"]], "format");
         /// Set format to configuration only of it's not 'auto',
         /// because we can have default format set in configuration.
-        if (format != "auto")
-            format = format;
+        if (format_ != "auto")
+            format = format_;
     }
 
     if (engine_args_to_idx.contains("structure"))
@@ -585,6 +585,30 @@ void StorageS3Configuration::addStructureAndFormatToArgsIfNeeded(
     }
 }
 
+ASTPtr StorageS3Configuration::createArgsWithAccessData() const
+{
+    auto arguments = std::make_shared<ASTExpressionList>();
+
+    arguments->children.push_back(std::make_shared<ASTLiteral>(url.uri_str));
+    if (auth_settings[S3AuthSetting::no_sign_request])
+    {
+        arguments->children.push_back(std::make_shared<ASTLiteral>("NOSIGN"));
+    }
+    else
+    {
+        arguments->children.push_back(std::make_shared<ASTLiteral>(auth_settings[S3AuthSetting::access_key_id].value));
+        arguments->children.push_back(std::make_shared<ASTLiteral>(auth_settings[S3AuthSetting::secret_access_key].value));
+        if (!auth_settings[S3AuthSetting::session_token].value.empty())
+            arguments->children.push_back(std::make_shared<ASTLiteral>(auth_settings[S3AuthSetting::session_token].value));
+        if (format != "auto")
+            arguments->children.push_back(std::make_shared<ASTLiteral>(format));
+        if (!compression_method.empty())
+            arguments->children.push_back(std::make_shared<ASTLiteral>(compression_method));
+    }
+
+    return arguments;
+}
+
 }
 
 #endif

src/Storages/ObjectStorage/S3/Configuration.h

@@ -94,6 +94,8 @@ class StorageS3Configuration : public StorageObjectStorage::Configuration
         ContextPtr context,
         bool with_structure) override;
 
+    ASTPtr createArgsWithAccessData() const override;
+
 private:
     void fromNamedCollection(const NamedCollection & collection, ContextPtr context) override;
     void fromAST(ASTs & args, ContextPtr context, bool with_structure) override;

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -60,6 +60,9 @@ String StorageObjectStorage::getPathSample(ContextPtr context)
     if (context->getSettingsRef()[Setting::use_hive_partitioning])
         local_distributed_processing = false;
 
+    if (!configuration->isArchive() && !configuration->isPathWithGlobs() && !local_distributed_processing)
+        return configuration->getPath();
+
     auto file_iterator = StorageObjectStorageSource::createFileIterator(
         configuration,
         query_settings,
@@ -72,9 +75,6 @@ String StorageObjectStorage::getPathSample(ContextPtr context)
         {} // file_progress_callback
     );
 
-    if (!configuration->isArchive() && !configuration->isPathWithGlobs() && !local_distributed_processing)
-        return configuration->getPath();
-
     if (auto file = file_iterator->next(0))
         return file->getPath();
     return "";
@@ -165,6 +165,13 @@ void StorageObjectStorage::Configuration::update(ObjectStoragePtr object_storage
 {
     IObjectStorage::ApplyNewSettingsOptions options{.allow_client_change = !isStaticConfiguration()};
     object_storage_ptr->applyNewSettings(context->getConfigRef(), getTypeName() + ".", context, options);
+    updated = true;
+}
+
+void StorageObjectStorage::Configuration::updateIfRequired(ObjectStoragePtr object_storage_ptr, ContextPtr local_context)
+{
+    if (!updated)
+        update(object_storage_ptr, local_context);
 }
 
 bool StorageObjectStorage::hasExternalDynamicMetadata() const

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -246,7 +246,13 @@ class StorageObjectStorage::Configuration
     String structure = "auto";
 
     virtual void update(ObjectStoragePtr object_storage, ContextPtr local_context);
+    void updateIfRequired(ObjectStoragePtr object_storage, ContextPtr local_context);
 
+    /// Create arguments for table function with path and access parameters
+    virtual ASTPtr createArgsWithAccessData() const
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method createArgsWithAccessData is not supported by storage {}", getEngineName());
+    }
 
 protected:
     virtual void fromNamedCollection(const NamedCollection & collection, ContextPtr context) = 0;
@@ -256,6 +262,7 @@ class StorageObjectStorage::Configuration
     void assertInitialized() const;
 
     bool initialized = false;
+    std::atomic<bool> updated = false;
     DataLakePartitionColumns partition_columns;
 
     bool allow_dynamic_metadata_for_data_lakes;