rename a few things

Author: arthurpassos

src/CMakeLists.txt

@@ -130,6 +130,7 @@ add_headers_and_sources(dbms Storages/ObjectStorage/Azure)
 add_headers_and_sources(dbms Storages/ObjectStorage/S3)
 add_headers_and_sources(dbms Storages/ObjectStorage/HDFS)
 add_headers_and_sources(dbms Storages/ObjectStorage/Local)
+add_headers_and_sources(dbms Storages/ObjectStorage/MergeTree)
 add_headers_and_sources(dbms Storages/ObjectStorage/DataLakes)
 add_headers_and_sources(dbms Storages/ObjectStorage/DataLakes/Iceberg)
 add_headers_and_sources(dbms Storages/ObjectStorage/DataLakes/DeltaLake)

src/Storages/ObjectStorage/MergeTree/StorageObjectStorageMergeTreePartImporterSink.cpp

@@ -0,0 +1,63 @@
+#include <Storages/ObjectStorage/MergeTree/StorageObjectStorageMergeTreePartImporterSink.h>
+
+namespace DB
+{
+
+StorageObjectStorageMergeTreePartImporterSink::StorageObjectStorageMergeTreePartImporterSink(
+        const DataPartPtr & part_,
+        const std::string & path_,
+        const ObjectStoragePtr & object_storage_,
+        const ConfigurationPtr & configuration_,
+        const std::optional<FormatSettings> & format_settings_,
+        const Block & sample_block_,
+        const std::function<void(MergeTreePartImportStats)> & part_log_,
+        const ContextPtr & context_)
+        : SinkToStorage(sample_block_)
+        , object_storage(object_storage_)
+        , configuration(configuration_)
+        , format_settings(format_settings_)
+        , sample_block(sample_block_)
+        , context(context_)
+        , part_log(part_log_)
+{
+    stats.part = part_;
+    stats.file_path = path_;
+    sink = std::make_shared<StorageObjectStorageSink>(
+        stats.file_path,
+        object_storage,
+        configuration,
+        format_settings,
+        sample_block,
+        context);
+}
+
+String StorageObjectStorageMergeTreePartImporterSink::getName() const
+{
+    return "StorageObjectStorageMergeTreePartImporterSink";
+}
+
+void StorageObjectStorageMergeTreePartImporterSink::consume(Chunk & chunk)
+{
+    sink->consume(chunk);
+
+    stats.read_bytes += chunk.bytes();
+    stats.read_rows += chunk.getNumRows();
+}
+
+void StorageObjectStorageMergeTreePartImporterSink::onFinish()
+{
+    sink->onFinish();
+    if (const auto object_metadata = object_storage->tryGetObjectMetadata(stats.file_path))
+    {
+        stats.bytes_on_disk = object_metadata->size_bytes;
+    }
+    part_log(stats);
+}
+
+void StorageObjectStorageMergeTreePartImporterSink::onException(std::exception_ptr exception)
+{
+    sink->onException(exception);
+    part_log(stats);
+}
+
+}

src/Storages/ObjectStorage/MergeTree/StorageObjectStorageMergeTreePartImporterSink.h

