add read columns

src/Core/MySQLClient.cpp

@@ -19,7 +19,7 @@ bool MySQLClient::connect()
 {
     if (connected)
     {
-        close();
+        disconnect();
     }
 
     socket = std::make_unique<Poco::Net::StreamSocket>();
@@ -33,7 +33,7 @@ bool MySQLClient::connect()
     return handshake();
 }
 
-void MySQLClient::close()
+void MySQLClient::disconnect()
 {
     in = nullptr;
     out = nullptr;
@@ -54,6 +54,7 @@ bool MySQLClient::handshake()
             "Only support " + mysql_native_password + " auth plugin name, but got " + handshake.auth_plugin_name,
             ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
     }
+    server_capability_flags = handshake.capability_flags;
 
     Native41 native41(password, handshake.auth_plugin_data);
     String auth_plugin_data = native41.getAuthPluginData();
@@ -62,21 +63,67 @@ bool MySQLClient::handshake()
         client_capability_flags, max_packet_size, charset_utf8, user, database, auth_plugin_data, mysql_native_password);
     packet_sender->sendPacket<HandshakeResponse>(handshake_response, true);
 
-    PacketResponse packet_response(handshake.capability_flags);
+    PacketResponse packet_response(server_capability_flags);
     packet_sender->receivePacket(packet_response);
-    if (packet_response.getType() != PACKET_ERR)
+    packet_sender->resetSequenceId();
-    {
+    if (packet_response.getType() == PACKET_ERR)
-        return true;
-    }
-    else
     {
         last_error = packet_response.err.error_message;
-        return false;
     }
+    return (packet_response.getType() != PACKET_ERR);
+}
+
+bool MySQLClient::ping()
+{
+    return writeCommand(Command::COM_PING, "");
+}
+
+bool MySQLClient::initdb(String db)
+{
+    return writeCommand(Command::COM_INIT_DB, db);
+}
+
+bool MySQLClient::query(String q)
+{
+    return writeCommand(Command::COM_QUERY, q);
 }
 
 String MySQLClient::error()
 {
     return last_error;
 }
+
+bool MySQLClient::writeCommand(char command, String query)
+{
+    bool ret = false;
+
+    WriteCommand write_command(command, query);
+    packet_sender->sendPacket<WriteCommand>(write_command, true);
+
+    PacketResponse packet_response(server_capability_flags);
+    packet_sender->receivePacket(packet_response);
+    switch (packet_response.getType())
+    {
+        case PACKET_ERR:
+            last_error = packet_response.err.error_message;
+            break;
+        case PACKET_OK:
+            ret = readColumns(packet_response.column_length);
+            break;
+        default:
+            break;
+    }
+    packet_sender->resetSequenceId();
+    return ret;
+}
+
+bool MySQLClient::readColumns(int column_length)
+{
+    for (auto i = 0; i < column_length; i++)
+    {
+        ColumnDefinition cd;
+        packet_sender->receivePacket(cd);
+    }
+    return true;
+}
 }

src/Core/MySQLClient.h

@@ -21,7 +21,10 @@ class MySQLClient
 public:
     MySQLClient(const String & host_, UInt16 port_, const String & user_, const String & password_, const String & database_);
     bool connect();
-    void close();
+    void disconnect();
+    bool ping();
+    bool initdb(String db);
+    bool query(String q);
     String error();
 
 private:
@@ -32,8 +35,9 @@ private:
     String database;
 
     bool connected = false;
-    UInt32 client_capability_flags = 0;
     String last_error;
+    UInt32 client_capability_flags = 0;
+    UInt32 server_capability_flags = 0;
 
     uint8_t seq = 0;
     UInt8 charset_utf8 = 33;
@@ -47,5 +51,27 @@ private:
     std::shared_ptr<PacketSender> packet_sender;
 
     bool handshake();
+    bool readColumns(int column_length);
+    bool writeCommand(char command, String query);
+};
+
+class WriteCommand : public WritePacket
+{
+public:
+    char command;
+    String query;
+
+    WriteCommand(char command_, String query_) : command(command_), query(query_) { }
+
+    size_t getPayloadSize() const override { return 1 + query.size(); }
+
+    void writePayloadImpl(WriteBuffer & buffer) const override
+    {
+        buffer.write(static_cast<char>(command));
+        if (!query.empty())
+        {
+            buffer.write(query.data(), query.size());
+        }
+    }
 };
 }

src/Core/MySQLProtocol.h

