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
|
|
|
};
|
|
|
|
|