@@ -0,0 +1,64 @@
+#pragma once
+
+#include <Interpreters/Context.h>
+#include <Storages/ObjectStorage/StorageObjectStorageSink.h>
+#include "Core/Settings.h"
+#include "Disks/ObjectStorages/IObjectStorage.h"
+#include "Disks/ObjectStorages/StoredObject.h"
+#include "Formats/FormatFactory.h"
+#include "IO/CompressionMethod.h"
+#include "Processors/Formats/IOutputFormat.h"
+#include "Storages/MergeTree/IMergeTreeDataPart.h"
+
+namespace DB
+{
+
+struct MergeTreePartImportStats
+{
+    ExecutionStatus status;
+    std::size_t bytes_on_disk = 0;
+    std::size_t read_rows = 0;
+    std::size_t read_bytes = 0;
+    std::string file_path = "";
+    DataPartPtr part = nullptr;
+};
+
+/*
+ * Wrapper around `StorageObjectsStorageSink` that takes care of accounting & metrics for partition export
+ */
+class StorageObjectStorageMergeTreePartImporterSink : public SinkToStorage
+{
+public:
+    using ConfigurationPtr = StorageObjectStorage::ConfigurationPtr;
+
+    StorageObjectStorageMergeTreePartImporterSink(
+        const DataPartPtr & part_,
+        const std::string & path_,
+        const ObjectStoragePtr & object_storage_,
+        const ConfigurationPtr & configuration_,
+        const std::optional<FormatSettings> & format_settings_,
+        const Block & sample_block_,
+        const std::function<void(MergeTreePartImportStats)> & part_log_,
+        const ContextPtr & context_);
+
+    String getName() const override;
+
+    void consume(Chunk & chunk) override;
+
+    void onFinish() override;
+
+    void onException(std::exception_ptr exception) override;
+
+private:
+    std::shared_ptr<StorageObjectStorageSink> sink;
+    ObjectStoragePtr object_storage;
+    ConfigurationPtr configuration;
+    std::optional<FormatSettings> format_settings;
+    Block sample_block;
+    ContextPtr context;
+    std::function<void(MergeTreePartImportStats)> part_log;
+
+    MergeTreePartImportStats stats;
+};
+
+}

src/Storages/ObjectStorage/MergeTree/StorageObjectStorageSinkMTPartImportDecorator.cpp

@@ -1,15 +1,15 @@
 #include <Core/ColumnWithTypeAndName.h>
 #include <Storages/MergeTree/IMergeTreeDataPart.h>
 #include <Storages/ObjectStorage/StorageObjectStorage.h>
-#include "MergeTree/StorageObjectStorageSinkMTPartImportDecorator.h"
+#include "MergeTree/StorageObjectStorageMergeTreePartImporterSink.h"
 
-#include <Common/logger_useful.h>
 #include <Core/Settings.h>
 #include <Formats/FormatFactory.h>
-#include <Parsers/ASTInsertQuery.h>
 #include <Formats/ReadSchemaUtils.h>
-#include <QueryPipeline/QueryPipelineBuilder.h>
 #include <Interpreters/Context.h>
+#include <Parsers/ASTInsertQuery.h>
+#include <QueryPipeline/QueryPipelineBuilder.h>
+#include <Common/logger_useful.h>
 
 #include <Processors/Sources/NullSource.h>
 #include <Processors/QueryPlan/QueryPlan.h>
@@ -663,7 +663,7 @@ void StorageObjectStorage::importMergeTreePartition(
         auto builder = plan_for_part.buildQueryPipeline(optimization_settings, pipeline_settings);
         auto pipeline = QueryPipelineBuilder::getPipeline(std::move(*builder));
 
-        auto sink = std::make_shared<StorageObjectStorageSinkMTPartImportDecorator>(
+        auto sink = std::make_shared<StorageObjectStorageMergeTreePartImporterSink>(
             data_part,
             file_path,
             object_storage,

src/Storages/ObjectStorage/MergeTree/StorageObjectStorageSinkMTPartImportDecorator.h

@@ -8,6 +8,7 @@ namespace DB
 {
 class StorageObjectStorageSink : public SinkToStorage
 {
+friend class StorageObjectStorageMergeTreePartImporterSink;
 public:
     using ConfigurationPtr = StorageObjectStorage::ConfigurationPtr;
 

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -48,7 +48,7 @@
 #include <Common/escapeForFileName.h>
 #include "Core/BackgroundSchedulePool.h"
 #include "Core/Names.h"
-#include "ObjectStorage/MergeTree/StorageObjectStorageSinkMTPartImportDecorator.h"
+#include "ObjectStorage/MergeTree/StorageObjectStorageMergeTreePartImporterSink.h"
 
 namespace DB
 {