@@ -141,6 +141,7 @@ enum ResponsePacketType
     PACKET_OK = 0x00,
     PACKET_ERR = 0xff,
     PACKET_EOF = 0xfe,
+    PACKET_LOCALINFILE = 0xfb,
 };
 
 // https://dev.mysql.com/doc/dev/mysql-server/latest/group__group__cs__column__definition__flags.html
@@ -422,6 +423,13 @@ public:
 
 uint64_t readLengthEncodedNumber(ReadBuffer & ss);
 
+inline void readLengthEncodedString(String & s, ReadBuffer & buffer)
+{
+    uint64_t len = readLengthEncodedNumber(buffer);
+    s.resize(len);
+    buffer.readStrict(reinterpret_cast<char *>(s.data()), len);
+}
+
 void writeLengthEncodedNumber(uint64_t x, WriteBuffer & buffer);
 
 inline void writeLengthEncodedString(const String & s, WriteBuffer & buffer)
@@ -453,7 +461,7 @@ public:
     String auth_plugin_name;
     String auth_plugin_data;
 
-    Handshake() = default;
+    Handshake() : connection_id(0x00), capability_flags(0x00), character_set(0x00), status_flags(0x00) { }
 
     Handshake(uint32_t capability_flags_, uint32_t connection_id_, String server_version_, String auth_plugin_name_, String auth_plugin_data_)
         : protocol_version(0xa)
@@ -489,43 +497,43 @@ public:
         writeChar(0x0, 1, buffer);
     }
 
-    void readPayloadImpl(ReadBuffer & buffer) override
+    void readPayloadImpl(ReadBuffer & payload) override
     {
-        buffer.readStrict(reinterpret_cast<char *>(&protocol_version), 1);
+        payload.readStrict(reinterpret_cast<char *>(&protocol_version), 1);
-        readNullTerminated(server_version, buffer);
+        readNullTerminated(server_version, payload);
-        buffer.readStrict(reinterpret_cast<char *>(&connection_id), 4);
+        payload.readStrict(reinterpret_cast<char *>(&connection_id), 4);
 
         auth_plugin_data.resize(AUTH_PLUGIN_DATA_PART_1_LENGTH);
-        buffer.readStrict(auth_plugin_data.data(), AUTH_PLUGIN_DATA_PART_1_LENGTH);
+        payload.readStrict(auth_plugin_data.data(), AUTH_PLUGIN_DATA_PART_1_LENGTH);
 
-        buffer.ignore(1);
+        payload.ignore(1);
-        buffer.readStrict(reinterpret_cast<char *>(&capability_flags), 2);
+        payload.readStrict(reinterpret_cast<char *>(&capability_flags), 2);
-        buffer.readStrict(reinterpret_cast<char *>(&character_set), 1);
+        payload.readStrict(reinterpret_cast<char *>(&character_set), 1);
-        buffer.readStrict(reinterpret_cast<char *>(&status_flags), 2);
+        payload.readStrict(reinterpret_cast<char *>(&status_flags), 2);
-        buffer.readStrict((reinterpret_cast<char *>(&capability_flags)) + 2, 2);
+        payload.readStrict((reinterpret_cast<char *>(&capability_flags)) + 2, 2);
 
         UInt8 auth_plugin_data_length = 0;
         if (capability_flags & MySQLProtocol::CLIENT_PLUGIN_AUTH)
         {
-            buffer.readStrict(reinterpret_cast<char *>(&auth_plugin_data_length), 1);
+            payload.readStrict(reinterpret_cast<char *>(&auth_plugin_data_length), 1);
         }
         else
         {
-            buffer.ignore(1);
+            payload.ignore(1);
         }
 
-        buffer.ignore(10);
+        payload.ignore(10);
         if (capability_flags & MySQLProtocol::CLIENT_SECURE_CONNECTION)
         {
             UInt8 part2_length = (SCRAMBLE_LENGTH - AUTH_PLUGIN_DATA_PART_1_LENGTH);
             auth_plugin_data.resize(SCRAMBLE_LENGTH);
-            buffer.readStrict(auth_plugin_data.data() + AUTH_PLUGIN_DATA_PART_1_LENGTH, part2_length);
+            payload.readStrict(auth_plugin_data.data() + AUTH_PLUGIN_DATA_PART_1_LENGTH, part2_length);
-            buffer.ignore(1);
+            payload.ignore(1);
         }
 
         if (capability_flags & MySQLProtocol::CLIENT_PLUGIN_AUTH)
         {
-            readNullTerminated(auth_plugin_name, buffer);
+            readNullTerminated(auth_plugin_name, payload);
         }
     }
 };
