From 35e7cfa80c57e2c7537cbb58a683fb7b479a3221 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Wed, 2 Nov 2022 21:15:56 +0530 Subject: [PATCH 1/4] Deprecate bigquery hook _bq_cast func --- airflow/providers/google/cloud/hooks/bigquery.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index c835ead045643..a7456456a8afa 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -2738,7 +2738,7 @@ def next(self) -> list | None: rows = query_results["rows"] for dict_row in rows: - typed_row = [_bq_cast(vs["v"], col_types[idx]) for idx, vs in enumerate(dict_row["f"])] + typed_row = [bq_cast(vs["v"], col_types[idx]) for idx, vs in enumerate(dict_row["f"])] self.buffer.append(typed_row) if not self.page_token: @@ -2844,6 +2844,18 @@ def _escape(s: str) -> str: def _bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str: + """ + Helper method that casts a BigQuery row to the appropriate data types. + This is useful because BigQuery returns all fields as strings. + """ + warnings.warn( + "This function is deprecated and will be Remove in future. Please use `bq_cast` function", + DeprecationWarning, + ) + return bq_cast(string_field=string_field, bq_type=bq_type) + + +def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str: """ Helper method that casts a BigQuery row to the appropriate data types. This is useful because BigQuery returns all fields as strings. From d1722d31bfa9c8218b1623c550fc936c66e73e40 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Sat, 5 Nov 2022 23:05:05 +0530 Subject: [PATCH 2/4] Fix it --- airflow/providers/google/cloud/hooks/bigquery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index a7456456a8afa..64521e36ff058 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -3080,7 +3080,7 @@ def get_records(self, query_results: dict[str, Any]) -> list[Any]: fields = query_results["schema"]["fields"] col_types = [field["type"] for field in fields] for dict_row in rows: - typed_row = [_bq_cast(vs["v"], col_types[idx]) for idx, vs in enumerate(dict_row["f"])] + typed_row = [bq_cast(vs["v"], col_types[idx]) for idx, vs in enumerate(dict_row["f"])] buffer.append(typed_row) return buffer From 5c7ef1579af5d3d4a05d83fa43488dd4de9ac8f6 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Tue, 8 Nov 2022 15:53:27 +0530 Subject: [PATCH 3/4] Remove DeprecationWarning message --- airflow/providers/google/cloud/hooks/bigquery.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index 64521e36ff058..3e5f718b7a374 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -2843,18 +2843,6 @@ def _escape(s: str) -> str: return e -def _bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str: - """ - Helper method that casts a BigQuery row to the appropriate data types. - This is useful because BigQuery returns all fields as strings. - """ - warnings.warn( - "This function is deprecated and will be Remove in future. Please use `bq_cast` function", - DeprecationWarning, - ) - return bq_cast(string_field=string_field, bq_type=bq_type) - - def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str: """ Helper method that casts a BigQuery row to the appropriate data types. From ab5894a2b44175d97082aae2be9cdaea771dbf49 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Wed, 9 Nov 2022 12:30:28 +0530 Subject: [PATCH 4/4] Move bq_cast in util module --- .../providers/google/cloud/hooks/bigquery.py | 20 +---------- .../providers/google/cloud/utils/bigquery.py | 36 +++++++++++++++++++ 2 files changed, 37 insertions(+), 19 deletions(-) create mode 100644 airflow/providers/google/cloud/utils/bigquery.py diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index 3e5f718b7a374..a3d19c4b9cc6c 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -57,6 +57,7 @@ from airflow.exceptions import AirflowException from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.google.cloud.utils.bigquery import bq_cast from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook, get_field from airflow.utils.helpers import convert_camel_to_snake @@ -2843,25 +2844,6 @@ def _escape(s: str) -> str: return e -def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str: - """ - Helper method that casts a BigQuery row to the appropriate data types. - This is useful because BigQuery returns all fields as strings. - """ - if string_field is None: - return None - elif bq_type == "INTEGER": - return int(string_field) - elif bq_type in ("FLOAT", "TIMESTAMP"): - return float(string_field) - elif bq_type == "BOOLEAN": - if string_field not in ["true", "false"]: - raise ValueError(f"{string_field} must have value 'true' or 'false'") - return string_field == "true" - else: - return string_field - - def split_tablename( table_input: str, default_project_id: str, var_name: str | None = None ) -> tuple[str, str, str]: diff --git a/airflow/providers/google/cloud/utils/bigquery.py b/airflow/providers/google/cloud/utils/bigquery.py new file mode 100644 index 0000000000000..03753c2423691 --- /dev/null +++ b/airflow/providers/google/cloud/utils/bigquery.py @@ -0,0 +1,36 @@ +# 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. +from __future__ import annotations + + +def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str: + """ + Helper method that casts a BigQuery row to the appropriate data types. + This is useful because BigQuery returns all fields as strings. + """ + if string_field is None: + return None + elif bq_type == "INTEGER": + return int(string_field) + elif bq_type in ("FLOAT", "TIMESTAMP"): + return float(string_field) + elif bq_type == "BOOLEAN": + if string_field not in ["true", "false"]: + raise ValueError(f"{string_field} must have value 'true' or 'false'") + return string_field == "true" + else: + return string_field