Refactoring, fix tests

programs/client/Client.cpp

@@ -92,6 +92,256 @@ namespace ErrorCodes
 }
 
 
+void Client::processParsedSingleQuery(std::optional<bool> echo_query)
+{
+    resetOutput();
+    client_exception.reset();
+    server_exception.reset();
+    have_error = false;
+    if (echo_query.value_or(echo_queries))
+    {
+        writeString(full_query, std_out);
+        writeChar('\n', std_out);
+        std_out.next();
+    }
+    if (is_interactive)
+    {
+        // Generate a new query_id
+        global_context->setCurrentQueryId("");
+        for (const auto & query_id_format : query_id_formats)
+        {
+            writeString(query_id_format.first, std_out);
+            writeString(fmt::format(query_id_format.second, fmt::arg("query_id", global_context->getCurrentQueryId())), std_out);
+            writeChar('\n', std_out);
+            std_out.next();
+        }
+    }
+    processed_rows = 0;
+    written_first_block = false;
+    progress_indication.resetProgress();
+    {
+        /// Temporarily apply query settings to context.
+        std::optional<Settings> old_settings;
+        SCOPE_EXIT_SAFE({
+            if (old_settings)
+                global_context->setSettings(*old_settings);
+        });
+        auto apply_query_settings = [&](const IAST & settings_ast)
+        {
+            if (!old_settings)
+                old_settings.emplace(global_context->getSettingsRef());
+            global_context->applySettingsChanges(settings_ast.as<ASTSetQuery>()->changes);
+        };
+        const auto * insert = parsed_query->as<ASTInsertQuery>();
+        if (insert && insert->settings_ast)
+            apply_query_settings(*insert->settings_ast);
+        /// FIXME: try to prettify this cast using `as<>()`
+        const auto * with_output = dynamic_cast<const ASTQueryWithOutput *>(parsed_query.get());
+        if (with_output && with_output->settings_ast)
+            apply_query_settings(*with_output->settings_ast);
+        if (!connection->checkConnected())
+            connect();
+        ASTPtr input_function;
+        if (insert && insert->select)
+            insert->tryFindInputFunction(input_function);
+        /// INSERT query for which data transfer is needed (not an INSERT SELECT or input()) is processed separately.
+        if (insert && (!insert->select || input_function) && !insert->watch)
+        {
+            if (input_function && insert->format.empty())
+                throw Exception("FORMAT must be specified for function input()", ErrorCodes::INVALID_USAGE_OF_INPUT);
+            executeInsertQuery();
+        }
+        else
+            executeOrdinaryQuery();
+    }
+    /// Do not change context (current DB, settings) in case of an exception.
+    if (!have_error)
+    {
+        if (const auto * set_query = parsed_query->as<ASTSetQuery>())
+        {
+            /// Save all changes in settings to avoid losing them if the connection is lost.
+            for (const auto & change : set_query->changes)
+            {
+                if (change.name == "profile")
+                    current_profile = change.value.safeGet<String>();
+                else
+                    global_context->applySettingChange(change);
+            }
+        }
+        if (const auto * use_query = parsed_query->as<ASTUseQuery>())
+        {
+            const String & new_database = use_query->database;
+            /// If the client initiates the reconnection, it takes the settings from the config.
+            config().setString("database", new_database);
+            /// If the connection initiates the reconnection, it uses its variable.
+            connection->setDefaultDatabase(new_database);
+        }
+    }
+    if (is_interactive)
+    {
+        std::cout << std::endl << processed_rows << " rows in set. Elapsed: " << progress_indication.elapsedSeconds() << " sec. ";
+        /// Write final progress if it makes sense to do so.
+        writeFinalProgress();
+        std::cout << std::endl << std::endl;
+    }
+    else if (print_time_to_stderr)
+    {
+        std::cerr << progress_indication.elapsedSeconds() << "\n";
+    }
+}
+
+void Client::processTextAsSingleQuery(const String & text_)
+{
+    full_query = text_;
+    /// Some parts of a query (result output and formatting) are executed
+    /// client-side. Thus we need to parse the query.
+    const char * begin = full_query.data();
+    parsed_query = parseQuery(begin, begin + full_query.size(), false);
+    if (!parsed_query)
+        return;
+    // An INSERT query may have the data that follow query text. Remove the
+    /// Send part of query without data, because data will be sent separately.
+    auto * insert = parsed_query->as<ASTInsertQuery>();
+    if (insert && insert->data)
+    {
+        query_to_execute = full_query.substr(0, insert->data - full_query.data());
+    }
+    else
+    {
+        query_to_execute = full_query;
+    }
+    processParsedSingleQuery();
+    if (have_error)
+    {
+        reportQueryError();
+    }
+}
+
+bool Client::processMultiQuery(const String & all_queries_text)
+{
+    // It makes sense not to base any control flow on this, so that it is
+    // the same in tests and in normal usage. The only difference is that in
+    // normal mode we ignore the test hints.
+    const bool test_mode = config().has("testmode");
+    {
+        /// disable logs if expects errors
+        TestHint test_hint(test_mode, all_queries_text);
+        if (test_hint.clientError() || test_hint.serverError())
+            processTextAsSingleQuery("SET send_logs_level = 'fatal'");
+    }
+
+    auto process_single_query = [&](const String & query)
+    {
+        // Now we know for sure where the query ends.
+        // Look for the hint in the text of query + insert data + trailing
+        // comments,
+        // e.g. insert into t format CSV 'a' -- { serverError 123 }.
+        // Use the updated query boundaries we just calculated.
+        TestHint test_hint(test_mode, query);
+        // Echo all queries if asked; makes for a more readable reference
+        // file.
+        auto echo_query = test_hint.echoQueries().value_or(echo_queries);
+        try
+        {
+            processParsedSingleQuery(echo_query);
+        }
+        catch (...)
+        {
+            // Surprisingly, this is a client error. A server error would
+            // have been reported w/o throwing (see onReceiveSeverException()).
+            client_exception = std::make_unique<Exception>(getCurrentExceptionMessage(true), getCurrentExceptionCode());
+            std::cerr << "\n\nReceived exception, setting have_error.\n\n";
+            have_error = true;
+        }
+        // Check whether the error (or its absence) matches the test hints
+        // (or their absence).
+        bool error_matches_hint = true;
+        if (have_error)
+        {
+            if (test_hint.serverError())
+            {
+                if (!server_exception)
+                {
+                    error_matches_hint = false;
+                    fmt::print(stderr, "Expected server error code '{}' but got no server error.\n", test_hint.serverError());
+                }
+                else if (server_exception->code() != test_hint.serverError())
+                {
+                    error_matches_hint = false;
+                    std::cerr << "Expected server error code: " << test_hint.serverError() << " but got: " << server_exception->code()
+                                << "." << std::endl;
+                }
+            }
+            if (test_hint.clientError())
+            {
+                if (!client_exception)
+                {
+                    error_matches_hint = false;
+                    fmt::print(stderr, "Expected client error code '{}' but got no client error.\n", test_hint.clientError());
+                }
+                else if (client_exception->code() != test_hint.clientError())
+                {
+                    error_matches_hint = false;
+                    fmt::print(
+                        stderr, "Expected client error code '{}' but got '{}'.\n", test_hint.clientError(), client_exception->code());
+                }
+            }
+            if (!test_hint.clientError() && !test_hint.serverError())
+            {
+                // No error was expected but it still occurred. This is the
+                // default case w/o test hint, doesn't need additional
+                // diagnostics.
+                error_matches_hint = false;
+            }
+        }
+        else
+        {
+            if (test_hint.clientError())
+            {
+                fmt::print(stderr, "The query succeeded but the client error '{}' was expected.\n", test_hint.clientError());
+                error_matches_hint = false;
+            }
+            if (test_hint.serverError())
+            {
+                fmt::print(stderr, "The query succeeded but the server error '{}' was expected.\n", test_hint.serverError());
+                error_matches_hint = false;
+            }
+        }
+        // If the error is expected, force reconnect and ignore it.
+        if (have_error && error_matches_hint)
+        {
+            client_exception.reset();
+            server_exception.reset();
+            have_error = false;
+            if (!connection->checkConnected())
+                connect();
+        }
+    };
+
+    auto process_parse_query_error = [&](const String & query, Exception & e)
+    {
+        // Try to find test hint for syntax error. We don't know where
+        // the query ends because we failed to parse it, so we consume
+        // the entire line.
+        TestHint hint(test_mode, query);
+        if (hint.serverError())
+        {
+            // Syntax errors are considered as client errors
+            e.addMessage("\nExpected server error '{}'.", hint.serverError());
+            throw;
+        }
+        if (hint.clientError() != e.code())
+        {
+            if (hint.clientError())
+                e.addMessage("\nExpected client error: " + std::to_string(hint.clientError()));
+            throw;
+        }
+    };
+
+    return processMultiQueryImpl(all_queries_text, process_single_query, process_parse_query_error);
+}
+
+
 static bool queryHasWithClause(const IAST * ast)
 {
     if (const auto * select = dynamic_cast<const ASTSelectQuery *>(ast); select && select->with())

programs/client/Client.h

@@ -25,6 +25,11 @@ protected:
             connect();
     }
 
