diff --git a/client/mysqltest.cc b/client/mysqltest.cc index a83463d..9d22658 100644 --- a/client/mysqltest.cc +++ b/client/mysqltest.cc @@ -434,6 +434,9 @@ enum enum_commands { Q_MOVE_FILE, Q_REMOVE_FILES_WILDCARD, Q_COPY_FILES_WILDCARD, 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, @@ -540,6 +543,9 @@ const char *command_names[]= "send_eval", "output", "reset_connection", + "query_attrs_add", + "query_attrs_delete", + "query_attrs_reset", 0 }; @@ -6973,6 +6979,53 @@ 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; @@ -10296,6 +10349,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 7b12124..d6a36b2 100644 --- a/include/my_command.h +++ b/include/my_command.h @@ -74,7 +74,12 @@ enum enum_server_command /* don't forget to update const char *command_name[] in sql_parse.cc */ /* Must be last */ - COM_END + COM_END, + /* + 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 13607e6..a5b741d 100644 --- a/include/mysql.h +++ b/include/mysql.h @@ -181,6 +181,8 @@ enum mysql_option MYSQL_OPT_SSL_CRL, MYSQL_OPT_SSL_CRLPATH, 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 7011f97..46d2e90 100644 --- a/include/mysql.h.pp +++ b/include/mysql.h.pp @@ -78,7 +78,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 { @@ -369,6 +370,8 @@ enum mysql_option MYSQL_OPT_SSL_CRL, MYSQL_OPT_SSL_CRLPATH, 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 f8e9b00..868f713 100644 --- a/include/mysql/com_data.h +++ b/include/mysql/com_data.h @@ -108,6 +108,8 @@ typedef struct st_com_query_data { const char *query; unsigned int length; + const char *query_attrs; + unsigned int query_attrs_length; } COM_QUERY_DATA; typedef struct st_com_field_list_data diff --git a/include/mysql/plugin_audit.h.pp b/include/mysql/plugin_audit.h.pp index 92b7391..78de7da 100644 --- a/include/mysql/plugin_audit.h.pp +++ b/include/mysql/plugin_audit.h.pp @@ -172,7 +172,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 51275b3..d627b53 100644 --- a/include/mysql/services.h.pp +++ b/include/mysql/services.h.pp @@ -110,6 +110,8 @@ typedef struct st_com_query_data { const char *query; unsigned int length; + const char *query_attrs; + unsigned int query_attrs_length; } COM_QUERY_DATA; typedef struct st_com_field_list_data { diff --git a/include/sql_common.h b/include/sql_common.h index b9b2d29..74ad044 100644 --- a/include/sql_common.h +++ b/include/sql_common.h @@ -121,6 +121,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 0000000..2d3c3f8 --- /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 0000000..fa5fddc --- /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 ec315c7..142f0ba 100644 --- a/sql-common/client.cc +++ b/sql-common/client.cc @@ -3393,41 +3393,55 @@ struct My_hash malloc_unordered_map hash{key_memory_mysql_options}; }; -uchar * -send_client_connect_attrs(MYSQL *mysql, uchar *buf) +static uchar * +encode_attrs(uchar *buf, struct My_hash *attrs, size_t attrs_length) { - /* 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); + buf= net_store_length(buf, attrs_length); - /* check if we have connection attributes */ - if (mysql->options.extension && - mysql->options.extension->connection_attributes) + /* check if we have attributes */ + if (attrs) + { + /* loop over and dump the connection attributes */ + for (const auto &key_and_value : attrs->hash) { - /* 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; + 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()); + /* 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()); - } + 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) + { + 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) { @@ -5594,6 +5608,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)); @@ -5813,7 +5828,26 @@ 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)); + } } @@ -6227,6 +6261,37 @@ mysql_options(MYSQL *mysql,enum mysql_option option, const void *arg) } } 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= @@ -6452,6 +6517,59 @@ mysql_get_option(MYSQL *mysql, enum mysql_option option, const void *arg) 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, @@ -6464,59 +6582,22 @@ mysql_options4(MYSQL *mysql,enum mysql_option 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/protocol_classic.cc b/sql/protocol_classic.cc index 521fe39..5cb2b9c 100644 --- a/sql/protocol_classic.cc +++ b/sql/protocol_classic.cc @@ -235,6 +235,7 @@ #include "mysql/com_data.h" #include "mysql/psi/mysql_socket.h" #include "mysql/psi/mysql_statement.h" +#include "mysql_com.h" #include "mysqld_error.h" #include "sql/auth/sql_security_ctx.h" #include "sql/field.h" @@ -1467,6 +1468,16 @@ 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); @@ -1537,7 +1548,7 @@ int Protocol_classic::get_command(COM_DATA *com_data, enum_server_command *cmd) *cmd= (enum enum_server_command) (uchar) input_raw_packet[0]; - if (*cmd >= COM_END) + if (*cmd >= COM_END && *cmd != COM_QUERY_ATTRS) *cmd= COM_END; // Wrong command DBUG_ASSERT(input_packet_length); @@ -1545,7 +1556,17 @@ int Protocol_classic::get_command(COM_DATA *com_data, enum_server_command *cmd) 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() diff --git a/sql/sql_class.h b/sql/sql_class.h index 324e36e..742250d 100644 --- a/sql/sql_class.h +++ b/sql/sql_class.h @@ -3951,7 +3951,16 @@ public: 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 45b68c7..03fbce9 100644 --- a/sql/sql_parse.cc +++ b/sql/sql_parse.cc @@ -1706,6 +1706,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)) @@ -2124,6 +2126,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;