diff --git a/airflow/providers/google/cloud/transfers/postgres_to_gcs.py b/airflow/providers/google/cloud/transfers/postgres_to_gcs.py index c233b5c4c9758..ce77d7b87beb8 100644 --- a/airflow/providers/google/cloud/transfers/postgres_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/postgres_to_gcs.py @@ -78,10 +78,10 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator): ui_color = '#a0e08c' type_map = { - 1114: 'TIMESTAMP', + 1114: 'DATETIME', 1184: 'TIMESTAMP', - 1082: 'TIMESTAMP', - 1083: 'TIMESTAMP', + 1082: 'DATE', + 1083: 'TIME', 1005: 'INTEGER', 1007: 'INTEGER', 1016: 'INTEGER', @@ -131,18 +131,24 @@ def field_to_bigquery(self, field) -> Dict[str, str]: def convert_type(self, value, schema_type): """ Takes a value from Postgres, 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. Times are converted to seconds. + JSON/Google Cloud Storage/BigQuery. + Timezone aware Datetime are converted to UTC seconds. + Unaware Datetime, Date and Time are converted to ISO formatted strings. + Decimals are converted to floats. """ - if isinstance(value, (datetime.datetime, datetime.date)): - return pendulum.parse(value.isoformat()).float_timestamp + if isinstance(value, datetime.datetime): + iso_format_value = value.isoformat() + if value.tzinfo is None: + return iso_format_value + return pendulum.parse(iso_format_value).float_timestamp + if isinstance(value, datetime.date): + return value.isoformat() if isinstance(value, datetime.time): formatted_time = time.strptime(str(value), "%H:%M:%S") - return int( - datetime.timedelta( - hours=formatted_time.tm_hour, minutes=formatted_time.tm_min, seconds=formatted_time.tm_sec - ).total_seconds() + time_delta = datetime.timedelta( + hours=formatted_time.tm_hour, minutes=formatted_time.tm_min, seconds=formatted_time.tm_sec ) + return str(time_delta) if isinstance(value, dict): return json.dumps(value) if isinstance(value, Decimal): diff --git a/tests/providers/google/cloud/transfers/test_postgres_to_gcs.py b/tests/providers/google/cloud/transfers/test_postgres_to_gcs.py index 564c5ead19033..d8621d8dbc449 100644 --- a/tests/providers/google/cloud/transfers/test_postgres_to_gcs.py +++ b/tests/providers/google/cloud/transfers/test_postgres_to_gcs.py @@ -15,11 +15,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - +import datetime import unittest from unittest.mock import patch import pytest +import pytz +from parameterized import parameterized from airflow.providers.google.cloud.transfers.postgres_to_gcs import PostgresToGCSOperator from airflow.providers.postgres.hooks.postgres import PostgresHook @@ -89,6 +91,28 @@ def _assert_uploaded_file_content(self, bucket, obj, tmp_filename, mime_type, gz with open(tmp_filename, 'rb') as file: assert b''.join(NDJSON_LINES) == file.read() + @parameterized.expand( + [ + ("string", "string"), + (32.9, 32.9), + (-2, -2), + (datetime.date(1970, 1, 2), "1970-01-02"), + (datetime.date(1000, 1, 2), "1000-01-02"), + (datetime.datetime(1970, 1, 1, 1, 0, tzinfo=None), "1970-01-01T01:00:00"), + ( + datetime.datetime(2022, 1, 1, 2, 0, tzinfo=pytz.UTC), + 1641002400.0, + ), + (datetime.time(hour=0, minute=0, second=0), "0:00:00"), + (datetime.time(hour=23, minute=59, second=59), "23:59:59"), + ] + ) + def test_convert_type(self, value, expected): + op = PostgresToGCSOperator( + task_id=TASK_ID, postgres_conn_id=POSTGRES_CONN_ID, sql=SQL, bucket=BUCKET, filename=FILENAME + ) + assert op.convert_type(value, None) == expected + @patch('airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook') def test_exec_success(self, gcs_hook_mock_class): """Test the execute function in case where the run is successful."""