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


[docs]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
[docs] 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
[docs] 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)