ClickHouse/src/Storages/StorageExecutable.cpp

115 lines
3.4 KiB
C++
Raw Normal View History

2021-04-14 17:51:55 +00:00
#include <Storages/StorageExecutable.h>
2021-08-25 19:30:22 +00:00
#include <filesystem>
#include <Common/ShellCommand.h>
#include <Core/Block.h>
2021-04-15 09:40:41 +00:00
#include <IO/ReadHelpers.h>
2021-08-25 19:30:22 +00:00
#include <Processors/Pipe.h>
2021-04-14 17:51:55 +00:00
#include <Interpreters/Context.h>
#include <Storages/StorageFactory.h>
2021-08-25 19:30:22 +00:00
#include <DataStreams/IBlockInputStream.h>
2021-08-24 19:38:42 +00:00
#include <DataStreams/ShellCommandSource.h>
2021-04-14 17:51:55 +00:00
namespace DB
{
2021-08-24 19:38:42 +00:00
2021-04-14 17:51:55 +00:00
namespace ErrorCodes
{
2021-08-25 19:30:22 +00:00
extern const int UNSUPPORTED_METHOD;
2021-08-28 19:47:59 +00:00
extern const int LOGICAL_ERROR;
2021-04-14 17:51:55 +00:00
}
StorageExecutable::StorageExecutable(
2021-04-15 09:40:41 +00:00
const StorageID & table_id_,
2021-08-25 19:30:22 +00:00
const String & script_name_,
2021-04-15 21:15:54 +00:00
const String & format_,
2021-08-25 19:30:22 +00:00
const std::vector<BlockInputStreamPtr> & inputs_,
2021-04-14 17:51:55 +00:00
const ColumnsDescription & columns,
2021-08-24 19:38:42 +00:00
const ConstraintsDescription & constraints)
2021-04-15 09:40:41 +00:00
: IStorage(table_id_)
2021-08-25 19:30:22 +00:00
, script_name(script_name_)
2021-04-15 21:15:54 +00:00
, format(format_)
2021-08-25 19:30:22 +00:00
, inputs(inputs_)
2021-08-24 19:38:42 +00:00
, log(&Poco::Logger::get("StorageExecutable"))
2021-04-14 17:51:55 +00:00
{
StorageInMemoryMetadata storage_metadata;
storage_metadata.setColumns(columns);
storage_metadata.setConstraints(constraints);
setInMemoryMetadata(storage_metadata);
}
Pipe StorageExecutable::read(
const Names & /*column_names*/,
const StorageMetadataPtr & metadata_snapshot,
SelectQueryInfo & /*query_info*/,
2021-08-25 19:30:22 +00:00
ContextPtr context,
2021-04-14 17:51:55 +00:00
QueryProcessingStage::Enum /*processed_stage*/,
size_t max_block_size,
2021-08-25 19:30:22 +00:00
unsigned /*threads*/)
2021-04-14 17:51:55 +00:00
{
2021-08-25 19:30:22 +00:00
auto user_scripts_path = context->getUserScriptsPath();
auto script_path = user_scripts_path + '/' + script_name;
if (!std::filesystem::exists(std::filesystem::path(script_path)))
throw Exception(ErrorCodes::UNSUPPORTED_METHOD,
"Executable file {} does not exists inside {}",
script_name,
user_scripts_path);
2021-08-24 19:38:42 +00:00
auto sample_block = metadata_snapshot->getSampleBlock();
2021-08-25 19:30:22 +00:00
ShellCommand::Config config(script_path);
2021-08-28 19:47:59 +00:00
for (size_t i = 1; i < inputs.size(); ++i)
config.write_fds.emplace_back(i + 2);
2021-08-25 19:30:22 +00:00
auto process = ShellCommand::execute(config);
2021-08-28 19:47:59 +00:00
std::vector<ShellCommandSource::SendDataTask> tasks;
tasks.reserve(inputs.size());
2021-08-25 19:30:22 +00:00
2021-08-28 19:47:59 +00:00
for (size_t i = 0; i < inputs.size(); ++i)
2021-08-25 19:30:22 +00:00
{
2021-08-28 19:47:59 +00:00
BlockInputStreamPtr input_stream = inputs[i];
WriteBufferFromFile * write_buffer;
if (i == 0)
{
write_buffer = &process->in;
}
else
{
auto descriptor = i + 2;
auto it = process->write_fds.find(descriptor);
if (it == process->write_fds.end())
throw Exception(ErrorCodes::LOGICAL_ERROR, "Process does not contain descriptor to write {}", descriptor);
write_buffer = &it->second;
}
ShellCommandSource::SendDataTask task = [input_stream, write_buffer, context, this]()
{
auto output_stream = context->getOutputStream(format, *write_buffer, input_stream->getHeader().cloneEmpty());
input_stream->readPrefix();
output_stream->writePrefix();
while (auto block = input_stream->read())
output_stream->write(block);
input_stream->readSuffix();
output_stream->writeSuffix();
output_stream->flush();
write_buffer->close();
};
tasks.emplace_back(std::move(task));
2021-08-25 19:30:22 +00:00
}
2021-08-28 19:47:59 +00:00
Pipe pipe(std::make_unique<ShellCommandSource>(context, format, sample_block, std::move(process), log, std::move(tasks), max_block_size));
return pipe;
2021-04-14 17:51:55 +00:00
}
2021-08-28 19:47:59 +00:00
2021-04-14 17:51:55 +00:00
};