ClickHouse/src/Storages/StorageURL.cpp

272 lines
8.5 KiB
C++
Raw Normal View History

2018-06-11 12:13:00 +00:00
#include <Storages/StorageFactory.h>
#include <Storages/StorageURL.h>
#include <Interpreters/Context.h>
#include <Interpreters/evaluateConstantExpression.h>
#include <Parsers/ASTLiteral.h>
2019-11-20 14:48:01 +00:00
#include <IO/ReadHelpers.h>
2018-06-11 12:13:00 +00:00
#include <IO/ReadWriteBufferFromHTTP.h>
#include <IO/WriteBufferFromHTTP.h>
2019-11-20 14:48:01 +00:00
#include <IO/WriteHelpers.h>
2018-06-11 12:13:00 +00:00
#include <Formats/FormatFactory.h>
#include <DataStreams/IBlockOutputStream.h>
#include <DataStreams/AddingDefaultsBlockInputStream.h>
2018-06-11 12:13:00 +00:00
#include <Poco/Net/HTTPRequest.h>
#include <Processors/Sources/SourceWithProgress.h>
#include <Processors/Pipe.h>
2018-06-11 12:13:00 +00:00
namespace DB
{
namespace ErrorCodes
{
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
}
2018-06-11 12:13:00 +00:00
2019-08-24 21:20:20 +00:00
IStorageURLBase::IStorageURLBase(
const Poco::URI & uri_,
const Context & context_,
2019-12-04 16:06:55 +00:00
const StorageID & table_id_,
2018-06-11 12:13:00 +00:00
const String & format_name_,
2019-08-24 21:20:20 +00:00
const ColumnsDescription & columns_,
const ConstraintsDescription & constraints_,
const String & compression_method_)
2019-12-04 16:06:55 +00:00
: IStorage(table_id_)
, uri(uri_)
, context_global(context_)
, compression_method(compression_method_)
, format_name(format_name_)
2018-06-11 12:13:00 +00:00
{
context_global.getRemoteHostFilter().checkURL(uri);
2020-06-19 15:39:41 +00:00
StorageInMemoryMetadata storage_metadata;
storage_metadata.setColumns(columns_);
storage_metadata.setConstraints(constraints_);
setInMemoryMetadata(storage_metadata);
2018-06-11 12:13:00 +00:00
}
namespace
{
class StorageURLSource : public SourceWithProgress
2018-06-11 12:13:00 +00:00
{
public:
StorageURLSource(const Poco::URI & uri,
const std::string & method,
std::function<void(std::ostream &)> callback,
2018-06-11 12:13:00 +00:00
const String & format,
String name_,
2018-06-11 12:13:00 +00:00
const Block & sample_block,
const Context & context,
2020-10-02 12:38:50 +00:00
const ColumnsDescription & columns,
2019-02-10 16:55:12 +00:00
UInt64 max_block_size,
const ConnectionTimeouts & timeouts,
const CompressionMethod compression_method)
: SourceWithProgress(sample_block), name(std::move(name_))
2018-06-11 12:13:00 +00:00
{
2020-08-28 01:21:08 +00:00
ReadWriteBufferFromHTTP::HTTPHeaderEntries header;
// Propagate OpenTelemetry trace context, if any, downstream.
2020-08-28 20:40:23 +00:00
const auto & client_info = context.getClientInfo();
2020-08-28 01:21:08 +00:00
if (client_info.opentelemetry_trace_id)
{
header.emplace_back("traceparent",
client_info.getOpenTelemetryTraceparentForChild());
if (!client_info.opentelemetry_tracestate.empty())
{
header.emplace_back("tracestate",
client_info.opentelemetry_tracestate);
}
}
read_buf = wrapReadBufferWithCompressionMethod(
std::make_unique<ReadWriteBufferFromHTTP>(
uri,
method,
std::move(callback),
timeouts,
context.getSettingsRef().max_http_get_redirects,
Poco::Net::HTTPBasicCredentials{},
DBMS_DEFAULT_BUFFER_SIZE,
2020-08-28 01:21:08 +00:00
header,
context.getRemoteHostFilter()),
compression_method);
2019-12-03 17:15:16 +00:00
reader = FormatFactory::instance().getInput(format, *read_buf, sample_block, context, max_block_size);
2020-10-02 12:38:50 +00:00
reader = std::make_shared<AddingDefaultsBlockInputStream>(reader, columns, context);
2018-06-11 12:13:00 +00:00
}
String getName() const override
{
return name;
}
Chunk generate() override
2018-06-11 12:13:00 +00:00
{
if (!reader)
return {};
2018-06-11 12:13:00 +00:00
if (!initialized)
reader->readPrefix();
2018-06-13 18:51:56 +00:00
2020-02-19 10:17:49 +00:00
initialized = true;
if (auto block = reader->read())
return Chunk(block.getColumns(), block.rows());
2018-06-11 12:13:00 +00:00
reader->readSuffix();
reader.reset();
return {};
2018-06-11 12:13:00 +00:00
}
private:
String name;
std::unique_ptr<ReadBuffer> read_buf;
2018-06-11 12:13:00 +00:00
BlockInputStreamPtr reader;
bool initialized = false;
2018-06-11 12:13:00 +00:00
};
}
2018-06-16 05:54:06 +00:00
2020-04-28 00:56:44 +00:00
StorageURLBlockOutputStream::StorageURLBlockOutputStream(const Poco::URI & uri,
const String & format,
const Block & sample_block_,
const Context & context,
const ConnectionTimeouts & timeouts,
const CompressionMethod compression_method)
: sample_block(sample_block_)
{
write_buf = wrapWriteBufferWithCompressionMethod(
std::make_unique<WriteBufferFromHTTP>(uri, Poco::Net::HTTPRequest::HTTP_POST, timeouts),
compression_method, 3);
writer = FormatFactory::instance().getOutput(format, *write_buf, sample_block, context);
}
2018-06-16 05:54:06 +00:00
2020-07-09 01:00:16 +00:00
void StorageURLBlockOutputStream::write(const Block & block)
{
writer->write(block);
}
void StorageURLBlockOutputStream::writePrefix()
{
writer->writePrefix();
}
void StorageURLBlockOutputStream::writeSuffix()
{
writer->writeSuffix();
writer->flush();
write_buf->finalize();
}
std::string IStorageURLBase::getReadMethod() const
{
return Poco::Net::HTTPRequest::HTTP_GET;
}
std::vector<std::pair<std::string, std::string>> IStorageURLBase::getReadURIParams(
const Names & /*column_names*/,
const StorageMetadataPtr & /*metadata_snapshot*/,
2018-06-11 12:13:00 +00:00
const SelectQueryInfo & /*query_info*/,
const Context & /*context*/,
2018-06-11 12:13:00 +00:00
QueryProcessingStage::Enum & /*processed_stage*/,
size_t /*max_block_size*/) const
{
return {};
}
std::function<void(std::ostream &)> IStorageURLBase::getReadPOSTDataCallback(
const Names & /*column_names*/,
const StorageMetadataPtr & /*metadata_snapshot*/,
2018-06-11 12:13:00 +00:00
const SelectQueryInfo & /*query_info*/,
const Context & /*context*/,
QueryProcessingStage::Enum & /*processed_stage*/,
size_t /*max_block_size*/) const
{
return nullptr;
}
2020-08-03 13:54:14 +00:00
Pipe IStorageURLBase::read(
const Names & column_names,
const StorageMetadataPtr & metadata_snapshot,
const SelectQueryInfo & query_info,
2018-06-11 12:13:00 +00:00
const Context & context,
2018-06-25 12:21:54 +00:00
QueryProcessingStage::Enum processed_stage,
size_t max_block_size,
2018-06-11 12:13:00 +00:00
unsigned /*num_streams*/)
{
auto request_uri = uri;
auto params = getReadURIParams(column_names, metadata_snapshot, query_info, context, processed_stage, max_block_size);
for (const auto & [param, value] : params)
request_uri.addQueryParameter(param, value);
2020-08-03 13:54:14 +00:00
return Pipe(std::make_shared<StorageURLSource>(
request_uri,
getReadMethod(),
getReadPOSTDataCallback(
column_names, metadata_snapshot, query_info,
context, processed_stage, max_block_size),
2018-06-11 12:13:00 +00:00
format_name,
getName(),
getHeaderBlock(column_names, metadata_snapshot),
2018-06-11 12:13:00 +00:00
context,
2020-10-02 12:38:50 +00:00
metadata_snapshot->getColumns(),
2018-06-11 12:13:00 +00:00
max_block_size,
ConnectionTimeouts::getHTTPTimeouts(context),
chooseCompressionMethod(request_uri.getPath(), compression_method)));
2018-06-11 12:13:00 +00:00
}
BlockOutputStreamPtr IStorageURLBase::write(const ASTPtr & /*query*/, const StorageMetadataPtr & metadata_snapshot, const Context & /*context*/)
2018-06-11 12:13:00 +00:00
{
return std::make_shared<StorageURLBlockOutputStream>(
uri, format_name, metadata_snapshot->getSampleBlock(), context_global,
ConnectionTimeouts::getHTTPTimeouts(context_global),
chooseCompressionMethod(uri.toString(), compression_method));
2018-06-11 12:13:00 +00:00
}
2018-06-16 05:54:06 +00:00
2018-06-11 12:13:00 +00:00
void registerStorageURL(StorageFactory & factory)
{
2018-08-24 00:07:25 +00:00
factory.registerStorage("URL", [](const StorageFactory::Arguments & args)
{
2018-06-11 12:13:00 +00:00
ASTs & engine_args = args.engine_args;
if (engine_args.size() != 2 && engine_args.size() != 3)
2018-06-11 12:13:00 +00:00
throw Exception(
"Storage URL requires 2 or 3 arguments: url, name of used format and optional compression method.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
2018-06-11 12:13:00 +00:00
engine_args[0] = evaluateConstantExpressionOrIdentifierAsLiteral(engine_args[0], args.local_context);
2019-03-15 17:09:14 +00:00
String url = engine_args[0]->as<ASTLiteral &>().value.safeGet<String>();
2018-06-11 12:13:00 +00:00
Poco::URI uri(url);
engine_args[1] = evaluateConstantExpressionOrIdentifierAsLiteral(engine_args[1], args.local_context);
2019-03-15 17:09:14 +00:00
String format_name = engine_args[1]->as<ASTLiteral &>().value.safeGet<String>();
2018-06-11 12:13:00 +00:00
String compression_method;
if (engine_args.size() == 3)
{
engine_args[2] = evaluateConstantExpressionOrIdentifierAsLiteral(engine_args[2], args.local_context);
compression_method = engine_args[2]->as<ASTLiteral &>().value.safeGet<String>();
} else compression_method = "auto";
return StorageURL::create(
uri,
2019-12-04 16:06:55 +00:00
args.table_id,
format_name,
args.columns, args.constraints, args.context,
compression_method);
},
{
.source_access_type = AccessType::URL,
2018-06-11 12:13:00 +00:00
});
}
}