diff --git a/dbms/src/Common/ErrorCodes.cpp b/dbms/src/Common/ErrorCodes.cpp index 513a26987e0428c2f8eb805c96133bbab82235c8..302a7bffd5e28c8d1ae4e7724c12e36ac9ccfa22 100644 --- a/dbms/src/Common/ErrorCodes.cpp +++ b/dbms/src/Common/ErrorCodes.cpp @@ -393,6 +393,7 @@ namespace ErrorCodes extern const int REPLICA_STATUS_CHANGED = 416; extern const int EXPECTED_ALL_OR_ANY = 417; extern const int UNKNOWN_JOIN_STRICTNESS = 418; + extern const int MULTIPLE_ASSIGNMENTS_TO_COLUMN = 419; extern const int KEEPER_EXCEPTION = 999; extern const int POCO_EXCEPTION = 1000; diff --git a/dbms/src/Parsers/ASTAlterQuery.cpp b/dbms/src/Parsers/ASTAlterQuery.cpp index cbf1c934b05687fb39389acd682ab1b322e8cd1b..7081b5122474ecddb34240a949a367cf1d94f84d 100644 --- a/dbms/src/Parsers/ASTAlterQuery.cpp +++ b/dbms/src/Parsers/ASTAlterQuery.cpp @@ -133,6 +133,14 @@ void ASTAlterCommand::formatImpl( settings.ostr << (settings.hilite ? hilite_keyword : "") << indent_str << "DELETE WHERE " << (settings.hilite ? hilite_none : ""); predicate->formatImpl(settings, state, frame); } + else if (type == ASTAlterCommand::UPDATE) + { + settings.ostr << (settings.hilite ? hilite_keyword : "") << indent_str << "UPDATE " << (settings.hilite ? hilite_none : ""); + update_assignments->formatImpl(settings, state, frame); + + settings.ostr << (settings.hilite ? hilite_keyword : "") << " WHERE " << (settings.hilite ? hilite_none : ""); + predicate->formatImpl(settings, state, frame); + } else throw Exception("Unexpected type of ALTER", ErrorCodes::UNEXPECTED_AST_STRUCTURE); } diff --git a/dbms/src/Parsers/ASTAlterQuery.h b/dbms/src/Parsers/ASTAlterQuery.h index 683d07808764ed2fc12666cd03e067517b934863..c79f9ba8b2f8974749510068fcf9e621e560d2bc 100644 --- a/dbms/src/Parsers/ASTAlterQuery.h +++ b/dbms/src/Parsers/ASTAlterQuery.h @@ -33,6 +33,7 @@ public: FREEZE_PARTITION, DELETE, + UPDATE, NO_TYPE, }; @@ -59,9 +60,12 @@ public: */ ASTPtr partition; - /// For DELETE WHERE: the predicate that filters the rows to delete. + /// For DELETE/UPDATE WHERE: the predicate that filters the rows to delete/update. ASTPtr predicate; + /// A list of expressions of the form `column = expr` for the UPDATE command. + ASTPtr update_assignments; + bool detach = false; /// true for DETACH PARTITION bool part = false; /// true for ATTACH PART diff --git a/dbms/src/Parsers/ASTAssignment.h b/dbms/src/Parsers/ASTAssignment.h new file mode 100644 index 0000000000000000000000000000000000000000..18bf46c171cc2d162d51fa69bfd4572384d6a11f --- /dev/null +++ b/dbms/src/Parsers/ASTAssignment.h @@ -0,0 +1,44 @@ +#pragma once + +#include + +namespace DB +{ + +/// Part of the ALTER UPDATE statement of the form: column = expr +class ASTAssignment : public IAST +{ +public: + String column_name; + ASTPtr expression; + + String getID() const override { return "Assignment_" + column_name; } + + ASTPtr clone() const override + { + auto res = std::make_shared(*this); + res->children.clear(); + + if (expression) + { + res->expression = expression->clone(); + res->children.push_back(res->expression); + } + + return res; + } + +protected: + void formatImpl(const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const override + { + settings.ostr << (settings.hilite ? hilite_identifier : ""); + settings.writeIdentifier(column_name); + settings.ostr << (settings.hilite ? hilite_none : ""); + + settings.ostr << (settings.hilite ? hilite_operator : "") << " = " << (settings.hilite ? hilite_none : ""); + + expression->formatImpl(settings, state, frame); + } +}; + +} diff --git a/dbms/src/Parsers/ASTColumnDeclaration.h b/dbms/src/Parsers/ASTColumnDeclaration.h index 57a1f7695d7558a5958fda0c692f95fa07b35a6d..308e9b665269ea4d8a06af1001515cb782707733 100644 --- a/dbms/src/Parsers/ASTColumnDeclaration.h +++ b/dbms/src/Parsers/ASTColumnDeclaration.h @@ -21,8 +21,6 @@ public: ASTPtr clone() const override { const auto res = std::make_shared(*this); - ASTPtr ptr{res}; - res->children.clear(); if (type) @@ -37,7 +35,7 @@ public: res->children.push_back(res->default_expression); } - return ptr; + return res; } protected: diff --git a/dbms/src/Parsers/ASTIdentifier.cpp b/dbms/src/Parsers/ASTIdentifier.cpp index 0670c3811a7a57d227f5aa1f3b21384114735784..e7032d24e8b99021a36eba688e779da0cbf7e750 100644 --- a/dbms/src/Parsers/ASTIdentifier.cpp +++ b/dbms/src/Parsers/ASTIdentifier.cpp @@ -11,11 +11,7 @@ void ASTIdentifier::formatImplWithoutAlias(const FormatSettings & settings, Form auto format_element = [&](const String & name) { settings.ostr << (settings.hilite ? hilite_identifier : ""); - - WriteBufferFromOStream wb(settings.ostr, 32); - settings.writeIdentifier(name, wb); - wb.next(); - + settings.writeIdentifier(name); settings.ostr << (settings.hilite ? hilite_none : ""); }; diff --git a/dbms/src/Parsers/ASTWithAlias.cpp b/dbms/src/Parsers/ASTWithAlias.cpp index ef261386cd08dfa3913e1053b03fa4fe7ef2e3f2..67a4401f9a5a5ea26b669e3e744634411fcaede6 100644 --- a/dbms/src/Parsers/ASTWithAlias.cpp +++ b/dbms/src/Parsers/ASTWithAlias.cpp @@ -9,11 +9,7 @@ namespace DB void ASTWithAlias::writeAlias(const String & name, const FormatSettings & settings) const { settings.ostr << (settings.hilite ? hilite_keyword : "") << " AS " << (settings.hilite ? hilite_alias : ""); - - WriteBufferFromOStream wb(settings.ostr, 32); - settings.writeIdentifier(name, wb); - wb.next(); - + settings.writeIdentifier(name); settings.ostr << (settings.hilite ? hilite_none : ""); } @@ -25,8 +21,7 @@ void ASTWithAlias::formatImpl(const FormatSettings & settings, FormatState & sta /// If we have previously output this node elsewhere in the query, now it is enough to output only the alias. if (!state.printed_asts_with_alias.emplace(frame.current_select, alias).second) { - WriteBufferFromOStream wb(settings.ostr, 32); - settings.writeIdentifier(alias, wb); + settings.writeIdentifier(alias); return; } } diff --git a/dbms/src/Parsers/IAST.cpp b/dbms/src/Parsers/IAST.cpp index e5d9b7fc29da79f307473fa15814d16f9abfcfd6..a3244e1abeb7ef4b82b5bd57b20854cfc7f35ed2 100644 --- a/dbms/src/Parsers/IAST.cpp +++ b/dbms/src/Parsers/IAST.cpp @@ -99,8 +99,10 @@ String IAST::getColumnName() const } -void IAST::FormatSettings::writeIdentifier(const String & name, WriteBuffer & out) const +void IAST::FormatSettings::writeIdentifier(const String & name) const { + WriteBufferFromOStream out(ostr, 32); + switch (identifier_quoting_style) { case IdentifierQuotingStyle::None: @@ -128,6 +130,8 @@ void IAST::FormatSettings::writeIdentifier(const String & name, WriteBuffer & ou break; } } + + out.next(); } } diff --git a/dbms/src/Parsers/IAST.h b/dbms/src/Parsers/IAST.h index b3d06abf584f699f00a1223b2c095fd3e85671ee..5135e08a9654b6f8c32f3904e3a180c5c3cabaf2 100644 --- a/dbms/src/Parsers/IAST.h +++ b/dbms/src/Parsers/IAST.h @@ -164,7 +164,7 @@ public: nl_or_ws = one_line ? ' ' : '\n'; } - void writeIdentifier(const String & name, WriteBuffer & out) const; + void writeIdentifier(const String & name) const; }; /// State. For example, a set of nodes can be remembered, which we already walk through. diff --git a/dbms/src/Parsers/ParserAlterQuery.cpp b/dbms/src/Parsers/ParserAlterQuery.cpp index c592cb66d7d5f1a2a024c0692befdc54a00d0e8c..56eaddb38eed554d48eed72cc5abc1c09950706b 100644 --- a/dbms/src/Parsers/ParserAlterQuery.cpp +++ b/dbms/src/Parsers/ParserAlterQuery.cpp @@ -1,11 +1,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include @@ -38,12 +40,17 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected ParserKeyword s_name("NAME"); ParserKeyword s_delete_where("DELETE WHERE"); + ParserKeyword s_update("UPDATE"); + ParserKeyword s_where("WHERE"); ParserCompoundIdentifier parser_name; ParserStringLiteral parser_string_literal; ParserCompoundColumnDeclaration parser_col_decl; ParserPartition parser_partition; - ParserExpression exp_elem; + ParserExpression parser_exp_elem; + ParserList parser_assignment_list( + std::make_unique(), std::make_unique(TokenType::Comma), + /* allow_empty = */ false); if (s_add_column.ignore(pos, expected)) { @@ -195,11 +202,24 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected } else if (s_delete_where.ignore(pos, expected)) { - if (!exp_elem.parse(pos, command->predicate, expected)) + if (!parser_exp_elem.parse(pos, command->predicate, expected)) return false; command->type = ASTAlterCommand::DELETE; } + else if (s_update.ignore(pos, expected)) + { + if (!parser_assignment_list.parse(pos, command->update_assignments, expected)) + return false; + + if (!s_where.ignore(pos, expected)) + return false; + + if (!parser_exp_elem.parse(pos, command->predicate, expected)) + return false; + + command->type = ASTAlterCommand::UPDATE; + } else return false; @@ -213,6 +233,8 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected command->children.push_back(command->partition); if (command->predicate) command->children.push_back(command->predicate); + if (command->update_assignments) + command->children.push_back(command->update_assignments); return true; } @@ -240,6 +262,33 @@ bool ParserAlterCommandList::parseImpl(Pos & pos, ASTPtr & node, Expected & expe } +bool ParserAssignment::parseImpl(Pos & pos, ASTPtr & node, Expected & expected) +{ + auto assignment = std::make_shared(); + node = assignment; + + ParserIdentifier p_identifier; + ParserToken s_equals(TokenType::Equals); + ParserExpression p_expression; + + ASTPtr column; + if (!p_identifier.parse(pos, column, expected)) + return false; + + if (!s_equals.ignore(pos, expected)) + return false; + + if (!p_expression.parse(pos, assignment->expression, expected)) + return false; + + assignment->column_name = typeid_cast(*column).name; + if (assignment->expression) + assignment->children.push_back(assignment->expression); + + return true; +} + + bool ParserAlterQuery::parseImpl(Pos & pos, ASTPtr & node, Expected & expected) { auto query = std::make_shared(); @@ -269,4 +318,5 @@ bool ParserAlterQuery::parseImpl(Pos & pos, ASTPtr & node, Expected & expected) return true; } + } diff --git a/dbms/src/Parsers/ParserAlterQuery.h b/dbms/src/Parsers/ParserAlterQuery.h index c758e0304b4a92e0d2ce4dac6b998b51977eabe9..46908ae135d8778321dc85833d35a9d5cda4d4c8 100644 --- a/dbms/src/Parsers/ParserAlterQuery.h +++ b/dbms/src/Parsers/ParserAlterQuery.h @@ -17,12 +17,13 @@ namespace DB * [FETCH PARTITION partition FROM ...] * [FREEZE PARTITION] * [DELETE WHERE ...] + * [UPDATE col_name = expr, ... WHERE ...] */ -class ParserAlterCommand : public IParserBase +class ParserAlterQuery : public IParserBase { protected: - const char * getName() const { return "ALTER command"; } + const char * getName() const { return "ALTER query"; } bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected); }; @@ -35,10 +36,19 @@ protected: }; -class ParserAlterQuery : public IParserBase +class ParserAlterCommand : public IParserBase { protected: - const char * getName() const { return "ALTER query"; } + const char * getName() const { return "ALTER command"; } + bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected); +}; + + +/// Part of the UPDATE command of the form: col_name = expr +class ParserAssignment : public IParserBase +{ +protected: + const char * getName() const { return "column assignment"; } bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected); }; diff --git a/dbms/src/Storages/MutationCommands.cpp b/dbms/src/Storages/MutationCommands.cpp index 38bacbfd9b6dde5611049b874d002cde1dd9f030..94bb77a48dc7c29c44e905a0ecb45244694a2ee6 100644 --- a/dbms/src/Storages/MutationCommands.cpp +++ b/dbms/src/Storages/MutationCommands.cpp @@ -7,6 +7,8 @@ #include #include #include +#include +#include #include @@ -16,6 +18,7 @@ namespace DB namespace ErrorCodes { extern const int UNKNOWN_MUTATION_COMMAND; + extern const int MULTIPLE_ASSIGNMENTS_TO_COLUMN; } std::optional MutationCommand::parse(ASTAlterCommand * command) @@ -28,6 +31,22 @@ std::optional MutationCommand::parse(ASTAlterCommand * command) res.predicate = command->predicate; return res; } + else if (command->type == ASTAlterCommand::UPDATE) + { + MutationCommand res; + res.ast = command->ptr(); + res.type = UPDATE; + res.predicate = command->predicate; + for (const ASTPtr & assignment_ast : command->update_assignments->children) + { + const auto & assignment = typeid_cast(*assignment_ast); + auto insertion = res.column_to_update_expression.emplace(assignment.column_name, assignment.expression); + if (!insertion.second) + throw Exception("Multiple assignments in the single statement to column `" + assignment.column_name + "`", + ErrorCodes::MULTIPLE_ASSIGNMENTS_TO_COLUMN); + } + return res; + } else return {}; } diff --git a/dbms/src/Storages/MutationCommands.h b/dbms/src/Storages/MutationCommands.h index 2d9d9fa5fde330a5cce2b50d3637c122ff5efb95..7a188f10553bbd52fcfe8a4216b4a0b52fdae81d 100644 --- a/dbms/src/Storages/MutationCommands.h +++ b/dbms/src/Storages/MutationCommands.h @@ -2,6 +2,7 @@ #include #include +#include namespace DB @@ -20,12 +21,15 @@ struct MutationCommand { EMPTY, /// Not used. DELETE, + UPDATE, }; Type type = EMPTY; ASTPtr predicate; + std::unordered_map column_to_update_expression; + static std::optional parse(ASTAlterCommand * command); };