mysql_hook — Airflow Documentation

# Source code for mysql_hook

    # -*- coding: utf-8 -*-
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    # http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.

    import MySQLdb
    import MySQLdb.cursors

    from airflow.hooks.dbapi_hook import DbApiHook

    class MySqlHook(DbApiHook):
        '''
        Interact with MySQL.

        You can specify charset in the extra field of your connection
        as ``{"charset": "utf8"}``. Also you can choose cursor as
        ``{"cursor": "SSCursor"}``. Refer to the MySQLdb.cursors for more details.
        '''

        conn_name_attr = 'mysql_conn_id'
        default_conn_name = 'mysql_default'
        supports_autocommit = True

        def get_conn(self):
            """
            Returns a mysql connection object
            """
            conn = self.get_connection(self.mysql_conn_id)
            conn_config = {
                "user": conn.login,
                "passwd": conn.password or ''
            }

            conn_config["host"] = conn.host or 'localhost'
            if not conn.port:
                conn_config["port"] = 3306
            else:
                conn_config["port"] = int(conn.port)

            conn_config["db"] = conn.schema or ''

            if conn.extra_dejson.get('charset', False):
                conn_config["charset"] = conn.extra_dejson["charset"]
                if (conn_config["charset"]).lower() == 'utf8' or\
                        (conn_config["charset"]).lower() == 'utf-8':
                    conn_config["use_unicode"] = True
            if conn.extra_dejson.get('cursor', False):
                if (conn.extra_dejson["cursor"]).lower() == 'sscursor':
                    conn_config["cursorclass"] = MySQLdb.cursors.SSCursor
                elif (conn.extra_dejson["cursor"]).lower() == 'dictcursor':
                    conn_config["cursorclass"] = MySQLdb.cursors.DictCursor
                elif (conn.extra_dejson["cursor"]).lower() == 'ssdictcursor':
                    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']
            if local_infile:
                conn_config["local_infile"] = 1
            conn = MySQLdb.connect(**conn_config)
            return conn

        def bulk_load(self, table, tmp_file):
            """
            Loads a tab-delimited file into a database table
            """
            conn = self.get_conn()
            cur = conn.cursor()
            cur.execute("""
                LOAD DATA LOCAL INFILE '{tmp_file}'
                INTO TABLE {table}
                """.format(**locals()))
            conn.commit()

        @staticmethod
        def _serialize_cell(cell, conn):
            """
            Returns the MySQL literal of the cell as a string.

            :param cell: The cell to insert into the table
            :type cell: object
            :param conn: The database connection
            :type conn: connection object
            :return: The serialized cell
            :rtype: str
            """

            return conn.literal(cell)