ClickHouse/src/Storages/StorageDelta.cpp

301 lines
9.5 KiB
C++
Raw Normal View History

2022-11-04 16:03:12 +00:00
#include "config.h"
2022-09-07 07:16:32 +00:00
#if USE_AWS_S3
2022-11-07 11:45:26 +00:00
#include <Storages/StorageDelta.h>
#include <Common/logger_useful.h>
2022-11-07 11:45:26 +00:00
#include <IO/ReadBufferFromS3.h>
#include <IO/ReadHelpers.h>
#include <IO/ReadSettings.h>
#include <IO/S3Common.h>
2022-11-07 11:45:26 +00:00
#include <Storages/ExternalDataSourceConfiguration.h>
#include <Storages/StorageFactory.h>
#include <Storages/checkAndGetLiteralArgument.h>
2022-11-03 17:28:41 +00:00
2022-11-07 11:45:26 +00:00
#include <Formats/FormatFactory.h>
2022-11-07 11:45:26 +00:00
#include <aws/core/auth/AWSCredentials.h>
#include <aws/s3/S3Client.h>
#include <aws/s3/model/ListObjectsV2Request.h>
2022-11-07 11:45:26 +00:00
#include <QueryPipeline/Pipe.h>
2022-11-07 11:45:26 +00:00
#include <fmt/format.h>
#include <fmt/ranges.h>
#include <ranges>
2022-09-20 10:24:47 +00:00
namespace DB
{
namespace ErrorCodes
{
extern const int S3_ERROR;
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
2022-09-20 10:24:47 +00:00
extern const int INCORRECT_DATA;
}
2022-09-21 16:14:51 +00:00
void DeltaLakeMetadata::setLastModifiedTime(const String & filename, uint64_t timestamp)
2022-09-06 15:01:34 +00:00
{
2022-09-19 15:23:07 +00:00
file_update_time[filename] = timestamp;
}
2022-09-19 15:23:07 +00:00
void DeltaLakeMetadata::remove(const String & filename, uint64_t /*timestamp */)
2022-09-06 15:01:34 +00:00
{
2022-09-20 10:24:47 +00:00
bool erase = file_update_time.erase(filename);
if (!erase)
throw Exception(ErrorCodes::INCORRECT_DATA, "Invalid table metadata, tried to remove {} before adding it", filename);
}
2022-09-06 15:01:34 +00:00
std::vector<String> DeltaLakeMetadata::ListCurrentFiles() &&
{
std::vector<String> keys;
keys.reserve(file_update_time.size());
2022-09-06 15:01:34 +00:00
for (auto && [k, _] : file_update_time)
keys.push_back(k);
return keys;
}
2022-11-04 20:49:24 +00:00
JsonMetadataGetter::JsonMetadataGetter(StorageS3::S3Configuration & configuration_, const String & table_path_, ContextPtr context)
: base_configuration(configuration_), table_path(table_path_)
2022-09-06 15:01:34 +00:00
{
2022-11-04 20:49:24 +00:00
Init(context);
2022-09-06 15:01:34 +00:00
}
2022-11-04 20:49:24 +00:00
void JsonMetadataGetter::Init(ContextPtr context)
2022-09-06 15:01:34 +00:00
{
auto keys = getJsonLogFiles();
2022-09-06 15:01:34 +00:00
2022-09-02 06:54:16 +00:00
// read data from every json log file
2022-09-06 15:01:34 +00:00
for (const String & key : keys)
{
2022-11-04 20:49:24 +00:00
auto buf = createS3ReadBuffer(key, context);
2022-09-02 06:54:16 +00:00
2022-11-07 11:45:26 +00:00
char c;
2022-09-20 14:16:27 +00:00
while (!buf->eof())
2022-09-28 11:21:32 +00:00
{
2022-11-07 11:45:26 +00:00
/// May be some invalid characters before json.
2022-09-28 11:21:32 +00:00
while (buf->peek(c) && c != '{')
buf->ignore();
2022-11-07 11:45:26 +00:00
2022-09-28 11:21:32 +00:00
if (buf->eof())
2022-09-21 16:14:51 +00:00
break;
String json_str;
readJSONObjectPossiblyInvalid(json_str, *buf);
2022-09-28 11:21:32 +00:00
2022-09-21 16:14:51 +00:00
if (json_str.empty())
2022-09-02 06:54:16 +00:00
continue;
2022-09-20 14:16:27 +00:00
const JSON json(json_str);
handleJSON(json);
}
}
}
2022-09-06 15:01:34 +00:00
std::vector<String> JsonMetadataGetter::getJsonLogFiles()
{
std::vector<String> keys;
const auto & client = base_configuration.client;
Aws::S3::Model::ListObjectsV2Request request;
Aws::S3::Model::ListObjectsV2Outcome outcome;
bool is_finished{false};
const auto bucket{base_configuration.uri.bucket};
request.SetBucket(bucket);
2022-11-04 20:49:24 +00:00
2022-11-07 11:45:26 +00:00
/// DeltaLake format stores all metadata json files in _delta_log directory
static constexpr auto deltalake_metadata_directory = "_delta_log";
request.SetPrefix(std::filesystem::path(table_path) / deltalake_metadata_directory);
while (!is_finished)
{
outcome = client->ListObjectsV2(request);
if (!outcome.IsSuccess())
throw Exception(
ErrorCodes::S3_ERROR,
"Could not list objects in bucket {} with key {}, S3 exception: {}, message: {}",
quoteString(bucket),
quoteString(table_path),
backQuote(outcome.GetError().GetExceptionName()),
quoteString(outcome.GetError().GetMessage()));
const auto & result_batch = outcome.GetResult().GetContents();
for (const auto & obj : result_batch)
{
const auto & filename = obj.GetKey();
2022-09-06 15:01:34 +00:00
2022-11-04 20:49:24 +00:00
// DeltaLake metadata files have json extension
2022-09-20 10:24:47 +00:00
if (std::filesystem::path(filename).extension() == ".json")
keys.push_back(filename);
}
2022-11-07 11:45:26 +00:00
/// Needed in case any more results are available
/// if so, we will continue reading, and not read keys that were already read
request.SetContinuationToken(outcome.GetResult().GetNextContinuationToken());
2022-11-04 20:51:44 +00:00
2022-11-04 20:49:24 +00:00
/// Set to false if all of the results were returned. Set to true if more keys
/// are available to return. If the number of results exceeds that specified by
/// MaxKeys, all of the results might not be returned
is_finished = !outcome.GetResult().GetIsTruncated();
}
return keys;
}
2022-11-04 20:49:24 +00:00
std::shared_ptr<ReadBuffer> JsonMetadataGetter::createS3ReadBuffer(const String & key, ContextPtr context)
2022-09-06 15:01:34 +00:00
{
2022-11-07 11:45:26 +00:00
/// TODO: add parallel downloads
2022-09-20 14:16:27 +00:00
return std::make_shared<ReadBufferFromS3>(
2022-09-06 15:01:34 +00:00
base_configuration.client,
base_configuration.uri.bucket,
key,
base_configuration.uri.version_id,
/* max single read retries */10,
2022-11-04 20:49:24 +00:00
context->getReadSettings());
}
2022-09-28 11:21:32 +00:00
void JsonMetadataGetter::handleJSON(const JSON & json)
{
2022-09-20 14:16:27 +00:00
if (json.has("add"))
{
auto path = json["add"]["path"].getString();
auto timestamp = json["add"]["modificationTime"].getInt();
2022-09-21 16:14:51 +00:00
metadata.setLastModifiedTime(path, timestamp);
2022-09-20 14:16:27 +00:00
}
else if (json.has("remove"))
{
auto path = json["remove"]["path"].getString();
auto timestamp = json["remove"]["deletionTimestamp"].getInt();
metadata.remove(path, timestamp);
}
}
StorageDelta::StorageDelta(
2022-11-03 17:28:41 +00:00
const StorageS3Configuration & configuration_,
const StorageID & table_id_,
2022-09-02 06:54:16 +00:00
ColumnsDescription columns_,
const ConstraintsDescription & constraints_,
const String & comment,
2022-11-03 17:28:41 +00:00
ContextPtr context_,
std::optional<FormatSettings> format_settings_)
: IStorage(table_id_)
, base_configuration{configuration_.url, configuration_.auth_settings, configuration_.rw_settings, configuration_.headers}
, log(&Poco::Logger::get("StorageDeltaLake (" + table_id_.table_name + ")"))
2022-11-03 17:28:41 +00:00
, table_path(base_configuration.uri.key)
{
2022-09-02 06:54:16 +00:00
StorageInMemoryMetadata storage_metadata;
2022-09-28 11:03:43 +00:00
StorageS3::updateS3Configuration(context_, base_configuration);
2022-09-06 15:01:34 +00:00
2022-11-04 20:49:24 +00:00
JsonMetadataGetter getter{base_configuration, table_path, context_};
2022-09-02 06:54:16 +00:00
auto keys = getter.getFiles();
auto new_uri = base_configuration.uri.uri.toString() + generateQueryFromKeys(std::move(keys));
LOG_DEBUG(log, "New uri: {}", new_uri);
LOG_DEBUG(log, "Table path: {}", table_path);
2022-11-03 18:17:39 +00:00
2022-11-03 17:28:41 +00:00
// set new url in configuration
StorageS3Configuration new_configuration;
new_configuration.url = new_uri;
new_configuration.auth_settings.access_key_id = configuration_.auth_settings.access_key_id;
2022-11-03 18:17:39 +00:00
new_configuration.auth_settings.secret_access_key = configuration_.auth_settings.secret_access_key;
2022-11-03 17:28:41 +00:00
new_configuration.format = configuration_.format;
2022-09-02 06:54:16 +00:00
if (columns_.empty())
{
2022-11-03 18:17:39 +00:00
columns_ = StorageS3::getTableStructureFromData(
new_configuration, /*distributed processing*/ false, format_settings_, context_, nullptr);
2022-09-02 06:54:16 +00:00
storage_metadata.setColumns(columns_);
}
else
2022-09-06 15:01:34 +00:00
storage_metadata.setColumns(columns_);
2022-09-02 06:54:16 +00:00
2022-11-03 18:17:39 +00:00
2022-09-02 06:54:16 +00:00
storage_metadata.setConstraints(constraints_);
storage_metadata.setComment(comment);
setInMemoryMetadata(storage_metadata);
s3engine = std::make_shared<StorageS3>(
2022-11-03 18:17:39 +00:00
new_configuration,
table_id_,
columns_,
constraints_,
comment,
context_,
format_settings_,
/* distributed_processing_ */ false,
nullptr);
}
Pipe StorageDelta::read(
const Names & column_names,
const StorageSnapshotPtr & storage_snapshot,
SelectQueryInfo & query_info,
ContextPtr context,
QueryProcessingStage::Enum processed_stage,
size_t max_block_size,
2022-11-04 16:03:12 +00:00
size_t num_streams)
{
2022-09-28 11:03:43 +00:00
StorageS3::updateS3Configuration(context, base_configuration);
return s3engine->read(column_names, storage_snapshot, query_info, context, processed_stage, max_block_size, num_streams);
}
String StorageDelta::generateQueryFromKeys(std::vector<String> && keys)
{
2022-11-04 20:49:24 +00:00
// DeltaLake store data parts in different files
// keys are filenames of parts
// for StorageS3 to read all parts we need format {key1,key2,key3,...keyn}
2022-09-20 10:24:47 +00:00
std::string new_query = fmt::format("{{{}}}", fmt::join(keys, ","));
return new_query;
}
void registerStorageDelta(StorageFactory & factory)
{
factory.registerStorage(
"DeltaLake",
[](const StorageFactory::Arguments & args)
{
auto & engine_args = args.engine_args;
if (engine_args.empty() || engine_args.size() < 3)
2022-09-28 11:21:32 +00:00
throw Exception(
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
"Storage DeltaLake requires 3 to 4 arguments: table_url, access_key, secret_access_key, [format]");
2022-11-03 17:28:41 +00:00
StorageS3Configuration configuration;
2022-11-03 18:17:39 +00:00
2022-11-03 17:28:41 +00:00
configuration.url = checkAndGetLiteralArgument<String>(engine_args[0], "url");
configuration.auth_settings.access_key_id = checkAndGetLiteralArgument<String>(engine_args[1], "access_key_id");
configuration.auth_settings.secret_access_key = checkAndGetLiteralArgument<String>(engine_args[2], "secret_access_key");
2022-11-03 18:17:39 +00:00
if (engine_args.size() == 4)
2022-11-03 17:28:41 +00:00
configuration.format = checkAndGetLiteralArgument<String>(engine_args[3], "format");
2022-11-07 11:45:26 +00:00
else
{
/// DeltaLake uses Parquet by default.
2022-11-03 17:28:41 +00:00
configuration.format = "Parquet";
2022-11-07 11:45:26 +00:00
}
return std::make_shared<StorageDelta>(
2022-11-03 18:17:39 +00:00
configuration, args.table_id, args.columns, args.constraints, args.comment, args.getContext(), std::nullopt);
},
{
.supports_settings = true,
.supports_schema_inference = true,
.source_access_type = AccessType::S3,
});
}
}
2022-09-07 07:16:32 +00:00
#endif