2021-03-19 21:49:18 +00:00
|
|
|
#include "Storages/StorageS3Distributed.h"
|
|
|
|
|
2021-04-10 14:58:29 +00:00
|
|
|
#if !defined(ARCADIA_BUILD)
|
2021-03-24 21:02:21 +00:00
|
|
|
#include <Common/config.h>
|
2021-04-10 14:58:29 +00:00
|
|
|
#endif
|
2021-03-24 21:02:21 +00:00
|
|
|
|
|
|
|
#if USE_AWS_S3
|
|
|
|
|
2021-03-22 17:12:31 +00:00
|
|
|
#include "Common/Exception.h"
|
2021-03-19 21:49:18 +00:00
|
|
|
#include <Common/Throttler.h>
|
|
|
|
#include "Client/Connection.h"
|
2021-03-22 23:06:10 +00:00
|
|
|
#include "Core/QueryProcessingStage.h"
|
2021-04-08 00:09:15 +00:00
|
|
|
#include <Core/UUID.h>
|
2021-03-19 21:49:18 +00:00
|
|
|
#include "DataStreams/RemoteBlockInputStream.h"
|
|
|
|
#include <Columns/ColumnsNumber.h>
|
|
|
|
#include <DataTypes/DataTypesNumber.h>
|
2021-03-22 17:12:31 +00:00
|
|
|
#include <DataTypes/DataTypeString.h>
|
|
|
|
#include <IO/ReadBufferFromS3.h>
|
|
|
|
#include <IO/ReadHelpers.h>
|
|
|
|
#include <IO/WriteBufferFromS3.h>
|
|
|
|
#include <IO/WriteHelpers.h>
|
2021-03-26 15:33:14 +00:00
|
|
|
#include <Interpreters/Context.h>
|
|
|
|
#include <Interpreters/getHeaderForProcessingStage.h>
|
|
|
|
#include <Interpreters/SelectQueryOptions.h>
|
|
|
|
#include <Interpreters/InterpreterSelectQuery.h>
|
|
|
|
#include <Interpreters/getTableExpressions.h>
|
2021-03-22 17:12:31 +00:00
|
|
|
#include <Formats/FormatFactory.h>
|
|
|
|
#include <DataStreams/IBlockOutputStream.h>
|
|
|
|
#include <DataStreams/AddingDefaultsBlockInputStream.h>
|
|
|
|
#include <DataStreams/narrowBlockInputStreams.h>
|
|
|
|
#include <Processors/Formats/InputStreamFromInputFormat.h>
|
2021-03-19 21:49:18 +00:00
|
|
|
#include <Processors/Pipe.h>
|
|
|
|
#include <Processors/Sources/SourceFromInputStream.h>
|
2021-04-10 14:58:29 +00:00
|
|
|
#include "Processors/Sources/SourceWithProgress.h"
|
2021-04-06 19:18:45 +00:00
|
|
|
#include <Processors/Sources/RemoteSource.h>
|
2021-03-19 21:49:18 +00:00
|
|
|
#include <Parsers/queryToString.h>
|
2021-03-24 18:36:31 +00:00
|
|
|
#include <Parsers/ASTTablesInSelectQuery.h>
|
2021-03-26 15:33:14 +00:00
|
|
|
#include <Storages/IStorage.h>
|
|
|
|
#include <Storages/SelectQueryInfo.h>
|
|
|
|
#include <common/logger_useful.h>
|
2021-03-19 21:49:18 +00:00
|
|
|
|
2021-03-26 15:33:14 +00:00
|
|
|
#include <aws/core/auth/AWSCredentials.h>
|
|
|
|
#include <aws/s3/S3Client.h>
|
|
|
|
#include <aws/s3/model/ListObjectsV2Request.h>
|
2021-03-22 17:12:31 +00:00
|
|
|
|
2021-03-23 17:58:29 +00:00
|
|
|
#include <ios>
|
2021-03-19 21:49:18 +00:00
|
|
|
#include <memory>
|
|
|
|
#include <string>
|
|
|
|
#include <thread>
|
2021-03-22 17:12:31 +00:00
|
|
|
#include <cassert>
|
|
|
|
|
2021-03-19 21:49:18 +00:00
|
|
|
namespace DB
|
|
|
|
{
|
|
|
|
|
|
|
|
|
|
|
|
namespace ErrorCodes
|
|
|
|
{
|
|
|
|
extern const int LOGICAL_ERROR;
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
StorageS3Distributed::StorageS3Distributed(
|
2021-04-06 19:18:45 +00:00
|
|
|
const String & filename_,
|
2021-03-22 17:12:31 +00:00
|
|
|
const String & access_key_id_,
|
|
|
|
const String & secret_access_key_,
|
|
|
|
const StorageID & table_id_,
|
|
|
|
String cluster_name_,
|
|
|
|
const String & format_name_,
|
|
|
|
UInt64 max_connections_,
|
|
|
|
const ColumnsDescription & columns_,
|
|
|
|
const ConstraintsDescription & constraints_,
|
|
|
|
const Context & context_,
|
|
|
|
const String & compression_method_)
|
2021-03-19 21:49:18 +00:00
|
|
|
: IStorage(table_id_)
|
2021-04-08 00:09:15 +00:00
|
|
|
, client_auth{S3::URI{Poco::URI{filename_}}, access_key_id_, secret_access_key_, max_connections_, {}, {}} /// Client and settings will be updated later
|
2021-04-06 19:18:45 +00:00
|
|
|
, filename(filename_)
|
2021-03-19 21:49:18 +00:00
|
|
|
, cluster_name(cluster_name_)
|
2021-03-22 17:12:31 +00:00
|
|
|
, cluster(context_.getCluster(cluster_name)->getClusterWithReplicasAsShards(context_.getSettings()))
|
|
|
|
, format_name(format_name_)
|
|
|
|
, compression_method(compression_method_)
|
2021-03-19 21:49:18 +00:00
|
|
|
{
|
|
|
|
StorageInMemoryMetadata storage_metadata;
|
2021-03-22 17:12:31 +00:00
|
|
|
storage_metadata.setColumns(columns_);
|
|
|
|
storage_metadata.setConstraints(constraints_);
|
2021-03-19 21:49:18 +00:00
|
|
|
setInMemoryMetadata(storage_metadata);
|
2021-04-08 00:09:15 +00:00
|
|
|
StorageS3::updateClientAndAuthSettings(context_, client_auth);
|
2021-03-19 21:49:18 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
Pipe StorageS3Distributed::read(
|
|
|
|
const Names & column_names,
|
|
|
|
const StorageMetadataPtr & metadata_snapshot,
|
|
|
|
SelectQueryInfo & query_info,
|
|
|
|
const Context & context,
|
2021-03-23 17:58:29 +00:00
|
|
|
QueryProcessingStage::Enum processed_stage,
|
2021-03-22 17:12:31 +00:00
|
|
|
size_t max_block_size,
|
2021-03-19 21:49:18 +00:00
|
|
|
unsigned /*num_streams*/)
|
|
|
|
{
|
2021-04-08 00:09:15 +00:00
|
|
|
StorageS3::updateClientAndAuthSettings(context, client_auth);
|
|
|
|
|
2021-03-19 21:49:18 +00:00
|
|
|
/// Secondary query, need to read from S3
|
2021-03-26 15:33:14 +00:00
|
|
|
if (context.getClientInfo().query_kind == ClientInfo::QueryKind::SECONDARY_QUERY)
|
2021-03-22 17:12:31 +00:00
|
|
|
{
|
|
|
|
bool need_path_column = false;
|
|
|
|
bool need_file_column = false;
|
|
|
|
for (const auto & column : column_names)
|
|
|
|
{
|
|
|
|
if (column == "_path")
|
|
|
|
need_path_column = true;
|
|
|
|
if (column == "_file")
|
|
|
|
need_file_column = true;
|
|
|
|
}
|
2021-04-08 00:09:15 +00:00
|
|
|
|
2021-04-10 02:21:18 +00:00
|
|
|
/// Save callback not to capture context by reference of copy it.
|
|
|
|
auto file_iterator = std::make_shared<StorageS3Source::IteratorWrapper>(
|
|
|
|
[callback = context.getReadTaskCallback()]() -> std::optional<String> {
|
|
|
|
return callback();
|
|
|
|
});
|
2021-04-08 00:09:15 +00:00
|
|
|
|
|
|
|
return Pipe(std::make_shared<StorageS3Source>(
|
|
|
|
need_path_column, need_file_column, format_name, getName(),
|
2021-04-07 14:43:34 +00:00
|
|
|
metadata_snapshot->getSampleBlock(), context,
|
|
|
|
metadata_snapshot->getColumns(), max_block_size,
|
2021-04-08 00:09:15 +00:00
|
|
|
compression_method,
|
|
|
|
client_auth.client,
|
|
|
|
client_auth.uri.bucket,
|
|
|
|
file_iterator
|
|
|
|
));
|
2021-03-19 21:49:18 +00:00
|
|
|
}
|
|
|
|
|
2021-03-26 15:33:14 +00:00
|
|
|
/// The code from here and below executes on initiator
|
2021-04-08 14:22:19 +00:00
|
|
|
S3::URI s3_uri(Poco::URI{filename});
|
|
|
|
StorageS3::updateClientAndAuthSettings(context, client_auth);
|
|
|
|
|
2021-04-10 02:21:18 +00:00
|
|
|
auto iterator = std::make_shared<StorageS3Source::DisclosedGlobIterator>(*client_auth.client, client_auth.uri);
|
|
|
|
auto callback = std::make_shared<StorageS3Source::IteratorWrapper>([iterator]() mutable -> std::optional<String>
|
2021-04-08 14:22:19 +00:00
|
|
|
{
|
2021-04-10 02:21:18 +00:00
|
|
|
return iterator->next();
|
|
|
|
});
|
2021-03-19 21:49:18 +00:00
|
|
|
|
2021-03-26 15:33:14 +00:00
|
|
|
/// Calculate the header. This is significant, because some columns could be thrown away in some cases like query with count(*)
|
2021-03-23 17:58:29 +00:00
|
|
|
Block header =
|
2021-04-06 11:51:16 +00:00
|
|
|
InterpreterSelectQuery(query_info.query, context, SelectQueryOptions(processed_stage).analyze()).getSampleBlock();
|
2021-03-23 17:58:29 +00:00
|
|
|
|
|
|
|
const Scalars & scalars = context.hasQueryContext() ? context.getQueryContext().getScalars() : Scalars{};
|
2021-03-19 21:49:18 +00:00
|
|
|
|
2021-03-24 18:36:31 +00:00
|
|
|
Pipes pipes;
|
|
|
|
connections.reserve(cluster->getShardCount());
|
|
|
|
|
2021-03-19 21:49:18 +00:00
|
|
|
for (const auto & replicas : cluster->getShardsAddresses()) {
|
|
|
|
/// There will be only one replica, because we consider each replica as a shard
|
|
|
|
for (const auto & node : replicas)
|
|
|
|
{
|
|
|
|
connections.emplace_back(std::make_shared<Connection>(
|
2021-04-06 19:18:45 +00:00
|
|
|
node.host_name, node.port, context.getGlobalContext().getCurrentDatabase(),
|
|
|
|
node.user, node.password, node.cluster, node.cluster_secret,
|
2021-04-07 14:43:34 +00:00
|
|
|
"S3DistributedInititiator",
|
|
|
|
node.compression,
|
|
|
|
node.secure
|
2021-03-19 21:49:18 +00:00
|
|
|
));
|
2021-04-06 19:18:45 +00:00
|
|
|
|
2021-04-08 14:22:19 +00:00
|
|
|
/// For unknown reason global context is passed to IStorage::read() method
|
|
|
|
/// So, task_identifier is passed as constructor argument. It is more obvious.
|
2021-04-06 19:18:45 +00:00
|
|
|
auto remote_query_executor = std::make_shared<RemoteQueryExecutor>(
|
2021-04-08 14:22:19 +00:00
|
|
|
*connections.back(), queryToString(query_info.query), header, context,
|
2021-04-10 02:21:18 +00:00
|
|
|
/*throttler=*/nullptr, scalars, Tables(), processed_stage, callback);
|
2021-04-06 19:18:45 +00:00
|
|
|
|
2021-04-07 14:43:34 +00:00
|
|
|
pipes.emplace_back(std::make_shared<RemoteSource>(remote_query_executor, false, false));
|
2021-03-19 21:49:18 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
metadata_snapshot->check(column_names, getVirtuals(), getStorageID());
|
|
|
|
return Pipe::unitePipes(std::move(pipes));
|
|
|
|
}
|
2021-04-07 14:43:34 +00:00
|
|
|
|
2021-04-08 14:22:19 +00:00
|
|
|
QueryProcessingStage::Enum StorageS3Distributed::getQueryProcessingStage(
|
|
|
|
const Context & context, QueryProcessingStage::Enum /*to_stage*/, SelectQueryInfo &) const
|
2021-04-07 14:43:34 +00:00
|
|
|
{
|
|
|
|
/// Initiator executes query on remote node.
|
|
|
|
if (context.getClientInfo().query_kind == ClientInfo::QueryKind::INITIAL_QUERY) {
|
|
|
|
return QueryProcessingStage::Enum::WithMergeableState;
|
|
|
|
}
|
|
|
|
/// Follower just reads the data.
|
|
|
|
return QueryProcessingStage::Enum::FetchColumns;
|
|
|
|
}
|
|
|
|
|
|
|
|
|
2021-03-19 21:49:18 +00:00
|
|
|
}
|
|
|
|
|
2021-03-24 21:02:21 +00:00
|
|
|
#endif
|