def test_file_splitting(self, gcs_hook_mock_class, mysql_hook_mock_class):
"""Test that ndjson is split by approx_max_file_size_bytes param."""
mysql_hook_mock = mysql_hook_mock_class.return_value
mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION
gcs_hook_mock = gcs_hook_mock_class.return_value
expected_upload = {
JSON_FILENAME.format(0): b''.join(NDJSON_LINES[:2]),
JSON_FILENAME.format(1): NDJSON_LINES[2],
}
def _assert_upload(bucket,
obj,
tmp_filename,
mime_type=None,
gzip=False):
self.assertEqual(BUCKET, bucket)
self.assertEqual('application/json', mime_type)
self.assertFalse(gzip)
with open(tmp_filename, 'rb') as file:
self.assertEqual(expected_upload[obj], file.read())
gcs_hook_mock.upload.side_effect = _assert_upload
op = MySQLToGCSOperator(
task_id=TASK_ID,
sql=SQL,
bucket=BUCKET,
filename=JSON_FILENAME,
approx_max_file_size_bytes=len(
expected_upload[JSON_FILENAME.format(0)]),
)
op.execute(None)
def test_schema_file_with_custom_schema(self, gcs_hook_mock_class,
mysql_hook_mock_class):
"""Test writing schema files with customized schema"""
mysql_hook_mock = mysql_hook_mock_class.return_value
mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION
gcs_hook_mock = gcs_hook_mock_class.return_value
def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip): # pylint: disable=unused-argument
if obj == SCHEMA_FILENAME:
self.assertFalse(gzip)
with open(tmp_filename, 'rb') as file:
self.assertEqual(b''.join(CUSTOM_SCHEMA_JSON), file.read())
gcs_hook_mock.upload.side_effect = _assert_upload
op = MySQLToGCSOperator(
task_id=TASK_ID,
sql=SQL,
bucket=BUCKET,
filename=JSON_FILENAME,
schema_filename=SCHEMA_FILENAME,
schema=SCHEMA,
)
op.execute(None)
# once for the file and once for the schema
self.assertEqual(2, gcs_hook_mock.upload.call_count)
def test_exec_success_json(self, gcs_hook_mock_class,
mysql_hook_mock_class):
"""Test successful run of execute function for JSON"""
op = MySQLToGCSOperator(task_id=TASK_ID,
mysql_conn_id=MYSQL_CONN_ID,
sql=SQL,
bucket=BUCKET,
filename=JSON_FILENAME)
mysql_hook_mock = mysql_hook_mock_class.return_value
mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION
gcs_hook_mock = gcs_hook_mock_class.return_value
def _assert_upload(bucket,
obj,
tmp_filename,
mime_type=None,
gzip=False):
self.assertEqual(BUCKET, bucket)
self.assertEqual(JSON_FILENAME.format(0), obj)
self.assertEqual('application/json', mime_type)
self.assertFalse(gzip)
with open(tmp_filename, 'rb') as file:
self.assertEqual(b''.join(NDJSON_LINES), file.read())
gcs_hook_mock.upload.side_effect = _assert_upload
op.execute(None)
mysql_hook_mock_class.assert_called_once_with(
mysql_conn_id=MYSQL_CONN_ID)
mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(
SQL)
def test_convert_type(self, value, schema_type, expected):
op = MySQLToGCSOperator(task_id=TASK_ID,
mysql_conn_id=MYSQL_CONN_ID,
sql=SQL,
bucket=BUCKET,
filename=JSON_FILENAME)
self.assertEqual(op.convert_type(value, schema_type), expected)
def test_exec_success_csv_with_delimiter(self, gcs_hook_mock_class, mysql_hook_mock_class):
"""Test successful run of execute function for CSV with a field delimiter"""
op = MySQLToGCSOperator(
task_id=TASK_ID,
mysql_conn_id=MYSQL_CONN_ID,
sql=SQL,
export_format='csv',
field_delimiter='|',
bucket=BUCKET,
filename=CSV_FILENAME,
)
mysql_hook_mock = mysql_hook_mock_class.return_value
mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION
gcs_hook_mock = gcs_hook_mock_class.return_value
def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False):
assert BUCKET == bucket
assert CSV_FILENAME.format(0) == obj
assert 'text/csv' == mime_type
assert not gzip
with open(tmp_filename, 'rb') as file:
assert b''.join(CSV_LINES_PIPE_DELIMITED) == file.read()
gcs_hook_mock.upload.side_effect = _assert_upload
op.execute(None)
mysql_hook_mock_class.assert_called_once_with(mysql_conn_id=MYSQL_CONN_ID)
mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL)
def test_execute_with_query_error(self, mock_gcs_hook, mock_mysql_hook):
mock_mysql_hook.return_value.get_conn.return_value.cursor.return_value.execute.side_effect = (
ProgrammingError
)
op = MySQLToGCSOperator(
task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME, schema_filename=SCHEMA_FILENAME
)
with pytest.raises(ProgrammingError):
op.execute(None)
def test_exec_success_csv_ensure_utc(self, gcs_hook_mock_class,
mysql_hook_mock_class):
"""Test successful run of execute function for CSV"""
op = MySQLToGCSOperator(
task_id=TASK_ID,
mysql_conn_id=MYSQL_CONN_ID,
sql=SQL,
export_format='CSV',
bucket=BUCKET,
filename=CSV_FILENAME,
ensure_utc=True,
)
mysql_hook_mock = mysql_hook_mock_class.return_value
mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION
gcs_hook_mock = gcs_hook_mock_class.return_value
def _assert_upload(bucket,
obj,
tmp_filename,
mime_type=None,
gzip=False):
self.assertEqual(BUCKET, bucket)
self.assertEqual(CSV_FILENAME.format(0), obj)
self.assertEqual('text/csv', mime_type)
self.assertFalse(gzip)
with open(tmp_filename, 'rb') as file:
self.assertEqual(b''.join(CSV_LINES), file.read())
gcs_hook_mock.upload.side_effect = _assert_upload
op.execute(None)
mysql_hook_mock_class.assert_called_once_with(
mysql_conn_id=MYSQL_CONN_ID)
mysql_hook_mock.get_conn().cursor().execute.assert_has_calls(
[mock.call(TZ_QUERY), mock.call(SQL)])
def test_init(self):
"""Test MySqlToGoogleCloudStorageOperator instance is properly initialized."""
op = MySQLToGCSOperator(task_id=TASK_ID,
sql=SQL,
bucket=BUCKET,
filename=JSON_FILENAME,
export_format='CSV',
field_delimiter='|')
self.assertEqual(op.task_id, TASK_ID)
self.assertEqual(op.sql, SQL)
self.assertEqual(op.bucket, BUCKET)
self.assertEqual(op.filename, JSON_FILENAME)
self.assertEqual(op.export_format, 'csv')
self.assertEqual(op.field_delimiter, '|')
def test_init(self):
"""Test MySqlToGoogleCloudStorageOperator instance is properly initialized."""
op = MySQLToGCSOperator(
task_id=TASK_ID,
sql=SQL,
bucket=BUCKET,
filename=JSON_FILENAME,
export_format='CSV',
field_delimiter='|',
)
assert op.task_id == TASK_ID
assert op.sql == SQL
assert op.bucket == BUCKET
assert op.filename == JSON_FILENAME
assert op.export_format == 'csv'
assert op.field_delimiter == '|'
# 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 os
from airflow import models
from airflow.providers.google.cloud.transfers.mysql_to_gcs import MySQLToGCSOperator
from airflow.utils import dates
GCS_BUCKET = os.environ.get("GCP_GCS_BUCKET", "example-airflow-mysql-gcs")
FILENAME = 'test_file'
SQL_QUERY = "SELECT * from test_table"
with models.DAG(
'example_mysql_to_gcs',
default_args=dict(start_date=dates.days_ago(1)),
schedule_interval=None,
tags=['example'],
) as dag:
# [START howto_operator_mysql_to_gcs]
upload = MySQLToGCSOperator(task_id='mysql_to_gcs',
sql=SQL_QUERY,
bucket=GCS_BUCKET,
filename=FILENAME,
export_format='csv')
# [END howto_operator_mysql_to_gcs]