ClickHouse/src/Storages/StorageSQLite.cpp

187 lines
5.8 KiB
C++
Raw Normal View History

2021-05-17 11:02:35 +00:00
#include "StorageSQLite.h"
2021-07-08 13:27:30 +00:00
#if USE_SQLITE
2021-10-02 07:13:14 +00:00
#include <base/range.h>
#include <base/logger_useful.h>
2021-08-11 16:44:34 +00:00
#include <DataStreams/SQLiteSource.h>
2021-07-14 13:13:37 +00:00
#include <Databases/SQLite/SQLiteUtils.h>
2021-05-17 11:02:35 +00:00
#include <DataTypes/DataTypeString.h>
2021-05-24 18:59:10 +00:00
#include <Interpreters/Context.h>
#include <Formats/FormatFactory.h>
2021-10-11 16:11:50 +00:00
#include <Processors/Formats/IOutputFormat.h>
2021-05-24 18:59:10 +00:00
#include <IO/Operators.h>
#include <IO/WriteHelpers.h>
2021-05-17 11:02:35 +00:00
#include <Interpreters/evaluateConstantExpression.h>
#include <Parsers/ASTLiteral.h>
2021-07-23 14:25:35 +00:00
#include <Processors/Sinks/SinkToStorage.h>
2021-05-17 11:02:35 +00:00
#include <Storages/StorageFactory.h>
#include <Storages/transformQueryForExternalDatabase.h>
2021-07-13 19:34:31 +00:00
#include <Common/filesystemHelpers.h>
2021-07-10 08:05:23 +00:00
2021-05-17 11:02:35 +00:00
namespace DB
{
2021-07-10 08:05:23 +00:00
2021-05-17 11:02:35 +00:00
namespace ErrorCodes
{
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
2021-07-10 11:00:59 +00:00
extern const int SQLITE_ENGINE_ERROR;
2021-05-17 11:02:35 +00:00
}
StorageSQLite::StorageSQLite(
const StorageID & table_id_,
2021-07-10 08:05:23 +00:00
SQLitePtr sqlite_db_,
2021-07-26 07:05:28 +00:00
const String & database_path_,
2021-05-17 11:02:35 +00:00
const String & remote_table_name_,
const ColumnsDescription & columns_,
const ConstraintsDescription & constraints_,
2021-06-01 20:52:12 +00:00
ContextPtr context_)
2021-05-17 11:02:35 +00:00
: IStorage(table_id_)
2021-05-24 18:59:10 +00:00
, WithContext(context_->getGlobalContext())
2021-05-17 11:02:35 +00:00
, remote_table_name(remote_table_name_)
2021-07-26 07:05:28 +00:00
, database_path(database_path_)
2021-05-17 11:02:35 +00:00
, global_context(context_)
2021-07-10 08:05:23 +00:00
, sqlite_db(sqlite_db_)
, log(&Poco::Logger::get("StorageSQLite (" + table_id_.table_name + ")"))
2021-05-17 11:02:35 +00:00
{
StorageInMemoryMetadata storage_metadata;
storage_metadata.setColumns(columns_);
storage_metadata.setConstraints(constraints_);
setInMemoryMetadata(storage_metadata);
}
2021-07-08 13:27:30 +00:00
2021-05-17 11:02:35 +00:00
Pipe StorageSQLite::read(
const Names & column_names,
const StorageMetadataPtr & metadata_snapshot,
SelectQueryInfo & query_info,
2021-05-24 18:59:10 +00:00
ContextPtr context_,
2021-05-17 11:02:35 +00:00
QueryProcessingStage::Enum,
size_t max_block_size,
unsigned int)
{
2021-07-26 07:05:28 +00:00
if (!sqlite_db)
sqlite_db = openSQLiteDB(database_path, getContext(), /* throw_on_error */true);
2021-05-17 11:02:35 +00:00
metadata_snapshot->check(column_names, getVirtuals(), getStorageID());
String query = transformQueryForExternalDatabase(
query_info,
metadata_snapshot->getColumns().getOrdinary(),
IdentifierQuotingStyle::DoubleQuotes,
2021-06-01 20:52:12 +00:00
"",
2021-05-17 11:02:35 +00:00
remote_table_name,
2021-05-24 18:59:10 +00:00
context_);
LOG_TRACE(log, "Query: {}", query);
2021-05-17 11:02:35 +00:00
Block sample_block;
for (const String & column_name : column_names)
{
auto column_data = metadata_snapshot->getColumns().getPhysical(column_name);
sample_block.insert({column_data.type, column_data.name});
}
2021-08-11 16:44:34 +00:00
return Pipe(std::make_shared<SQLiteSource>(sqlite_db, query, sample_block, max_block_size));
2021-05-17 11:02:35 +00:00
}
2021-07-08 13:27:30 +00:00
2021-07-23 14:25:35 +00:00
class SQLiteSink : public SinkToStorage
2021-05-24 18:59:10 +00:00
{
public:
2021-07-23 14:25:35 +00:00
explicit SQLiteSink(
2021-05-24 18:59:10 +00:00
const StorageSQLite & storage_,
const StorageMetadataPtr & metadata_snapshot_,
2021-07-10 08:05:23 +00:00
StorageSQLite::SQLitePtr sqlite_db_,
const String & remote_table_name_)
2021-07-26 10:08:40 +00:00
: SinkToStorage(metadata_snapshot_->getSampleBlock())
2021-07-23 14:25:35 +00:00
, storage{storage_}
2021-05-24 18:59:10 +00:00
, metadata_snapshot(metadata_snapshot_)
2021-07-10 08:05:23 +00:00
, sqlite_db(sqlite_db_)
2021-05-24 18:59:10 +00:00
, remote_table_name(remote_table_name_)
{
}
2021-07-23 19:33:59 +00:00
String getName() const override { return "SQLiteSink"; }
2021-07-23 14:25:35 +00:00
void consume(Chunk chunk) override
2021-05-24 18:59:10 +00:00
{
auto block = getHeader().cloneWithColumns(chunk.getColumns());
2021-05-24 18:59:10 +00:00
WriteBufferFromOwnString sqlbuf;
2021-07-10 08:05:23 +00:00
2021-05-24 18:59:10 +00:00
sqlbuf << "INSERT INTO ";
2021-06-01 20:52:12 +00:00
sqlbuf << doubleQuoteString(remote_table_name);
2021-05-24 18:59:10 +00:00
sqlbuf << " (";
for (auto it = block.begin(); it != block.end(); ++it)
{
if (it != block.begin())
sqlbuf << ", ";
2021-06-01 20:52:12 +00:00
sqlbuf << quoteString(it->name);
2021-05-24 18:59:10 +00:00
}
sqlbuf << ") VALUES ";
2021-10-11 16:11:50 +00:00
auto writer = FormatFactory::instance().getOutputFormat("Values", sqlbuf, metadata_snapshot->getSampleBlock(), storage.getContext());
2021-05-24 18:59:10 +00:00
writer->write(block);
sqlbuf << ";";
char * err_message = nullptr;
2021-07-10 08:05:23 +00:00
int status = sqlite3_exec(sqlite_db.get(), sqlbuf.str().c_str(), nullptr, nullptr, &err_message);
2021-05-24 18:59:10 +00:00
if (status != SQLITE_OK)
{
String err_msg(err_message);
sqlite3_free(err_message);
2021-07-10 11:00:59 +00:00
throw Exception(ErrorCodes::SQLITE_ENGINE_ERROR,
"Failed to execute sqlite INSERT query. Status: {}. Message: {}",
status, err_msg);
2021-05-24 18:59:10 +00:00
}
}
private:
const StorageSQLite & storage;
StorageMetadataPtr metadata_snapshot;
2021-07-10 08:05:23 +00:00
StorageSQLite::SQLitePtr sqlite_db;
String remote_table_name;
2021-05-24 18:59:10 +00:00
};
2021-07-08 13:27:30 +00:00
2021-07-23 14:25:35 +00:00
SinkToStoragePtr StorageSQLite::write(const ASTPtr & /* query */, const StorageMetadataPtr & metadata_snapshot, ContextPtr)
2021-05-24 18:59:10 +00:00
{
2021-07-26 07:05:28 +00:00
if (!sqlite_db)
sqlite_db = openSQLiteDB(database_path, getContext(), /* throw_on_error */true);
2021-07-23 14:25:35 +00:00
return std::make_shared<SQLiteSink>(*this, metadata_snapshot, sqlite_db, remote_table_name);
2021-05-24 18:59:10 +00:00
}
2021-07-10 08:05:23 +00:00
2021-05-17 11:02:35 +00:00
void registerStorageSQLite(StorageFactory & factory)
{
2021-07-10 08:05:23 +00:00
factory.registerStorage("SQLite", [](const StorageFactory::Arguments & args) -> StoragePtr
{
ASTs & engine_args = args.engine_args;
if (engine_args.size() != 2)
throw Exception("SQLite database requires 2 arguments: database path, table name",
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
for (auto & engine_arg : engine_args)
engine_arg = evaluateConstantExpressionOrIdentifierAsLiteral(engine_arg, args.getLocalContext());
const auto database_path = engine_args[0]->as<ASTLiteral &>().value.safeGet<String>();
const auto table_name = engine_args[1]->as<ASTLiteral &>().value.safeGet<String>();
2021-07-26 07:05:28 +00:00
auto sqlite_db = openSQLiteDB(database_path, args.getContext(), /* throw_on_error */!args.attach);
2021-07-13 12:38:37 +00:00
2021-07-26 07:05:28 +00:00
return StorageSQLite::create(args.table_id, sqlite_db, database_path,
2021-07-10 08:05:23 +00:00
table_name, args.columns, args.constraints, args.getContext());
},
{
.source_access_type = AccessType::SQLITE,
});
2021-05-17 11:02:35 +00:00
}
}
2021-07-08 13:27:30 +00:00
#endif