Optionally set null marker in csv exports in BaseSQLToGCSOperator (#11409)

jmcarp · web-flow · commit bd204bb91b4b · 2020-10-11T11:48:54.000+02:00
diff --git a/airflow/providers/google/cloud/transfers/sql_to_gcs.py b/airflow/providers/google/cloud/transfers/sql_to_gcs.py
@@ -57,6 +57,8 @@ class BaseSQLToGCSOperator(BaseOperator):
     :type export_format: str
     :param field_delimiter: The delimiter to be used for CSV files.
     :type field_delimiter: str
+    :param null_marker: The null marker to be used for CSV files.
+    :type null_marker: str
     :param gzip: Option to compress file for upload (does not apply to schemas).
     :type gzip: bool
     :param schema: The schema to use, if any. Should be a list of dict or
@@ -109,6 +111,7 @@ def __init__(
         approx_max_file_size_bytes=1900000000,
         export_format='json',
         field_delimiter=',',
+        null_marker=None,
         gzip=False,
         schema=None,
         parameters=None,
@@ -136,6 +139,7 @@ def __init__(
         self.approx_max_file_size_bytes = approx_max_file_size_bytes
         self.export_format = export_format.lower()
         self.field_delimiter = field_delimiter
+        self.null_marker = null_marker
         self.gzip = gzip
         self.schema = schema
         self.parameters = parameters
@@ -204,6 +208,8 @@ def _write_local_data_files(self, cursor):
             row = self.convert_types(schema, col_type_dict, row)
 
             if self.export_format == 'csv':
+                if self.null_marker is not None:
+                    row = [value if value is not None else self.null_marker for value in row]
                 csv_writer.writerow(row)
             else:
                 row_dict = dict(zip(schema, row))
diff --git a/tests/providers/google/cloud/transfers/test_sql_to_gcs.py b/tests/providers/google/cloud/transfers/test_sql_to_gcs.py
@@ -159,3 +159,26 @@ def test_exec(
         mock_flush.assert_called_once()
         mock_upload.assert_called_once_with(BUCKET, FILENAME, TMP_FILE_NAME, mime_type=APP_JSON, gzip=False)
         mock_close.assert_called_once()
+
+        # Test null marker
+        cursor_mock.__iter__ = Mock(return_value=iter(INPUT_DATA))
+        mock_convert_type.return_value = None
+
+        operator = DummySQLToGCSOperator(
+            sql=SQL,
+            bucket=BUCKET,
+            filename=FILENAME,
+            task_id=TASK_ID,
+            export_format="csv",
+            null_marker="NULL",
+        )
+        operator.execute(context=dict())
+
+        mock_writerow.assert_has_calls(
+            [
+                mock.call(COLUMNS),
+                mock.call(["NULL", "NULL", "NULL"]),
+                mock.call(["NULL", "NULL", "NULL"]),
+                mock.call(["NULL", "NULL", "NULL"]),
+            ]
+        )
