From dfee4d2bb0acf983f092f250b2414db39c4e8dfd Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Mon, 8 May 2017 15:51:11 +0300 Subject: [PATCH] cql3: Fix partition key bind indices for prepared statements Fix the CQL front-end to populate the partition key bind index array in result message prepared metadata, which is needed for CQL binary protocol v4 to function correctly. Fixes #2355. Message-Id: <1494247871-3148-1-git-send-email-penberg@scylladb.com> --- cql3/statements/batch_statement.cc | 17 ++++++++++++++++- cql3/statements/modification_statement.cc | 4 +++- cql3/statements/parsed_statement.cc | 13 +++++++------ cql3/statements/prepared_statement.hh | 7 ++++--- cql3/statements/select_statement.cc | 4 +++- cql3/variable_specifications.hh | 20 ++++++++++++++++++++ transport/messages/result_message.hh | 3 +-- 7 files changed, 54 insertions(+), 14 deletions(-) diff --git a/cql3/statements/batch_statement.cc b/cql3/statements/batch_statement.cc index ce3a68e4d2..0053f8415d 100644 --- a/cql3/statements/batch_statement.cc +++ b/cql3/statements/batch_statement.cc @@ -414,8 +414,18 @@ std::unique_ptr batch_statement::prepare(database& db, cql_stats& stats) { auto&& bound_names = get_bound_variables(); + stdx::optional first_ks; + stdx::optional first_cf; + bool have_multiple_cfs = false; + std::vector> statements; for (auto&& parsed : _parsed_statements) { + if (!first_ks) { + first_ks = parsed->keyspace(); + first_cf = parsed->column_family(); + } else { + have_multiple_cfs = first_ks.value() != parsed->keyspace() || first_cf.value() != parsed->column_family(); + } statements.push_back(parsed->prepare(db, bound_names, stats)); } @@ -425,8 +435,13 @@ batch_statement::prepare(database& db, cql_stats& stats) { cql3::statements::batch_statement batch_statement_(bound_names->size(), _type, std::move(statements), std::move(prep_attrs), stats); batch_statement_.validate(); + std::vector partition_key_bind_indices; + if (!have_multiple_cfs && batch_statement_.get_statements().size() > 0) { + partition_key_bind_indices = bound_names->get_partition_key_bind_indexes(batch_statement_.get_statements()[0]->s); + } return std::make_unique(make_shared(std::move(batch_statement_)), - bound_names->get_specifications()); + bound_names->get_specifications(), + std::move(partition_key_bind_indices)); } } diff --git a/cql3/statements/modification_statement.cc b/cql3/statements/modification_statement.cc index 3657a93b53..846897ec8c 100644 --- a/cql3/statements/modification_statement.cc +++ b/cql3/statements/modification_statement.cc @@ -516,9 +516,11 @@ namespace raw { std::unique_ptr modification_statement::modification_statement::prepare(database& db, cql_stats& stats) { + schema_ptr schema = validation::validate_column_family(db, keyspace(), column_family()); auto bound_names = get_bound_variables(); auto statement = prepare(db, bound_names, stats); - return std::make_unique(std::move(statement), *bound_names); + auto partition_key_bind_indices = bound_names->get_partition_key_bind_indexes(schema); + return std::make_unique(std::move(statement), *bound_names, std::move(partition_key_bind_indices)); } ::shared_ptr diff --git a/cql3/statements/parsed_statement.cc b/cql3/statements/parsed_statement.cc index 298515f923..3a13920572 100644 --- a/cql3/statements/parsed_statement.cc +++ b/cql3/statements/parsed_statement.cc @@ -67,21 +67,22 @@ bool parsed_statement::uses_function(const sstring& ks_name, const sstring& func } -prepared_statement::prepared_statement(::shared_ptr statement_, std::vector<::shared_ptr> bound_names_) +prepared_statement::prepared_statement(::shared_ptr statement_, std::vector<::shared_ptr> bound_names_, std::vector partition_key_bind_indices) : statement(std::move(statement_)) , bound_names(std::move(bound_names_)) + , partition_key_bind_indices(std::move(partition_key_bind_indices)) { } -prepared_statement::prepared_statement(::shared_ptr statement_, const variable_specifications& names) - : prepared_statement(statement_, names.get_specifications()) +prepared_statement::prepared_statement(::shared_ptr statement_, const variable_specifications& names, const std::vector& partition_key_bind_indices) + : prepared_statement(statement_, names.get_specifications(), partition_key_bind_indices) { } -prepared_statement::prepared_statement(::shared_ptr statement_, variable_specifications&& names) - : prepared_statement(statement_, std::move(names).get_specifications()) +prepared_statement::prepared_statement(::shared_ptr statement_, variable_specifications&& names, std::vector&& partition_key_bind_indices) + : prepared_statement(statement_, std::move(names).get_specifications(), std::move(partition_key_bind_indices)) { } prepared_statement::prepared_statement(::shared_ptr&& statement_) - : prepared_statement(statement_, std::vector<::shared_ptr>()) + : prepared_statement(statement_, std::vector<::shared_ptr>(), std::vector()) { } } diff --git a/cql3/statements/prepared_statement.hh b/cql3/statements/prepared_statement.hh index fef202d69d..1856534a4c 100644 --- a/cql3/statements/prepared_statement.hh +++ b/cql3/statements/prepared_statement.hh @@ -71,12 +71,13 @@ public: sstring raw_cql_statement; const ::shared_ptr statement; const std::vector<::shared_ptr> bound_names; + std::vector partition_key_bind_indices; - prepared_statement(::shared_ptr statement_, std::vector<::shared_ptr> bound_names_); + prepared_statement(::shared_ptr statement_, std::vector<::shared_ptr> bound_names_, std::vector partition_key_bind_indices); - prepared_statement(::shared_ptr statement_, const variable_specifications& names); + prepared_statement(::shared_ptr statement_, const variable_specifications& names, const std::vector& partition_key_bind_indices); - prepared_statement(::shared_ptr statement_, variable_specifications&& names); + prepared_statement(::shared_ptr statement_, variable_specifications&& names, std::vector&& partition_key_bind_indices); prepared_statement(::shared_ptr&& statement_); diff --git a/cql3/statements/select_statement.cc b/cql3/statements/select_statement.cc index 22ff56db40..cd6be4fffa 100644 --- a/cql3/statements/select_statement.cc +++ b/cql3/statements/select_statement.cc @@ -460,7 +460,9 @@ std::unique_ptr select_statement::prepare(database& db, cql_ prepare_limit(db, bound_names), stats); - return std::make_unique(std::move(stmt), std::move(*bound_names)); + auto partition_key_bind_indices = bound_names->get_partition_key_bind_indexes(schema); + + return std::make_unique(std::move(stmt), std::move(*bound_names), std::move(partition_key_bind_indices)); } ::shared_ptr diff --git a/cql3/variable_specifications.hh b/cql3/variable_specifications.hh index 8dad138edc..48e78daa3b 100644 --- a/cql3/variable_specifications.hh +++ b/cql3/variable_specifications.hh @@ -80,6 +80,26 @@ public: return std::move(_specs); } + std::vector get_partition_key_bind_indexes(schema_ptr schema) const { + auto count = schema->partition_key_columns().size(); + std::vector partition_key_positions(count, uint16_t(0)); + std::vector set(count, false); + for (size_t i = 0; i < _specs.size(); i++) { + auto& target_column = _specs[i]; + const auto* cdef = schema->get_column_definition(target_column->name->name()); + if (cdef && cdef->is_partition_key()) { + partition_key_positions[cdef->position()] = i; + set[cdef->position()] = true; + } + } + for (bool b : set) { + if (!b) { + return {}; + } + } + return partition_key_positions; + } + void add(int32_t bind_index, ::shared_ptr spec) { auto name = _variable_names[bind_index]; // Use the user name, if there is one diff --git a/transport/messages/result_message.hh b/transport/messages/result_message.hh index 3ed9d4589a..3ed4ff66fd 100644 --- a/transport/messages/result_message.hh +++ b/transport/messages/result_message.hh @@ -43,8 +43,7 @@ private: protected: prepared(cql3::statements::prepared_statement::checked_weak_ptr prepared) : _prepared(std::move(prepared)) - // FIXME: Populate partition key bind indices for prepared_metadata. - , _metadata{::make_shared(_prepared->bound_names, std::vector())} + , _metadata{::make_shared(_prepared->bound_names, _prepared->partition_key_bind_indices)} , _result_metadata{extract_result_metadata(_prepared->statement)} { } public: