ClickHouse/src/Storages/StorageMySQL.cpp

310 lines
10 KiB
C++
Raw Normal View History

#include "StorageMySQL.h"
#if USE_MYSQL
#include <Storages/StorageFactory.h>
2017-12-26 21:34:06 +00:00
#include <Storages/transformQueryForExternalDatabase.h>
2021-08-11 16:44:34 +00:00
#include <Formats/MySQLSource.h>
#include <Interpreters/evaluateConstantExpression.h>
#include <Core/Settings.h>
#include <Interpreters/Context.h>
#include <DataTypes/DataTypeString.h>
#include <DataStreams/IBlockOutputStream.h>
#include <Formats/FormatFactory.h>
#include <Common/parseAddress.h>
#include <IO/Operators.h>
#include <IO/WriteHelpers.h>
#include <Parsers/ASTLiteral.h>
#include <Parsers/ASTCreateQuery.h>
#include <mysqlxx/Transaction.h>
#include <Processors/Sources/SourceFromInputStream.h>
2021-07-23 19:33:59 +00:00
#include <Processors/Sinks/SinkToStorage.h>
#include <Processors/Pipe.h>
#include <Common/parseRemoteDescription.h>
namespace DB
{
namespace ErrorCodes
{
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
2018-05-14 11:00:22 +00:00
extern const int BAD_ARGUMENTS;
}
2019-12-15 06:34:43 +00:00
static String backQuoteMySQL(const String & x)
{
String res(x.size(), '\0');
{
WriteBufferFromString wb(res);
writeBackQuotedStringMySQL(x, wb);
}
return res;
}
StorageMySQL::StorageMySQL(
2019-12-04 16:06:55 +00:00
const StorageID & table_id_,
mysqlxx::PoolWithFailover && pool_,
2019-08-03 11:02:40 +00:00
const std::string & remote_database_name_,
const std::string & remote_table_name_,
const bool replace_query_,
const std::string & on_duplicate_clause_,
const ColumnsDescription & columns_,
2019-08-24 21:20:20 +00:00
const ConstraintsDescription & constraints_,
2021-04-23 12:18:23 +00:00
const String & comment,
ContextPtr context_,
const MySQLSettings & mysql_settings_)
2019-12-04 16:06:55 +00:00
: IStorage(table_id_)
, WithContext(context_->getGlobalContext())
2019-08-03 11:02:40 +00:00
, remote_database_name(remote_database_name_)
, remote_table_name(remote_table_name_)
, replace_query{replace_query_}
, on_duplicate_clause{on_duplicate_clause_}
, mysql_settings(mysql_settings_)
2021-03-27 14:35:44 +00:00
, pool(std::make_shared<mysqlxx::PoolWithFailover>(pool_))
{
2020-06-19 15:39:41 +00:00
StorageInMemoryMetadata storage_metadata;
storage_metadata.setColumns(columns_);
storage_metadata.setConstraints(constraints_);
2021-04-23 12:18:23 +00:00
storage_metadata.setComment(comment);
2020-06-19 15:39:41 +00:00
setInMemoryMetadata(storage_metadata);
}
2020-08-03 13:54:14 +00:00
Pipe StorageMySQL::read(
2019-08-03 11:02:40 +00:00
const Names & column_names_,
const StorageMetadataPtr & metadata_snapshot,
SelectQueryInfo & query_info_,
ContextPtr context_,
QueryProcessingStage::Enum /*processed_stage*/,
size_t /*max_block_size*/,
2017-12-26 21:34:06 +00:00
unsigned)
{
2020-06-19 17:17:13 +00:00
metadata_snapshot->check(column_names_, getVirtuals(), getStorageID());
String query = transformQueryForExternalDatabase(
query_info_,
metadata_snapshot->getColumns().getOrdinary(),
IdentifierQuotingStyle::BackticksMySQL,
remote_database_name,
remote_table_name,
context_);
Block sample_block;
2019-08-03 11:02:40 +00:00
for (const String & column_name : column_names_)
{
auto column_data = metadata_snapshot->getColumns().getPhysical(column_name);
WhichDataType which(column_data.type);
/// Convert enum to string.
if (which.isEnum())
column_data.type = std::make_shared<DataTypeString>();
sample_block.insert({ column_data.type, column_data.name });
}
StreamSettings mysql_input_stream_settings(context_->getSettingsRef(),
mysql_settings.connection_auto_close);
return Pipe(std::make_shared<MySQLWithFailoverSource>(pool, query, sample_block, mysql_input_stream_settings));
}
2017-12-26 21:34:06 +00:00
2021-07-23 19:33:59 +00:00
class StorageMySQLSink : public SinkToStorage
{
public:
2021-07-23 19:33:59 +00:00
explicit StorageMySQLSink(
const StorageMySQL & storage_,
const StorageMetadataPtr & metadata_snapshot_,
2019-08-03 11:02:40 +00:00
const std::string & remote_database_name_,
const std::string & remote_table_name_,
const mysqlxx::PoolWithFailover::Entry & entry_,
2018-05-14 11:00:22 +00:00
const size_t & mysql_max_rows_to_insert)
2021-07-26 10:08:40 +00:00
: SinkToStorage(metadata_snapshot_->getSampleBlock())
2021-07-23 19:33:59 +00:00
, storage{storage_}
, metadata_snapshot{metadata_snapshot_}
2019-08-03 11:02:40 +00:00
, remote_database_name{remote_database_name_}
, remote_table_name{remote_table_name_}
, entry{entry_}
, max_batch_rows{mysql_max_rows_to_insert}
{
}
2021-07-23 19:33:59 +00:00
String getName() const override { return "StorageMySQLSink"; }
2018-05-14 11:00:22 +00:00
2021-07-23 19:33:59 +00:00
void consume(Chunk chunk) override
2018-05-14 11:00:22 +00:00
{
2021-07-23 19:33:59 +00:00
auto block = getPort().getHeader().cloneWithColumns(chunk.detachColumns());
2018-05-14 11:00:22 +00:00
auto blocks = splitBlocks(block, max_batch_rows);
mysqlxx::Transaction trans(entry);
try
{
for (const Block & batch_data : blocks)
{
writeBlockData(batch_data);
}
trans.commit();
}
catch (...)
2018-05-14 11:00:22 +00:00
{
trans.rollback();
throw;
}
}
void writeBlockData(const Block & block)
{
WriteBufferFromOwnString sqlbuf;
sqlbuf << (storage.replace_query ? "REPLACE" : "INSERT") << " INTO ";
if (!remote_database_name.empty())
sqlbuf << backQuoteMySQL(remote_database_name) << ".";
sqlbuf << backQuoteMySQL(remote_table_name);
2018-11-24 01:48:06 +00:00
sqlbuf << " (" << dumpNamesWithBackQuote(block) << ") VALUES ";
2018-05-14 11:00:22 +00:00
auto writer = FormatFactory::instance().getOutputStream("Values", sqlbuf, metadata_snapshot->getSampleBlock(), storage.getContext());
2018-05-14 11:00:22 +00:00
writer->write(block);
if (!storage.on_duplicate_clause.empty())
sqlbuf << " ON DUPLICATE KEY " << storage.on_duplicate_clause;
sqlbuf << ";";
auto query = this->entry->query(sqlbuf.str());
query.execute();
}
Blocks splitBlocks(const Block & block, const size_t & max_rows) const
{
/// Avoid Excessive copy when block is small enough
if (block.rows() <= max_rows)
return Blocks{std::move(block)};
2021-06-28 17:02:22 +00:00
const size_t split_block_size = ceil(block.rows() * 1.0 / max_rows);
Blocks split_blocks(split_block_size);
2018-05-14 11:00:22 +00:00
2021-06-28 17:02:22 +00:00
for (size_t idx = 0; idx < split_block_size; ++idx)
split_blocks[idx] = block.cloneEmpty();
2018-05-14 11:00:22 +00:00
const size_t columns = block.columns();
const size_t rows = block.rows();
size_t offsets = 0;
2019-02-10 15:17:45 +00:00
UInt64 limits = max_batch_rows;
2021-06-28 17:02:22 +00:00
for (size_t idx = 0; idx < split_block_size; ++idx)
2018-05-14 11:00:22 +00:00
{
/// For last batch, limits should be the remain size
2021-06-28 17:02:22 +00:00
if (idx == split_block_size - 1) limits = rows - offsets;
2018-05-14 11:00:22 +00:00
for (size_t col_idx = 0; col_idx < columns; ++col_idx)
{
2021-06-28 17:02:22 +00:00
split_blocks[idx].getByPosition(col_idx).column = block.getByPosition(col_idx).column->cut(offsets, limits);
2018-05-14 11:00:22 +00:00
}
offsets += max_batch_rows;
}
2021-06-28 17:02:22 +00:00
return split_blocks;
2018-05-14 11:00:22 +00:00
}
2020-03-18 00:57:00 +00:00
static std::string dumpNamesWithBackQuote(const Block & block)
2018-05-14 11:00:22 +00:00
{
WriteBufferFromOwnString out;
for (auto it = block.begin(); it != block.end(); ++it)
{
if (it != block.begin())
out << ", ";
out << backQuoteMySQL(it->name);
2018-05-14 11:00:22 +00:00
}
return out.str();
}
private:
2018-05-14 11:00:22 +00:00
const StorageMySQL & storage;
StorageMetadataPtr metadata_snapshot;
2018-05-14 11:00:22 +00:00
std::string remote_database_name;
std::string remote_table_name;
mysqlxx::PoolWithFailover::Entry entry;
size_t max_batch_rows;
};
2021-07-23 19:33:59 +00:00
SinkToStoragePtr StorageMySQL::write(const ASTPtr & /*query*/, const StorageMetadataPtr & metadata_snapshot, ContextPtr local_context)
{
2021-07-23 19:33:59 +00:00
return std::make_shared<StorageMySQLSink>(
*this,
metadata_snapshot,
remote_database_name,
remote_table_name,
pool->get(),
local_context->getSettingsRef().mysql_max_rows_to_insert);
}
void registerStorageMySQL(StorageFactory & factory)
{
factory.registerStorage("MySQL", [](const StorageFactory::Arguments & args)
{
ASTs & engine_args = args.engine_args;
if (engine_args.size() < 5 || engine_args.size() > 7)
throw Exception(
"Storage MySQL requires 5-7 parameters: MySQL('host:port' (or 'addresses_pattern'), database, table, 'user', 'password'[, replace_query, 'on_duplicate_clause']).",
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
2020-03-09 01:22:33 +00:00
for (auto & engine_arg : engine_args)
engine_arg = evaluateConstantExpressionOrIdentifierAsLiteral(engine_arg, args.getLocalContext());
/// 3306 is the default MySQL port.
const String & host_port = engine_args[0]->as<ASTLiteral &>().value.safeGet<String>();
2019-03-15 17:09:14 +00:00
const String & remote_database = engine_args[1]->as<ASTLiteral &>().value.safeGet<String>();
const String & remote_table = engine_args[2]->as<ASTLiteral &>().value.safeGet<String>();
const String & username = engine_args[3]->as<ASTLiteral &>().value.safeGet<String>();
const String & password = engine_args[4]->as<ASTLiteral &>().value.safeGet<String>();
size_t max_addresses = args.getContext()->getSettingsRef().glob_expansion_max_elements;
/// TODO: move some arguments from the arguments to the SETTINGS.
MySQLSettings mysql_settings;
if (args.storage_def->settings)
{
mysql_settings.loadFromQuery(*args.storage_def);
}
if (!mysql_settings.connection_pool_size)
throw Exception("connection_pool_size cannot be zero.", ErrorCodes::BAD_ARGUMENTS);
2021-04-01 11:30:20 +00:00
auto addresses = parseRemoteDescriptionForExternalDatabase(host_port, max_addresses, 3306);
mysqlxx::PoolWithFailover pool(remote_database, addresses,
username, password,
MYSQLXX_POOL_WITH_FAILOVER_DEFAULT_START_CONNECTIONS,
mysql_settings.connection_pool_size,
mysql_settings.connection_max_tries);
bool replace_query = false;
std::string on_duplicate_clause;
2018-05-14 11:00:22 +00:00
if (engine_args.size() >= 6)
2019-03-15 17:09:14 +00:00
replace_query = engine_args[5]->as<ASTLiteral &>().value.safeGet<UInt64>();
2018-05-14 11:00:22 +00:00
if (engine_args.size() == 7)
2019-03-15 17:09:14 +00:00
on_duplicate_clause = engine_args[6]->as<ASTLiteral &>().value.safeGet<String>();
2018-05-14 11:00:22 +00:00
if (replace_query && !on_duplicate_clause.empty())
throw Exception(
"Only one of 'replace_query' and 'on_duplicate_clause' can be specified, or none of them",
ErrorCodes::BAD_ARGUMENTS);
return StorageMySQL::create(
2019-12-04 16:06:55 +00:00
args.table_id,
std::move(pool),
remote_database,
remote_table,
replace_query,
on_duplicate_clause,
args.columns,
2019-08-24 21:20:20 +00:00
args.constraints,
2021-04-23 12:18:23 +00:00
args.comment,
args.getContext(),
mysql_settings);
},
{
.supports_settings = true,
.source_access_type = AccessType::MYSQL,
});
}
}
#endif