#include "RedisBlockInputStream.h" #include #include #include #include #include #include #include #include #include #include #include #include "DictionaryStructure.h" namespace DB { namespace ErrorCodes { extern const int TYPE_MISMATCH; extern const int LOGICAL_ERROR; extern const int NUMBER_OF_COLUMNS_DOESNT_MATCH; extern const int INTERNAL_REDIS_ERROR; extern const int UNKNOWN_TYPE; } RedisBlockInputStream::RedisBlockInputStream( const std::shared_ptr & client_, const RedisArray & keys_, const RedisStorageType & storage_type_, const DB::Block & sample_block, const size_t max_block_size_) : client(client_), keys(keys_), storage_type(storage_type_), max_block_size{max_block_size_} { description.init(sample_block); } RedisBlockInputStream::~RedisBlockInputStream() = default; namespace { using ValueType = ExternalResultDescription::ValueType; template inline void insert(IColumn & column, const String & string_value) { assert_cast &>(column).insertValue(parse(string_value)); } void insertValue(IColumn & column, const ValueType type, const Poco::Redis::BulkString & bulk_string) { if (bulk_string.isNull()) throw Exception(ErrorCodes::TYPE_MISMATCH, "Type mismatch, expected not Null String"); const String & string_value = bulk_string.value(); switch (type) { case ValueType::vtUInt8: insert(column, string_value); break; case ValueType::vtUInt16: insert(column, string_value); break; case ValueType::vtUInt32: insert(column, string_value); break; case ValueType::vtUInt64: insert(column, string_value); break; case ValueType::vtInt8: insert(column, string_value); break; case ValueType::vtInt16: insert(column, string_value); break; case ValueType::vtInt32: insert(column, string_value); break; case ValueType::vtInt64: insert(column, string_value); break; case ValueType::vtFloat32: insert(column, string_value); break; case ValueType::vtFloat64: insert(column, string_value); break; case ValueType::vtString: assert_cast(column).insert(parse(string_value)); break; case ValueType::vtDate: assert_cast(column).insertValue(parse(string_value).getDayNum()); break; case ValueType::vtDateTime: { ReadBufferFromString in(string_value); time_t time = 0; readDateTimeText(time, in); if (time < 0) time = 0; assert_cast(column).insertValue(time); break; } case ValueType::vtUUID: assert_cast(column).insertValue(parse(string_value)); break; default: throw Exception(ErrorCodes::UNKNOWN_TYPE, "Value of unsupported type: {}", column.getName()); } } } Block RedisBlockInputStream::readImpl() { if (keys.isNull() || description.sample_block.rows() == 0 || cursor >= keys.size()) all_read = true; if (all_read) return {}; const size_t size = description.sample_block.columns(); MutableColumns columns(size); for (size_t i = 0; i < size; ++i) columns[i] = description.sample_block.getByPosition(i).column->cloneEmpty(); const auto insert_value_by_idx = [this, &columns](size_t idx, const auto & value) { if (description.types[idx].second) { ColumnNullable & column_nullable = static_cast(*columns[idx]); insertValue(column_nullable.getNestedColumn(), description.types[idx].first, value); column_nullable.getNullMapData().emplace_back(0); } else insertValue(*columns[idx], description.types[idx].first, value); }; if (storage_type == RedisStorageType::HASH_MAP) { size_t num_rows = 0; for (; cursor < keys.size(); ++cursor) { const auto & keys_array = keys.get(cursor); if (keys_array.size() < 2) { throw Exception(ErrorCodes::LOGICAL_ERROR, "Too low keys in request to source: {}, expected 2 or more", DB::toString(keys_array.size())); } if (num_rows + keys_array.size() - 1 > max_block_size) break; Poco::Redis::Command command_for_values("HMGET"); for (const auto & elem : keys_array) command_for_values.addRedisType(elem); auto values = client->execute(command_for_values); if (keys_array.size() != values.size() + 1) // 'HMGET' primary_key secondary_keys throw Exception(ErrorCodes::NUMBER_OF_COLUMNS_DOESNT_MATCH, "Inconsistent sizes of keys and values in Redis request"); const auto & primary_key = keys_array.get(0); for (size_t i = 0; i < values.size(); ++i) { const auto & secondary_key = keys_array.get(i + 1); const auto & value = values.get(i); /// null string means 'no value for requested key' if (!value.isNull()) { insert_value_by_idx(0, primary_key); insert_value_by_idx(1, secondary_key); insert_value_by_idx(2, value); ++num_rows; } } } } else { Poco::Redis::Command command_for_values("MGET"); size_t need_values = std::min(max_block_size, keys.size() - cursor); for (size_t i = 0; i < need_values; ++i) command_for_values.add(keys.get(cursor + i)); auto values = client->execute(command_for_values); if (values.size() != need_values) throw Exception(ErrorCodes::INTERNAL_REDIS_ERROR, "Inconsistent sizes of keys and values in Redis request"); for (size_t i = 0; i < values.size(); ++i) { const auto & key = keys.get(cursor + i); const auto & value = values.get(i); /// Null string means 'no value for requested key' if (!value.isNull()) { insert_value_by_idx(0, key); insert_value_by_idx(1, value); } } cursor += need_values; } return description.sample_block.cloneWithColumns(std::move(columns)); } }