diff --git a/client/mysqltest.cc b/client/mysqltest.cc index 23453f58820..8d48e7bed48 100644 --- a/client/mysqltest.cc +++ b/client/mysqltest.cc @@ -483,6 +483,9 @@ enum enum_commands { Q_SEND_EVAL, Q_OUTPUT, /* redirect output to a file */ Q_RESET_CONNECTION, + Q_QUERY_ATTRS_ADD, + Q_QUERY_ATTRS_DELETE, + Q_QUERY_ATTRS_RESET, Q_UNKNOWN, /* Unknown command. */ Q_COMMENT, /* Comments, ignored. */ Q_COMMENT_WITH_COMMAND, @@ -518,7 +521,8 @@ const char *command_names[] = { "list_files", "list_files_write_file", "list_files_append_file", "send_shutdown", "shutdown_server", "result_format", "move_file", "remove_files_wildcard", "copy_files_wildcard", "send_eval", "output", - "reset_connection", + "reset_connection", "query_attrs_add", "query_attrs_delete", + "query_attrs_reset", 0}; @@ -6566,6 +6570,45 @@ static void do_reset_connection() { DBUG_VOID_RETURN; } +void do_query_attrs_add(struct st_command *command) { + int error; + static DYNAMIC_STRING key; + static DYNAMIC_STRING value; + const struct command_arg query_attrs_args[] = { + {"key", ARG_STRING, true, &key, "Key for query attributes"}, + {"value", ARG_STRING, true, &value, "Value for query attributes"}}; + DBUG_ENTER("do_query_attrs_add"); + + check_command_args(command, command->first_argument, query_attrs_args, + sizeof(query_attrs_args) / sizeof(struct command_arg), + ' '); + + error = mysql_options4(&cur_con->mysql, MYSQL_OPT_QUERY_ATTR_ADD, key.str, + value.str); + handle_command_error(command, error); + dynstr_free(&key); + dynstr_free(&value); + DBUG_VOID_RETURN; +} + +void do_query_attrs_delete(struct st_command *command) { + int error; + static DYNAMIC_STRING key; + const struct command_arg query_attrs_args[] = { + {"key", ARG_STRING, true, &key, "Key for query attributes"}, + }; + DBUG_ENTER("do_query_attrs_delete"); + + check_command_args(command, command->first_argument, query_attrs_args, + sizeof(query_attrs_args) / sizeof(struct command_arg), + ' '); + + error = mysql_options(&cur_con->mysql, MYSQL_OPT_QUERY_ATTR_DELETE, key.str); + handle_command_error(command, error); + dynstr_free(&key); + DBUG_VOID_RETURN; +} + bool match_delimiter(int c, const char *delim, size_t length) { uint i; char tmp[MAX_DELIMITER_LENGTH]; @@ -9579,6 +9622,15 @@ int main(int argc, char **argv) { dynstr_free(&ds_to_file); break; } + case Q_QUERY_ATTRS_ADD: + do_query_attrs_add(command); + break; + case Q_QUERY_ATTRS_DELETE: + do_query_attrs_delete(command); + break; + case Q_QUERY_ATTRS_RESET: + mysql_options(&cur_con->mysql, MYSQL_OPT_QUERY_ATTR_RESET, 0); + break; default: processed = 0; diff --git a/include/my_command.h b/include/my_command.h index 00e6c805428..198e93b07c3 100644 --- a/include/my_command.h +++ b/include/my_command.h @@ -95,7 +95,12 @@ enum enum_server_command { /* don't forget to update const char *command_name[] in sql_parse.cc */ /* Must be last */ - COM_END /**< Not a real command. Refused. */ + COM_END, /**< Not a real command. Refused. */ + /* + The following are Facebook specific commands. They are put at the top end + to avoid conflicting with upstream. + */ + COM_QUERY_ATTRS = 255, }; #endif /* _mysql_command_h */ diff --git a/include/mysql.h b/include/mysql.h index 726ede2903d..29f1bdd6329 100644 --- a/include/mysql.h +++ b/include/mysql.h @@ -198,6 +198,9 @@ enum mysql_option { MYSQL_OPT_CONNECT_ATTR_RESET, MYSQL_OPT_CONNECT_ATTR_ADD, MYSQL_OPT_CONNECT_ATTR_DELETE, + MYSQL_OPT_QUERY_ATTR_RESET, + MYSQL_OPT_QUERY_ATTR_ADD, + MYSQL_OPT_QUERY_ATTR_DELETE, MYSQL_SERVER_PUBLIC_KEY, MYSQL_ENABLE_CLEARTEXT_PLUGIN, MYSQL_OPT_CAN_HANDLE_EXPIRED_PASSWORDS, diff --git a/include/mysql.h.pp b/include/mysql.h.pp index 7d1262669ac..4b784edbaf0 100644 --- a/include/mysql.h.pp +++ b/include/mysql.h.pp @@ -82,7 +82,8 @@ enum enum_server_command { COM_DAEMON, COM_BINLOG_DUMP_GTID, COM_RESET_CONNECTION, - COM_END + COM_END, + COM_QUERY_ATTRS = 255, }; enum SERVER_STATUS_flags_enum { SERVER_STATUS_IN_TRANS = 1, @@ -447,6 +448,9 @@ enum mysql_option { MYSQL_OPT_CONNECT_ATTR_RESET, MYSQL_OPT_CONNECT_ATTR_ADD, MYSQL_OPT_CONNECT_ATTR_DELETE, + MYSQL_OPT_QUERY_ATTR_RESET, + MYSQL_OPT_QUERY_ATTR_ADD, + MYSQL_OPT_QUERY_ATTR_DELETE, MYSQL_SERVER_PUBLIC_KEY, MYSQL_ENABLE_CLEARTEXT_PLUGIN, MYSQL_OPT_CAN_HANDLE_EXPIRED_PASSWORDS, diff --git a/include/mysql/com_data.h b/include/mysql/com_data.h index e64b6eb0baf..75983e15ca6 100644 --- a/include/mysql/com_data.h +++ b/include/mysql/com_data.h @@ -94,6 +94,8 @@ struct COM_STMT_RESET_DATA { struct COM_QUERY_DATA { const char *query; unsigned int length; + const char *query_attrs; + unsigned int query_attrs_length; }; struct COM_FIELD_LIST_DATA { diff --git a/include/mysql/plugin_audit.h.pp b/include/mysql/plugin_audit.h.pp index 1e1987e4d05..b5e44a7dda3 100644 --- a/include/mysql/plugin_audit.h.pp +++ b/include/mysql/plugin_audit.h.pp @@ -167,7 +167,8 @@ enum enum_server_command { COM_DAEMON, COM_BINLOG_DUMP_GTID, COM_RESET_CONNECTION, - COM_END + COM_END, + COM_QUERY_ATTRS = 255, }; #include "my_sqlcommand.h" enum enum_sql_command { diff --git a/include/mysql/services.h.pp b/include/mysql/services.h.pp index 56dd8d02865..cbfd22fa405 100644 --- a/include/mysql/services.h.pp +++ b/include/mysql/services.h.pp @@ -50,6 +50,8 @@ struct COM_STMT_RESET_DATA { struct COM_QUERY_DATA { const char *query; unsigned int length; + const char *query_attrs; + unsigned int query_attrs_length; }; struct COM_FIELD_LIST_DATA { unsigned char *table_name; diff --git a/include/sql_common.h b/include/sql_common.h index b01a2250025..8cb496fff91 100644 --- a/include/sql_common.h +++ b/include/sql_common.h @@ -117,6 +117,8 @@ struct st_mysql_options_extention { struct My_hash *connection_attributes; char *server_public_key_path; size_t connection_attributes_length; + struct My_hash *query_attributes; + size_t query_attributes_length; bool enable_cleartext_plugin; bool get_server_public_key; /* Former ssl_enforce */ char *tls_version; /* TLS version option */ diff --git a/mysql-test/r/query_attrs.result b/mysql-test/r/query_attrs.result new file mode 100644 index 00000000000..2d3c3f8a9fc --- /dev/null +++ b/mysql-test/r/query_attrs.result @@ -0,0 +1,18 @@ +create database db1; +use db1; +select 1; +1 +1 +select 1; +1 +1 +select 1; +1 +1 +select 1; +1 +1 +select 1; +1 +1 +drop database db1; diff --git a/mysql-test/t/query_attrs.test b/mysql-test/t/query_attrs.test new file mode 100644 index 00000000000..fa5fddcf102 --- /dev/null +++ b/mysql-test/t/query_attrs.test @@ -0,0 +1,25 @@ +create database db1; +use db1; + +connect (con,localhost,root,,db1); + +query_attrs_add a b; +select 1; + +query_attrs_add b a; +select 1; + +connection default; +select 1; + +connection con; +query_attrs_delete a; +select 1; + +query_attrs_reset; +select 1; + +connection default; +disconnect con; + +drop database db1; diff --git a/sql-common/client.cc b/sql-common/client.cc index 6e75b77cdf4..cf53c1c38a7 100644 --- a/sql-common/client.cc +++ b/sql-common/client.cc @@ -1362,7 +1362,7 @@ net_async_status cli_advanced_command_nonblocking( can happen if a client sends a query but does not reap the result before attempting to close the connection. */ - DBUG_ASSERT(command <= COM_END); + DBUG_ASSERT(command <= COM_END || command == COM_QUERY_ATTRS); net_clear(&mysql->net, (command != COM_QUIT)); mysql->async_send_command_status = NET_ASYNC_SEND_COMMAND_WRITE_COMMAND; } @@ -3757,35 +3757,49 @@ struct My_hash { malloc_unordered_map hash{key_memory_mysql_options}; }; +static uchar *encode_attrs(uchar *buf, struct My_hash *attrs, + size_t attrs_length) { + buf = net_store_length(buf, attrs_length); + + /* check if we have attributes */ + if (attrs) { + /* loop over and dump the connection attributes */ + for (const auto &key_and_value : attrs->hash) { + const string &key = key_and_value.first; + const string &value = key_and_value.second; + + /* we can't have zero length keys */ + DBUG_ASSERT(!key.empty()); + + buf = write_length_encoded_string3(buf, key.data(), key.size()); + buf = write_length_encoded_string3(buf, value.data(), value.size()); + } + } + return buf; +} + uchar *send_client_connect_attrs(MYSQL *mysql, uchar *buf) { /* check if the server supports connection attributes */ if (mysql->server_capabilities & CLIENT_CONNECT_ATTRS) { - /* Always store the length if the client supports it */ - buf = net_store_length( - buf, mysql->options.extension - ? mysql->options.extension->connection_attributes_length - : 0); - - /* check if we have connection attributes */ - if (mysql->options.extension && - mysql->options.extension->connection_attributes) { - /* loop over and dump the connection attributes */ - for (const auto &key_and_value : - mysql->options.extension->connection_attributes->hash) { - const string &key = key_and_value.first; - const string &value = key_and_value.second; - - /* we can't have zero length keys */ - DBUG_ASSERT(!key.empty()); - - buf = write_length_encoded_string3(buf, key.data(), key.size()); - buf = write_length_encoded_string3(buf, value.data(), value.size()); - } - } + if (mysql->options.extension) + buf = + encode_attrs(buf, mysql->options.extension->connection_attributes, + mysql->options.extension->connection_attributes_length); + else + buf = encode_attrs(buf, nullptr, 0); } return buf; } +uchar *send_client_query_attrs(MYSQL *mysql, uchar *buf) { + if (mysql->options.extension) + buf = encode_attrs(buf, mysql->options.extension->query_attributes, + mysql->options.extension->query_attributes_length); + else + buf = encode_attrs(buf, nullptr, 0); + return buf; +} + static size_t get_length_store_length(size_t length) { /* as defined in net_store_length */ #define MAX_VARIABLE_STRING_LENGTH 9 @@ -6759,6 +6773,7 @@ void mysql_close_free_options(MYSQL *mysql) { my_free(mysql->options.extension->default_auth); my_free(mysql->options.extension->server_public_key_path); delete mysql->options.extension->connection_attributes; + delete mysql->options.extension->query_attributes; my_free(mysql->options.extension); } memset(&mysql->options, 0, sizeof(mysql->options)); @@ -7100,7 +7115,26 @@ int STDCALL mysql_send_query(MYSQL *mysql, const char *query, ulong length) { if ((info = STATE_DATA(mysql))) free_state_change_info(static_cast(mysql->extension)); - DBUG_RETURN(simple_command(mysql, COM_QUERY, (uchar *)query, length, 1)); + size_t query_attrs_len = + mysql->options.extension + ? mysql->options.extension->query_attributes_length + : 0; + if (query_attrs_len > 0) { + bool ret; + uchar *buf = (uchar *)my_malloc(PSI_NOT_INSTRUMENTED, query_attrs_len + 9, + MYF(MY_WME | MY_ZEROFILL)); + + uchar *end = send_client_query_attrs(mysql, buf); + query_attrs_len = end - buf; + + ret = (*mysql->methods->advanced_command)(mysql, COM_QUERY_ATTRS, buf, + query_attrs_len, (uchar *)query, + length, 1, NULL); + my_free(buf); + DBUG_RETURN(ret); + } else { + DBUG_RETURN(simple_command(mysql, COM_QUERY, (uchar *)query, length, 1)); + } } net_async_status STDCALL mysql_send_query_nonblocking(MYSQL *mysql, @@ -7692,6 +7726,33 @@ int STDCALL mysql_options(MYSQL *mysql, enum mysql_option option, } } break; + case MYSQL_OPT_QUERY_ATTR_RESET: + ENSURE_EXTENSIONS_PRESENT(&mysql->options); + if (mysql->options.extension->query_attributes) { + delete mysql->options.extension->query_attributes; + mysql->options.extension->query_attributes = NULL; + mysql->options.extension->query_attributes_length = 0; + } + break; + case MYSQL_OPT_QUERY_ATTR_DELETE: + ENSURE_EXTENSIONS_PRESENT(&mysql->options); + if (mysql->options.extension->query_attributes) { + string key = arg ? pointer_cast(arg) : ""; + + if (!key.empty()) { + auto it = mysql->options.extension->query_attributes->hash.find(key); + if (it != mysql->options.extension->query_attributes->hash.end()) { + const string &key = it->first; + const string &value = it->second; + mysql->options.extension->query_attributes_length -= + get_length_store_length(key.size()) + key.size() + + get_length_store_length(value.size()) + value.size(); + + mysql->options.extension->query_attributes->hash.erase(it); + } + } + } + break; case MYSQL_ENABLE_CLEARTEXT_PLUGIN: ENSURE_EXTENSIONS_PRESENT(&mysql->options); mysql->options.extension->enable_cleartext_plugin = @@ -7935,6 +7996,53 @@ int STDCALL mysql_get_option(MYSQL *mysql, enum mysql_option option, DBUG_RETURN(0); } +static int add_attributes(MYSQL *mysql, const void *arg1, const void *arg2, + struct My_hash **attrs, size_t *attrs_length) { + DBUG_ENTER(__func__); + const char *key = static_cast(arg1); + const char *value = static_cast(arg2); + size_t key_len = arg1 ? strlen(key) : 0; + size_t value_len = arg2 ? strlen(value) : 0; + size_t attr_storage_length = key_len + value_len; + + /* we can't have a zero length key */ + if (!key_len) { + set_mysql_error(mysql, CR_INVALID_PARAMETER_NO, unknown_sqlstate); + DBUG_RETURN(1); + } + + /* calculate the total storage length of the attribute */ + attr_storage_length += get_length_store_length(key_len); + attr_storage_length += get_length_store_length(value_len); + + /* + Throw and error if the maximum combined length of the attribute value + will be greater than the maximum that we can safely transmit. + */ + if (attr_storage_length + *attrs_length > + MAX_CONNECTION_ATTR_STORAGE_LENGTH) { + set_mysql_error(mysql, CR_INVALID_PARAMETER_NO, unknown_sqlstate); + DBUG_RETURN(1); + } + + if (!*attrs) { + *attrs = new (std::nothrow) My_hash(); + if (!*attrs) { + set_mysql_error(mysql, CR_OUT_OF_MEMORY, unknown_sqlstate); + DBUG_RETURN(1); + } + } + if (!(*attrs)->hash.emplace(key, value).second) { + /* can't insert the value */ + set_mysql_error(mysql, CR_DUPLICATE_CONNECTION_ATTR, unknown_sqlstate); + DBUG_RETURN(1); + } + + *attrs_length += attr_storage_length; + + DBUG_RETURN(0); +} + int STDCALL mysql_options4(MYSQL *mysql, enum mysql_option option, const void *arg1, const void *arg2) { DBUG_ENTER("mysql_option"); @@ -7942,54 +8050,22 @@ int STDCALL mysql_options4(MYSQL *mysql, enum mysql_option option, switch (option) { case MYSQL_OPT_CONNECT_ATTR_ADD: { - const char *key = static_cast(arg1); - const char *value = static_cast(arg2); - size_t key_len = arg1 ? strlen(key) : 0; - size_t value_len = arg2 ? strlen(value) : 0; - size_t attr_storage_length = key_len + value_len; - - /* we can't have a zero length key */ - if (!key_len) { - set_mysql_error(mysql, CR_INVALID_PARAMETER_NO, unknown_sqlstate); - DBUG_RETURN(1); - } - - /* calculate the total storage length of the attribute */ - attr_storage_length += get_length_store_length(key_len); - attr_storage_length += get_length_store_length(value_len); - ENSURE_EXTENSIONS_PRESENT(&mysql->options); - - /* - Throw and error if the maximum combined length of the attribute value - will be greater than the maximum that we can safely transmit. - */ - if (attr_storage_length + - mysql->options.extension->connection_attributes_length > - MAX_CONNECTION_ATTR_STORAGE_LENGTH) { - set_mysql_error(mysql, CR_INVALID_PARAMETER_NO, unknown_sqlstate); + if (add_attributes( + mysql, arg1, arg2, + &mysql->options.extension->connection_attributes, + &mysql->options.extension->connection_attributes_length)) { DBUG_RETURN(1); } - - if (!mysql->options.extension->connection_attributes) { - mysql->options.extension->connection_attributes = - new (std::nothrow) My_hash(); - if (!mysql->options.extension->connection_attributes) { - set_mysql_error(mysql, CR_OUT_OF_MEMORY, unknown_sqlstate); - DBUG_RETURN(1); - } - } - if (!mysql->options.extension->connection_attributes->hash - .emplace(key, value) - .second) { - /* can't insert the value */ - set_mysql_error(mysql, CR_DUPLICATE_CONNECTION_ATTR, unknown_sqlstate); + break; + } + case MYSQL_OPT_QUERY_ATTR_ADD: { + ENSURE_EXTENSIONS_PRESENT(&mysql->options); + if (add_attributes(mysql, arg1, arg2, + &mysql->options.extension->query_attributes, + &mysql->options.extension->query_attributes_length)) { DBUG_RETURN(1); } - - mysql->options.extension->connection_attributes_length += - attr_storage_length; - break; } diff --git a/sql-common/net_serv.cc b/sql-common/net_serv.cc index 35786226ce4..8f02cb6b04c 100644 --- a/sql-common/net_serv.cc +++ b/sql-common/net_serv.cc @@ -462,7 +462,7 @@ static int begin_packet_write_state(NET *net, uchar command, header_len += NET_HEADER_SIZE + COMP_HEADER_SIZE; } size_t total_len = packet_len + prefix_len; - bool include_command = (command < COM_END); + bool include_command = (command != COM_END); if (include_command) { ++total_len; } diff --git a/sql/protocol_classic.cc b/sql/protocol_classic.cc index 2b9e6a50cfe..68b0c12365a 100644 --- a/sql/protocol_classic.cc +++ b/sql/protocol_classic.cc @@ -445,6 +445,7 @@ #include "my_time.h" #include "mysql/com_data.h" #include "mysql/psi/mysql_socket.h" +#include "mysql_com.h" #include "mysqld_error.h" #include "sql/field.h" #include "sql/item.h" @@ -2723,6 +2724,15 @@ bool Protocol_classic::parse_packet(union COM_DATA *data, data->com_stmt_reset.stmt_id = uint4korr(input_raw_packet); break; } + case COM_QUERY_ATTRS: { + uchar *pos = input_raw_packet; + data->com_query.query_attrs_length = net_field_length_ll((uchar **)&pos); + data->com_query.query_attrs = reinterpret_cast(pos); + pos += data->com_query.query_attrs_length; + data->com_query.query = reinterpret_cast(pos); + data->com_query.length = input_raw_packet + input_packet_length - pos; + break; + } case COM_QUERY: { data->com_query.query = reinterpret_cast(input_raw_packet); data->com_query.length = input_packet_length; @@ -2788,14 +2798,24 @@ int Protocol_classic::get_command(COM_DATA *com_data, *cmd = (enum enum_server_command)(uchar)input_raw_packet[0]; - if (*cmd >= COM_END) *cmd = COM_END; // Wrong command + if (*cmd >= COM_END && *cmd != COM_QUERY_ATTRS) + *cmd = COM_END; // Wrong command DBUG_ASSERT(input_packet_length); // Skip 'command' input_packet_length--; input_raw_packet++; - return parse_packet(com_data, *cmd); + int rc = parse_packet(com_data, *cmd); + // Here we pretend that we just received a COM_QUERY requested. This will + // make tracking in perfschema easier, since it won't split our workload + // according to whether query attributes were sent or not. Ideally, if this + // were upstreamed, we would be able to change the COM_QUERY packet + // structure via client capabilities. + if (*cmd == COM_QUERY_ATTRS) { + *cmd = COM_QUERY; + } + return rc; } uint Protocol_classic::get_rw_status() { return m_thd->net.reading_or_writing; } diff --git a/sql/sql_class.h b/sql/sql_class.h index 1ef9c8c3cc5..77bdb065553 100644 --- a/sql/sql_class.h +++ b/sql/sql_class.h @@ -3700,7 +3700,14 @@ class THD : public MDL_context_owner, void mark_transaction_to_rollback(bool all); + void set_query_attrs(const LEX_CSTRING &arg) { query_attrs_string = arg; } + void reset_query_attrs() { set_query_attrs(LEX_CSTRING()); } + inline const char *query_attrs() const { return query_attrs_string.str; } + inline uint32 query_attrs_length() const { return query_attrs_string.length; } + private: + LEX_CSTRING query_attrs_string; + /** The current internal error handler for this thread, or NULL. */ Internal_error_handler *m_internal_handler; diff --git a/sql/sql_parse.cc b/sql/sql_parse.cc index 86178eb5961..63705103fe2 100644 --- a/sql/sql_parse.cc +++ b/sql/sql_parse.cc @@ -1605,6 +1605,8 @@ bool dispatch_command(THD *thd, const COM_DATA *com_data, DBUG_ASSERT(thd->m_digest == NULL); thd->m_digest = &thd->m_digest_state; thd->m_digest->reset(thd->m_token_array, max_digest_length); + thd->set_query_attrs({com_data->com_query.query_attrs, + com_data->com_query.query_attrs_length}); if (alloc_query(thd, com_data->com_query.query, com_data->com_query.length)) @@ -2004,6 +2006,7 @@ done: THD_STAGE_INFO(thd, stage_cleaning_up); thd->reset_query(); + thd->reset_query_attrs(); thd->set_command(COM_SLEEP); thd->proc_info = 0; thd->lex->sql_command = SQLCOM_END;