[AIRFLOW-3268] Better handling of extras field in MySQL connection (apache#4113)

Author: potiuk

airflow/hooks/mysql_hook.py

@@ -19,6 +19,8 @@
 
 import MySQLdb
 import MySQLdb.cursors
+import json
+import six
 
 from airflow.hooks.dbapi_hook import DbApiHook
 
@@ -87,7 +89,13 @@ def get_conn(self):
                 conn_config["cursorclass"] = MySQLdb.cursors.SSDictCursor
         local_infile = conn.extra_dejson.get('local_infile', False)
         if conn.extra_dejson.get('ssl', False):
-            conn_config['ssl'] = conn.extra_dejson['ssl']
+            # SSL parameter for MySQL has to be a dictionary and in case
+            # of extra/dejson we can get string if extra is passed via
+            # URL parameters
+            dejson_ssl = conn.extra_dejson['ssl']
+            if isinstance(dejson_ssl, six.string_types):
+                dejson_ssl = json.loads(dejson_ssl)
+            conn_config['ssl'] = dejson_ssl
         if conn.extra_dejson.get('unix_socket'):
             conn_config['unix_socket'] = conn.extra_dejson['unix_socket']
         if local_infile:

docs/howto/manage-connections.rst

@@ -133,8 +133,9 @@ Scopes (comma separated)
         Scopes are ignored when using application default credentials. See
         issue `AIRFLOW-2522
         <https://issues.apache.org/jira/browse/AIRFLOW-2522>`_.
+
 MySQL
-~~~~~~~~~~~~~~~~~~~~~
+~~~~~
 The MySQL connect type allows to connect with MySQL database.
 
 Configuring the Connection
@@ -152,7 +153,61 @@ Password (required)
     Specify the password to connect.    
 
 Extra (optional)
-    Specify the charset. Example: {"charset": "utf8"}
-    
+    Specify the extra parameters (as json dictionary) that can be used in mysql
+    connection. The following parameters are supported:
+
+    * **charset**: specify charset of the connection
+    * **cursor**: one of "sscursor", "dictcursor, "ssdictcursor" - specifies cursor class to be
+      used
+    * **local_infile**: controls MySQL's LOCAL capability (permitting local data loading by
+      clients). See `MySQLdb docs <https://mysqlclient.readthedocs.io/user_guide.html>`_
+      for details.
+    * **unix_socket**: UNIX socket used instead of the default socket
+    * **ssl**: Dictionary of SSL parameters that control connecting using SSL (those
+      parameters are server specific and should contain "ca", "cert", "key", "capath",
+      "cipher" parameters. See
+      `MySQLdb docs <https://mysqlclient.readthedocs.io/user_guide.html>`_ for details.
+      Note that in order to be useful in URL notation, this parameter might also be
+      a string where the SSL dictionary is a string-encoded JSON dictionary.
+
+    Example "extras" field:
+
+    .. code-block:: json
+
+       {
+          "charset": "utf8",
+          "cursorclass": "sscursor",
+          "local_infile": true,
+          "unix_socket": "/var/socket",
+          "ssl": {
+            "cert": "/tmp/client-cert.pem",
+            "ca": "/tmp/server-ca.pem'",
+            "key": "/tmp/client-key.pem"
+          }
+       }
+
+    or
+
+    .. code-block:: json
+
+       {
+          "charset": "utf8",
+          "cursorclass": "sscursor",
+          "local_infile": true,
+          "unix_socket": "/var/socket",
+          "ssl": "{\"cert\": \"/tmp/client-cert.pem\", \"ca\": \"/tmp/server-ca.pem\", \"key\": \"/tmp/client-key.pem\"}"
+       }
+
+    When specifying the connection as URI (in AIRFLOW_CONN_* variable) you should specify it
+    following the standard syntax of DB connections, where extras as passed as parameters
+    of the URI (note that all components of the URI should be URL-encoded).
+
+    For example:
+
+    .. code-block:: bash
+
+       mysql://mysql_user:XXXXXXXXXXXX@1.1.1.1:3306/mysqldb?ssl=%7B%22cert%22%3A+%22%2Ftmp%2Fclient-cert.pem%22%2C+%22ca%22%3A+%22%2Ftmp%2Fserver-ca.pem%22%2C+%22key%22%3A+%22%2Ftmp%2Fclient-key.pem%22%7D
+
     .. note::
-        If encounter UnicodeDecodeError while working with MySQL connection check the charset defined is matched to the database charset.
+        If encounter UnicodeDecodeError while working with MySQL connection check
+        the charset defined is matched to the database charset.

tests/hooks/test_mysql_hook.py

@@ -27,6 +27,12 @@
 from airflow import models
 from airflow.hooks.mysql_hook import MySqlHook
 
+SSL_DICT = {
+    'cert': '/tmp/client-cert.pem',
+    'ca': '/tmp/server-ca.pem',
+    'key': '/tmp/client-key.pem'
+}
+
 
 class TestMySqlHookConn(unittest.TestCase):
 
@@ -101,6 +107,24 @@ def test_get_con_unix_socket(self, mock_connect):
         self.assertEqual(args, ())
         self.assertEqual(kwargs['unix_socket'], '/tmp/socket')
 
+    @mock.patch('airflow.hooks.mysql_hook.MySQLdb.connect')
+    def test_get_conn_ssl_as_dictionary(self, mock_connect):
+        self.connection.extra = json.dumps({'ssl': SSL_DICT})
+        self.db_hook.get_conn()
+        mock_connect.assert_called_once()
+        args, kwargs = mock_connect.call_args
+        self.assertEqual(args, ())
+        self.assertEqual(kwargs['ssl'], SSL_DICT)
+
+    @mock.patch('airflow.hooks.mysql_hook.MySQLdb.connect')
+    def test_get_conn_ssl_as_string(self, mock_connect):
+        self.connection.extra = json.dumps({'ssl': json.dumps(SSL_DICT)})
+        self.db_hook.get_conn()
+        mock_connect.assert_called_once()
+        args, kwargs = mock_connect.call_args
+        self.assertEqual(args, ())
+        self.assertEqual(kwargs['ssl'], SSL_DICT)
+
 
 class TestMySqlHook(unittest.TestCase):