@@ -649,9 +657,7 @@ public:
 
         if (capability_flags & CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA)
         {
-            auto len = readLengthEncodedNumber(payload);
+            readLengthEncodedString(auth_response, payload);
-            auth_response.resize(len);
-            payload.readStrict(auth_response.data(), len);
         }
         else if (capability_flags & CLIENT_SECURE_CONNECTION)
         {
@@ -837,14 +843,10 @@ public:
 
         if (capabilities & CLIENT_SESSION_TRACK)
         {
-            auto len = readLengthEncodedNumber(payload);
+            readLengthEncodedString(info, payload);
-            info.resize(len);
-            payload.readStrict(info.data(), len);
             if (status_flags & SERVER_SESSION_STATE_CHANGED)
             {
-                len = readLengthEncodedNumber(payload);
+                readLengthEncodedString(session_state_changes, payload);
-                session_state_changes.resize(len);
-                payload.readStrict(session_state_changes.data(), len);
             }
         }
         else
@@ -920,12 +922,16 @@ public:
         assert(header == 0xff);
 
         payload.readStrict(reinterpret_cast<char *>(&error_code), 2);
-        payload.ignore(1);
 
+        /// SQL State [optional: # + 5bytes string]
+        UInt8 sharp = static_cast<unsigned char>(*payload.position());
+        if (sharp == 0x23)
+        {
+            payload.ignore(1);
             sql_state.resize(5);
             payload.readStrict(reinterpret_cast<char *>(sql_state.data()), 5);
-
+        }
-        readNullTerminated(error_message, payload);
+        readString(error_message, payload);
     }
 };
 
@@ -936,12 +942,13 @@ public:
     OK_Packet ok;
     ERR_Packet err;
     EOF_Packet eof;
+    UInt64 column_length = 0;
 
     PacketResponse(UInt32 server_capability_flags_) : ok(OK_Packet(server_capability_flags_)) { }
 
     void readPayloadImpl(ReadBuffer & payload) override
     {
-        UInt8 header = static_cast<unsigned char>(*payload.position());
+        UInt16 header = static_cast<unsigned char>(*payload.position());
         switch (header)
         {
             case PACKET_OK:
@@ -956,6 +963,11 @@ public:
                 packetType = PACKET_EOF;
                 eof.readPayloadImpl(payload);
                 break;
+            case PACKET_LOCALINFILE:
+                packetType = PACKET_LOCALINFILE;
+                break;
+            default:
+                column_length = readLengthEncodedNumber(payload);
         }
     }
 
@@ -965,8 +977,9 @@ private:
     ResponsePacketType packetType = PACKET_OK;
 };
 