+    bool processMultiQuery(const String & all_queries_text) override;
+
+    void processParsedSingleQuery(std::optional<bool> echo_query = {});
+    void processTextAsSingleQuery(const String & text_);
+
     bool processMultiQueryFromFile(const String & file) override
     {
         connection->setDefaultDatabase(connection_parameters.default_database);
@@ -34,6 +39,8 @@ protected:
         return processMultiQuery(text);
     }
 
+    bool validateParsedOptions() override { return true; }
+
     std::vector<String> loadWarningMessages();
 
     void loadSuggestionDataIfPossible() override;

programs/local/LocalServer.h

@@ -63,6 +63,15 @@ protected:
         return getInitialCreateTableQuery();
     }
 
+    bool processMultiQuery(const String & all_queries_text) override
+    {
+        auto process_single_query = [&](const String & query)
+        {
+            prepareAndExecuteQuery(query);
+        };
+        return processMultiQueryImpl(all_queries_text, process_single_query);
+    }
+
     void executeParsedQueryImpl() override;
 
     void reportQueryError() const override {}

src/Client/ClientBase.cpp

@@ -210,12 +210,15 @@ void ClientBase::prepareAndExecuteQuery(const String & query)
 
 void ClientBase::executeParsedQuery(std::optional<bool> echo_query_, bool report_error)
 {
+    resetOutput();
+    client_exception.reset();
+    server_exception.reset();
+
     have_error = false;
     processed_rows = 0;
     written_first_block = false;
     progress_indication.resetProgress();
 
-    resetOutput();
     outputQueryInfo(echo_query_.value_or(echo_queries));
 
     executeParsedQueryImpl();
@@ -237,22 +240,15 @@ void ClientBase::executeParsedQuery(std::optional<bool> echo_query_, bool report
 }
 
 
-bool ClientBase::processMultiQuery(const String & all_queries_text)
+bool ClientBase::processMultiQueryImpl(const String & all_queries_text,
+                                       std::function<void(const String &)> process_single_query,
+                                       std::function<void(const String &, Exception &)> process_parse_query_error)
 {
-    // It makes sense not to base any control flow on this, so that it is
-    // the same in tests and in normal usage. The only difference is that in
-    // normal mode we ignore the test hints.
-    const bool test_mode = config().has("testmode");
-
-    {
-        /// disable logs if expects errors
-        TestHint test_hint(test_mode, all_queries_text);
-        if (test_hint.clientError() || test_hint.serverError())
-            prepareAndExecuteQuery("SET send_logs_level = 'fatal'");
-    }
-
-    bool echo_query = echo_queries;
 
+    /// Several queries separated by ';'.
+    /// INSERT data is ended by the end of line, not ';'.
+    /// An exception is VALUES format where we also support semicolon in
+    /// addition to end of line.
     const char * this_query_begin = all_queries_text.data();
     const char * all_queries_end = all_queries_text.data() + all_queries_text.size();
 
@@ -278,10 +274,11 @@ bool ClientBase::processMultiQuery(const String & all_queries_text)
             Tokens tokens(this_query_begin, all_queries_end);
             IParser::Pos token_iterator(tokens, global_context->getSettingsRef().max_parser_depth);
             if (!token_iterator.isValid())
+            {
                 break;
+            }
         }
 
-        // Try to parse the query.
         const char * this_query_end = this_query_begin;
         try
         {
@@ -289,27 +286,9 @@ bool ClientBase::processMultiQuery(const String & all_queries_text)
         }
         catch (Exception & e)
         {
-            // Try to find test hint for syntax error. We don't know where
-            // the query ends because we failed to parse it, so we consume
-            // the entire line.
             this_query_end = find_first_symbols<'\n'>(this_query_end, all_queries_end);
-            TestHint hint(test_mode, String(this_query_begin, this_query_end - this_query_begin));
-
-            if (hint.serverError())
-            {
-                // Syntax errors are considered as client errors
-                e.addMessage("\nExpected server error '{}'.", hint.serverError());
-                throw;
-            }
-
-            if (hint.clientError() != e.code())
-            {
-                if (hint.clientError())
-                    e.addMessage("\nExpected client error: " + std::to_string(hint.clientError()));
-                throw;
-            }
-
-            /// It's expected syntax error, skip the line.
+            process_parse_query_error(String(this_query_begin, this_query_end - this_query_begin), e);
+            /// It's expected syntax error, skip the line
             this_query_begin = this_query_end;
             continue;
         }
@@ -320,11 +299,9 @@ bool ClientBase::processMultiQuery(const String & all_queries_text)
             {
                 Tokens tokens(this_query_begin, all_queries_end);
                 IParser::Pos token_iterator(tokens, global_context->getSettingsRef().max_parser_depth);
-
                 while (token_iterator->type != TokenType::Semicolon && token_iterator.isValid())
                     ++token_iterator;
                 this_query_begin = token_iterator->end;
-
                 continue;
             }
 
@@ -365,86 +342,43 @@ bool ClientBase::processMultiQuery(const String & all_queries_text)
         // full_query is the query + inline INSERT data + trailing comments
         // (the latter is our best guess for now).
         full_query = all_queries_text.substr(this_query_begin - all_queries_text.data(), this_query_end - this_query_begin);
-
         if (query_fuzzer_runs)
         {
             if (!processWithFuzzing(full_query))
                 return false;
-
             this_query_begin = this_query_end;
             continue;
         }
 
-        // Now we know for sure where the query ends.
-        // Look for the hint in the text of query + insert data + trailing
-        // comments,
-        // e.g. insert into t format CSV 'a' -- { serverError 123 }.
-        // Use the updated query boundaries we just calculated.
-        TestHint test_hint(test_mode, std::string(this_query_begin, this_query_end - this_query_begin));
+        process_single_query(String(this_query_begin, this_query_end - this_query_begin));
 
-        // Echo all queries if asked; makes for a more readable reference
-        // file.
-        echo_query = test_hint.echoQueries().value_or(echo_query);
-
-        try
-        {
-            executeParsedQuery(echo_query, false);
-        }
-        catch (...)
-        {
-            // Surprisingly, this is a client error. A server error would
-            // have been reported w/o throwing (see onReceiveSeverException()).
-            client_exception = std::make_unique<Exception>(getCurrentExceptionMessage(true), getCurrentExceptionCode());
-            have_error = true;
-        }
-
-        // For INSERTs with inline data: use the end of inline data as reported by the format
-        // parser (it is saved in sendData()). This allows us to handle queries like:
-        // insert into t values (1); select 1, where the inline data is delimited by semicolon
-        // and not by a newline.
-        /// TODO: Better way
+        // For INSERTs with inline data: use the end of inline data as
+        // reported by the format parser (it is saved in sendData()).
+        // This allows us to handle queries like:
+        //   insert into t values (1); select 1
+        // , where the inline data is delimited by semicolon and not by a
+        // newline.
         if (insert_ast && insert_ast->data)
         {
             this_query_end = insert_ast->end;
             adjustQueryEnd(this_query_end, all_queries_end, global_context->getSettingsRef().max_parser_depth);
         }
 
-        // Check whether the error (or its absence) matches the test hints
-        // (or their absence).
-        bool error_matches_hint = checkErrorMatchesHints(test_hint, have_error);
-
-        // If the error is expected, force reconnect and ignore it.
-        if (have_error && error_matches_hint)
-        {
-            client_exception.reset();
-            server_exception.reset();
-            have_error = false;
-
-            reconnectIfNeeded();
-        }
-
         // Report error.
         if (have_error)
-        {
             reportQueryError();
-        }
 
         // Stop processing queries if needed.
         if (have_error && !ignore_error)
         {
             if (is_interactive)
-            {
                 break;
-            }
             else
-            {
                 return false;
-            }
         }
 
         this_query_begin = this_query_end;
     }
-
     return true;
 }
 
@@ -703,7 +637,7 @@ void ClientBase::init(int argc, char ** argv)
     /// Don't parse options with Poco library, we prefer neat boost::program_options.
     stopOptionsProcessing();
 
-    Arguments common_arguments{}; /// 0th argument is ignored.
+    Arguments common_arguments{""}; /// 0th argument is ignored.
     std::vector<Arguments> external_tables_arguments;
     readArguments(argc, argv, common_arguments, external_tables_arguments);
 
@@ -719,10 +653,9 @@ void ClientBase::init(int argc, char ** argv)
 
     /// Parse main commandline options.
     po::parsed_options parsed = po::command_line_parser(common_arguments).options(options_description.main_description.value()).run();
-
-    //auto unrecognized_options = po::collect_unrecognized(parsed.options, po::collect_unrecognized_mode::include_positional);
-    //if (!unrecognized_options.empty())
-    //    throw Exception(ErrorCodes::UNRECOGNIZED_ARGUMENTS, "Unrecognized option '{}'", unrecognized_options[0]);
+    auto unrecognized_options = po::collect_unrecognized(parsed.options, po::collect_unrecognized_mode::include_positional);
+    if (validateParsedOptions() && unrecognized_options.size() > 1)
+        throw Exception(ErrorCodes::UNRECOGNIZED_ARGUMENTS, "Unrecognized option '{}'", unrecognized_options[1]);
 
     po::variables_map options;
     po::store(parsed, options);

src/Client/ClientBase.h

@@ -62,6 +62,7 @@ protected:
 
     virtual void processConfig() = 0;
 
+    virtual bool validateParsedOptions() { return false; }
 
     void runInteractive();
 
@@ -73,7 +74,10 @@ protected:
      * - INSERT data is ended by the end of line, not ';'.
      * - An exception is VALUES format where we also support semicolon in addition to end of line.
     **/
-    bool processMultiQuery(const String & all_queries_text);
+    virtual bool processMultiQuery(const String & all_queries_text) = 0;
+    bool processMultiQueryImpl(const String & all_queries_text,
+                               std::function<void(const String &)> process_single_query,
+                               std::function<void(const String &, Exception &)> process_parse_query_error = {});
 
     /// Process single file (with queries) from non-interactive mode.
     virtual bool processMultiQueryFromFile(const String & file) = 0;