From b3b7c9148bec44383a8ea561bdb1fee122f6a64e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20van=20Eeden?= Date: Sat, 9 May 2015 11:30:35 +0200 Subject: [PATCH 1/2] Initial support for connection attributes. --- lib/mysql/connector/connection.py | 1 + lib/mysql/connector/protocol.py | 20 +++++++++++++++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/lib/mysql/connector/connection.py b/lib/mysql/connector/connection.py index e7b49c2..0daca7e 100644 --- a/lib/mysql/connector/connection.py +++ b/lib/mysql/connector/connection.py @@ -113,6 +113,7 @@ def _do_handshake(self): if handshake['capabilities'] & ClientFlag.PLUGIN_AUTH: self.set_client_flags([ClientFlag.PLUGIN_AUTH]) + self.set_client_flags([ClientFlag.CONNECT_ARGS]) self._handshake = handshake def _do_auth(self, username=None, password=None, database=None, diff --git a/lib/mysql/connector/protocol.py b/lib/mysql/connector/protocol.py index 5091814..1e9229f 100644 --- a/lib/mysql/connector/protocol.py +++ b/lib/mysql/connector/protocol.py @@ -73,7 +73,7 @@ def _auth_response(self, client_flags, username, password, database, def make_auth(self, handshake, username=None, password=None, database=None, charset=33, client_flags=0, max_allowed_packet=1073741824, ssl_enabled=False, - auth_plugin=None): + auth_plugin=None, connattrs=None): """Make a MySQL Authentication packet""" try: @@ -105,8 +105,26 @@ def make_auth(self, handshake, username=None, password=None, database=None, if client_flags & ClientFlag.PLUGIN_AUTH: packet += auth_plugin.encode('utf8') + b'x00' + if connattrs is None: + connattrs = {'_client_name': 'MySQL Connector/Python'} + if client_flags & ClientFlag.CONNECT_ARGS: + packet += self.make_connattrs(connattrs) + return packet + def make_connattrs(self, connattrs): + """Encode the connection attributes""" + if not isinstance(connattrs, dict): + raise ValueError('connattrs must be of type dict') + + connattrs_len = sum([len(x) + len(connattrs[x]) for x in connattrs]) + len(connattrs.keys()) + len(connattrs.values()) + connattrs_packet = struct.pack(' Date: Sat, 9 May 2015 11:59:26 +0200 Subject: [PATCH 2/2] Support custom connectin attributes --- lib/mysql/connector/connection.py | 8 +++++--- lib/mysql/connector/constants.py | 1 + lib/mysql/connector/protocol.py | 17 ++++++++++++----- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/lib/mysql/connector/connection.py b/lib/mysql/connector/connection.py index 0daca7e..bb625dd 100644 --- a/lib/mysql/connector/connection.py +++ b/lib/mysql/connector/connection.py @@ -90,6 +90,7 @@ def __init__(self, *args, **kwargs): self._ssl_active = False self._auth_plugin = None self._pool_config_version = None + self._connattrs = None if len(kwargs) > 0: self.connect(**kwargs) @@ -117,7 +118,7 @@ def _do_handshake(self): self._handshake = handshake def _do_auth(self, username=None, password=None, database=None, - client_flags=0, charset=33, ssl_options=None): + client_flags=0, charset=33, ssl_options=None, connattrs=None): """Authenticate with the MySQL server Authentication happens in two parts. We first send a response to the @@ -140,7 +141,8 @@ def _do_auth(self, username=None, password=None, database=None, username=username, password=password, database=database, charset=charset, client_flags=client_flags, ssl_enabled=self._ssl_active, - auth_plugin=self._auth_plugin) + auth_plugin=self._auth_plugin, + connattrs=connattrs) self._socket.send(packet) self._auth_switch_request(username, password) @@ -211,7 +213,7 @@ def _open_connection(self): self._do_handshake() self._do_auth(self._user, self._password, self._database, self._client_flags, self._charset_id, - self._ssl) + self._ssl, self._connattrs) self.set_converter_class(self._converter_class) if self._client_flags & ClientFlag.COMPRESS: self._socket.recv = self._socket.recv_compressed diff --git a/lib/mysql/connector/constants.py b/lib/mysql/connector/constants.py index 79dc178..d631017 100644 --- a/lib/mysql/connector/constants.py +++ b/lib/mysql/connector/constants.py @@ -64,6 +64,7 @@ 'auth_plugin': None, 'allow_local_infile': True, 'consume_results': False, + 'connattrs': None, } CNX_POOL_ARGS = ('pool_name', 'pool_size', 'pool_reset_session') diff --git a/lib/mysql/connector/protocol.py b/lib/mysql/connector/protocol.py index 1e9229f..b99946d 100644 --- a/lib/mysql/connector/protocol.py +++ b/lib/mysql/connector/protocol.py @@ -30,7 +30,7 @@ from .constants import ( FieldFlag, ServerCmd, FieldType, ClientFlag, MAX_MYSQL_TABLE_COLUMNS) -from . import errors, utils +from . import errors, utils, version from .authentication import get_auth_plugin from .catch23 import PY2, struct_unpack @@ -105,8 +105,6 @@ def make_auth(self, handshake, username=None, password=None, database=None, if client_flags & ClientFlag.PLUGIN_AUTH: packet += auth_plugin.encode('utf8') + b'x00' - if connattrs is None: - connattrs = {'_client_name': 'MySQL Connector/Python'} if client_flags & ClientFlag.CONNECT_ARGS: packet += self.make_connattrs(connattrs) @@ -114,8 +112,17 @@ def make_auth(self, handshake, username=None, password=None, database=None, def make_connattrs(self, connattrs): """Encode the connection attributes""" - if not isinstance(connattrs, dict): - raise ValueError('connattrs must be of type dict') + + default_connattrs = {'_client_name': 'MySQL Connector/Python', + '_client_version': version.VERSION_TEXT} + + if isinstance(connattrs, dict): + default_connattrs.update(connattrs) + connattrs = default_connattrs + elif connattrs is None: + connattrs = default_connattrs + else: + raise ValueError('connattrs must be of type dict or None') connattrs_len = sum([len(x) + len(connattrs[x]) for x in connattrs]) + len(connattrs.keys()) + len(connattrs.values()) connattrs_packet = struct.pack('