-class ColumnDefinition : public WritePacket
+class ColumnDefinition : public WritePacket, public ReadPacket
 {
+public:
     String schema;
     String table;
     String org_table;
@@ -978,7 +991,9 @@ class ColumnDefinition : public WritePacket
     ColumnType column_type;
     uint16_t flags;
     uint8_t decimals = 0x00;
-public:
+
+    ColumnDefinition() : character_set(0x00), column_length(0), column_type(MYSQL_TYPE_DECIMAL), flags(0x00) { }
+
     ColumnDefinition(
         String schema_,
         String table_,
@@ -1009,7 +1024,6 @@ public:
     {
     }
 
-protected:
     size_t getPayloadSize() const override
     {
         return 13 + getLengthEncodedStringSize("def") + getLengthEncodedStringSize(schema) + getLengthEncodedStringSize(table) + getLengthEncodedStringSize(org_table) + \
@@ -1032,6 +1046,25 @@ protected:
         buffer.write(reinterpret_cast<const char *>(&decimals), 2);
         writeChar(0x0, 2, buffer);
     }
+
+    void readPayloadImpl(ReadBuffer & payload) override
+    {
+        String def;
+        readLengthEncodedString(def, payload);
+        assert(def == "def");
+        readLengthEncodedString(schema, payload);
+        readLengthEncodedString(table, payload);
+        readLengthEncodedString(org_table, payload);
+        readLengthEncodedString(name, payload);
+        readLengthEncodedString(org_name, payload);
+        next_length = readLengthEncodedNumber(payload);
+        payload.readStrict(reinterpret_cast<char *>(&character_set), 2);
+        payload.readStrict(reinterpret_cast<char *>(&column_length), 4);
+        payload.readStrict(reinterpret_cast<char *>(&column_type), 1);
+        payload.readStrict(reinterpret_cast<char *>(&flags), 2);
+        payload.readStrict(reinterpret_cast<char *>(&decimals), 2);
+        payload.ignore(2);
+    }
 };
 
 class ComFieldList : public LimitedReadPacket

src/Core/tests/mysql_protocol.cpp

@@ -1,5 +1,6 @@
 #include <string>
 
+#include <Core/MySQLClient.h>
 #include <Core/MySQLProtocol.h>
 #include <IO/ReadBufferFromString.h>
 #include <IO/WriteBufferFromString.h>
@@ -11,15 +12,7 @@ int main(int, char **)
     using namespace MySQLProtocol;
     using namespace MySQLProtocol::Authentication;
 
-    /*
+
-    UInt16 port = 9001;
-    String host = "127.0.0.1", user = "default", password = "123";
-    MySQLClient client(host, port, user, password, "");
-    if (!client.connect())
-    {
-        std::cerr << "Connect Error: " << client.error() << std::endl;
-    }
-     */
     String user = "default";
     String password = "123";
     String database = "";
@@ -90,9 +83,9 @@ int main(int, char **)
         server.writePayloadImpl(out0);
 
         // 2. Client reads packet
-        ReadBufferFromString in1(s0);
+        ReadBufferFromString in0(s0);
         PacketResponse client(server_capability_flags);
-        client.readPayloadImpl(in1);
+        client.readPayloadImpl(in0);
 
         // Check
         ASSERT(client.getType() == PACKET_OK);
@@ -110,16 +103,16 @@ int main(int, char **)
         server.writePayloadImpl(out0);
 
         // 2. Client reads packet
-        ReadBufferFromString in1(s0);
+        ReadBufferFromString in0(s0);
         PacketResponse client(server_capability_flags);
-        client.readPayloadImpl(in1);
+        client.readPayloadImpl(in0);
 
         // Check
         ASSERT(client.getType() == PACKET_ERR);
         ASSERT(client.err.header == server.header);
         ASSERT(client.err.error_code == server.error_code);
         ASSERT(client.err.sql_state == server.sql_state);
-        ASSERT(client.err.error_message == server.error_message);
+        ASSERT(client.err.error_message.data() == server.error_message);
     }
 
     /// EOF Packet
@@ -131,9 +124,9 @@ int main(int, char **)
         server.writePayloadImpl(out0);
 
         // 2. Client reads packet
-        ReadBufferFromString in1(s0);
+        ReadBufferFromString in0(s0);
         PacketResponse client(server_capability_flags);
-        client.readPayloadImpl(in1);
+        client.readPayloadImpl(in0);
 
         // Check
         ASSERT(client.getType() == PACKET_EOF);
@@ -142,5 +135,57 @@ int main(int, char **)
         ASSERT(client.eof.status_flags == server.status_flags);
     }
 
+    /// ColumnDefinition Packet
+    {
+        // 1. Server writes packet
+        std::string s0;
+        WriteBufferFromString out0(s0);
+        ColumnDefinition server("schema", "tbl", "org_tbl", "name", "org_name", 33, 0x00, MYSQL_TYPE_STRING, 0x00, 0x00);
+        server.writePayloadImpl(out0);
+
+        // 2. Client reads packet
+        ReadBufferFromString in0(s0);
+        ColumnDefinition client;
+        client.readPayloadImpl(in0);
+
+        // Check
+        ASSERT(client.column_type == server.column_type);
+        ASSERT(client.column_length == server.column_length);
+        ASSERT(client.next_length == server.next_length);
+        ASSERT(client.character_set == server.character_set);
+        ASSERT(client.decimals == server.decimals);
+        ASSERT(client.name == server.name);
+        ASSERT(client.org_name == server.org_name);
+        ASSERT(client.table == server.table);
+        ASSERT(client.org_table == server.org_table);
+        ASSERT(client.schema == server.schema);
+    }
+
+    {
+        MySQLClient client1("127.0.0.1", 9001, "default", "123", "");
+        if (!client1.connect())
+        {
+            std::cerr << "Connect Error: " << client1.error() << std::endl;
+            return 1;
+        }
+
+        if (!client1.ping())
+        {
+            std::cerr << "Connect Error: " << client1.error() << std::endl;
+            return 1;
+        }
+
+        if (!client1.initdb("default"))
+        {
+            std::cerr << "Connect Error: " << client1.error() << std::endl;
+            return 1;
+        }
+
+        if (!client1.query("select 1"))
+        {
+            std::cerr << "Connect Error: " << client1.error() << std::endl;
+            return 1;
+        }
+    }
     return 0;
 }