This is an automated email from the ASF dual-hosted git repository. potiuk pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new 65f3b18fc1 Sql to gcs with exclude columns (#23695) 65f3b18fc1 is described below commit 65f3b18fc1142c0d23e715fa1a98f21662df9584 Author: 서재권(Data Platform) <90180644+jaegwon...@users.noreply.github.com> AuthorDate: Sun May 22 21:50:49 2022 +0900 Sql to gcs with exclude columns (#23695) --- .../providers/google/cloud/transfers/sql_to_gcs.py | 16 ++++++++++-- tests/providers/google/cloud/transfers/temp-file | Bin 3739 -> 0 bytes .../google/cloud/transfers/test_sql_to_gcs.py | 29 +++++++++++++++++++++ 3 files changed, 43 insertions(+), 2 deletions(-) diff --git a/airflow/providers/google/cloud/transfers/sql_to_gcs.py b/airflow/providers/google/cloud/transfers/sql_to_gcs.py index 0df46bc777..46e1ad505d 100644 --- a/airflow/providers/google/cloud/transfers/sql_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/sql_to_gcs.py @@ -71,6 +71,7 @@ class BaseSQLToGCSOperator(BaseOperator): If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated). + :param exclude_columns: set of columns to exclude from transmission """ template_fields: Sequence[str] = ( @@ -103,9 +104,13 @@ class BaseSQLToGCSOperator(BaseOperator): gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None, impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + exclude_columns=None, **kwargs, ) -> None: super().__init__(**kwargs) + if exclude_columns is None: + exclude_columns = set() + self.sql = sql self.bucket = bucket self.filename = filename @@ -120,6 +125,7 @@ class BaseSQLToGCSOperator(BaseOperator): self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to self.impersonation_chain = impersonation_chain + self.exclude_columns = exclude_columns def execute(self, context: 'Context'): self.log.info("Executing query") @@ -165,7 +171,9 @@ class BaseSQLToGCSOperator(BaseOperator): names in GCS, and values are file handles to local files that contain the data for the GCS objects. """ - schema = list(map(lambda schema_tuple: schema_tuple[0], cursor.description)) + org_schema = list(map(lambda schema_tuple: schema_tuple[0], cursor.description)) + schema = [column for column in org_schema if column not in self.exclude_columns] + col_type_dict = self._get_col_type_dict() file_no = 0 @@ -314,7 +322,11 @@ class BaseSQLToGCSOperator(BaseOperator): schema = self.schema else: self.log.info("Starts generating schema") - schema = [self.field_to_bigquery(field) for field in cursor.description] + schema = [ + self.field_to_bigquery(field) + for field in cursor.description + if field[0] not in self.exclude_columns + ] if isinstance(schema, list): schema = json.dumps(schema, sort_keys=True) diff --git a/tests/providers/google/cloud/transfers/temp-file b/tests/providers/google/cloud/transfers/temp-file deleted file mode 100644 index d2282fc46c..0000000000 Binary files a/tests/providers/google/cloud/transfers/temp-file and /dev/null differ diff --git a/tests/providers/google/cloud/transfers/test_sql_to_gcs.py b/tests/providers/google/cloud/transfers/test_sql_to_gcs.py index 525e04bd0e..824ab8ff31 100644 --- a/tests/providers/google/cloud/transfers/test_sql_to_gcs.py +++ b/tests/providers/google/cloud/transfers/test_sql_to_gcs.py @@ -61,6 +61,12 @@ APP_JSON = "application/json" OUTPUT_DF = pd.DataFrame([['convert_type_return_value'] * 3] * 3, columns=COLUMNS) +EXCLUDE_COLUMNS = set('column_c') +NEW_COLUMNS = [c for c in COLUMNS if c not in EXCLUDE_COLUMNS] +OUTPUT_DF_WITH_EXCLUDE_COLUMNS = pd.DataFrame( + [['convert_type_return_value'] * len(NEW_COLUMNS)] * 3, columns=NEW_COLUMNS +) + class DummySQLToGCSOperator(BaseSQLToGCSOperator): def field_to_bigquery(self, field) -> Dict[str, str]: @@ -287,3 +293,26 @@ class TestBaseSQLToGCSOperator(unittest.TestCase): file.flush() df = pd.read_parquet(file.name) assert df.equals(OUTPUT_DF) + + def test__write_local_data_files_json_with_exclude_columns(self): + op = DummySQLToGCSOperator( + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + task_id=TASK_ID, + schema_filename=SCHEMA_FILE, + export_format="json", + gzip=False, + schema=SCHEMA, + gcp_conn_id='google_cloud_default', + exclude_columns=EXCLUDE_COLUMNS, + ) + cursor = MagicMock() + cursor.__iter__.return_value = INPUT_DATA + cursor.description = CURSOR_DESCRIPTION + + files = op._write_local_data_files(cursor) + file = next(files)['file_handle'] + file.flush() + df = pd.read_json(file.name, orient='records', lines=True) + assert df.equals(OUTPUT_DF_WITH_EXCLUDE_COLUMNS)