Skip to content

Commit

Permalink
Honor schema type for MySQL to GCS data pre-process (#8090)
Browse files Browse the repository at this point in the history
  • Loading branch information
whynick1 authored Apr 10, 2020
1 parent c9c336c commit bb5e403
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 12 deletions.
38 changes: 26 additions & 12 deletions airflow/providers/google/cloud/operators/mysql_to_gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,10 +102,16 @@ def field_to_bigquery(self, field):
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
JSON/Google Cloud Storage/BigQuery.
* Datetimes are converted to UTC seconds.
* Decimals are converted to floats.
* Dates are converted to ISO formatted string if given schema_type is
DATE, or UTC seconds otherwise.
* Binary type fields are converted to integer if given schema_type is
INTEGER, or encoded with base64 otherwise. Imported BYTES data must
be base64-encoded according to BigQuery documentation:
https://cloud.google.com/bigquery/data-types
:param value: MySQLdb column value
:type value: Any
Expand All @@ -114,12 +120,20 @@ def convert_type(self, value, schema_type):
"""
if value is None:
return value
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 isinstance(value, bytes) or schema_type == "BYTES":
return base64.standard_b64encode(value).decode('ascii')
if isinstance(value, datetime):
value = calendar.timegm(value.timetuple())
elif isinstance(value, timedelta):
value = value.total_seconds()
elif isinstance(value, Decimal):
value = float(value)
elif isinstance(value, date):
if schema_type == "DATE":
value = value.isoformat()
else:
value = calendar.timegm(value.timetuple())
elif isinstance(value, bytes):
if schema_type == "INTEGER":
value = int.from_bytes(value, "big")
else:
value = base64.standard_b64encode(value).decode('ascii')
return value
2 changes: 2 additions & 0 deletions tests/providers/google/cloud/operators/test_mysql_to_gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,9 +92,11 @@ def test_init(self):
@parameterized.expand([
("string", None, "string"),
(datetime.date(1970, 1, 2), None, 86400),
(datetime.date(1970, 1, 2), "DATE", "1970-01-02"),
(datetime.datetime(1970, 1, 1, 1, 0), None, 3600),
(decimal.Decimal(5), None, 5),
(b"bytes", "BYTES", "Ynl0ZXM="),
(b"\x00\x01", "INTEGER", 1),
(None, "BYTES", None)
])
def test_convert_type(self, value, schema_type, expected):
Expand Down

0 comments on commit bb5e403

Please sign in to comment.