ClickHouse/dbms/Dictionaries/ExecutableDictionarySource.cpp

243 lines
8.0 KiB
C++
Raw Normal View History

#include "ExecutableDictionarySource.h"
#include <future>
#include <thread>
2019-03-01 16:38:52 +00:00
#include <ext/scope_guard.h>
#include <DataStreams/IBlockOutputStream.h>
#include <DataStreams/OwningBlockInputStream.h>
#include <Interpreters/Context.h>
2019-03-01 16:38:52 +00:00
#include <IO/WriteHelpers.h>
#include <Common/ShellCommand.h>
#include <Common/ThreadPool.h>
#include <common/logger_useful.h>
2019-03-01 16:38:52 +00:00
#include <common/LocalDateTime.h>
#include "DictionarySourceFactory.h"
#include "DictionarySourceHelpers.h"
#include "DictionaryStructure.h"
2019-12-15 06:34:43 +00:00
#include "registerDictionaries.h"
2016-11-15 19:51:06 +00:00
namespace DB
{
2019-02-10 16:55:12 +00:00
static const UInt64 max_block_size = 8192;
2016-12-08 02:49:04 +00:00
namespace ErrorCodes
{
2020-02-25 18:02:41 +00:00
extern const int LOGICAL_ERROR;
extern const int DICTIONARY_ACCESS_DENIED;
}
2016-12-08 02:49:04 +00:00
namespace
{
/// Owns ShellCommand and calls wait for it.
class ShellCommandOwningBlockInputStream : public OwningBlockInputStream<ShellCommand>
{
public:
ShellCommandOwningBlockInputStream(const BlockInputStreamPtr & impl, std::unique_ptr<ShellCommand> own_)
: OwningBlockInputStream(std::move(impl), std::move(own_))
{
}
void readSuffix() override
{
OwningBlockInputStream<ShellCommand>::readSuffix();
own->wait();
}
};
}
2017-10-12 19:35:01 +00:00
ExecutableDictionarySource::ExecutableDictionarySource(
const DictionaryStructure & dict_struct_,
const Poco::Util::AbstractConfiguration & config,
const std::string & config_prefix,
2019-08-03 11:02:40 +00:00
Block & sample_block_,
const Context & context_)
: log(&Logger::get("ExecutableDictionarySource"))
, dict_struct{dict_struct_}
, command{config.getString(config_prefix + ".command")}
, update_field{config.getString(config_prefix + ".update_field", "")}
, format{config.getString(config_prefix + ".format")}
2019-08-03 11:02:40 +00:00
, sample_block{sample_block_}
, context(context_)
2016-11-15 19:51:06 +00:00
{
}
ExecutableDictionarySource::ExecutableDictionarySource(const ExecutableDictionarySource & other)
: log(&Logger::get("ExecutableDictionarySource"))
, update_time{other.update_time}
, dict_struct{other.dict_struct}
, command{other.command}
, update_field{other.update_field}
, format{other.format}
, sample_block{other.sample_block}
, context(other.context)
2016-11-15 19:51:06 +00:00
{
}
BlockInputStreamPtr ExecutableDictionarySource::loadAll()
{
LOG_TRACE(log, "loadAll " + toString());
auto process = ShellCommand::execute(command);
2017-05-25 19:21:57 +00:00
auto input_stream = context.getInputFormat(format, process->out, sample_block, max_block_size);
return std::make_shared<ShellCommandOwningBlockInputStream>(input_stream, std::move(process));
2016-11-27 13:01:35 +00:00
}
BlockInputStreamPtr ExecutableDictionarySource::loadUpdatedAll()
{
2019-03-01 16:38:52 +00:00
time_t new_update_time = time(nullptr);
SCOPE_EXIT(update_time = new_update_time);
std::string command_with_update_field = command;
if (update_time)
command_with_update_field += " " + update_field + " " + DB::toString(LocalDateTime(update_time - 1));
LOG_TRACE(log, "loadUpdatedAll " + command_with_update_field);
auto process = ShellCommand::execute(command_with_update_field);
auto input_stream = context.getInputFormat(format, process->out, sample_block, max_block_size);
return std::make_shared<ShellCommandOwningBlockInputStream>(input_stream, std::move(process));
}
namespace
{
/** A stream, that also runs and waits for background thread
* (that will feed data into pipe to be read from the other side of the pipe).
*/
class BlockInputStreamWithBackgroundThread final : public IBlockInputStream
{
public:
BlockInputStreamWithBackgroundThread(
const BlockInputStreamPtr & stream_, std::unique_ptr<ShellCommand> && command_, std::packaged_task<void()> && task_)
: stream{stream_}, command{std::move(command_)}, task(std::move(task_)), thread([this] {
task();
command->in.close();
})
{
children.push_back(stream);
}
~BlockInputStreamWithBackgroundThread() override
{
if (thread.joinable())
{
try
{
readSuffix();
}
catch (...)
{
tryLogCurrentException(__PRETTY_FUNCTION__);
}
}
}
Block getHeader() const override { return stream->getHeader(); }
private:
Block readImpl() override { return stream->read(); }
void readSuffix() override
{
IBlockInputStream::readSuffix();
if (!wait_called)
{
wait_called = true;
command->wait();
}
thread.join();
/// To rethrow an exception, if any.
task.get_future().get();
}
String getName() const override { return "WithBackgroundThread"; }
BlockInputStreamPtr stream;
std::unique_ptr<ShellCommand> command;
std::packaged_task<void()> task;
ThreadFromGlobalPool thread;
bool wait_called = false;
};
}
2016-11-27 13:01:35 +00:00
BlockInputStreamPtr ExecutableDictionarySource::loadIds(const std::vector<UInt64> & ids)
{
2017-05-25 18:24:47 +00:00
LOG_TRACE(log, "loadIds " << toString() << " size = " << ids.size());
auto process = ShellCommand::execute(command);
2017-05-25 18:24:47 +00:00
2017-05-25 19:21:57 +00:00
auto output_stream = context.getOutputFormat(format, process->in, sample_block);
auto input_stream = context.getInputFormat(format, process->out, sample_block, max_block_size);
return std::make_shared<BlockInputStreamWithBackgroundThread>(
input_stream, std::move(process), std::packaged_task<void()>([output_stream, &ids]() mutable { formatIDs(output_stream, ids); }));
2016-11-15 19:51:06 +00:00
}
BlockInputStreamPtr ExecutableDictionarySource::loadKeys(const Columns & key_columns, const std::vector<size_t> & requested_rows)
2016-11-15 19:51:06 +00:00
{
2017-05-25 18:24:47 +00:00
LOG_TRACE(log, "loadKeys " << toString() << " size = " << requested_rows.size());
auto process = ShellCommand::execute(command);
2016-11-18 01:48:13 +00:00
2017-05-25 19:21:57 +00:00
auto output_stream = context.getOutputFormat(format, process->in, sample_block);
auto input_stream = context.getInputFormat(format, process->out, sample_block, max_block_size);
return std::make_shared<BlockInputStreamWithBackgroundThread>(
2018-12-10 15:50:58 +00:00
input_stream, std::move(process), std::packaged_task<void()>([output_stream, key_columns, &requested_rows, this]() mutable
{
formatKeys(dict_struct, output_stream, key_columns, requested_rows);
}));
2016-11-15 19:51:06 +00:00
}
bool ExecutableDictionarySource::isModified() const
{
return true;
2016-11-15 19:51:06 +00:00
}
bool ExecutableDictionarySource::supportsSelectiveLoad() const
{
return true;
2016-11-15 19:51:06 +00:00
}
bool ExecutableDictionarySource::hasUpdateField() const
{
2020-03-09 03:14:24 +00:00
return !update_field.empty();
}
2016-11-15 19:51:06 +00:00
DictionarySourcePtr ExecutableDictionarySource::clone() const
{
return std::make_unique<ExecutableDictionarySource>(*this);
2016-11-15 19:51:06 +00:00
}
std::string ExecutableDictionarySource::toString() const
{
return "Executable: " + command;
2016-11-15 19:51:06 +00:00
}
void registerDictionarySourceExecutable(DictionarySourceFactory & factory)
{
auto create_table_source = [=](const DictionaryStructure & dict_struct,
const Poco::Util::AbstractConfiguration & config,
const std::string & config_prefix,
Block & sample_block,
const Context & context,
bool check_config) -> DictionarySourcePtr
{
if (dict_struct.has_expressions)
throw Exception{"Dictionary source of type `executable` does not support attribute expressions", ErrorCodes::LOGICAL_ERROR};
/// Executable dictionaries may execute arbitrary commands.
/// It's OK for dictionaries created by administrator from xml-file, but
/// maybe dangerous for dictionaries created from DDL-queries.
if (check_config)
throw Exception("Dictionaries with Executable dictionary source is not allowed", ErrorCodes::DICTIONARY_ACCESS_DENIED);
return std::make_unique<ExecutableDictionarySource>(
dict_struct, config, config_prefix + ".executable",
sample_block, context);
};
factory.registerSource("executable", create_table_source);
}
2016-11-15 19:51:06 +00:00
}