From f6788fc6600ba6c23ede027050487ac81ef3659b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Mar=C3=ADn?= <git@rmr.ninja>
Date: Thu, 12 Aug 2021 11:29:50 +0200
Subject: [PATCH] Mysql handler: Move format check to the handler

---
 src/Formats/FormatFactory.cpp                         | 11 +++--------
 src/Server/MySQLHandler.cpp                           |  5 ++++-
 .../0_stateless/01176_mysql_client_interactive.expect |  8 ++++++++
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/Formats/FormatFactory.cpp b/src/Formats/FormatFactory.cpp
index cd55c66fbdf..9b701816a2c 100644
--- a/src/Formats/FormatFactory.cpp
+++ b/src/Formats/FormatFactory.cpp
@@ -5,7 +5,6 @@
 #include <Interpreters/Context.h>
 #include <Core/Settings.h>
 #include <DataStreams/MaterializingBlockOutputStream.h>
-#include <DataStreams/NativeBlockInputStream.h>
 #include <Formats/FormatSettings.h>
 #include <Processors/Formats/IRowInputFormat.h>
 #include <Processors/Formats/IRowOutputFormat.h>
@@ -214,8 +213,8 @@ BlockOutputStreamPtr FormatFactory::getOutputStreamParallelIfPossible(
     bool parallel_formatting = settings.output_format_parallel_formatting;
     auto format_settings = _format_settings ? *_format_settings : getFormatSettings(context);
 
-    if (output_getter && parallel_formatting && getCreators(name).supports_parallel_formatting && !settings.output_format_json_array_of_rows
-        && !format_settings.mysql_wire.sequence_id)
+    if (output_getter && parallel_formatting && getCreators(name).supports_parallel_formatting
+        && !settings.output_format_json_array_of_rows)
     {
         auto formatter_creator = [output_getter, sample, callback, format_settings]
             (WriteBuffer & output) -> OutputFormatPtr
@@ -315,7 +314,7 @@ OutputFormatPtr FormatFactory::getOutputFormatParallelIfPossible(
     const Settings & settings = context->getSettingsRef();
 
     if (settings.output_format_parallel_formatting && getCreators(name).supports_parallel_formatting
-        && !settings.output_format_json_array_of_rows && !format_settings.mysql_wire.sequence_id)
+        && !settings.output_format_json_array_of_rows)
     {
         auto formatter_creator = [output_getter, sample, callback, format_settings]
         (WriteBuffer & output) -> OutputFormatPtr
@@ -353,10 +352,6 @@ OutputFormatPtr FormatFactory::getOutputFormat(
 
     auto format_settings = _format_settings ? *_format_settings : getFormatSettings(context);
 
-    /// If we're handling MySQL protocol connection right now then MySQLWire is only allowed output format.
-    if (format_settings.mysql_wire.sequence_id && (name != "MySQLWire"))
-        throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "MySQL protocol does not support custom output formats");
-
     /** TODO: Materialization is needed, because formats can use the functions `IDataType`,
       *  which only work with full columns.
       */
diff --git a/src/Server/MySQLHandler.cpp b/src/Server/MySQLHandler.cpp
index 375f248d939..b05f96bbfe1 100644
--- a/src/Server/MySQLHandler.cpp
+++ b/src/Server/MySQLHandler.cpp
@@ -57,6 +57,7 @@ namespace ErrorCodes
     extern const int NOT_IMPLEMENTED;
     extern const int MYSQL_CLIENT_INSUFFICIENT_CAPABILITIES;
     extern const int SUPPORT_IS_DISABLED;
+    extern const int UNSUPPORTED_METHOD;
 }
 
 
@@ -352,8 +353,10 @@ void MySQLHandler::comQuery(ReadBuffer & payload)
         format_settings.mysql_wire.max_packet_size = max_packet_size;
         format_settings.mysql_wire.sequence_id = &sequence_id;
 
-        auto set_result_details = [&with_output](const String &, const String &, const String &, const String &)
+        auto set_result_details = [&with_output](const String &, const String &, const String &format, const String &)
         {
+            if (format != "MySQLWire")
+                throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "MySQL protocol does not support custom output formats");
             with_output = true;
         };
 
diff --git a/tests/queries/0_stateless/01176_mysql_client_interactive.expect b/tests/queries/0_stateless/01176_mysql_client_interactive.expect
index 6f0cea48e76..8f56d047e8e 100755
--- a/tests/queries/0_stateless/01176_mysql_client_interactive.expect
+++ b/tests/queries/0_stateless/01176_mysql_client_interactive.expect
@@ -34,6 +34,14 @@ expect "ERROR 395 (00000): Code: 395"
 send -- "select * from system.one format TSV;\r"
 expect "ERROR 1 (00000): Code: 1"
 
+send -- "select * from system.one format JSON;\r"
+expect "ERROR 1 (00000): Code: 1"
+
+send -- "select * from system.one format MySQLWire;\r"
+expect "| dummy |"
+expect "|     0 |"
+expect "1 row in set"
+
 send -- "select count(number), sum(number) from numbers(10);\r"
 expect "+---------------+-------------+"
 expect "| count(number) | sum(number) |"