Source code for airflow.contrib.operators.mysql_to_gcs

# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""
MySQL to GCS operator.
"""

import base64
import calendar
import sys


from datetime import date, datetime, timedelta
from decimal import Decimal

from MySQLdb.constants import FIELD_TYPE

from airflow.hooks.mysql_hook import MySqlHook
from airflow.utils.decorators import apply_defaults
from airflow.contrib.operators.sql_to_gcs import BaseSQLToGoogleCloudStorageOperator

[docs]PY3 = sys.version_info[0] == 3
[docs]class MySqlToGoogleCloudStorageOperator(BaseSQLToGoogleCloudStorageOperator): """Copy data from MySQL to Google cloud storage in JSON or CSV format. :param mysql_conn_id: Reference to a specific MySQL hook. :type mysql_conn_id: str :param ensure_utc: Ensure TIMESTAMP columns exported as UTC. If set to `False`, TIMESTAMP columns will be exported using the MySQL server's default timezone. :type ensure_utc: bool """
[docs] ui_color = '#a0e08c'
[docs] type_map = { FIELD_TYPE.BIT: 'INTEGER', FIELD_TYPE.DATETIME: 'TIMESTAMP', FIELD_TYPE.DATE: 'TIMESTAMP', FIELD_TYPE.DECIMAL: 'FLOAT', FIELD_TYPE.NEWDECIMAL: 'FLOAT', FIELD_TYPE.DOUBLE: 'FLOAT', FIELD_TYPE.FLOAT: 'FLOAT', FIELD_TYPE.INT24: 'INTEGER', FIELD_TYPE.LONG: 'INTEGER', FIELD_TYPE.LONGLONG: 'INTEGER', FIELD_TYPE.SHORT: 'INTEGER', FIELD_TYPE.TIME: 'TIME', FIELD_TYPE.TIMESTAMP: 'TIMESTAMP', FIELD_TYPE.TINY: 'INTEGER', FIELD_TYPE.YEAR: 'INTEGER',
} @apply_defaults def __init__(self, mysql_conn_id='mysql_default', ensure_utc=False, *args, **kwargs): super(MySqlToGoogleCloudStorageOperator, self).__init__(*args, **kwargs) self.mysql_conn_id = mysql_conn_id self.ensure_utc = ensure_utc
[docs] def query(self): """ Queries mysql and returns a cursor to the results. """ mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id) conn = mysql.get_conn() cursor = conn.cursor() if self.ensure_utc: # Ensure TIMESTAMP results are in UTC tz_query = "SET time_zone = '+00:00'" self.log.info('Executing: %s', tz_query) cursor.execute(tz_query) self.log.info('Executing: %s', self.sql) cursor.execute(self.sql) return cursor
[docs] def field_to_bigquery(self, field): field_type = self.type_map.get(field[1], "STRING") # Always allow TIMESTAMP to be nullable. MySQLdb returns None types # for required fields because some MySQL timestamps can't be # represented by Python's datetime (e.g. 0000-00-00 00:00:00). field_mode = "NULLABLE" if field[6] or field_type == "TIMESTAMP" else "REQUIRED" return { 'name': field[0], 'type': field_type, 'mode': field_mode,
}
[docs] def convert_type(self, value, schema_type): """ Takes a value from MySQLdb, and converts it to a value that's safe for JSON/Google cloud storage/BigQuery. Dates are converted to UTC seconds. Decimals are converted to floats. Binary type fields are encoded with base64, as imported BYTES data must be base64-encoded according to Bigquery SQL date type documentation: https://cloud.google.com/bigquery/data-types :param value: MySQLdb column value :type value: Any :param schema_type: BigQuery data type :type schema_type: str """ if isinstance(value, (datetime, date)): return calendar.timegm(value.timetuple()) if isinstance(value, timedelta): return value.total_seconds() if isinstance(value, Decimal): return float(value) if schema_type == "BYTES": col_val = base64.standard_b64encode(value) if PY3: col_val = col_val.decode('ascii') return col_val return value